| from __future__ import annotations |
|
|
| from collections import defaultdict |
| from typing import Dict, Iterable, List, Optional, Sequence, Tuple, Union |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| MaskType = Union[np.ndarray, torch.Tensor] |
|
|
|
|
| def _to_numpy_mask(mask: MaskType) -> np.ndarray: |
| """ |
| Convert assorted mask formats to a 2D numpy boolean array. |
| """ |
| if isinstance(mask, torch.Tensor): |
| mask_np = mask.detach().cpu().numpy() |
| else: |
| mask_np = np.asarray(mask) |
|
|
| |
| while mask_np.ndim > 2 and mask_np.shape[0] == 1: |
| mask_np = np.squeeze(mask_np, axis=0) |
| if mask_np.ndim > 2 and mask_np.shape[-1] == 1: |
| mask_np = np.squeeze(mask_np, axis=-1) |
|
|
| if mask_np.ndim != 2: |
| raise ValueError(f"Expected mask to be 2D after squeezing, got shape {mask_np.shape}") |
|
|
| return mask_np.astype(bool) |
|
|
|
|
| def _mask_to_bbox(mask: np.ndarray) -> Optional[Tuple[int, int, int, int]]: |
| """ |
| Compute a bounding box for a 2D boolean mask. |
| """ |
| if not mask.any(): |
| return None |
| rows, cols = np.nonzero(mask) |
| y_min, y_max = rows.min(), rows.max() |
| x_min, x_max = cols.min(), cols.max() |
| return x_min, y_min, x_max, y_max |
|
|
|
|
| def flatten_segments_for_batch( |
| video_id: int, |
| segments: Dict[int, Dict[int, MaskType]], |
| bbox_min_dim: int = 5, |
| ) -> Dict[str, List]: |
| """ |
| Flatten nested segmentation data into batched lists suitable for predicate |
| models or downstream visualizations. Mirrors the notebook helper but is |
| robust to differing mask dtypes/shapes. |
| """ |
| batched_object_ids: List[Tuple[int, int, int]] = [] |
| batched_masks: List[np.ndarray] = [] |
| batched_bboxes: List[Tuple[int, int, int, int]] = [] |
| frame_pairs: List[Tuple[int, int, Tuple[int, int]]] = [] |
|
|
| for frame_id, frame_objects in segments.items(): |
| valid_objects: List[int] = [] |
| for object_id, raw_mask in frame_objects.items(): |
| mask = _to_numpy_mask(raw_mask) |
| bbox = _mask_to_bbox(mask) |
| if bbox is None: |
| continue |
|
|
| x_min, y_min, x_max, y_max = bbox |
| if abs(y_max - y_min) < bbox_min_dim or abs(x_max - x_min) < bbox_min_dim: |
| continue |
|
|
| valid_objects.append(object_id) |
| batched_object_ids.append((video_id, frame_id, object_id)) |
| batched_masks.append(mask) |
| batched_bboxes.append(bbox) |
|
|
| for i in valid_objects: |
| for j in valid_objects: |
| if i == j: |
| continue |
| frame_pairs.append((video_id, frame_id, (i, j))) |
|
|
| return { |
| "object_ids": batched_object_ids, |
| "masks": batched_masks, |
| "bboxes": batched_bboxes, |
| "pairs": frame_pairs, |
| } |
|
|
|
|
| def extract_valid_object_pairs( |
| batched_object_ids: Sequence[Tuple[int, int, int]], |
| interested_object_pairs: Optional[Iterable[Tuple[int, int]]] = None, |
| ) -> List[Tuple[int, int, Tuple[int, int]]]: |
| """ |
| Filter object pairs per frame. If `interested_object_pairs` is provided, only |
| emit those combinations when both objects are present; otherwise emit all |
| permutations (i, j) with i != j for each frame. |
| """ |
| frame_to_objects: Dict[Tuple[int, int], set] = defaultdict(set) |
| for vid, fid, oid in batched_object_ids: |
| frame_to_objects[(vid, fid)].add(oid) |
|
|
| interested = ( |
| list(interested_object_pairs) |
| if interested_object_pairs is not None |
| else None |
| ) |
|
|
| valid_pairs: List[Tuple[int, int, Tuple[int, int]]] = [] |
| for (vid, fid), object_ids in frame_to_objects.items(): |
| if interested: |
| for src, dst in interested: |
| if src in object_ids and dst in object_ids: |
| valid_pairs.append((vid, fid, (src, dst))) |
| else: |
| for src in object_ids: |
| for dst in object_ids: |
| if src == dst: |
| continue |
| valid_pairs.append((vid, fid, (src, dst))) |
|
|
| return valid_pairs |
|
|
