Update app.py to use local model files

app.py

@@ -2,33 +2,37 @@ import gradio as gr
 from ultralytics import YOLO
 import cv2
 import numpy as np
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image
 import os
-import requests
-from io import BytesIO
 
-# Load the model with error handling
+# Load the model with proper error handling
 def load_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:
+    model_paths = [
+        'best_model.pt',
+        'tree_disease_detector.pt',
+        './best_model.pt',
+        './tree_disease_detector.pt'
+    ]
+    
+    # Try to load from local files first
+    for path in model_paths:
+        if os.path.exists(path):
             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)"
+                print(f"Loading model from {path}")
+                model = YOLO(path)
+                return model, f"Tree Disease Detection Model ({path})"
+            except Exception as e:
+                print(f"Error loading {path}: {e}")
+                continue
+    
+    # Fallback to standard YOLOv8s
+    try:
+        print("Loading standard YOLOv8s model...")
+        model = YOLO('yolov8s.pt')
+        return model, "Standard YOLOv8s Model (Fallback)"
+    except Exception as e:
+        print(f"Error loading YOLOv8s: {e}")
+        return None, "No model available"
 
 # Load model and get status
 model, model_status = load_model()
@@ -36,12 +40,20 @@ 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"""
     
+    if model is None:
+        return image, "Error: No model available"
+    
     # Convert PIL image to numpy array
     image_np = np.array(image)
     
     # Run inference
     results = model(image_np, conf=conf_threshold, iou=iou_threshold)
     
+    # 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)
+    
     # Extract detections
     detections = []
     for result in results:
@@ -51,17 +63,14 @@ 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' if model_status.startswith("Custom") else 'object'
+                    'class': 'unhealthy'
                 }
                 detections.append(detection)
     
-    # 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
-    if model_status.startswith("Custom"):
+    is_custom_model = "Tree Disease Detection Model" in model_status
+    
+    if is_custom_model:
         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"
@@ -75,10 +84,11 @@ def detect_tree_disease(image, conf_threshold=0.25, iou_threshold=0.45):
     
     return annotated_img, summary
 
-# Create example images
+# Create example images (tree images)
 example_images = [
-    ["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
+    ["https://images.pexels.com/photos/1632790/pexels-photo-1632790.jpeg", 0.25, 0.45],
+    ["https://images.pexels.com/photos/38537/woodland-road-falling-leaf-natural-38537.jpeg", 0.25, 0.45],
+    ["https://upload.wikimedia.org/wikipedia/commons/thumb/e/eb/Ash_Tree_-_geograph.org.uk_-_590710.jpg/640px-Ash_Tree_-_geograph.org.uk_-_590710.jpg", 0.25, 0.45],
 ]
 
 # Create Gradio interface
@@ -92,10 +102,9 @@ with gr.Blocks(title="Tree Disease Detection") as demo:
     **Current Model**: {model_status}
     """)
     
-    if not model_status.startswith("Custom"):
+    if "Fallback" in model_status:
         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.
+        ⚠️ **Note**: Using a fallback model. Detection will work but won't be specific to tree diseases.
         """)
     
     with gr.Row():
@@ -134,7 +143,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=False,  # Disable caching to avoid initialization issues
+        cache_examples=False,
     )
     
     gr.Markdown("""
@@ -144,6 +153,8 @@ with gr.Blocks(title="Tree Disease Detection") as demo:
     - **Dataset**: [PDT Dataset](https://huggingface.co/datasets/qwer0213/PDT_dataset)
     - **mAP50**: 0.933
     - **mAP50-95**: 0.659
+    - **Precision**: 0.878
+    - **Recall**: 0.863
     - **Classes**: 1 (unhealthy trees)
     
     ## Usage Tips
@@ -158,4 +169,5 @@ with gr.Blocks(title="Tree Disease Detection") as demo:
     """)
 
 # Launch the app
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()