app.py

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()