import gradio as gr 
import torch 
from torchvision import transforms 
from models.experimental import attempt_load
from utils.general import non_max_suppression, plot_one_box
import numpy as np 

# Load YOLOv7 model 
weights_path = "cattle.pt" 
config_path = "yolov7.yaml" 

# Initialize your YOLOv7 model here 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = attempt_load(weights_path, map_location=device)

# Define function to detect objects using YOLOv7 
def detect_objects(image): 
    global net 

    # Perform any necessary preprocessing on the image 
    transform = transforms.Compose([ 
        transforms.Resize((416, 416)), # Resize image to expected input size 
        transforms.ToTensor(), # Convert image to tensor 
    ]) 
    image_tensor = transform(image).unsqueeze(0) 

    # Perform object detection 
    with torch.no_grad(): 
        # Forward pass through the network 
        outs = net(image_tensor) 
        
        # Apply non-maximum suppression
        pred = non_max_suppression(outs, conf_thres=0.4, iou_thres=0.5, classes=None, agnostic=False)
        
        # Process detection results and draw bounding boxes
        object_count = 0
        for i, det in enumerate(pred):
            # Skip if no detections
            if len(det):
                # Increment object count
                object_count += len(det)
                # Loop over the detections
                for *xyxy, conf, cls in reversed(det):
                    # Draw bounding box
                    plot_one_box(xyxy, image, label=f'{conf:.2f}', color=(255, 0, 0), line_thickness=3)
                    
    return image, object_count 

# Create Gradio interface 
inputs = gr.inputs.Image(label="Upload Image or Video") 
outputs = [gr.outputs.Image(label="Output Image with Objects Detected"), gr.outputs.Text(label="Object Count")] 
gr.Interface(detect_objects, inputs, outputs, title="YOLOv7 Object Detection").launch()