Update app.py

app.py

@@ -8,33 +8,32 @@ from ultralytics import YOLO
 # Load YOLO model
 model = YOLO("best.pt").to("cpu")
 
-def detect_and_crop(image):
-    """Detect objects using YOLO and crop them from the image."""
+def predict(image):
+    """Detect objects using YOLO and draw bounding boxes on the image."""
     image_np = np.array(image)
+    
+    # Perform inference
     results = model(image_np, conf=0.85, device='cpu')
     
-    cropped_images = {}
+    # Draw bounding boxes
     for result in results:
         for box in result.boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])
             class_name = model.names[int(box.cls[0])]
-            cropped = image_np[y1:y2, x1:x2]
-            cropped_images[class_name] = Image.fromarray(cropped)
+            
+            # Draw rectangle and label
+            cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.putText(image_np, class_name, (x1, y1 - 10), 
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
     
-    return cropped_images
-
-def predict(image):
-    """Process image: detect objects and crop them."""
-    cropped_images = detect_and_crop(image)
-    return list(cropped_images.values()) if cropped_images else image
+    return Image.fromarray(image_np)
 
 # Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs="image",
-    # outputs="image",
-    outputs=gr.Gallery(), 
-    title="License Field Detection & Cropping"
+    outputs="image",
+    title="YOLO Object Detection"
 )
 
 iface.launch()