Update app.py

app.py

@@ -1,41 +1,38 @@
+
 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/resolve/main/1/best.pt')
-last_model = YOLO('https://huggingface.co/poudel/yolov8-cargo-package-counter/resolve/main/1/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)
-
+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)
-
-    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")
+    
+    # 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, 
+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()
+             description="Upload an image to count the number of packages detected.").launch()