import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import json

# Load your custom YOLO model
model = YOLO("fentanyl_oft.pt")

def detect_keypoints(image):
    """
    Run YOLO inference and return keypoints data
    """
    try:
        # Convert PIL Image to numpy array
        if isinstance(image, Image.Image):
            image_np = np.array(image)
            image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
        else:
            image_cv2 = image
        
        # Run inference with your exact parameters
        results = model.predict(
            source=image_cv2,
            conf=0.05,
            iou=0.7,
            max_det=1,
            imgsz=1440,
            device='cpu',
            verbose=False
        )
        
        keypoints_data = []
        
        for result in results:
            if result.keypoints is not None and len(result.keypoints.data) > 0:
                # Extract keypoints from the first detection
                kpts = result.keypoints.data[0]  # Shape: [num_keypoints, 3] (x, y, confidence)
                
                h, w = image_cv2.shape[:2]
                
                for i, (x, y, conf) in enumerate(kpts):
                    if conf > 0.3:  # Confidence threshold for individual keypoints
                        keypoints_data.append({
                            "id": int(i),
                            "x": float(x),
                            "y": float(y), 
                            "confidence": float(conf)
                        })
        
        return {
            "success": True,
            "keypoints": keypoints_data,
            "image_width": image_cv2.shape[1],
            "image_height": image_cv2.shape[0],
            "num_keypoints": len(keypoints_data)
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "keypoints": []
        }

# Create simple Gradio interface
interface = gr.Interface(
    fn=detect_keypoints,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="YOLO Keypoint Detection",
    description="Upload an image to detect keypoints using custom YOLO model"
)

if __name__ == "__main__":
    interface.launch()