import os
import zipfile
from ultralytics import YOLO
from PIL import Image
import gradio as gr

# Unzip the model file if it's not already unzipped
zip_file = "yolov8_tuned_with_weightsv4.zip"
model_file = "yolov8_tuned_with_weightsv4.pt"

if not os.path.exists(model_file):
    with zipfile.ZipFile(zip_file, "r") as zip_ref:
        zip_ref.extractall()  # Extract to current directory

# Load the YOLOv8 model
model = YOLO(model_file)

# Define the class names
class_names = {0: "Dumbbell", 1: "Kettlebell"}  # Adjust based on your dataset

# Define prediction function
def predict(image):
    # Perform inference
    results = model.predict(image)
    output_image = results[0].plot()  # Annotated image with predictions

    # Prepare detection details
    detections = [
        {
            "class": class_names.get(int(result.cls), "Unknown"),  # Map class index to name
            "confidence": float(result.conf),
            "bbox": result.xyxy.tolist(),
        }
        for result in results[0].boxes
    ]

    return output_image, detections

# Create Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an Image"),  # File upload button
    outputs=[
        gr.Image(type="pil", label="Predicted Image"),  # Display annotated image
        gr.JSON(label="Detection Details"),  # Show detection results as JSON
    ],
    title="YOLOv8 Dumbbell and Kettlebell Detection",
    description="Upload an image to detect dumbbells and kettlebells using YOLOv8.",
)

# Launch the app
if __name__ == "__main__":
    interface.launch()