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):
    # Preprocessing: Convert the colour space to RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # print("converted the colour to RGB.")
    
    # Make prediction 
    results = model(image)
    #print("ran the model")

    # Postprocessing: Convert the colour space back to BGR
    annotated_img = results[0].plot()
    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)
    # print("converted the colour to BGR.")
    
    return annotated_img
    
# Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(sources=["webcam"], type="numpy"),  # Accepts image input
    outputs="image" # Customize based on your output format
)

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