app.py

import gradio as gr
from ultralytics import YOLO

# Load YOLOv11 models from Hugging Face Hub
best_model = YOLO('https://huggingface.co/poudel/yolov8-cargo-package-counter/blob/main/best.pt')
last_model = YOLO('https://huggingface.co/poudel/yolov8-cargo-package-counter/blob/main/last.pt')

# Function to detect and count packages using the selected model
def count_packages(image, model_choice):
    # Choose the model based on user input
    if model_choice == "Best Model":
        model = best_model
    else:
        model = last_model

    # Run inference using the selected model
    results = model(image)

    # Get the number of detected boxes (packages)
    package_count = len(results[0].boxes)

    return package_count

# Gradio interface with model selection dropdown
inputs = [
    gr.Image(type="pil", label="Upload an Image"),
    gr.Radio(["Best Model", "Last Model"], label="Choose Model")
]

outputs = gr.Textbox(label="Package Count")

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