Update app.py

app.py

@@ -1,38 +1,48 @@
-import gradio as gr
-import torch
-import cv2
-import numpy as np
-from PIL import Image
-from pathlib import Path
-import subprocess
-import os
-
-# ✅ 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 (placeholder logic below)
-model_path = Path("best.pt")
-assert model_path.exists(), "best.pt model not found. Upload your trained YOLOv12 model."
-
-# Simulate a dummy model detection for placeholder purposes
-def detect_damage(image: Image.Image):
-    img = np.array(image.convert("RGB"))
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    std_dev = np.std(gray)
-    is_damaged = std_dev > 25  # Replace with actual model logic
-    return {"is_damaged": is_damaged}
-
-# 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 a YOLOv12 model."
-)
-
-demo.launch()
+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
+
+# ✅ 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"
+model = YOLO(str(model_path))
+
+# ✅ Detection function
+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)
+            if "apple_damaged" in label.lower():
+                is_damaged = True
+                break
+
+    return {"is_damaged": is_damaged}
+
+# ✅ Gradio UI
+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."
+)
+
+demo.launch()