Update app.py

app.py

@@ -1,111 +1,80 @@
 import gradio as gr
 from transformers import AutoImageProcessor, AutoModelForImageClassification
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 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
+import os
 
 # Configure Logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
 # 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)
+try:
+    processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
+    model = AutoModelForImageClassification.from_pretrained(model_name)
+    logging.info("✅ Model and processor loaded successfully.")
+except Exception as e:
+    logging.error(f"❌ Failed to load model: {str(e)}")
+    raise RuntimeError("Failed to load the model. Please check the logs for details.")
 
-# FastAPI Setup
-app = FastAPI()
+# Load or Create Disease Treatment Database
+disease_treatments_file = "disease_treatments.json"
 
-# 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.",
+if not os.path.exists(disease_treatments_file):
+    logging.warning("⚠️ Treatment database file not found. Creating a default one.")
+    default_treatments = {
+        "Powdery Mildew": "Use fungicides like sulfur or neem oil.",
+        "Leaf Blight": "Apply copper-based fungicides.",
+        "Rust": "Use resistant varieties.",
+        "Healthy": "No disease detected!",
     }
+    with open(disease_treatments_file, "w") as file:
+        json.dump(default_treatments, file)
+    logging.info("✅ Created default 'disease_treatments.json' file.")
+
+# Load Treatments
+with open(disease_treatments_file, "r") as file:
+    disease_treatments = json.load(file)
+    logging.info("✅ Treatment database loaded successfully.")
 
-# Input Validation for Image Size
+# Image Validation
 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.")
+        raise ValueError("⚠️ Image is too small. Please upload an image of at least 64x64 pixels.")
     return image
 
-# Define Prediction Function
+# Prediction Function
 def predict(image):
     try:
-        image = validate_image(Image.fromarray(np.uint8(image)).convert("RGB"))
+        image = Image.fromarray(np.uint8(image)).convert("RGB")
+        validate_image(image)
+        
         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}
+        
+        treatment = disease_treatments.get(predicted_label, "No treatment information available for this disease.")
+        return f"Predicted Disease: {predicted_label}\nTreatment: {treatment}"
+    
     except Exception as e:
         logging.error(f"Prediction failed: {str(e)}")
-        return {"error": f"Prediction failed: {str(e)}"}
+        return 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 or take a photo to identify plant diseases.",
+    fn=predict,
+    inputs=gr.Image(type="numpy", label="Upload or capture plant image"),
+    outputs=gr.Textbox(label="Result"),
+    title="Plant Disease Detector",
+    description="Upload a plant leaf image to detect diseases and get treatment suggestions.",
+    allow_flagging="never",
 )
 
-# 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__":
-    threading.Thread(target=run_fastapi, daemon=True).start()
-    iface.launch(server_name="0.0.0.0", server_port=GRADIO_PORT, share=True)
+# Launch Gradio App
+iface.launch(share=True)  # No fixed port to avoid conflicts