Spaces:
Running on Zero
Running on Zero
| """Bounding-box and image transforms for accident detection. | |
| The original Parquet annotations are COCO-style xywh in absolute pixels for | |
| 640x640 images. This module provides conversions to the xyxy format required | |
| by torchvision detection models, plus resize-aware box rescaling. | |
| """ | |
| from __future__ import annotations | |
| from typing import Any | |
| import numpy as np | |
| import torch | |
| import torchvision.transforms.functional as F | |
| from PIL import Image | |
| class AccidentObject: | |
| """Container for one annotation with conversion helpers.""" | |
| def __init__(self, bbox_xywh: list[float], category_id: int): | |
| self.bbox_xywh = [float(v) for v in bbox_xywh] | |
| self.category_id = int(category_id) | |
| def to_xyxy(self) -> list[float]: | |
| """Convert COCO [x, y, w, h] to [x_min, y_min, x_max, y_max].""" | |
| x, y, w, h = self.bbox_xywh | |
| return [x, y, x + w, y + h] | |
| def area(self) -> float: | |
| return self.bbox_xywh[2] * self.bbox_xywh[3] | |
| def is_valid(self, image_width: int, image_height: int) -> bool: | |
| """Check that the box has positive area and lies inside the image.""" | |
| x, y, w, h = self.bbox_xywh | |
| if w <= 0 or h <= 0: | |
| return False | |
| if x < 0 or y < 0: | |
| return False | |
| if x + w > image_width or y + h > image_height: | |
| return False | |
| return True | |
| def to_dict(self) -> dict[str, Any]: | |
| return { | |
| "bbox_xywh": self.bbox_xywh, | |
| "bbox_xyxy": self.to_xyxy(), | |
| "category_id": self.category_id, | |
| "area": self.area, | |
| } | |
| def xywh_to_xyxy(boxes_xywh: np.ndarray) -> np.ndarray: | |
| """Vectorized COCO xywh -> xyxy conversion.""" | |
| boxes_xywh = np.asarray(boxes_xywh, dtype=np.float32) | |
| if boxes_xywh.size == 0: | |
| return np.zeros((0, 4), dtype=np.float32) | |
| xyxy = np.copy(boxes_xywh) | |
| xyxy[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2] | |
| xyxy[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3] | |
| return xyxy | |
| def xyxy_to_xywh(boxes_xyxy: np.ndarray) -> np.ndarray: | |
| """Vectorized xyxy -> COCO xywh conversion.""" | |
| boxes_xyxy = np.asarray(boxes_xyxy, dtype=np.float32) | |
| if boxes_xyxy.size == 0: | |
| return np.zeros((0, 4), dtype=np.float32) | |
| xywh = np.copy(boxes_xyxy) | |
| xywh[:, 2] = boxes_xyxy[:, 2] - boxes_xyxy[:, 0] | |
| xywh[:, 3] = boxes_xyxy[:, 3] - boxes_xyxy[:, 1] | |
| return xywh | |
| def scale_boxes( | |
| boxes_xyxy: np.ndarray, | |
| orig_size: tuple[int, int], | |
| new_size: tuple[int, int], | |
| ) -> np.ndarray: | |
| """Rescale absolute xyxy boxes after image resize. | |
| Args: | |
| boxes_xyxy: array of shape (N, 4) in original absolute coordinates. | |
| orig_size: (width, height) of the source image. | |
| new_size: (width, height) of the resized image. | |
| Returns: | |
| Rescaled boxes in the new image coordinate system. | |
| """ | |
| boxes_xyxy = np.asarray(boxes_xyxy, dtype=np.float32) | |
| if boxes_xyxy.size == 0: | |
| return boxes_xyxy | |
| orig_w, orig_h = orig_size | |
| new_w, new_h = new_size | |
| sx = new_w / orig_w | |
| sy = new_h / orig_h | |
| scaled = np.copy(boxes_xyxy) | |
| scaled[:, [0, 2]] *= sx | |
| scaled[:, [1, 3]] *= sy | |
| return scaled | |
| class ResizeTransform: | |
| """Resize image and rescale bounding boxes deterministically.""" | |
| def __init__(self, target_size: tuple[int, int] = (320, 320)): | |
| self.target_size = target_size # (width, height) | |
| def __call__( | |
| self, | |
| image: Image.Image, | |
| boxes_xyxy: np.ndarray, | |
| ) -> tuple[torch.Tensor, np.ndarray]: | |
| orig_size = image.size # (width, height) | |
| resized = F.resize(image, self.target_size[::-1]) # PIL expects (h, w) | |
| tensor = F.pil_to_tensor(resized).float() / 255.0 | |
| scaled = scale_boxes(boxes_xyxy, orig_size, self.target_size) | |
| return tensor, scaled | |
| def clamp_boxes(boxes_xyxy: np.ndarray, width: int, height: int) -> np.ndarray: | |
| """Clamp xyxy boxes to image bounds and ensure x2 > x1, y2 > y1.""" | |
| boxes = np.copy(boxes_xyxy) | |
| boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, width) | |
| boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, height) | |
| boxes[:, 2] = np.maximum(boxes[:, 2], boxes[:, 0] + 1e-3) | |
| boxes[:, 3] = np.maximum(boxes[:, 3], boxes[:, 1] + 1e-3) | |
| return boxes | |