Upload app.py

app.py

@@ -41,15 +41,6 @@ class TrackingAlgorithm:
     SORT = "SORT (2016)"
 
 
-TRACKERS = [None, TrackingAlgorithm.BYTETRACK, TrackingAlgorithm.DEEPSORT, TrackingAlgorithm.SORT]
-VIDEO_EXAMPLES = [
-    {"path": "./examples/videos/dogs_running.mp4", "label": "Local Video", "tracker": None, "classes": "all"},
-    {"path": "./examples/videos/traffic.mp4", "label": "Local Video", "tracker": TrackingAlgorithm.BYTETRACK,
-     "classes": "car, truck, bus"},
-    {"path": "./examples/videos/fast_and_furious.mp4", "label": "Local Video", "tracker": None, "classes": "all"},
-    {"path": "./examples/videos/break_dance.mp4", "label": "Local Video", "tracker": None, "classes": "all"},
-]
-
 # Create a color palette for visualization
 # These hex color codes define different colors for tracking different objects
 color = sv.ColorPalette.from_hex([
@@ -72,17 +63,16 @@ def get_model_and_processor(checkpoint: str):
 
 @spaces.GPU(duration=20)
 def detect_objects(
-        checkpoint: str,
         images: List[np.ndarray] | np.ndarray,
-        confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
         target_size: Optional[Tuple[int, int]] = None,
-        batch_size: int = BATCH_SIZE,
-        classes: Optional[List[str]] = None,
+        batch_size: int = BATCH_SIZE
 ):
+    checkpoint = "ustc-community/dfine-xlarge-obj2coco"
+    confidence_threshold = 0.3
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model, image_processor = get_model_and_processor(checkpoint)
     model = model.to(device)
-
+    classes = ["Airplane", "Drone", "Helicopter", "Satellite", "Quadcopter", "Vehicle"]
     if classes is not None:
         wrong_classes = [cls for cls in classes if cls not in model.config.label2id]
         if wrong_classes:
@@ -92,7 +82,7 @@ def detect_objects(
         keep_ids = None
 
     if isinstance(images, np.ndarray) and images.ndim == 4:
-        images = [x for x in images]  # split video array into list of images
+        images = [x for x in images]
 
     batches = [images[i:i + batch_size] for i in range(0, len(images), batch_size)]
 
@@ -164,7 +154,8 @@ def read_video_k_frames(video_path: str, k: int, read_every_i_frame: int = 1):
     return frames
 
 
-def get_tracker(tracker: str, fps: float):
+def get_tracker(fps: float):
+    tracker = TrackingAlgorithm.BYTETRACK
     if tracker == TrackingAlgorithm.SORT:
         return trackers.SORTTracker(frame_rate=fps)
     elif tracker == TrackingAlgorithm.DEEPSORT:
@@ -272,33 +263,7 @@ def create_video_inputs() -> List[gr.components.Component]:
             interactive=True,
             format="mp4",  # Ensure MP4 format
             elem_classes="input-component",
-        ),
-        gr.Dropdown(
-            choices=CHECKPOINTS,
-            label="Select Model Checkpoint",
-            value=DEFAULT_CHECKPOINT,
-            elem_classes="input-component",
-        ),
-        gr.Dropdown(
-            choices=TRACKERS,
-            label="Select Tracker (Optional)",
-            value=None,
-            elem_classes="input-component",
-        ),
-        gr.TextArea(
-            label="Specify Class Names to Detect (comma separated)",
-            value="all",
-            lines=1,
-            elem_classes="input-component",
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=DEFAULT_CONFIDENCE_THRESHOLD,
-            step=0.1,
-            label="Confidence Threshold",
-            elem_classes="input-component",
-        ),
+        )
     ]
 
 
@@ -329,7 +294,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             with gr.Row():
                 with gr.Column(scale=1, min_width=300):
                     with gr.Group():
-                        video_input, video_checkpoint, video_tracker, video_classes, video_confidence_threshold = create_video_inputs()
+                        video_input = create_video_inputs()
                         video_detect_button, video_clear_button = create_button_row()
                 with gr.Column(scale=2):
                     video_output = gr.Video(
@@ -337,44 +302,19 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         format="mp4",  # Explicit MP4 format
                         elem_classes="output-component",
                     )
-
-            gr.Examples(
-                examples=[
-                    [example["path"], DEFAULT_CHECKPOINT, example["tracker"], example["classes"],
-                     DEFAULT_CONFIDENCE_THRESHOLD]
-                    for example in VIDEO_EXAMPLES
-                ],
-                inputs=[video_input, video_checkpoint, video_tracker, video_classes, video_confidence_threshold],
-                outputs=[video_output],
-                fn=process_video,
-                cache_examples=False,
-                label="Select a video example to populate inputs",
-            )
-
-    # Video clear button
     video_clear_button.click(
         fn=lambda: (
             None,
-            DEFAULT_CHECKPOINT,
-            None,
-            "all",
-            DEFAULT_CONFIDENCE_THRESHOLD,
             None,
         ),
         outputs=[
             video_input,
-            video_checkpoint,
-            video_tracker,
-            video_classes,
-            video_confidence_threshold,
             video_output,
         ],
     )
-
-    # Video detect button
     video_detect_button.click(
         fn=process_video,
-        inputs=[video_input, video_checkpoint, video_tracker, video_classes, video_confidence_threshold],
+        inputs=[video_input],
         outputs=[video_output],
     )