done

app.py

@@ -47,7 +47,7 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-def run_tracking(video_file, yolo_model, reid_model, tracking_method, conf_threshold):
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, target_classes, conf_threshold):
     """Run object tracking on the uploaded video."""
     try:
         # Create temporary workspace
@@ -74,6 +74,12 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, conf_thres
                 "--exist-ok"
             ]
             
+            # Add class filtering if specific classes are selected
+            if target_classes and target_classes != ["all"]:
+                classes_arg = ",".join(str(YOLO_CLASSES.index(cls)) for cls in target_classes if cls in YOLO_CLASSES)
+                if classes_arg:
+                    cmd.extend(["--classes", classes_arg])
+            
             # Special handling for OcSort
             if tracking_method == "ocsort":
                 cmd.append("--per-class")
@@ -151,19 +157,20 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, conf_thres
         return None, f"Error: {str(e)}"
 
 # Define the Gradio interface
-def process_video(video_path, yolo_model, reid_model, tracking_method, conf_threshold):
+def process_video(video_path, yolo_model, reid_model, tracking_method, target_classes, conf_threshold):
     # Validate inputs
     if not video_path:
         return None, "Please upload a video file"
     
     print(f"Processing video: {video_path}")
-    print(f"Parameters: model={yolo_model}, reid={reid_model}, tracker={tracking_method}, conf={conf_threshold}")
+    print(f"Parameters: model={yolo_model}, reid={reid_model}, tracker={tracking_method}, classes={target_classes}, conf={conf_threshold}")
     
     output_path, status = run_tracking(
         video_path,
         yolo_model,
         reid_model,
         tracking_method,
+        target_classes,
         conf_threshold
     )
     
@@ -181,6 +188,28 @@ yolo_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"]
 reid_models = ["osnet_x0_25_msmt17.pt"]
 tracking_methods = ["bytetrack", "botsort", "ocsort", "strongsort"]
 
+# YOLO COCO class names
+YOLO_CLASSES = [
+    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
+    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
+    "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie",
+    "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon",
+    "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut",
+    "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
+    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book",
+    "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
+]
+
+# Common object groups for convenience
+COMMON_CLASSES = {
+    "all": ["all"],
+    "people": ["person"],
+    "vehicles": ["bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat"],
+    "animals": ["bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe"],
+    "common_objects": ["backpack", "umbrella", "handbag", "tie", "suitcase", "cell phone", "laptop", "book"]
+}
+
 # Ensure dependencies and apply patches at startup
 ensure_dependencies()
 apply_patches()
@@ -191,7 +220,7 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
     gr.Markdown("Upload a video file to detect and track objects. Processing may take a few minutes depending on video length.")
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             input_video = gr.Video(label="Input Video", sources=["upload"])
             
             with gr.Group():
@@ -210,6 +239,22 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
                     value="bytetrack", 
                     label="Tracking Method"
                 )
+                
+                # Object category selection
+                class_category = gr.Radio(
+                    choices=list(COMMON_CLASSES.keys()),
+                    value="all",
+                    label="Object Category"
+                )
+                
+                # Individual class selection
+                target_classes = gr.CheckboxGroup(
+                    choices=YOLO_CLASSES,
+                    value=["all"],
+                    label="Target Classes",
+                    interactive=True
+                )
+                
                 conf_threshold = gr.Slider(
                     minimum=0.1, 
                     maximum=0.9, 
@@ -220,13 +265,22 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
                 
             process_btn = gr.Button("Process Video", variant="primary")
         
-        with gr.Column():
+        with gr.Column(scale=1):
             output_video = gr.Video(label="Output Video with Tracking")
             status_text = gr.Textbox(label="Status", value="Ready to process video")
     
+    # Update the target classes based on the selected category
+    def update_classes(category):
+        if category == "all":
+            return ["all"]
+        else:
+            return COMMON_CLASSES[category]
+    
+    class_category.change(fn=update_classes, inputs=class_category, outputs=target_classes)
+    
     process_btn.click(
         fn=process_video,
-        inputs=[input_video, yolo_model, reid_model, tracking_method, conf_threshold],
+        inputs=[input_video, yolo_model, reid_model, tracking_method, target_classes, conf_threshold],
         outputs=[output_video, status_text]
     )