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