|
|
@@ -0,0 +1,576 @@
|
|
|
+# Copyright (c) 2025 D-Robotics Corporation
|
|
|
+#
|
|
|
+# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
+# you may not use this file except in compliance with the License.
|
|
|
+# You may obtain a copy of the License at
|
|
|
+#
|
|
|
+# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
+#
|
|
|
+# Unless required by applicable law or agreed to in writing, software
|
|
|
+# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
+# See the License for the specific language governing permissions and
|
|
|
+# limitations under the License.
|
|
|
+
|
|
|
+# flake8: noqa: E501
|
|
|
+
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
+from hbm_runtime import QuantParams
|
|
|
+from scipy.special import softmax
|
|
|
+
|
|
|
+
|
|
|
+def sigmoid(x: np.ndarray) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Compute the sigmoid activation function.
|
|
|
+ @param x Input NumPy array.
|
|
|
+ @return NumPy array after applying sigmoid function element-wise.
|
|
|
+ """
|
|
|
+ return 1.0 / (1.0 + cv2.exp(-x))
|
|
|
+
|
|
|
+
|
|
|
+def recover_to_original_size(img: np.ndarray,
|
|
|
+ orig_w: int,
|
|
|
+ orig_h: int,
|
|
|
+ resize_type: int = 1) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Restore resized image back to original size.
|
|
|
+ @details Supports direct resize or reverse letterbox removal.
|
|
|
+ @param img Input image of shape (H, W, C).
|
|
|
+ @param orig_w Original image width.
|
|
|
+ @param orig_h Original image height.
|
|
|
+ @param resize_type Resize type used before: 0 (direct) or 1 (letterbox).
|
|
|
+ @return Resized image of shape (orig_h, orig_w, C).
|
|
|
+ """
|
|
|
+ h, w = img.shape[:2] # current size after preprocess
|
|
|
+
|
|
|
+ if resize_type == 0:
|
|
|
+ # Resize directly to original dimensions
|
|
|
+ img_resized = cv2.resize(img, (orig_w, orig_h),
|
|
|
+ interpolation=cv2.INTER_NEAREST)
|
|
|
+ elif resize_type == 1:
|
|
|
+ # Remove padding and resize back from letterbox
|
|
|
+ scale = min(h / orig_h, w / orig_w)
|
|
|
+ new_w, new_h = int(orig_w * scale), int(orig_h * scale)
|
|
|
+
|
|
|
+ pad_w = w - new_w
|
|
|
+ pad_h = h - new_h
|
|
|
+ left = pad_w // 2
|
|
|
+ top = pad_h // 2
|
|
|
+
|
|
|
+ # Crop out the letterbox padding
|
|
|
+ cropped = img[top:top + new_h, left:left + new_w]
|
|
|
+
|
|
|
+ # Resize cropped region to original size
|
|
|
+ img_resized = cv2.resize(cropped, (orig_w, orig_h),
|
|
|
+ interpolation=cv2.INTER_NEAREST)
|
|
|
+ else:
|
|
|
+ raise ValueError(f"Invalid resize_type: {resize_type}, must be 0 or 1")
|
|
|
+
|
|
|
+ return img_resized
|
|
|
+
|
|
|
+
|
|
|
+def print_topk_predictions(output: np.ndarray,
|
|
|
+ idx2label: dict,
|
|
|
+ topk: int = 5) -> None:
|
|
|
+ """
|
|
|
+ @brief Print top-k classification predictions.
|
|
|
+ @details Uses softmax to compute probability and selects top-k.
|
|
|
+ @param output Raw logits as NumPy array (shape: [num_classes]).
|
|
|
+ @param idx2label Dictionary mapping class indices to labels.
|
|
|
+ @param topk Number of top predictions to display.
|
|
|
+ @return None
|
|
|
+ """
|
|
|
+ # Softmax with stability adjustment
|
|
|
+ exp_logits = np.exp(output - np.max(output))
|
|
|
+ probabilities = exp_logits / np.sum(exp_logits)
|
|
|
+
|
|
|
+ # Top-k indices
|
|
|
+ topk_idx = np.argsort(probabilities)[-topk:][::-1]
|
|
|
+ topk_prob = probabilities[topk_idx]
|
|
|
+
|
|
|
+ print(f"Top-{topk} Predictions:")
|
|
|
+ for i in range(topk):
|
|
|
+ idx = topk_idx[i]
|
|
|
+ prob = topk_prob[i]
|
|
|
+ label = idx2label[idx] if idx2label and idx in idx2label else f"Class {idx}"
|
|
|
+ print(f"{label}: {prob:.4f}")
|
|
|
+
|
|
|
+
|
|
|
+def dequantize_tensor(q_tensor: np.ndarray, quant_info: QuantParams) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Dequantize a quantized tensor to floating-point values.
|
|
|
+ @details Supports both per-tensor and per-channel dequantization based on quant_info.
|
|
|
+ @param q_tensor Quantized tensor (e.g., int8 or uint8).
|
|
|
+ @param quant_info Quantization parameters (scale, zero_point, axis, type).
|
|
|
+ @return Dequantized tensor (float32).
|
|
|
+ """
|
|
|
+ if quant_info.quant_type != 1: # 1 indicates linear scale quantization
|
|
|
+ return q_tensor
|
|
|
+
|
|
|
+ if quant_info.scale.ndim == 0 or q_tensor.ndim == 1 or quant_info.scale.size == 1:
|
|
|
+ # Per-tensor dequantization
|
|
|
+ return (q_tensor.astype(np.float32) - quant_info.zero_point.astype(np.float32)) * quant_info.scale
|
|
|
+ else:
|
|
|
+ # Per-channel dequantization
|
|
|
+ shape = [1] * q_tensor.ndim
|
|
|
+ shape[quant_info.axis] = -1
|
|
|
+ scale = quant_info.scale.reshape(shape)
|
|
|
+ zero_point = quant_info.zero_point.reshape(shape)
|
|
|
+ return (q_tensor.astype(np.float32) - zero_point.astype(np.float32)) * scale
|
|
|
+
|
|
|
+
|
|
|
+def dequantize_outputs(outputs: dict, quan_infos: dict) -> dict:
|
|
|
+ """
|
|
|
+ @brief Dequantize a dictionary of quantized model outputs.
|
|
|
+ @param outputs Dictionary of quantized output tensors.
|
|
|
+ @param quan_infos Dictionary of quantization parameters per output.
|
|
|
+ @return Dictionary of dequantized float32 outputs.
|
|
|
+ """
|
|
|
+ fp32_outputs = {}
|
|
|
+ for name, output in outputs.items():
|
|
|
+ quant_info = quan_infos[name]
|
|
|
+ fp32_outputs[name] = dequantize_tensor(output, quant_info)
|
|
|
+ return fp32_outputs
|
|
|
+
|
|
|
+
|
|
|
+def scale_coords_back(xyxy: np.ndarray,
|
|
|
+ img_w: int,
|
|
|
+ img_h: int,
|
|
|
+ input_w: int,
|
|
|
+ input_h: int,
|
|
|
+ resize_type: int = 1) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Map coordinates from resized image back to original image scale.
|
|
|
+ @param xyxy Bounding boxes (N, 4) in resized image.
|
|
|
+ @param img_w Original image width.
|
|
|
+ @param img_h Original image height.
|
|
|
+ @param input_w Network input width.
|
|
|
+ @param input_h Network input height.
|
|
|
+ @param resize_type Resize strategy: 0 (resize), 1 (letterbox).
|
|
|
+ @return Bounding boxes rescaled to original image dimensions.
|
|
|
+ """
|
|
|
+ if resize_type == 0:
|
|
|
+ # Direct resize
|
|
|
+ scale_x = img_w / input_w
|
|
|
+ scale_y = img_h / input_h
|
|
|
+ xyxy[:, [0, 2]] *= scale_x
|
|
|
+ xyxy[:, [1, 3]] *= scale_y
|
|
|
+ elif resize_type == 1:
|
|
|
+ # Letterbox resize
|
|
|
+ scale = min(input_w / img_w, input_h / img_h)
|
|
|
+ pad_w = (input_w - img_w * scale) / 2
|
|
|
+ pad_h = (input_h - img_h * scale) / 2
|
|
|
+ xyxy[:, [0, 2]] = (xyxy[:, [0, 2]] - pad_w) / scale
|
|
|
+ xyxy[:, [1, 3]] = (xyxy[:, [1, 3]] - pad_h) / scale
|
|
|
+ else:
|
|
|
+ raise ValueError("resize_type must be 0 (resize) or 1 (letterbox)")
|
|
|
+
|
|
|
+ # Clamp coordinates within valid image bounds
|
|
|
+ xyxy[:, [0, 2]] = np.clip(xyxy[:, [0, 2]], 0, img_w)
|
|
|
+ xyxy[:, [1, 3]] = np.clip(xyxy[:, [1, 3]], 0, img_h)
|
|
|
+
|
|
|
+ return xyxy
|
|
|
+
|
|
|
+
|
|
|
+def NMS(xyxy: np.ndarray,
|
|
|
+ score: np.ndarray,
|
|
|
+ cls: np.ndarray,
|
|
|
+ iou_thresh: float = 0.45) -> list:
|
|
|
+ """
|
|
|
+ @brief Perform class-wise Non-Maximum Suppression (NMS).
|
|
|
+ @details Keeps boxes with highest scores and removes overlaps above IoU threshold.
|
|
|
+ @param xyxy Bounding boxes (N, 4).
|
|
|
+ @param score Confidence scores (N,).
|
|
|
+ @param cls Class IDs for each box (N,).
|
|
|
+ @param iou_thresh IoU threshold for suppression.
|
|
|
+ @return List of indices to keep.
|
|
|
+ """
|
|
|
+ keep = []
|
|
|
+ for c in np.unique(cls):
|
|
|
+ idx = np.where(cls == c)[0]
|
|
|
+ x1, y1, x2, y2 = xyxy[idx].T
|
|
|
+ area = (x2 - x1) * (y2 - y1)
|
|
|
+ order = score[idx].argsort()[::-1] # Sort by descending score
|
|
|
+
|
|
|
+ while order.size > 0:
|
|
|
+ i = order[0]
|
|
|
+ keep.append(idx[i])
|
|
|
+ # Compute IoU with remaining boxes
|
|
|
+ xx1 = np.maximum(x1[i], x1[order[1:]])
|
|
|
+ yy1 = np.maximum(y1[i], y1[order[1:]])
|
|
|
+ xx2 = np.minimum(x2[i], x2[order[1:]])
|
|
|
+ yy2 = np.minimum(y2[i], y2[order[1:]])
|
|
|
+ inter = np.clip(xx2 - xx1, 0, None) * np.clip(yy2 - yy1, 0, None)
|
|
|
+ iou = inter / (area[i] + area[order[1:]] - inter + 1e-9)
|
|
|
+
|
|
|
+ # Keep boxes with IoU below threshold
|
|
|
+ order = order[1:][iou < iou_thresh]
|
|
|
+
|
|
|
+ return keep
|
|
|
+
|
|
|
+
|
|
|
+def xywh_to_xyxy(xywh: np.ndarray) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Convert bounding boxes from (x_center, y_center, w, h) to (x1, y1, x2, y2).
|
|
|
+ @param xywh (N, 4) array in [center_x, center_y, width, height] format.
|
|
|
+ @return (N, 4) array in [x1, y1, x2, y2] format.
|
|
|
+ """
|
|
|
+ x1y1 = xywh[:, :2] - xywh[:, 2:] / 2
|
|
|
+ x2y2 = xywh[:, :2] + xywh[:, 2:] / 2
|
|
|
+ return np.hstack([x1y1, x2y2])
|
|
|
+
|
|
|
+
|
|
|
+def filter_classification(cls_output: np.ndarray, conf_thres_raw: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Filter classification outputs using raw confidence threshold.
|
|
|
+ @param cls_output Classification logits of shape (N, C).
|
|
|
+ @param conf_thres_raw Threshold applied to max logit (before sigmoid).
|
|
|
+ @return Tuple of:
|
|
|
+ - scores: Sigmoid confidence scores of selected predictions
|
|
|
+ - ids: Class indices of selected predictions
|
|
|
+ - valid_indices: Original indices of selected predictions
|
|
|
+ """
|
|
|
+ cls_output = cls_output.reshape(-1, cls_output.shape[-1])
|
|
|
+ max_scores = np.max(cls_output, axis=1)
|
|
|
+ valid_indices = np.flatnonzero(max_scores >= conf_thres_raw)
|
|
|
+ ids = np.argmax(cls_output[valid_indices], axis=1)
|
|
|
+ # Apply sigmoid
|
|
|
+ scores = 1 / (1 + np.exp(-max_scores[valid_indices]))
|
|
|
+ return scores, ids, valid_indices
|
|
|
+
|
|
|
+
|
|
|
+def filter_mces(mces_output: np.ndarray, valid_indices: np.ndarray) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Extract MCES features from selected predictions.
|
|
|
+ @param mces_output.
|
|
|
+ @param valid_indices Indices of valid predictions.
|
|
|
+ @return Filtered MCES tensor of shape (K, D), K = len(valid_indices).
|
|
|
+ """
|
|
|
+ mces_output = mces_output.reshape(-1, mces_output.shape[-1])
|
|
|
+ mces = mces_output[valid_indices, :]
|
|
|
+ return mces
|
|
|
+
|
|
|
+
|
|
|
+def filter_predictions(pred: np.ndarray, score_thres: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Filter detection predictions by confidence threshold.
|
|
|
+ @param pred Tensor of shape (N, 5 + C): [x, y, w, h, obj_conf, class_probs...].
|
|
|
+ @param score_thres Threshold on (obj_conf * class_conf).
|
|
|
+ @return Tuple of:
|
|
|
+ - xyxy: Filtered bounding boxes (Nf, 4)
|
|
|
+ - score: Filtered scores (Nf,)
|
|
|
+ - cls: Class indices (Nf,)
|
|
|
+ """
|
|
|
+ xywh = pred[:, :4]
|
|
|
+
|
|
|
+ # Combine objectness and class scores
|
|
|
+ conf_all = pred[:, 4:5] * pred[:, 5:]
|
|
|
+ cls = conf_all.argmax(axis=1)
|
|
|
+ score = conf_all[np.arange(len(pred)), cls]
|
|
|
+ mask = score > score_thres
|
|
|
+ xyxy = xywh_to_xyxy(xywh[mask])
|
|
|
+ return xyxy, score[mask], cls[mask]
|
|
|
+
|
|
|
+
|
|
|
+def gen_anchor(grid_size: int) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Generate anchor center positions on a square grid.
|
|
|
+ @param grid_size Size of the square grid (e.g., 80 for 80x80).
|
|
|
+ @return (N, 2) array of anchor coordinates [x, y].
|
|
|
+ """
|
|
|
+ x = np.tile(np.linspace(0.5, grid_size - 0.5, grid_size), reps=grid_size)
|
|
|
+ y = np.repeat(np.linspace(0.5, grid_size - 0.5, grid_size), grid_size)
|
|
|
+ return np.stack([x, y], axis=1)
|
|
|
+
|
|
|
+
|
|
|
+def decode_boxes(boxes_output: np.ndarray,
|
|
|
+ valid_indices: np.ndarray,
|
|
|
+ grid_size: int,
|
|
|
+ stride: int,
|
|
|
+ weights_static: np.ndarray) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Decode bounding boxes from distributional predictions.
|
|
|
+ @param boxes_output Tensor of shape (N, 4 * 16).
|
|
|
+ @param valid_indices Indices of valid predictions.
|
|
|
+ @param grid_size Feature map grid size.
|
|
|
+ @param stride Downsampling factor.
|
|
|
+ @param weights_static Discrete location weights (e.g., 0~15).
|
|
|
+ @return Decoded bounding boxes in xyxy format (M, 4).
|
|
|
+ """
|
|
|
+ bboxes = boxes_output.reshape(-1, boxes_output.shape[-1])
|
|
|
+ bboxes_float32 = bboxes[valid_indices]
|
|
|
+ # Softmax over 16 bins per LTRB side and apply expectation
|
|
|
+ ltrb = np.sum(softmax(bboxes_float32.reshape(-1, 4, 16), axis=2) *
|
|
|
+ weights_static, axis=2)
|
|
|
+ anchor = gen_anchor(grid_size)[valid_indices]
|
|
|
+ x1y1 = anchor - ltrb[:, 0:2]
|
|
|
+ x2y2 = anchor + ltrb[:, 2:4]
|
|
|
+ return np.hstack([x1y1, x2y2]) * stride
|
|
|
+
|
|
|
+
|
|
|
+def decode_masks(mces: np.ndarray,
|
|
|
+ boxes: np.ndarray,
|
|
|
+ protos: np.ndarray,
|
|
|
+ input_w: int,
|
|
|
+ input_h: int,
|
|
|
+ mask_w: int,
|
|
|
+ mask_h: int,
|
|
|
+ mask_thresh: float = 0.5) -> list[np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Decode instance segmentation masks.
|
|
|
+ @param mces Mask coefficients for each detection (M, C).
|
|
|
+ @param boxes Bounding boxes (M, 4).
|
|
|
+ @param protos Mask prototype feature map (H, W, C).
|
|
|
+ @param input_w Width of the input image.
|
|
|
+ @param input_h Height of the input image.
|
|
|
+ @param mask_w Width of the mask proto.
|
|
|
+ @param mask_h Height of the mask proto.
|
|
|
+ @param mask_thresh Threshold to binarize masks.
|
|
|
+ @return List of (H, W) binary mask arrays.
|
|
|
+ """
|
|
|
+ masks = []
|
|
|
+ x_scale = mask_w / input_w
|
|
|
+ y_scale = mask_h / input_h
|
|
|
+
|
|
|
+ for i, (x1, y1, x2, y2) in enumerate(boxes):
|
|
|
+ # Crop proto features using scaled coordinates
|
|
|
+ x1_corp = int(x1 * x_scale)
|
|
|
+ y1_corp = int(y1 * y_scale)
|
|
|
+ x2_corp = int(x2 * x_scale)
|
|
|
+ y2_corp = int(y2 * y_scale)
|
|
|
+
|
|
|
+ proto_crop = protos[y1_corp:y2_corp, x1_corp:x2_corp, :] # (H, W, C)
|
|
|
+ mc = mces[i]
|
|
|
+ # Linear combination and thresholding
|
|
|
+ mask = (np.sum(proto_crop * mc[np.newaxis, np.newaxis, :], axis=2)
|
|
|
+ > mask_thresh).astype(np.uint8)
|
|
|
+ masks.append(mask)
|
|
|
+
|
|
|
+ return masks
|
|
|
+
|
|
|
+
|
|
|
+def decode_kpts(kpts_output: np.ndarray,
|
|
|
+ valid_indices: np.ndarray,
|
|
|
+ grid_size: int,
|
|
|
+ stride: int,
|
|
|
+ anchor: np.ndarray = None) -> tuple[np.ndarray, np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Decode keypoint coordinates from model output.
|
|
|
+ @param kpts_output Keypoint tensor of shape (N, 17*3).
|
|
|
+ @param valid_indices Indices of valid predictions.
|
|
|
+ @param grid_size Size of feature map grid.
|
|
|
+ @param stride Downsampling factor (e.g., 8, 16, 32).
|
|
|
+ @param anchor Optional anchor points. If None, generated automatically.
|
|
|
+ @return Tuple:
|
|
|
+ - kpts_xy: (M, 17, 2) pixel coordinates of keypoints.
|
|
|
+ - kpts_score: (M, 17, 1) keypoint confidence scores.
|
|
|
+ """
|
|
|
+ kpts_output = kpts_output.reshape(-1, kpts_output.shape[-1])
|
|
|
+ kpts = kpts_output[valid_indices].reshape(-1, 17, 3) # (M, 17, 3)
|
|
|
+
|
|
|
+ if anchor is None:
|
|
|
+ anchor = gen_anchor(grid_size)[valid_indices] # (M, 2)
|
|
|
+
|
|
|
+ # Decode x, y using anchor and stride
|
|
|
+ kpts_xy = (kpts[:, :, :2] * 2.0 + (anchor[:, None, :] - 0.5)) * stride
|
|
|
+
|
|
|
+ # Extract score without activation (or apply sigmoid optionally)
|
|
|
+ kpts_score = kpts[:, :, 2:3]
|
|
|
+
|
|
|
+ return kpts_xy, kpts_score
|
|
|
+
|
|
|
+
|
|
|
+def decode_layer(feat: np.ndarray,
|
|
|
+ stride: int,
|
|
|
+ anchor: np.ndarray,
|
|
|
+ classes_num: int = 80) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Decode a single feature layer from detection head.
|
|
|
+ @param feat Raw model output tensor of shape (1, na, h, w, c).
|
|
|
+ @param stride Stride of the feature layer.
|
|
|
+ @param anchor Anchor sizes for this layer (na, 2).
|
|
|
+ @param classes_num Number of output classes.
|
|
|
+ @return Decoded prediction array of shape (N, 5 + num_classes).
|
|
|
+ """
|
|
|
+ _, _, h, w, _ = feat.shape # h/w: feature map size
|
|
|
+
|
|
|
+ # Create coordinate grid of shape (1, 1, h, w, 2)
|
|
|
+ grid_y, grid_x = np.mgrid[0:h, 0:w]
|
|
|
+ grid = np.stack((grid_x, grid_y), axis=-1)[None, None]
|
|
|
+
|
|
|
+ # batch sigmoid
|
|
|
+ feat_sig = sigmoid(feat[..., :5 + classes_num])
|
|
|
+
|
|
|
+ # Decode center offsets (dx, dy) and size (dw, dh)
|
|
|
+ dxdy = feat_sig[..., :2]
|
|
|
+ dwdh = feat_sig[..., 2:4]
|
|
|
+ obj = feat_sig[..., 4:5]
|
|
|
+ cls = feat_sig[..., 5:]
|
|
|
+
|
|
|
+ # Compute center coordinates in original image scale
|
|
|
+ xy = (dxdy * 2. - 0.5 + grid) * stride
|
|
|
+
|
|
|
+ # Compute width/height from anchor sizes
|
|
|
+ wh = (dwdh * 2.) ** 2 * anchor[:, None, None, :]
|
|
|
+
|
|
|
+ # Construct final output tensor (xywh + obj + class scores)
|
|
|
+ out = np.empty((*xy.shape[:-1], 5 + classes_num), dtype=np.float32)
|
|
|
+ out[..., 0:2] = xy
|
|
|
+ out[..., 2:4] = wh
|
|
|
+ out[..., 4:5] = obj
|
|
|
+ out[..., 5:] = cls
|
|
|
+
|
|
|
+ return out.reshape(-1, 5 + classes_num)
|
|
|
+
|
|
|
+
|
|
|
+def decode_outputs(output_names: list[str],
|
|
|
+ fp32_outputs: dict[str, np.ndarray],
|
|
|
+ strides: list[int],
|
|
|
+ anchors: list[np.ndarray],
|
|
|
+ classes_num: int = 80) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Decode all feature maps from model output.
|
|
|
+ @param output_names List of output tensor names.
|
|
|
+ @param fp32_outputs Dict of decoded tensors from model.
|
|
|
+ @param strides Stride values for each output head.
|
|
|
+ @param anchors Anchor arrays for each head.
|
|
|
+ @param classes_num Number of output classes.
|
|
|
+ @return Concatenated prediction tensor of shape (N, 5 + classes).
|
|
|
+ """
|
|
|
+ decoded = []
|
|
|
+ for i, key in enumerate(output_names):
|
|
|
+ out = fp32_outputs[key]
|
|
|
+ h, w = out.shape[1:3]
|
|
|
+ # Reshape and transpose to (1, na, h, w, c)
|
|
|
+ feat = out.reshape(1, h, w, 3, 5 + classes_num).transpose(0, 3, 1, 2, 4)
|
|
|
+ decoded.append(decode_layer(feat, strides[i], anchors[i], classes_num))
|
|
|
+ return np.concatenate(decoded, axis=0)
|
|
|
+
|
|
|
+
|
|
|
+def get_bounding_boxes(dilated_polys: list[np.ndarray], min_area: float) -> list[np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Extract minimum area bounding boxes from polygon contours.
|
|
|
+ @param dilated_polys List of polygon contours. Each element is a NumPy array of shape (N, 1, 2).
|
|
|
+ @param min_area Minimum area threshold to filter small boxes.
|
|
|
+ @return List of bounding boxes. Each is a NumPy array of shape (4, 2), type int.
|
|
|
+ """
|
|
|
+ boxes_list = []
|
|
|
+ for cnt in dilated_polys:
|
|
|
+ if cv2.contourArea(cnt) < min_area:
|
|
|
+ continue # Skip small contours
|
|
|
+ rect = cv2.minAreaRect(cnt)
|
|
|
+ box = cv2.boxPoints(rect).astype(np.int_)
|
|
|
+ boxes_list.append(box)
|
|
|
+ return boxes_list
|
|
|
+
|
|
|
+
|
|
|
+def resize_masks_to_boxes(masks: list[np.ndarray],
|
|
|
+ boxes: list[tuple[float, float, float, float]],
|
|
|
+ img_w: int, img_h: int,
|
|
|
+ interpolation: int = cv2.INTER_LANCZOS4,
|
|
|
+ do_morph: bool = True) -> list[np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Resize binary masks to fit inside their corresponding bounding boxes.
|
|
|
+ @param masks List of binary mask arrays of shape (H, W), dtype=uint8.
|
|
|
+ @param boxes List of bounding boxes in (x1, y1, x2, y2) format.
|
|
|
+ @param img_w Width of the original image.
|
|
|
+ @param img_h Height of the original image.
|
|
|
+ @param interpolation OpenCV interpolation method used for resizing.
|
|
|
+ @param do_morph Whether to apply morphological open to smooth the mask.
|
|
|
+ @return List of resized binary masks cropped to box size.
|
|
|
+ """
|
|
|
+ resized_masks = []
|
|
|
+ for mask, (x1, y1, x2, y2) in zip(masks, boxes):
|
|
|
+ # Clamp coordinates to image bounds
|
|
|
+ x1, y1 = max(int(x1), 0), max(int(y1), 0)
|
|
|
+ x2, y2 = min(int(x2), img_w), min(int(y2), img_h)
|
|
|
+
|
|
|
+ target_w = max(x2 - x1, 1)
|
|
|
+ target_h = max(y2 - y1, 1)
|
|
|
+
|
|
|
+ resized = cv2.resize(mask, (target_w, target_h), interpolation=interpolation)
|
|
|
+
|
|
|
+ if do_morph:
|
|
|
+ # Apply morphological filtering
|
|
|
+ resized = cv2.morphologyEx(resized, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
|
|
|
+
|
|
|
+ resized_masks.append(resized)
|
|
|
+
|
|
|
+ return resized_masks
|
|
|
+
|
|
|
+
|
|
|
+def scale_keypoints_to_original_image(kpts_xy: np.ndarray,
|
|
|
+ kpts_score: np.ndarray,
|
|
|
+ boxes: list[tuple[float, float, float, float]],
|
|
|
+ img_w: int, img_h: int,
|
|
|
+ input_w: int, input_h: int,
|
|
|
+ resize_type: int = 1) -> tuple[np.ndarray, np.ndarray]:
|
|
|
+ """
|
|
|
+ @brief Scale keypoints back to original image coordinates.
|
|
|
+ @param kpts_xy Keypoint coordinates of shape (M, 17, 2), float32.
|
|
|
+ @param kpts_score Keypoint scores of shape (M, 17, 1), float32.
|
|
|
+ @param boxes List of bounding boxes, not used here.
|
|
|
+ @param img_w Width of the original image.
|
|
|
+ @param img_h Height of the original image.
|
|
|
+ @param input_w Width of model input.
|
|
|
+ @param input_h Height of model input.
|
|
|
+ @param resize_type 0 = direct resize, 1 = letterbox resize.
|
|
|
+ @return Tuple of (scaled keypoints, scores), both NumPy arrays.
|
|
|
+ """
|
|
|
+ scaled_kpts = kpts_xy.copy()
|
|
|
+
|
|
|
+ if resize_type == 0:
|
|
|
+ scale_x = img_w / input_w
|
|
|
+ scale_y = img_h / input_h
|
|
|
+ scaled_kpts[..., 0] *= scale_x
|
|
|
+ scaled_kpts[..., 1] *= scale_y
|
|
|
+
|
|
|
+ elif resize_type == 1:
|
|
|
+ scale = min(input_w / img_w, input_h / img_h)
|
|
|
+ pad_w = (input_w - img_w * scale) / 2
|
|
|
+ pad_h = (input_h - img_h * scale) / 2
|
|
|
+ scaled_kpts[..., 0] = (scaled_kpts[..., 0] - pad_w) / scale
|
|
|
+ scaled_kpts[..., 1] = (scaled_kpts[..., 1] - pad_h) / scale
|
|
|
+
|
|
|
+ else:
|
|
|
+ raise ValueError("resize_type must be 0 or 1")
|
|
|
+
|
|
|
+ # Clip to image bounds
|
|
|
+ scaled_kpts[..., 0] = np.clip(scaled_kpts[..., 0], 0, img_w)
|
|
|
+ scaled_kpts[..., 1] = np.clip(scaled_kpts[..., 1], 0, img_h)
|
|
|
+
|
|
|
+ return scaled_kpts, kpts_score
|
|
|
+
|
|
|
+
|
|
|
+def crop_and_rotate_image(img: np.ndarray, box: np.ndarray) -> np.ndarray:
|
|
|
+ """
|
|
|
+ @brief Crop and rotate a region from the image using a rotated bounding box.
|
|
|
+ @param img Input image array of shape (H, W, C), dtype=uint8.
|
|
|
+ @param box Bounding box as 4-point array of shape (4, 2).
|
|
|
+ @return Cropped and rotated region image as a NumPy array.
|
|
|
+ """
|
|
|
+ rect = cv2.minAreaRect(box)
|
|
|
+ box = cv2.boxPoints(rect).astype(np.intp)
|
|
|
+ width = int(rect[1][0])
|
|
|
+ height = int(rect[1][1])
|
|
|
+ angle = rect[2]
|
|
|
+
|
|
|
+ src_pts = box.astype("float32")
|
|
|
+ dst_pts = np.array([[0, height - 1],
|
|
|
+ [0, 0],
|
|
|
+ [width - 1, 0],
|
|
|
+ [width - 1, height - 1]], dtype="float32")
|
|
|
+
|
|
|
+ # Apply perspective transform
|
|
|
+ M = cv2.getPerspectiveTransform(src_pts, dst_pts)
|
|
|
+ warped = cv2.warpPerspective(img, M, (width, height))
|
|
|
+
|
|
|
+ # Rotate if angle is large
|
|
|
+ if angle >= 45:
|
|
|
+ rotated = cv2.rotate(warped, cv2.ROTATE_90_CLOCKWISE)
|
|
|
+ else:
|
|
|
+ rotated = warped
|
|
|
+
|
|
|
+ print("width:", rotated.shape[1], "height:", rotated.shape[0])
|
|
|
+ return rotated
|
