app

app.py

@@ -3,16 +3,35 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
+import os
 import requests
 from io import BytesIO
 
-# Load the model
+# Load the model with error handling
 def load_model():
-    # Replace with your model path on HuggingFace
-    model = YOLO('https://huggingface.co/IsmatS/crop_desease_detection/resolve/main/best.pt')
-    return model
+    try:
+        # Try to load from HuggingFace first
+        model = YOLO('IsmatS/crop_desease_detection')
+        return model, "Custom Tree Disease Detection Model"
+    except:
+        try:
+            # Try direct URL
+            model = YOLO('https://huggingface.co/IsmatS/crop_desease_detection/resolve/main/best.pt')
+            return model, "Custom Tree Disease Detection Model"
+        except:
+            try:
+                # Try local file if exists
+                if os.path.exists('best.pt'):
+                    model = YOLO('best.pt')
+                    return model, "Custom Tree Disease Detection Model (Local)"
+            except:
+                # Fallback to standard YOLOv8s
+                print("Loading standard YOLOv8s model as fallback...")
+                model = YOLO('yolov8s.pt')
+                return model, "Standard YOLOv8s Model (Fallback)"
 
-model = load_model()
+# Load model and get status
+model, model_status = load_model()
 
 def detect_tree_disease(image, conf_threshold=0.25, iou_threshold=0.45):
     """Detect unhealthy trees in the uploaded image"""
@@ -32,75 +51,71 @@ def detect_tree_disease(image, conf_threshold=0.25, iou_threshold=0.45):
                 detection = {
                     'confidence': float(box.conf[0]),
                     'bbox': box.xyxy[0].tolist(),
-                    'class': 'unhealthy'
+                    'class': 'unhealthy' if model_status.startswith("Custom") else 'object'
                 }
                 detections.append(detection)
     
-    # Draw annotations on the image
-    annotated_img = image.copy()
-    draw = ImageDraw.Draw(annotated_img)
-    
-    # Try to use a default font, fall back to PIL default if not available
-    try:
-        font = ImageFont.truetype("arial.ttf", 20)
-    except:
-        font = ImageFont.load_default()
-    
-    for det in detections:
-        bbox = det['bbox']
-        conf = det['confidence']
-        
-        # Draw bounding box
-        draw.rectangle(bbox, outline="red", width=3)
-        
-        # Draw label with confidence
-        label = f"Unhealthy: {conf:.2f}"
-        text_bbox = draw.textbbox((bbox[0], bbox[1] - 25), label, font=font)
-        draw.rectangle(text_bbox, fill="red")
-        draw.text((bbox[0], bbox[1] - 25), label, fill="white", font=font)
+    # Get annotated image directly from results
+    annotated_img = results[0].plot()
+    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
+    annotated_img = Image.fromarray(annotated_img)
     
     # Create detection summary
-    summary = f"Detected {len(detections)} unhealthy tree(s)\n\n"
-    for i, det in enumerate(detections, 1):
-        summary += f"Tree {i}: Confidence {det['confidence']:.2f}\n"
+    if model_status.startswith("Custom"):
+        summary = f"Detected {len(detections)} unhealthy tree(s)\n\n"
+        for i, det in enumerate(detections, 1):
+            summary += f"Tree {i}: Confidence {det['confidence']:.2f}\n"
+    else:
+        summary = f"Using {model_status}\n"
+        summary += f"Detected {len(detections)} object(s)\n\n"
+        for i, det in enumerate(detections, 1):
+            summary += f"Object {i}: Confidence {det['confidence']:.2f}\n"
+    
+    summary += f"\nModel Status: {model_status}"
     
     return annotated_img, summary
 
 # Create example images
 example_images = [
-    ["https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/zidane.jpg", 0.25, 0.45],
-    # Add your own example images here
+    ["https://hips.hearstapps.com/hmg-prod/images/gettyimages-1841066-1691513468.jpg", 0.25, 0.45],  # Tree image
+    ["https://www.fs.usda.gov/Internet/FSE_MEDIA/fseprd1115588.jpg", 0.25, 0.45],  # Another tree
 ]
 
 # Create Gradio interface
 with gr.Blocks(title="Tree Disease Detection") as demo:
-    gr.Markdown("""
+    gr.Markdown(f"""
     # 🌳 Tree Disease Detection with YOLOv8
     
     This model detects unhealthy/diseased trees in aerial UAV imagery. 
     Upload an image or use one of the examples below to detect diseased trees.
     
-    **Model**: YOLOv8s trained on PDT (Pests and Diseases Tree) dataset
+    **Current Model**: {model_status}
     """)
     
+    if not model_status.startswith("Custom"):
+        gr.Markdown("""
+        ⚠️ **Note**: Currently using a fallback model. The specialized tree disease model is being updated.
+        Detection will work but won't be specific to tree diseases.
+        """)
+    
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil", label="Upload Image")
             conf_threshold = gr.Slider(
-                minimum=0.0, 
-                maximum=1.0, 
-                value=0.25, 
+                minimum=0.0,
+                maximum=1.0,
+                value=0.25,
                 step=0.05,
                 label="Confidence Threshold"
             )
             iou_threshold = gr.Slider(
-                minimum=0.0, 
-                maximum=1.0, 
-                value=0.45, 
+                minimum=0.0,
+                maximum=1.0,
+                value=0.45,
                 step=0.05,
                 label="IoU Threshold"
             )
-            detect_button = gr.Button("Detect Tree Disease")
+            detect_button = gr.Button("Detect Tree Disease", variant="primary")
         
         with gr.Column():
             output_image = gr.Image(type="pil", label="Detection Results")
@@ -119,7 +134,7 @@ with gr.Blocks(title="Tree Disease Detection") as demo:
         inputs=[input_image, conf_threshold, iou_threshold],
         outputs=[output_image, detection_summary],
         fn=detect_tree_disease,
-        cache_examples=True,
+        cache_examples=False,  # Disable caching to avoid initialization issues
     )
     
     gr.Markdown("""
@@ -143,4 +158,4 @@ with gr.Blocks(title="Tree Disease Detection") as demo:
     """)
 
 # Launch the app
-demo.launch()
+demo.launch()