import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np

# Load the YOLOv8 model
best_model = YOLO('https://huggingface.co/poudel/yolov8-cargo-package-counter/resolve/main/best.pt')

# Function to detect and count packages
def count_packages(image):
    # Convert the PIL image to a NumPy array
    image = np.array(image)
    
    # Run inference using the loaded YOLO model
    results = best_model(image)
    
    # Get the number of detected boxes (packages)
    package_count = len(results[0].boxes)
    
    # Annotate the image with the detected boxes
    annotated_image = results[0].plot()
    
    # Return the package count and the annotated image
    return package_count, annotated_image

# Gradio interface
inputs = gr.Image(type="pil", label="Upload an Image")
outputs = [gr.Textbox(label="Package Count"),
           gr.Image(type="numpy", label="Annotated Image")]

# Launch the Gradio app
gr.Interface(fn=count_packages,
             inputs=inputs,
             outputs=outputs,
             title="Cargo Package Counting App",
             description="Upload an image to count the number of packages detected.").launch()