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

# Load your model
model = YOLO("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
    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)
    # return output  # or post-process the results as needed
    results = model(image)
    # print(type(results))
    # print(results)
    annotated_img = results[0].plot()
    return annotated_img

# Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="webcam"),  # Accepts image input
    outputs="image"  # Customize based on your output format
)

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