Add auto data-labeling pipeline

auto-data-labeling/annotate_dataset.py

@@ -0,0 +1,70 @@
+import json
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from transformers import Sam3Model, Sam3Processor
+from pathlib import Path
+import time
+from annotate_frame import process_frame
+
+def process_dataset():
+
+    # Setup paths for annotations and exports
+    input_dir = Path("./Test_Images")
+    output_dir = Path("./Test_JSON")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    json_path = output_dir / "raw_predictions.json"
+
+    # Filter out non-valid image formats
+    valid_extensions = {".png", ".jpg", ".jpeg"}
+    img_paths = [p for p in input_dir.iterdir() if p.suffix.lower() in valid_extensions]
+
+    # Load model weights ONCE
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+    print(f"Using device: {device}")
+
+    print("Loading model weights into memory...")
+    load_start = time.time()
+
+    model = Sam3Model.from_pretrained("facebook/sam3").to(device)
+    processor = Sam3Processor.from_pretrained("facebook/sam3")
+
+    load_end = time.time()
+    print(f"Model loaded in: {load_end - load_start:.2f} seconds\n")
+
+    prompts = ["vertical black pole"] # "horizontal black pole"
+
+    print(f"Running sequential processing for {[p for p in prompts]}")
+
+    preds_dct = {}
+
+    inference_start = time.time()
+
+    # Loop through each image
+    for idx, path in enumerate(img_paths, 1):
+        print(f"\n --- [{idx}/{len(img_paths)}] Processing {path.name} ---")
+
+        image = Image.open(path).convert("RGB")
+
+        keypoints = process_frame(image, processor, model, prompts, device)
+
+        preds_dct[path.name] = keypoints
+
+        print(f"--- Done with image {idx} ---")
+
+    inference_end = time.time()
+
+    print(f"\nSaving raw predictions to {json_path}")
+    with open(json_path, "w") as f:
+        json.dump(preds_dct, f, indent=4)
+
+    print(f"Dataset processing complete, ran in {inference_end - inference_start:.2f} seconds")
+    return json_path
+
+if __name__ == "__main__":
+    process_dataset()

auto-data-labeling/annotate_frame.py

