Refactor repository

app.py

@@ -5,7 +5,6 @@ import uuid
 import logging
 
 import torch
-import spaces
 import trackers
 import numpy as np
 import gradio as gr
@@ -13,12 +12,9 @@ import imageio.v3 as iio
 import supervision as sv
 
 from pathlib import Path
-from functools import lru_cache
 from typing import List, Optional, Tuple
 
 from PIL import Image
-from transformers import AutoModelForObjectDetection, AutoImageProcessor
-from transformers.image_utils import load_image
 from pipeline import build_pipeline
 from utils import cfg, load_config, load_onnx_model
 
@@ -28,26 +24,19 @@ DETECTORS = {
     "yolo8n-640": 'downloads/yolo8n-640.onnx',
     "yolo8n-416": 'downloads/yolo8n-416.onnx',
 }
-DEFAULT_CHECKPOINT = list(DETECTORS.keys())[0]
+DEFAULT_DETECTOR = list(DETECTORS.keys())[0]
 DEFAULT_CONFIDENCE_THRESHOLD = 0.6
 
-TORCH_DTYPE = torch.float32
 
 # Image
 IMAGE_EXAMPLES = [
-    {"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,
-        "url": "https://live.staticflickr.com/65535/33021460783_1646d43c54_b.jpg",
-        "label": "Flickr Image",
-    },
+    {"path": "./examples/images/forest.jpg", "label": "Local Image"},
+    {"path": "./examples/images/river.jpg", "label": "Local Image"},
+    {"path": "./examples/images/ocean.jpg", "label": "Local Image"},
 ]
 
 # Video
 MAX_NUM_FRAMES = 250
-BATCH_SIZE = 4
 ALLOWED_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov"}
 VIDEO_OUTPUT_DIR = Path("static/videos")
 VIDEO_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
@@ -59,10 +48,8 @@ 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/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"},
+    {"path": "./examples/videos/sea.mp4", "label": "Local Video", "tracker": TrackingAlgorithm.BYTETRACK, "classes": "Person, Boat"},
+    {"path": "./examples/videos/forest.mp4", "label": "Local Video", "tracker": TrackingAlgorithm.BYTETRACK, "classes": "LightVehicle, Person, Boat"},
 ]
 
 
@@ -92,102 +79,64 @@ def detect_objects(
     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,
 ):
-
-    # device = "cuda" if torch.cuda.is_available() else "cpu"
-    # 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')
-    # config.detector.thresholds.confidence = confidence_threshold
     config.defrost()
-    config.detector.thresholds.confidence = confidence_threshold
+    config.detector.thresholds.confidence = float(confidence_threshold)
     config.freeze()
     pipeline = get_pipeline(config, onnx_path)
-    detector_category_mapping = pipeline.detector.get_category_mapping()
+    id2label = pipeline.detector.get_category_mapping()
+    label2id = {v: k for k, v in pipeline.detector.get_category_mapping().items()}
     if classes is not None:
-        wrong_classes = [cls for cls in classes if cls not in detector_category_mapping]
+        wrong_classes = [cls for cls in classes if cls not in label2id]
         if wrong_classes:
             gr.Warning(f"Classes not found in model config: {wrong_classes}")
-        keep_ids = [detector_category_mapping[cls] for cls in classes if cls in detector_category_mapping]
+        keep_ids = [label2id[cls] for cls in classes if cls in label2id]
     else:
         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
 
-    batches = [images[i:i + batch_size] for i in range(0, len(images), batch_size)]
-
     results = []
-    for batch in tqdm.tqdm(batches, desc="Processing frames"):
-
-        # preprocess images
-        # inputs = image_processor(images=batch, return_tensors="pt")
-        # inputs = inputs.to(device).to(TORCH_DTYPE)
-
-        # # forward pass
-        # with torch.no_grad():
-        #     outputs = model(**inputs)
+    for img in tqdm.tqdm(images, desc="Processing frames"):
+        output_ = pipeline(img)
+        output_reshaped = {
+        "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,
+        }
+        results.append(output_reshaped)
+        if target_size:
+            # Resize boxes to target size
+            scale_x = target_size[0] / img.shape[1]
+            scale_y = target_size[1] / img.shape[0]
+            output_reshaped["boxes"][:, [0, 2]] *= scale_x
+            output_reshaped["boxes"][:, [1, 3]] *= scale_y
         
