Update app.py

app.py

@@ -5,7 +5,7 @@ from ultralytics import ASSETS, YOLO
 model = YOLO("yolo12x.pt")
 
 def predict_image(img, conf_threshold, iou_threshold):
-    """Predicts persons in an image and returns the image with detections and count."""
+    """Predicts persons and cars in an image and returns the image with detections and counts."""
     results = model.predict(
         source=img,
         conf=conf_threshold,
@@ -13,16 +13,29 @@ def predict_image(img, conf_threshold, iou_threshold):
         show_labels=True,
         show_conf=True,
         imgsz=640,
-        classes=[0]
+        classes=[0, 2]  # 0 for person, 2 for car
     )
-
+    
     for r in results:
         im_array = r.plot()
         im = Image.fromarray(im_array[..., ::-1])
-    
-    person_count = len(results[0].boxes) if results[0].boxes is not None else 0
-
-    return im, f"Number of persons detected: {person_count}"
+        
+        # Count persons and cars separately
+        person_count = 0
+        car_count = 0
+        
+        if results[0].boxes is not None:
+            for box in results[0].boxes:
+                class_id = int(box.cls[0])
+                if class_id == 0:  # person
+                    person_count += 1
+                elif class_id == 2:  # car
+                    car_count += 1
+        
+        total_count = person_count + car_count
+        count_text = f"Persons: {person_count} | Cars: {car_count} | Total: {total_count}"
+        
+    return im, count_text
 
 iface = gr.Interface(
     fn=predict_image,
@@ -33,10 +46,10 @@ iface = gr.Interface(
     ],
     outputs=[
         gr.Image(type="pil", label="Result"),
-        gr.Textbox(label="Person Count")
+        gr.Textbox(label="Detection Count")
     ],
-    title="Image Person Detection",
-    description="Upload images to detect persons and get a count",
+    title="Person and Car Detection",
+    description="Upload images to detect persons and cars with individual counts",
 )
 
 if __name__ == "__main__":