Ready for testing

.gitignore

@@ -1,6 +1,3 @@
-# Virtual environment
-pesar_env/
-
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -137,6 +134,8 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+pisar_env/
+
 
 # Spyder project settings
 .spyderproject

README.md

@@ -1,6 +1,6 @@
 ---
-title: PeSAR
-emoji: 📡 
+title: PiSAR
+emoji: 🚁 
 colorFrom: green
 colorTo: yellow
 sdk: gradio
@@ -8,7 +8,7 @@ sdk_version: 5.29.0
 app_file: app.py
 pinned: false
 license: agpl-3.0
-short_description: 'Perception for Aerial Search and Rescue'
+short_description: 'Pipelines for Aerial Search and Rescue'
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -24,32 +24,19 @@ from utils import cfg, load_config, load_onnx_model
 
 
 # Configuration constants
-CHECKPOINTS = [
-    "yolo8n-cpu",
-    # "ustc-community/dfine-medium-obj2coco",
-    # "ustc-community/dfine-medium-coco",    
-    # "ustc-community/dfine-medium-obj365",
-    # "ustc-community/dfine-nano-coco",
-    # "ustc-community/dfine-small-coco",
-    # "ustc-community/dfine-large-coco",
-    # "ustc-community/dfine-xlarge-coco",
-    # "ustc-community/dfine-small-obj365",
-    # "ustc-community/dfine-large-obj365",
-    # "ustc-community/dfine-xlarge-obj365",
-    # "ustc-community/dfine-small-obj2coco",
-    # "ustc-community/dfine-large-obj2coco-e25",
-    # "ustc-community/dfine-xlarge-obj2coco",
-]
-DEFAULT_CHECKPOINT = CHECKPOINTS[0]
-DEFAULT_CONFIDENCE_THRESHOLD = 0.3
+DETECTORS = {
+    "yolo8n-640": 'downloads/yolo8n-640.onnx',
+    "yolo8n-416": 'downloads/yolo8n-416.onnx',
+}
+DEFAULT_CHECKPOINT = list(DETECTORS.keys())[0]
+DEFAULT_CONFIDENCE_THRESHOLD = 0.6
 
 TORCH_DTYPE = torch.float32
 
 # Image
 IMAGE_EXAMPLES = [
-    {"path": "./examples/images/road.jpg", "use_url": False, "url": "", "label": "Local Image"},
-    {"path": "./examples/images/bay.jpg", "use_url": False, "url": "", "label": "Local Image"},
-    {"path": "./examples/images/airport.jpg", "use_url": False, "url": "", "label": "Local Image"},
+    {"path": "./examples/images/forest.jpg", "use_url": False, "url": "", "label": "Local Image"},
+    {"path": "./examples/images/coast.jpg", "use_url": False, "url": "", "label": "Local Image"},
     {
         "path": None,
         "use_url": True,
@@ -73,7 +60,7 @@ class TrackingAlgorithm:
 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/forest.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"},
 ]
@@ -93,16 +80,15 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-@lru_cache(maxsize=3)
-def get_model_and_processor(checkpoint: str):
-    model = AutoModelForObjectDetection.from_pretrained(checkpoint, torch_dtype=TORCH_DTYPE)
-    image_processor = AutoImageProcessor.from_pretrained(checkpoint)
-    return model, image_processor
+def get_pipeline(config: dict, onnx_path: str):
+    pipeline = build_pipeline(config)
+    load_onnx_model(pipeline.detector, onnx_path) 
+    return pipeline
 
 
-# @spaces.GPU(duration=20)
 def detect_objects(
-    checkpoint: str,
+    config: dict,
+    onnx_path: str,
     images: List[np.ndarray] | np.ndarray, 
     confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
     target_size: Optional[Tuple[int, int]] = None,
@@ -114,15 +100,20 @@ def detect_objects(
     # model, image_processor = get_model_and_processor(checkpoint)
     # model = model.to(device)
     
-    load_config(cfg, f'configs/{checkpoint}.yaml')
-    pipeline = build_pipeline(cfg.pipeline)
-    load_onnx_model(pipeline.detector, 'downloads/yolo8n-416.onnx')
-
+    # load_config(cfg, f'configs/{checkpoint}.yaml')
+    # pipeline = build_pipeline(cfg.pipeline)
+    # load_onnx_model(pipeline.detector, 'downloads/yolo8n-416.onnx')
+    # config.detector.thresholds.confidence = confidence_threshold
+    config.defrost()
+    config.detector.thresholds.confidence = confidence_threshold
+    config.freeze()
+    pipeline = get_pipeline(config, onnx_path)
+    detector_category_mapping = pipeline.detector.get_category_mapping()
     if classes is not None:
-        wrong_classes = [cls for cls in classes if cls not in model.config.label2id]
+        wrong_classes = [cls for cls in classes if cls not in detector_category_mapping]
         if wrong_classes:
             gr.Warning(f"Classes not found in model config: {wrong_classes}")
-        keep_ids = [model.config.label2id[cls] for cls in classes if cls in model.config.label2id]
+        keep_ids = [detector_category_mapping[cls] for cls in classes if cls in detector_category_mapping]
     else:
         keep_ids = None
 
@@ -142,16 +133,16 @@ def detect_objects(
         # with torch.no_grad():
         #     outputs = model(**inputs)
         
-        outputs = []
+        batch_results = []
         for i in range(len(batch)):
             img = batch[i]
             output_ = pipeline(img)
-            output_1 = {
+            output = {
             "scores": torch.from_numpy(output_.confidence) if isinstance(output_.confidence, np.ndarray) else output_.confidence,
             "labels": torch.from_numpy(output_.class_id) if isinstance(output_.class_id, np.ndarray) else output_.class_id,
             "boxes": torch.from_numpy(output_.xyxy) if isinstance(output_.xyxy, np.ndarray) else output_.xyxy,
             }
-            outputs.append(output_1)
+            batch_results.append(output)
         
 
         # postprocess outputs
@@ -163,16 +154,15 @@ def detect_objects(
         # batch_results = image_processor.post_process_object_detection(
         #     outputs, target_sizes=target_sizes, threshold=confidence_threshold
         # )
-        batch_results = outputs
 
         results.extend(batch_results)
 
-    # move results to cpu
-    for i, result in enumerate(results):
-        results[i] = {k: v.cpu() for k, v in result.items()}
-        if keep_ids is not None:
-            keep = torch.isin(results[i]["labels"], torch.tensor(keep_ids))
-            results[i] = {k: v[keep] for k, v in results[i].items()}
+    # # move results to cpu
+    # for i, result in enumerate(results):
+    #     results[i] = {k: v.cpu() for k, v in result.items()}
+    #     if keep_ids is not None:
+    #         keep = torch.isin(results[i]["labels"], torch.tensor(keep_ids))
+    #         results[i] = {k: v[keep] for k, v in results[i].items()}
     
     # return results, model.config.id2label
     return results, pipeline.detector.get_category_mapping()
@@ -193,9 +183,11 @@ def process_image(
 
     if url:
         image = load_image(url)
-
+    
+    load_config(cfg, f'configs/{checkpoint}.yaml')
     results, id2label = detect_objects(
-        checkpoint=checkpoint,
+        config=cfg.pipeline,
+        onnx_path=DETECTORS[checkpoint],
         images=[np.array(image)],
         confidence_threshold=confidence_threshold,
     )
@@ -313,9 +305,11 @@ def process_video(
     else:
         classes_list = None
 
+    load_config(cfg, f'configs/{checkpoint}.yaml')
     results, id2label = detect_objects(
+        config=cfg.pipeline,
+        onnx_path=DETECTORS[checkpoint],
         images=np.array(frames),
-        checkpoint=checkpoint,
         confidence_threshold=confidence_threshold,
         target_size=(target_height, target_width),
         classes=classes_list,
@@ -368,7 +362,7 @@ def create_image_inputs() -> List[gr.components.Component]:
             elem_classes="input-component",
         ),
         gr.Dropdown(
-            choices=CHECKPOINTS,
+            choices=list(DETECTORS.keys()),
             label="Select Model Checkpoint",
             value=DEFAULT_CHECKPOINT,
             elem_classes="input-component",
@@ -394,7 +388,7 @@ def create_video_inputs() -> List[gr.components.Component]:
             elem_classes="input-component",
         ),
         gr.Dropdown(
-            choices=CHECKPOINTS,
+            choices=list(DETECTORS.keys()),
             label="Select Model Checkpoint",
             value=DEFAULT_CHECKPOINT,
             elem_classes="input-component",

configs/yolo8n-640.yaml

@@ -0,0 +1,15 @@
+# YOLOv8n + ByteTrack Configuration
+pipeline: 
+  detector:
+    model: yolov8n
+    categories: ['LightVehicle', 'Person', 'Building', 'UPole', 'Boat', 'Bike', 'Container', 'Truck', 'Gastank', 'Digger', 'Solarpanels', 'Bus']
+    thresholds:
+      confidence: 0.6
+      iou: 0.4
+    slicing:
+      overlap: 0.2
+    device: cpu
+
+  tracker:
+    algorithm: dummytrack
+    

examples/images/road.jpg → downloads/yolo8n-640.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fe183df95295eb7a5a0e55810c1c5e46412eeb6f99029267a8cd45f4437280af
-size 5325543
+oid sha256:346359fb811146dc1793406cbad1ac5f88dddce172276f24beae789e85d5efba
+size 12259807

examples/videos/dogs_running.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:043c1a28bddcd9beeddbab946c38d1a91de2a18da3ae114d307cd01b30b64ca2
-size 22119043

examples/videos/fast_and_furious.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5980eada9d80c65b4da5b536427ccf8ff8ea2707ee3e4aa52fb2c4e1b1979dae
-size 16872922

examples/videos/{break_dance.mp4 → forest.mp4}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:747f9c2f9d19e4955603e1a13b69663187882d4c6a8fbcad18ddbd04ee792d4d
-size 1972564
+oid sha256:a6e1f240ad106045504afbd114e79387cceb11877af61f759f5371c548b6dab6
+size 14931145

examples/videos/traffic.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:71908c136bba6b50b9071fb2015553f651c91a7ee857924f33616c046011aaed
-size 8591523