Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from ultralytics import YOLO
-from PIL import Image
 
 # Load YOLOv8 models from Hugging Face Hub
 best_model = YOLO('https://huggingface.co/poudel/yolov8-cargo-package-counter/resolve/main/best.pt')
@@ -16,11 +15,11 @@ def count_packages(image, model_choice):
 
     # 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 = [
@@ -28,8 +27,7 @@ inputs = [
     gr.Radio(["Best Model", "Last Model"], label="Choose Model")
 ]
 
-outputs = [gr.Textbox(label="Package Count"),
-           gr.Image(type="pil", label="Annotated Image")]
+outputs = gr.Textbox(label="Package Count")
 
 # Launch the Gradio app
 gr.Interface(fn=count_packages, 
@@ -37,4 +35,4 @@ gr.Interface(fn=count_packages,
              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()
+             live=True).launch()