Update app.py

app.py

@@ -4,7 +4,7 @@ from PIL import Image
 import numpy as np
 import cv2
 
-# Load the YOLOv8 model (make sure the path is correct and model is downloaded)
+# Load the YOLOv8 model (ensure the model file is correct and accessible)
 try:
     model = YOLO('yolov8n.pt')
     print("Model loaded successfully.")
@@ -12,20 +12,25 @@ except Exception as e:
     print(f"Error loading model: {e}")
 
 def identify_disease(image):
-    # Convert the image to RGB if it's not
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
+    # Convert the image to RGB if it's not already in RGB
+    try:
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        print("Image converted to RGB successfully.")
+    except Exception as e:
+        print(f"Error converting image to RGB: {e}")
+        return image, [{"Disease": "Error", "Confidence": "N/A"}]
     
     # Perform inference
     try:
         results = model(image)
         predictions = results[0]
-        print("Inference completed.")
+        print("Model inference completed successfully.")
     except Exception as e:
-        print(f"Error during inference: {e}")
-        return image, [{"Disease": "Error", "Confidence": "N/A"}]
+        print(f"Error during model inference: {e}")
+        return image, [{"Disease": "Error during inference", "Confidence": "N/A"}]
     
-    # Check if there are any detections
+    # Check for detections
     if len(predictions.boxes) == 0:
         print("No detections found.")
         annotated_image = np.array(image)
@@ -34,29 +39,31 @@ def identify_disease(image):
         annotated_image = Image.fromarray(annotated_image)
         return annotated_image, [{"Disease": "None", "Confidence": "N/A"}]
     
-    # Extract predictions
-    boxes = predictions.boxes
-    labels = boxes.cls.cpu().numpy() if boxes.cls is not None else []
-    scores = boxes.conf.cpu().numpy() if boxes.conf is not None else []
-    class_names = model.names
-
-    # Annotate image with bounding boxes and labels
-    annotated_image = np.array(image)
-    for box, label, score in zip(boxes.xyxy.cpu().numpy(), labels, scores):
-        x1, y1, x2, y2 = map(int, box)
-        class_name = class_names[int(label)]
-        confidence = f"{score * 100:.2f}%"
-        annotated_image = cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        annotated_image = cv2.putText(annotated_image, f"{class_name} {confidence}", (x1, y1 - 10),
-                                      cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-    
-    # Convert annotated image back to PIL format
-    annotated_image = Image.fromarray(annotated_image)
-    
-    # Prepare results for display
-    results_list = [{"Disease": class_names[int(label)], "Confidence": f"{score * 100:.2f}%"} for label, score in zip(labels, scores)]
+    # Extract predictions and annotate image
+    try:
+        boxes = predictions.boxes
+        labels = boxes.cls.cpu().numpy() if boxes.cls is not None else []
+        scores = boxes.conf.cpu().numpy() if boxes.conf is not None else []
+        class_names = model.names
+        
+        annotated_image = np.array(image)
+        for box, label, score in zip(boxes.xyxy.cpu().numpy(), labels, scores):
+            x1, y1, x2, y2 = map(int, box)
+            class_name = class_names[int(label)]
+            confidence = f"{score * 100:.2f}%"
+            annotated_image = cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            annotated_image = cv2.putText(annotated_image, f"{class_name} {confidence}", (x1, y1 - 10),
+                                          cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        annotated_image = Image.fromarray(annotated_image)
+        print("Image annotation completed.")
+        
+        # Prepare results list for output
+        results_list = [{"Disease": class_names[int(label)], "Confidence": f"{score * 100:.2f}%"} for label, score in zip(labels, scores)]
+        return annotated_image, results_list
     
-    return annotated_image, results_list
+    except Exception as e:
+        print(f"Error during annotation: {e}")
+        return image, [{"Disease": "Error during annotation", "Confidence": "N/A"}]
 
 # Define Gradio interface with updated syntax
 interface = gr.Interface(