Update app.py

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, target_classes, conf_threshold):
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold):
     """Run object tracking on the uploaded video."""
     try:
         # Create temporary workspace
@@ -74,11 +74,18 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, target_cla
                 "--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])
+            # Add class filtering if specific classes are provided
+            if class_ids and class_ids.strip():
+                # Parse the comma-separated class IDs
+                try:
+                    # Split by comma and convert to integers to validate
+                    class_list = [int(c.strip()) for c in class_ids.split(",") if c.strip()]
+                    # Add each class ID as a separate argument
+                    if class_list:
+                        cmd.append("--classes")
+                        cmd.extend(str(c) for c in class_list)
+                except ValueError:
+                    return None, "Invalid class IDs. Please enter comma-separated numbers (e.g., '0,1,2')."
             
             # Special handling for OcSort
             if tracking_method == "ocsort":
@@ -157,20 +164,20 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, target_cla
         return None, f"Error: {str(e)}"
 
 # Define the Gradio interface
-def process_video(video_path, yolo_model, reid_model, tracking_method, target_classes, conf_threshold):
+def process_video(video_path, yolo_model, reid_model, tracking_method, class_ids, 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}, classes={target_classes}, conf={conf_threshold}")
+    print(f"Parameters: model={yolo_model}, reid={reid_model}, tracker={tracking_method}, classes={class_ids}, conf={conf_threshold}")
     
     output_path, status = run_tracking(
         video_path,
         yolo_model,
         reid_model,
         tracking_method,
-        target_classes,
+        class_ids,
         conf_threshold
     )
     
@@ -188,28 +195,6 @@ 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()
@@ -219,6 +204,28 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
     gr.Markdown("# 🚀 YOLO Object Tracking")
     gr.Markdown("Upload a video file to detect and track objects. Processing may take a few minutes depending on video length.")
     
+    # Add class reference information
+    with gr.Accordion("YOLO Class Reference", open=False):
+        gr.Markdown("""
+        # YOLO Class IDs Reference
+        
+        Enter the class IDs as comma-separated numbers in the "Target Classes" field.
+        Leave empty to track all classes.
+        
+        ## Common Class IDs:
+        - 0: person
+        - 1: bicycle
+        - 2: car
+        - 3: motorcycle
+        - 5: bus
+        - 7: truck
+        - 16: dog
+        - 17: horse
+        - 67: cell phone
+        
+        [See full COCO class list here](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datasets/coco128.yaml)
+        """)
+    
     with gr.Row():
         with gr.Column(scale=1):
             input_video = gr.Video(label="Input Video", sources=["upload"])
@@ -240,19 +247,11 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
                     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
+                # Class ID input field
+                class_ids = gr.Textbox(
+                    value="",
+                    label="Target Classes (comma-separated IDs, e.g. '0,2,5', leave empty for all classes)",
+                    placeholder="e.g. 0,2,5"
                 )
                 
                 conf_threshold = gr.Slider(
@@ -269,18 +268,9 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
             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, target_classes, conf_threshold],
+        inputs=[input_video, yolo_model, reid_model, tracking_method, class_ids, conf_threshold],
         outputs=[output_video, status_text]
     )