-        batch_results = []
-        for i in range(len(batch)):
-            img = batch[i]
-            output_ = pipeline(img)
-            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,
-            }
-            batch_results.append(output)
         
 
-        # postprocess outputs
-        # if target_size:
-        #     target_sizes = [target_size] * len(batch)
-        # else:
-        #     target_sizes = [(image.shape[0], image.shape[1]) for image in batch]
-
-        # batch_results = image_processor.post_process_object_detection(
-        #     outputs, target_sizes=target_sizes, threshold=confidence_threshold
-        # )
-
-        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()}
+    for i, result in enumerate(results):
+        results[i] = {k: v 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()
 
 
 def process_image(
-    checkpoint: str = DEFAULT_CHECKPOINT,
+    model: str = DEFAULT_DETECTOR,
     image: Optional[Image.Image] = None,
-    url: Optional[str] = None,
-    use_url: bool = False,
     confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
 ):
-    if not use_url:
-        url = None
-
-    if (image is None) ^ bool(url):
-        raise ValueError(f"Either image or url must be provided, but not both.")
-
-    if url:
-        image = load_image(url)
     
-    load_config(cfg, f'configs/{checkpoint}.yaml')
+    load_config(cfg, f'configs/{model}.yaml')
     results, id2label = detect_objects(
         config=cfg.pipeline,
-        onnx_path=DETECTORS[checkpoint],
+        onnx_path=DETECTORS[model],
         images=[np.array(image)],
         confidence_threshold=confidence_threshold,
     )
@@ -297,11 +246,10 @@ def process_video(
 
     box_annotator = sv.BoxAnnotator(color, color_lookup=color_lookup, thickness=1)
     label_annotator = sv.LabelAnnotator(color, color_lookup=color_lookup, text_scale=0.5)
-    trace_annotator = sv.TraceAnnotator(color, color_lookup=color_lookup, thickness=1, trace_length=100)
 
     # preprocess classes
     if classes != "all":
-        classes_list = [cls.strip().lower() for cls in classes.split(",")]
+        classes_list = [cls.strip() for cls in classes.split(",")]
     else:
         classes_list = None
 
@@ -328,7 +276,6 @@ def process_video(
             labels = [f"#{tracker_id} {id2label[class_id]}" for class_id, tracker_id in zip(detections.class_id, detections.tracker_id)]
             annotated_frame = box_annotator.annotate(scene=frame, detections=detections)
             annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
-            annotated_frame = trace_annotator.annotate(scene=annotated_frame, detections=detections)
             annotated_frames.append(annotated_frame)
         
     else:
@@ -354,17 +301,10 @@ def create_image_inputs() -> List[gr.components.Component]:
             interactive=True,
             elem_classes="input-component",
         ),
-        gr.Checkbox(label="Use Image URL Instead", value=False),
-        gr.Textbox(
-            label="Image URL",
-            placeholder="https://example.com/image.jpg",
-            visible=False,
-            elem_classes="input-component",
-        ),
         gr.Dropdown(
             choices=list(DETECTORS.keys()),
             label="Select Model Checkpoint",
-            value=DEFAULT_CHECKPOINT,
+            value=DEFAULT_DETECTOR,
             elem_classes="input-component",
         ),
         gr.Slider(
@@ -390,7 +330,7 @@ def create_video_inputs() -> List[gr.components.Component]:
         gr.Dropdown(
             choices=list(DETECTORS.keys()),
             label="Select Model Checkpoint",
-            value=DEFAULT_CHECKPOINT,
+            value=DEFAULT_DETECTOR,
             elem_classes="input-component",
         ),
         gr.Dropdown(
@@ -434,6 +374,8 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
          - **Image** and **Video** modes are supported.
          - Select a model and adjust the confidence threshold to see detections!
          - On video mode, you can enable tracking powered by [Supervision](https://github.com/roboflow/supervision) and [Trackers](https://github.com/roboflow/trackers) from Roboflow.
+
+        For more details and source code, visit the [GitHub Repository](https://github.com/eadali/PiSAR).
         """,
         elem_classes="header-text",
     )
@@ -445,8 +387,6 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
                     with gr.Group():
                         (
                             image_input,
-                            use_url,
-                            url_input,
                             image_model_checkpoint,
                             image_confidence_threshold,
                         ) = create_image_inputs()
@@ -461,10 +401,8 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
             gr.Examples(
                 examples=[
                     [
-                        DEFAULT_CHECKPOINT,
+                        DEFAULT_DETECTOR,
                         example["path"],
-                        example["url"],
-                        example["use_url"],
                         DEFAULT_CONFIDENCE_THRESHOLD,
                     ]
                     for example in IMAGE_EXAMPLES
@@ -472,8 +410,6 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
                 inputs=[
                     image_model_checkpoint,
                     image_input,
-                    url_input,
-                    use_url,
                     image_confidence_threshold,
                 ],
                 outputs=[image_output],
@@ -501,7 +437,7 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
 
             gr.Examples(
                 examples=[
-                    [example["path"], DEFAULT_CHECKPOINT, example["tracker"], example["classes"], DEFAULT_CONFIDENCE_THRESHOLD]
+                    [example["path"], DEFAULT_DETECTOR, example["tracker"], example["classes"], DEFAULT_CONFIDENCE_THRESHOLD]
                     for example in VIDEO_EXAMPLES
                 ],
                 inputs=[video_input, video_checkpoint, video_tracker, video_classes, video_confidence_threshold],
@@ -511,27 +447,19 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
                 label="Select a video example to populate inputs",
             )
 
-    # Dynamic visibility for URL input
-    use_url.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_url,
-        outputs=url_input,
-    )
-
     # Image clear button
     image_clear_button.click(
         fn=lambda: (
             None,
             False,
             "",
-            DEFAULT_CHECKPOINT,
+            DEFAULT_DETECTOR,
             DEFAULT_CONFIDENCE_THRESHOLD,
             None,
         ),
         outputs=[
             image_input,
-            use_url,
-            url_input,
+
             image_model_checkpoint,
             image_confidence_threshold,
             image_output,
@@ -542,7 +470,7 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     video_clear_button.click(
         fn=lambda: (
             None,
-            DEFAULT_CHECKPOINT,
+            DEFAULT_DETECTOR,
             None,
             "all",
             DEFAULT_CONFIDENCE_THRESHOLD,
@@ -564,8 +492,6 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
         inputs=[
             image_model_checkpoint,
             image_input,
-            url_input,
-            use_url,
             image_confidence_threshold,
         ],
         outputs=[image_output],

examples/videos/sea.mp4

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

examples/videos/seaa.mp4

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

utils/config.py

@@ -2,29 +2,23 @@ from .yacs import CfgNode
 
 cfg = CfgNode(new_allowed=True)
 cfg.save_dir = "./"
-# common params for NETWORK
-cfg.model = CfgNode(new_allowed=True)
-cfg.model.arch = CfgNode(new_allowed=True)
-cfg.model.arch.backbone = CfgNode(new_allowed=True)
-cfg.model.arch.fpn = CfgNode(new_allowed=True)
-cfg.model.arch.head = CfgNode(new_allowed=True)
-
-# DATASET related params
-cfg.data = CfgNode(new_allowed=True)
-cfg.data.train = CfgNode(new_allowed=True)
-cfg.data.val = CfgNode(new_allowed=True)
-cfg.device = CfgNode(new_allowed=True)
-cfg.device.precision = 32
-# train
-cfg.schedule = CfgNode(new_allowed=True)
-
-# logger
-cfg.log = CfgNode()
-cfg.log.interval = 50
-
-# testing
-cfg.test = CfgNode()
-# size of images for each device
+cfg.pipeline = CfgNode(new_allowed=True)
+
+# Detector config
+cfg.pipeline.detector = CfgNode(new_allowed=True)
+cfg.pipeline.detector.model = "yolov8n"
+cfg.pipeline.detector.categories = []
+cfg.pipeline.detector.thresholds = CfgNode(new_allowed=True)
+cfg.pipeline.detector.thresholds.confidence = 0.6
+cfg.pipeline.detector.thresholds.iou = 0.4
+cfg.pipeline.detector.slicing = CfgNode(new_allowed=True)
+cfg.pipeline.detector.slicing.overlap = 0.2
+cfg.pipeline.detector.device = "cpu"
+
+# Tracker config
+cfg.pipeline.tracker = CfgNode(new_allowed=True)
+cfg.pipeline.tracker.algorithm = "dummytrack"
+
 
 
 def load_config(cfg, args_cfg):