app.py

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)
    print("converted the colour to RGB.")
    
    # 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)
    print("ran the model")
    annotated_img = results[0].plot()
    print("got annotated img")
    print("type annotated img:", type(annotated_img))
    
    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()