import gradio as gr
import torch
from PIL import Image
import torchvision.transforms as T
from ultralytics import YOLO
import cv2
import numpy as np

# Load the PT model
model = YOLO("Model_IV.pt")
checkpoint = torch.load("Model_IV.pt")

# Define preprocessing
transform = T.Compose([
    T.Resize((224, 224)),  # Adjust to your model's input size
    T.ToTensor(),
])

def predict(image):
    # Preprocess the image by converting the colour space to RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
    
    # # Make prediction
    # with torch.no_grad():
    #     output = model(img_tensor)
    
    # Process output (adjust based on your model's format)
    results = model(image)
    annotated_img = results[0].plot()
    return annotated_img
    
# Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(sources=["webcam"], type="pil"),  # Accepts image input
    outputs="image" # Customize based on your output format
)

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