Update app.py

app.py

@@ -1,13 +1,111 @@
 import gradio as gr
-from server import predict
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image, UnidentifiedImageError
+import torch
+import numpy as np
+from fastapi import FastAPI, UploadFile, File, HTTPException
+import uvicorn
+import io
+import json
+import threading
+import os
+import logging
+import signal
+import sys
 
+# Configure Logging
+logging.basicConfig(level=logging.INFO)
+
+# Load Model & Processor
+model_name = "linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification"
+processor = AutoImageProcessor.from_pretrained(model_name, use_fast=False)  # Avoid fast mode issue
+model = AutoModelForImageClassification.from_pretrained(model_name)
+
+# FastAPI Setup
+app = FastAPI()
+
+# Load Disease Treatment Database Dynamically
+disease_treatments = {}
+try:
+    with open("disease_treatments.json", "r") as file:
+        disease_treatments = json.load(file)
+except FileNotFoundError:
+    logging.warning("Treatment database file not found. Using default treatments.")
+    disease_treatments = {
+        "Powdery Mildew": "Use fungicides like sulfur or neem oil. Improve air circulation.",
+        "Leaf Blight": "Apply copper-based fungicides and ensure proper plant spacing.",
+        "Rust": "Use resistant varieties and apply organic sulfur fungicide.",
+        "Healthy": "No disease detected! Keep maintaining proper watering and soil health.",
+    }
+
+# Input Validation for Image Size
+def validate_image(image):
+    if image.size[0] < 64 or image.size[1] < 64:
+        raise ValueError("Image is too small. Please upload an image with dimensions at least 64x64.")
+    return image
+
+# Define Prediction Function
+def predict(image):
+    try:
+        image = validate_image(Image.fromarray(np.uint8(image)).convert("RGB"))
+        inputs = processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+        predicted_label = model.config.id2label[predicted_class_idx]
+        treatment = disease_treatments.get(predicted_label, "No treatment information available.")
+        return {"Disease": predicted_label, "Treatment": treatment}
+    except Exception as e:
+        logging.error(f"Prediction failed: {str(e)}")
+        return {"error": f"Prediction failed: {str(e)}"}
+
+# Gradio Interface
 iface = gr.Interface(
     fn=lambda img: (lambda pred: f"Predicted Disease: {pred.get('Disease', 'Error')}\nRecommended Treatment: {pred.get('Treatment', 'N/A')}")(predict(img)),
     inputs=gr.Image(type="numpy"),
     outputs="text",
     title="🌿 Plant Disease Detector",
-    description="Upload an image to identify plant diseases.",
+    description="Upload an image or take a photo to identify plant diseases.",
 )
 
+# FastAPI Endpoint
+@app.post("/predict")
+async def api_predict(file: UploadFile = File(...)):
+    logging.info(f"Received file: {file.filename}")
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        image_array = np.array(image)
+        prediction = predict(image_array)
+        logging.info(f"Prediction successful: {prediction}")
+        return {"prediction": prediction}
+    except UnidentifiedImageError:
+        logging.error("Invalid image file uploaded.")
+        raise HTTPException(status_code=400, detail="Invalid image file. Please upload a valid image.")
+    except ValueError as ve:
+        logging.error(f"ValueError during prediction: {str(ve)}")
+        raise HTTPException(status_code=400, detail=f"Invalid input: {str(ve)}")
+    except Exception as e:
+        logging.error(f"Internal server error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+
+# Configurable Ports
+FASTAPI_PORT = int(os.getenv("FASTAPI_PORT", 7860))
+GRADIO_PORT = int(os.getenv("GRADIO_PORT", 7862))  # Changed from 7861 to 7862
+
+# Graceful Shutdown
+def shutdown(signum, frame):
+    logging.info("Shutting down servers...")
+    sys.exit(0)
+
+signal.signal(signal.SIGINT, shutdown)
+signal.signal(signal.SIGTERM, shutdown)
+
+# Run FastAPI & Gradio Together
+def run_fastapi():
+    uvicorn.run(app, host="0.0.0.0", port=FASTAPI_PORT)
+
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7861, share=True)
+    threading.Thread(target=run_fastapi, daemon=True).start()
+    iface.launch(server_name="0.0.0.0", server_port=GRADIO_PORT, share=True)