Update app.py

app.py

@@ -1,48 +1,37 @@
 import gradio as gr
-import torch
-import cv2
 import numpy as np
 from PIL import Image
-from pathlib import Path
-import subprocess
-import os
 from ultralytics import YOLO
+from pathlib import Path
 
-# ✅ Install flash_attn from local wheel (must be uploaded to the space)
-whl_file = "flash_attn-2.7.3+cu11torch2.2cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
-if os.path.exists(whl_file):
-    subprocess.run(["pip", "install", whl_file])
-else:
-    print("⚠️ .whl file for flash_attn not found. Please upload it to the space.")
-
-# ✅ Load YOLOv12 model
-model_path = Path("runs/detect/train7/weights/best.pt")
-assert model_path.exists(), "Model not found at runs/detect/train7/weights/best.pt"
+# ✅ Load YOLOv12 model (make sure this path matches your repo)
+model_path = Path("best.pt")  # Or change to "weights/best.pt" if uploaded in subfolder
+assert model_path.exists(), f"Model not found at {model_path}"
 model = YOLO(str(model_path))
 
-# ✅ Detection function
+# ✅ Damage detection using YOLOv12
 def detect_damage(image: Image.Image):
     results = model(image)
-
+    
     is_damaged = False
     for result in results:
         for box in result.boxes:
             class_id = int(box.cls[0])
             label = model.names[class_id]
-            print("Detected:", label)
+            print("Detected label:", label)
             if "apple_damaged" in label.lower():
                 is_damaged = True
                 break
 
     return {"is_damaged": is_damaged}
 
-# ✅ Gradio UI
+# ✅ Gradio interface
 demo = gr.Interface(
     fn=detect_damage,
     inputs=gr.Image(type="pil"),
     outputs=gr.JSON(label="Detection Result"),
     title="🍎 Apple Damage Detector",
-    description="Upload an image of an apple to detect whether it is damaged using your trained YOLOv12 model."
+    description="Upload an apple image to detect whether it is damaged using a YOLOv12 model."
 )
 
 demo.launch()