| from __future__ import annotations |
|
|
| from typing import Iterable, List, Sequence, Tuple |
|
|
| import cv2 |
| import numpy as np |
|
|
|
|
| |
| CLASS_NAMES: Tuple[str, ...] = ( |
| "plane", "ship", "storage tank", "baseball diamond", "tennis court", |
| "basketball court", "ground track field", "harbor", "bridge", |
| "large vehicle", "small vehicle", "helicopter", "roundabout", |
| "soccer ball field", "swimming pool", |
| ) |
|
|
| |
| CLASS_COLORS: Tuple[Tuple[int, int, int], ...] = ( |
| (0, 215, 255), (0, 140, 255), (255, 200, 0), (255, 120, 200), |
| (90, 230, 0), (255, 90, 140), (200, 200, 0), (128, 0, 255), |
| (0, 255, 180), (255, 0, 0), (0, 0, 255), (60, 76, 231), |
| (255, 255, 0), (0, 165, 255), (147, 20, 255), |
| ) |
|
|
| LABEL_FG = (255, 255, 255) |
| FONT = cv2.FONT_HERSHEY_SIMPLEX |
|
|
|
|
| def letterbox( |
| image: np.ndarray, |
| input_size: int = 1024, |
| pad_value: int = 114, |
| ) -> Tuple[np.ndarray, float, int, int]: |
| """Ultralytics-style letterbox: keep ratio, centered gray padding.""" |
| height, width = image.shape[:2] |
| ratio = min(input_size / width, input_size / height) |
| resized_w, resized_h = round(width * ratio), round(height * ratio) |
| resized = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR) |
| pad_w, pad_h = input_size - resized_w, input_size - resized_h |
| left, top = pad_w // 2, pad_h // 2 |
| canvas = cv2.copyMakeBorder( |
| resized, |
| top, |
| pad_h - top, |
| left, |
| pad_w - left, |
| cv2.BORDER_CONSTANT, |
| value=(pad_value, pad_value, pad_value), |
| ) |
| return canvas, ratio, left, top |
|
|
|
|
| def preprocess_image( |
| image: np.ndarray, |
| input_size: int = 1024, |
| pad_value: int = 114, |
| ) -> Tuple[np.ndarray, float, int, int]: |
| canvas, ratio, pad_x, pad_y = letterbox(image, input_size, pad_value) |
| rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB) |
| tensor = (rgb.astype(np.float32) / 255.0)[None, :] |
| return tensor, ratio, pad_x, pad_y |
|
|
|
|
| def xywhr2xyxyxyxy(boxes: np.ndarray) -> np.ndarray: |
| """Convert [cx, cy, w, h, angle] boxes to (N, 4, 2) corner points.""" |
| boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 5) |
| center = boxes[:, :2] |
| w, h, angle = boxes[:, 2:3], boxes[:, 3:4], boxes[:, 4:5] |
| cos, sin = np.cos(angle), np.sin(angle) |
| vec1 = np.concatenate([w / 2 * cos, w / 2 * sin], axis=-1) |
| vec2 = np.concatenate([-h / 2 * sin, h / 2 * cos], axis=-1) |
| pt1 = center + vec1 + vec2 |
| pt2 = center + vec1 - vec2 |
| pt3 = center - vec1 - vec2 |
| pt4 = center - vec1 + vec2 |
| return np.stack([pt1, pt2, pt3, pt4], axis=1) |
|
|
|
|
| def _covariance_matrix(boxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """Gaussian covariance components (a, b, c) for xywhr boxes.""" |
| w2 = boxes[:, 2] ** 2 / 12.0 |
| h2 = boxes[:, 3] ** 2 / 12.0 |
| angle = boxes[:, 4] |
| cos2 = np.cos(angle) ** 2 |
| sin2 = np.sin(angle) ** 2 |
| a = w2 * cos2 + h2 * sin2 |
| b = w2 * sin2 + h2 * cos2 |
| c = (w2 - h2) * np.cos(angle) * np.sin(angle) |
| return a, b, c |
|
|
|
|
| def batch_probiou(obb1: np.ndarray, obb2: np.ndarray, eps: float = 1e-7) -> np.ndarray: |
| """Probabilistic IoU between two sets of xywhr boxes -> shape (N, M). |
| |
| Mirrors ``ultralytics.utils.metrics.batch_probiou`` so on-device NMS |
| matches the official OBB post-processing. |
| """ |
| obb1 = np.asarray(obb1, dtype=np.float32).reshape(-1, 5) |
| obb2 = np.asarray(obb2, dtype=np.float32).reshape(-1, 5) |
| x1, y1 = obb1[:, 0:1], obb1[:, 1:2] |
| x2, y2 = obb2[None, :, 0], obb2[None, :, 1] |
| a1, b1, c1 = (v[:, None] for v in _covariance_matrix(obb1)) |
| a2, b2, c2 = (v[None, :] for v in _covariance_matrix(obb2)) |
|
|
| denom = (a1 + a2) * (b1 + b2) - (c1 + c2) ** 2 + eps |
| t1 = (((a1 + a2) * (y1 - y2) ** 2 + (b1 + b2) * (x1 - x2) ** 2) / denom) * 0.25 |
| t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / denom) * 0.5 |
| det = np.sqrt( |
| np.clip(a1 * b1 - c1 ** 2, 0, None) * np.clip(a2 * b2 - c2 ** 2, 0, None) |
| ) |
| t3 = np.log(((a1 + a2) * (b1 + b2) - (c1 + c2) ** 2) / (4 * det + eps) + eps) * 0.5 |
| bd = np.clip(t1 + t2 + t3, eps, 100.0) |
| hd = np.sqrt(1.0 - np.exp(-bd) + eps) |
| return 1.0 - hd |
|
|
|
|
| def obb_nms(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> List[int]: |
| """Greedy NMS over xywhr boxes using probabilistic IoU.""" |
| if len(boxes) == 0: |
| return [] |
| order = scores.argsort()[::-1] |
| keep: List[int] = [] |
| while order.size: |
| i = int(order[0]) |
| keep.append(i) |
| if order.size == 1: |
| break |
| rest = order[1:] |
| iou = batch_probiou(boxes[i:i + 1], boxes[rest])[0] |
| order = rest[iou <= iou_thres] |
| return keep |
|
|
|
|
| def obb_postprocess( |
| prediction: np.ndarray, |
| ratio: float, |
| pad_x: int, |
| pad_y: int, |
| original_shape: Sequence[int], |
| conf_thres: float = 0.25, |
| iou_thres: float = 0.45, |
| max_det: int = 300, |
| class_names: Sequence[str] = CLASS_NAMES, |
| class_agnostic: bool = False, |
| ) -> List[dict]: |
| """Decode output0 `[1, 4+nc+1, anchors]` into oriented boxes. |
| |
| Channel order matches the Ultralytics OBB head: xywh, class scores, angle. |
| """ |
| prediction = np.asarray(prediction, dtype=np.float32) |
| if prediction.ndim == 3: |
| if prediction.shape[0] != 1: |
| raise ValueError(f"Only batch size 1 is supported, got {prediction.shape}") |
| prediction = prediction[0] |
| if prediction.shape[0] < prediction.shape[1]: |
| prediction = prediction.T |
| if prediction.ndim != 2 or prediction.shape[1] < 6: |
| raise ValueError(f"Unexpected prediction shape: {prediction.shape}") |
|
|
| nc = prediction.shape[1] - 5 |
| class_scores = prediction[:, 4:4 + nc] |
| class_ids = class_scores.argmax(axis=1) |
| scores = class_scores[np.arange(class_scores.shape[0]), class_ids] |
| mask = scores >= conf_thres |
| if not np.any(mask): |
| return [] |
|
|
| xywh = prediction[mask, :4] |
| angle = prediction[mask, 4 + nc] |
| scores = scores[mask] |
| class_ids = class_ids[mask] |
| rboxes = np.concatenate([xywh, angle[:, None]], axis=1) |
|
|
| |
| nms_boxes = rboxes.copy() |
| if not class_agnostic: |
| offsets = class_ids.astype(np.float32) * 7680.0 |
| nms_boxes[:, 0] += offsets |
| nms_boxes[:, 1] += offsets |
| keep = obb_nms(nms_boxes, scores, iou_thres)[:max_det] |
|
|
| boxes = rboxes[keep].copy() |
| boxes[:, 0] = (boxes[:, 0] - pad_x) / ratio |
| boxes[:, 1] = (boxes[:, 1] - pad_y) / ratio |
| boxes[:, 2] = boxes[:, 2] / ratio |
| boxes[:, 3] = boxes[:, 3] / ratio |
| corners = xywhr2xyxyxyxy(boxes) |
| corners[..., 0].clip(0, original_shape[1], out=corners[..., 0]) |
| corners[..., 1].clip(0, original_shape[0], out=corners[..., 1]) |
|
|
| detections: List[dict] = [] |
| for row, idx in enumerate(keep): |
| cls_id = int(class_ids[idx]) |
| detections.append({ |
| "box": boxes[row].astype(np.float32), |
| "corners": corners[row].astype(np.float32), |
| "score": float(scores[idx]), |
| "class_id": cls_id, |
| "class_name": class_names[cls_id] if cls_id < len(class_names) else str(cls_id), |
| }) |
| return detections |
|
|
|
|
| def draw_label(result: np.ndarray, text: str, x: int, bottom: int, color) -> None: |
| """Filled label above ``(x, bottom)``; drops inside the image when it clips.""" |
| scale, thickness, pad = 0.5, 1, 4 |
| (text_w, text_h), _ = cv2.getTextSize(text, FONT, scale, thickness) |
| height = text_h + 2 * pad |
| top = bottom - height |
| if top < 0: |
| top = bottom |
| right = min(result.shape[1] - 1, x + text_w + 2 * pad) |
| cv2.rectangle(result, (x, top), (right, top + height), color, -1, cv2.LINE_AA) |
| cv2.putText(result, text, (x + pad, top + text_h + pad), FONT, scale, |
| LABEL_FG, thickness, cv2.LINE_AA) |
|
|
|
|
| def draw_obb_res(image: np.ndarray, detections: Iterable[dict]) -> np.ndarray: |
| result = image.astype(np.uint8).copy() |
| for i, det in enumerate(detections, start=1): |
| corners = np.asarray(det["corners"], dtype=np.float32) |
| score = float(det["score"]) |
| name = det["class_name"] |
| color = CLASS_COLORS[det["class_id"] % len(CLASS_COLORS)] |
| points = np.round(corners).astype(np.int32).reshape(-1, 1, 2) |
| cv2.polylines(result, [points], True, color, 2, cv2.LINE_AA) |
| top_x, top_y = np.round(corners[np.argmin(corners[:, 1])]).astype(int) |
| draw_label(result, f"{name} {score:.2f}", int(top_x), int(top_y), color) |
| print(i, det["class_id"], name, [round(float(v), 2) for v in det["box"]], score) |
| return result |
|
|
