Blablablab/audio-classification
0
1"""2CV Export Utilities3 4Shared helper functions for computer vision export formats (COCO, YOLO, VOC).5"""6 7from typing import Dict, List, Tuple, Any, Optional8import logging9 10logger = logging.getLogger(__name__)11 12 13def build_category_mapping(annotations: List[dict], schemas: List[dict]) -> Dict[str, int]:14 """15 Build a mapping from label names to integer category IDs.16 17 Extracts labels from image_annotation schemas first (preserving config order),18 then discovers any additional labels from annotations.19 20 Args:21 annotations: List of annotation records22 schemas: List of annotation_scheme config dicts23 24 Returns:25 Dict mapping label name -> integer ID (starting from 1 for COCO, 0-indexed for YOLO)26 """27 labels = []28 seen = set()29 30 # First, collect labels from schema configs (preserves defined order)31 for schema in schemas:32 if schema.get("annotation_type") == "image_annotation":33 for label_def in schema.get("labels", []):34 name = label_def if isinstance(label_def, str) else label_def.get("name", "")35 if name and name not in seen:36 labels.append(name)37 seen.add(name)38 39 # Then discover any labels in annotation data not already in config40 for ann in annotations:41 for schema_name, img_annotations in ann.get("image_annotations", {}).items():42 if not isinstance(img_annotations, list):43 continue44 for obj in img_annotations:45 label = obj.get("label", "")46 if label and label not in seen:47 labels.append(label)48 seen.add(label)49 50 return {name: idx for idx, name in enumerate(labels)}51 52 53def polygon_to_bbox(points: List[List[float]]) -> Tuple[float, float, float, float]:54 """55 Compute axis-aligned bounding box from a polygon.56 57 Args:58 points: List of [x, y] coordinate pairs59 60 Returns:61 Tuple of (x_min, y_min, width, height)62 """63 if not points:64 return (0, 0, 0, 0)65 66 xs = [p[0] for p in points]67 ys = [p[1] for p in points]68 x_min = min(xs)69 y_min = min(ys)70 return (x_min, y_min, max(xs) - x_min, max(ys) - y_min)71 72 73def polygon_area(points: List[List[float]]) -> float:74 """75 Compute the area of a polygon using the shoelace formula.76 77 Args:78 points: List of [x, y] coordinate pairs79 80 Returns:81 Absolute area of the polygon82 """83 n = len(points)84 if n < 3:85 return 0.086 area = 0.087 for i in range(n):88 j = (i + 1) % n89 area += points[i][0] * points[j][1]90 area -= points[j][0] * points[i][1]91 return abs(area) / 2.092 93 94def normalize_bbox(x: float, y: float, w: float, h: float,95 img_w: float, img_h: float) -> Tuple[float, float, float, float]:96 """97 Normalize bounding box coordinates to [0, 1] range.98 99 Args:100 x, y: Top-left corner coordinates101 w, h: Width and height102 img_w, img_h: Image dimensions103 104 Returns:105 Tuple of (center_x, center_y, width, height) normalized to [0, 1]106 """107 if img_w <= 0 or img_h <= 0:108 return (0, 0, 0, 0)109 cx = max(0.0, min(1.0, (x + w / 2) / img_w))110 cy = max(0.0, min(1.0, (y + h / 2) / img_h))111 nw = max(0.0, min(1.0, w / img_w))112 nh = max(0.0, min(1.0, h / img_h))113 return (cx, cy, nw, nh)114 115 116def flatten_polygon(points: List[List[float]]) -> List[float]:117 """118 Flatten a list of [x, y] points into a flat coordinate list [x1, y1, x2, y2, ...].119 120 This is the format used by COCO segmentation.121 122 Args:123 points: List of [x, y] coordinate pairs124 125 Returns:126 Flat list of coordinates127 """128 result = []129 for p in points:130 result.extend(p[:2])131 return result132 133 134def extract_image_annotations(annotation: dict) -> List[Tuple[str, List[dict]]]:135 """136 Extract image annotation objects from an annotation record.137 138 Args:139 annotation: Single annotation record with image_annotations field140 141 Returns:142 List of (schema_name, annotation_objects) tuples143 """144 results = []145 for schema_name, objects in annotation.get("image_annotations", {}).items():146 if isinstance(objects, list) and objects:147 results.append((schema_name, objects))148 return results149 150 151def get_image_dimensions(item: dict, default_width: int = 0,152 default_height: int = 0) -> Tuple[int, int]:153 """154 Extract image dimensions from item metadata.155 156 Checks common field names for image width/height.157 158 Args:159 item: Item data dict160 default_width: Fallback width161 default_height: Fallback height162 163 Returns:164 Tuple of (width, height)165 """166 # Check common field patterns167 width = default_width168 for w_key in ("image_width", "width", "img_width", "w"):169 if w_key in item:170 try:171 width = int(item[w_key])172 except (ValueError, TypeError):173 pass174 break175 176 height = default_height177 for h_key in ("image_height", "height", "img_height", "h"):178 if h_key in item:179 try:180 height = int(item[h_key])181 except (ValueError, TypeError):182 pass183 break184 185 return (width, height)186 187 188def get_image_filename(item: dict) -> Optional[str]:189 """190 Extract image filename from item data.191 192 Args:193 item: Item data dict194 195 Returns:196 Image filename/path string or None197 """198 for key in ("image", "image_path", "image_url", "file_name", "filename", "img"):199 if key in item and item[key]:200 return str(item[key])201 return None202 203 204# ---------------------------------------------------------------------------205# RLE mask utilities (Potato RLE <-> COCO RLE conversion)206# ---------------------------------------------------------------------------207 208 209def decode_rle(rle: dict, width: int, height: int) -> List[int]:210 """211 Decode Potato RLE-encoded mask to a flat binary array (row-major order).212 213 Potato RLE stores counts alternating between 0-pixels and 1-pixels,214 starting with 0s, in row-major (left-to-right, top-to-bottom) order.215 216 Args:217 rle: Dict with 'counts' (list of ints) and 'size' [height, width]218 width: Image width219 height: Image height220 221 Returns:222 Flat list of 0/1 values in row-major order223 """224 counts = rle.get("counts", [])225 total = width * height226 mask = [0] * total227 pos = 0228 val = 0229 for count in counts:230 for _ in range(count):231 if pos < total:232 mask[pos] = val233 pos += 1234 val = 1 - val235 return mask236 237 238def rle_bbox(mask: List[int], width: int, height: int) -> List[float]:239 """240 Compute axis-aligned bounding box [x, y, w, h] from a flat binary mask.241 242 Args:243 mask: Flat list of 0/1 values (row-major)244 width: Image width245 height: Image height246 247 Returns:248 [x_min, y_min, bbox_width, bbox_height] or [0, 0, 0, 0] if empty249 """250 x_min, y_min = width, height251 x_max, y_max = -1, -1252 for i, val in enumerate(mask):253 if val:254 y = i // width255 x = i % width256 if x < x_min:257 x_min = x258 if x > x_max:259 x_max = x260 if y < y_min:261 y_min = y262 if y > y_max:263 y_max = y264 if x_max < 0:265 return [0, 0, 0, 0]266 return [float(x_min), float(y_min),267 float(x_max - x_min + 1), float(y_max - y_min + 1)]268 269 270def rle_area(mask: List[int]) -> int:271 """272 Compute mask area as the count of foreground pixels.273 274 Args:275 mask: Flat list of 0/1 values276 277 Returns:278 Number of 1-pixels279 """280 return sum(mask)281 282 283def _column_major_rle_counts(mask_2d: List[List[int]], height: int,284 width: int) -> List[int]:285 """286 Read a 2D mask in column-major order and compute RLE counts.287 288 Counts alternate between 0-pixels and 1-pixels, starting with 0s.289 290 Args:291 mask_2d: 2D list [height][width] of 0/1 values292 height: Image height293 width: Image width294 295 Returns:296 List of integer run counts in column-major order297 """298 counts: List[int] = []299 current_val = 0300 current_run = 0301 302 for x in range(width):303 for y in range(height):304 pixel = mask_2d[y][x]305 if pixel == current_val:306 current_run += 1307 else:308 counts.append(current_run)309 current_val = 1 - current_val310 current_run = 1311 counts.append(current_run)312 return counts313 314 315def _encode_coco_rle_string(counts: List[int]) -> str:316 """317 Encode RLE integer counts as a COCO compressed ASCII string.318 319 Implements the exact algorithm from pycocotools maskApi.c rleToString():320 - Delta encoding for i > 2: x = counts[i] - counts[i-2]321 - Each value encoded as 6-bit groups (5 data bits + 1 continuation bit)322 - Each group offset by 48 to produce printable ASCII323 - Signed values supported via arithmetic right shift324 325 Args:326 counts: List of integer run counts327 328 Returns:329 Encoded ASCII string330 """331 chars = []332 for i, cnt in enumerate(counts):333 # Delta encoding: for i > 2, encode difference from counts[i-2]334 x = cnt - counts[i - 2] if i > 2 else cnt335 while True:336 c = x & 0x1F337 x >>= 5338 # If bit 4 set, sign bit is 1 → more groups unless x is all-ones (-1)339 # If bit 4 clear, sign bit is 0 → more groups unless x is all-zeros (0)340 if c & 0x10:341 more = (x != -1)342 else:343 more = (x != 0)344 if more:345 c |= 0x20346 chars.append(chr(c + 48))347 if not more:348 break349 return "".join(chars)350 351 352def rle_to_coco_rle(rle: dict, width: int, height: int) -> Dict[str, Any]:353 """354 Convert Potato RLE to COCO RLE format.355 356 Potato RLE is row-major; COCO RLE is column-major with compressed357 ASCII string encoding.358 359 Args:360 rle: Potato RLE dict with 'counts' and 'size'361 width: Image width362 height: Image height363 364 Returns:365 COCO RLE dict {"counts": "encoded_string", "size": [height, width]}366 """367 # Decode to flat row-major mask368 flat = decode_rle(rle, width, height)369 370 # Reshape to 2D371 mask_2d = []372 for y in range(height):373 row = flat[y * width:(y + 1) * width]374 mask_2d.append(row)375 376 # Compute column-major RLE counts377 col_counts = _column_major_rle_counts(mask_2d, height, width)378 379 # Encode as COCO compressed string380 encoded = _encode_coco_rle_string(col_counts)381 382 return {"counts": encoded, "size": [height, width]}383