@@ -0,0 +1,119 @@
+import time
+import torch
+import cv2
+import numpy as np
+from utils import (
+    keep_largest_component,
+    get_labeled_corners,
+)
+
+confidence_threshold = 0.5
+post_ratio_thresh = 0.5
+
+def process_frame(image, processor, model, prompts, device):
+    inference_start = time.time()
+
+    # Setup empty lists to accumulate results across all prompts
+    all_masks = []
+    all_boxes = []
+    all_scores = []
+
+    # Loop through each prompt
+    for prompt in prompts:
+        print(f"  -> Processing: '{prompt}'")
+        inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        results = processor.post_process_instance_segmentation(
+            outputs,
+            threshold=confidence_threshold,
+            mask_threshold=confidence_threshold,
+            target_sizes=[image.size[::-1]], 
+        )[0]
+
+        # Only append if the model actually found a matching mask for this specific prompt
+        if len(results["masks"]) > 0:
+            masks = results["masks"]
+            boxes = results["boxes"]
+            scores = results["scores"]
+
+            areas = []
+            # Convert tensor masks to numpy arrays for OpenCV
+            masks_np = masks.cpu().numpy().astype(np.uint8)
+
+            for m in masks_np:
+                # Scale 0/1 binary mask to 0/255 for cv2
+                m_255 = m * 255
+                contours, _ = cv2.findContours(m_255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+                if contours:
+                    # Find the largest contour area within this specific mask
+                    max_area = max([cv2.contourArea(c) for c in contours])
+                    areas.append(max_area)
+                else:
+                    areas.append(0.0)
+
+            # Sort indices by area in descending order
+            sorted_indices = np.argsort(areas)[::-1].tolist()
+
+            valid_indices = []
+            if len(sorted_indices) > 0:
+
+                # Always keep the largest area (the primary post)
+                largest_idx = sorted_indices[0]
+                largest_area = areas[largest_idx]
+
+                if largest_area > 0:
+                    valid_indices.append(largest_idx)
+
+                # Evaluate the second largest if it exists
+                if len(sorted_indices) > 1:
+                    second_idx = sorted_indices[1]
+                    second_area = areas[second_idx]
+
+                    # Check if the second post meets the 0.7 ratio threshold
+                    if (second_area / largest_area) >= post_ratio_thresh:
+                        valid_indices.append(second_idx)
+                    else:
+                        print(f"    -> Discarding second post. Ratio: {(second_area/largest_area):.2f} < {post_ratio_thresh}")
+
+            # Filter the tensors
+            filtered_masks = masks[valid_indices]
+            filtered_boxes = boxes[valid_indices]
+            filtered_scores = scores[valid_indices]
+
+            cleaned_masks = []
+            cleaned_boxes = []
+
+            for mask, box in zip(filtered_masks, filtered_boxes):
+                c_mask, c_box = keep_largest_component(mask, box)
+                cleaned_masks.append(c_mask)
+                cleaned_boxes.append(c_box)
+
+            filtered_masks = torch.stack(cleaned_masks)
+            filtered_boxes = torch.stack(cleaned_boxes)
+
+            all_masks.append(filtered_masks)
+            all_boxes.append(filtered_boxes)
+            all_scores.append(filtered_scores)
+
+    inference_end = time.time()
+    elapsed_time = inference_end - inference_start
+
+    print(f"Total sequential execution time: {elapsed_time:.4f} seconds")
+
+    # Combine all collected data and draw
+    if all_masks:
+        combined_boxes = torch.cat(all_boxes, dim=0)
+        combined_scores = torch.cat(all_scores, dim=0)
+
+        print(f"Found {len(torch.cat(all_masks, dim=0))} instance(s) combined")
+        print(f"Scores: {combined_scores.tolist()}")
+        print("Keypoints stored sucessfully")
+
+        return get_labeled_corners(combined_boxes, image.width)
+    else:
+        print("No instances found for any of the provided prompts.")
+        return None

auto-data-labeling/annotation_review.py

@@ -0,0 +1,103 @@
+import gradio as gr
+import json
+from pathlib import Path
+from PIL import Image, ImageDraw
+from utils import (
+    draw_JSON_kpts,
+    export_to_yolo, 
+)
+
+# Fetch JSON predictions
+json_path = Path("./Test_JSON/raw_predictions.json")
+image_dir = Path("./Test_Images")
+
+with open(json_path, "r") as f:
+    raw_preds = json.load(f)
+
+img_fn = list(raw_preds.keys())
+
+with gr.Blocks(title="Underwater Gate Pose Annotator") as app:
+
+    # State variable to keep track of index
+    curr_idx = gr.State(0)
+
+    # Image display
+    with gr.Row():
+        img_display = gr.Image(type="pil", interactive=True, label="Current Gate Frame")
+
+    # Corner selections for adjustment
+    with gr.Row():
+        corner_selector = gr.Radio(
+            choices=["TL", "TR", "BL", "BR"], 
+            value="TL", 
+            label="Corner to Adjust", 
+            interactive=True
+        )
+
+    # Accept button
+    with gr.Row():
+        btn_accept = gr.Button("Accept & Export")
+
+    # ---------------------------------------
+    #           HELPER FUNCTIONS
+    # ---------------------------------------
+
+    # Display each set of predictions on its corresponding image
+    def render_image(index):
+        if index >= len(img_fn):
+            return Image.open("alldone.png").convert("RGB")
+        filename = img_fn[index]
+        img_path = image_dir / filename
+        kpts = raw_preds[filename]
+        return draw_JSON_kpts(img_path, kpts)
+
+    # Logic for updating keypoints
+    def update_kp(index, corner, evt: gr.SelectData):
+        x, y = evt.index
+
+        filename = img_fn[index]
+
+        raw_preds[filename][corner] = [x, y]
+
+        return render_image(index)
+
+    # Accepting the corner points
+    def accept_and_next(index):
+        if index >= len(img_fn):
+            return index, render_image(index)
+
+        filename = img_fn[index]
+        points = raw_preds[filename]
+
+        img_path = image_dir / filename
+        with Image.open(img_path) as img:
+            img_width, img_height = img.size
+
+        export_to_yolo(filename, points, img_width, img_height)
+
+        # Move onto next image
+        new_index = index + 1
+        new_img = render_image(new_index)
+
+        return new_index, new_img
+
+    # ---------------------------------------
+    #            BUTTON CONFIGS
+    # ---------------------------------------
+
+    app.load(fn=render_image, inputs=curr_idx, outputs=img_display)
+
+    img_display.select(
+        fn=update_kp,
+        inputs=[curr_idx, corner_selector],
+        outputs=img_display,
+    )
+
+    btn_accept.click(
+        fn=accept_and_next,
+        inputs=curr_idx,
+        outputs=[curr_idx, img_display]
+    )
+
+if __name__ == "__main__":
+    app.launch()