Shaun ZhouandClaude Opus 5
af22dc639c
- yolo: images/ 与同层 labels/ 的路径约定、`class_id cx cy w h` 归一化行解析 （兼容分割多边形取外接矩形）、data.yaml 读写；缺失或空 txt 即无缺陷样本 - film_defect: 抛出图像、框、图像级标签；label_mode 决定多热向量还是互斥下标 - transforms: 图框同步的 Resize / Letterbox / CenterCrop / RandomResizedCrop / 翻转 / ColorJitter / GaussianBlur / PerImageStandardize - target_encoder: 高斯半径与 heatmap / wh / offset / mask 监督信号编码 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
112 lines
3.8 KiB
Python
112 lines
3.8 KiB
Python
"""把框标注编码成 CenterHead 需要的监督信号。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def gaussian_radius(height: float, width: float, min_overlap: float = 0.7) -> float:
|
|
"""CenterNet 的半径公式：保证半径内的框与 GT 的 IoU >= min_overlap。"""
|
|
a1 = 1.0
|
|
b1 = height + width
|
|
c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
|
|
r1 = (b1 - np.sqrt(max(b1**2 - 4 * a1 * c1, 0.0))) / (2 * a1)
|
|
|
|
a2 = 4.0
|
|
b2 = 2 * (height + width)
|
|
c2 = (1 - min_overlap) * width * height
|
|
r2 = (b2 - np.sqrt(max(b2**2 - 4 * a2 * c2, 0.0))) / (2 * a2)
|
|
|
|
a3 = 4 * min_overlap
|
|
b3 = -2 * min_overlap * (height + width)
|
|
c3 = (min_overlap - 1) * width * height
|
|
r3 = (-b3 + np.sqrt(max(b3**2 - 4 * a3 * c3, 0.0))) / (2 * a3)
|
|
|
|
return max(0.0, min(r1, r2, r3))
|
|
|
|
|
|
def draw_gaussian(heatmap: np.ndarray, center: tuple[int, int], radius: int) -> None:
|
|
"""在 heatmap 上叠加二维高斯（取 max，避免相邻目标互相削弱）。"""
|
|
diameter = 2 * radius + 1
|
|
sigma = diameter / 6.0
|
|
grid = np.arange(-radius, radius + 1, dtype=np.float32)
|
|
gaussian = np.exp(-(grid[None, :] ** 2 + grid[:, None] ** 2) / (2 * sigma**2))
|
|
|
|
cx, cy = center
|
|
height, width = heatmap.shape
|
|
left, right = min(cx, radius), min(width - cx, radius + 1)
|
|
top, bottom = min(cy, radius), min(height - cy, radius + 1)
|
|
if right <= -left or bottom <= -top:
|
|
return
|
|
|
|
masked_heatmap = heatmap[cy - top : cy + bottom, cx - left : cx + right]
|
|
masked_gaussian = gaussian[radius - top : radius + bottom, radius - left : radius + right]
|
|
np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap)
|
|
|
|
|
|
class CenterTargetEncoder:
|
|
"""生成 heatmap / wh / offset / mask 四个监督张量。
|
|
|
|
Args:
|
|
num_classes: 类别数。
|
|
stride: 模型输出相对输入的下采样倍数。
|
|
max_objects: 单图最多编码的目标数。
|
|
min_overlap: 高斯半径的 IoU 阈值。
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_classes: int,
|
|
stride: int,
|
|
max_objects: int = 64,
|
|
min_overlap: float = 0.7,
|
|
) -> None:
|
|
self.num_classes = num_classes
|
|
self.stride = stride
|
|
self.max_objects = max_objects
|
|
self.min_overlap = min_overlap
|
|
|
|
def __call__(
|
|
self,
|
|
boxes: np.ndarray,
|
|
labels: np.ndarray,
|
|
image_size: tuple[int, int],
|
|
) -> dict[str, torch.Tensor]:
|
|
img_h, img_w = image_size
|
|
out_h, out_w = img_h // self.stride, img_w // self.stride
|
|
|
|
heatmap = np.zeros((self.num_classes, out_h, out_w), dtype=np.float32)
|
|
wh = np.zeros((2, out_h, out_w), dtype=np.float32)
|
|
offset = np.zeros((2, out_h, out_w), dtype=np.float32)
|
|
mask = np.zeros((out_h, out_w), dtype=np.float32)
|
|
|
|
count = 0
|
|
for box, label in zip(np.asarray(boxes, dtype=np.float32), np.asarray(labels), strict=True):
|
|
if count >= self.max_objects:
|
|
break
|
|
x1, y1, x2, y2 = box / self.stride
|
|
box_w, box_h = x2 - x1, y2 - y1
|
|
if box_w <= 0 or box_h <= 0:
|
|
continue
|
|
|
|
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
|
|
cx_int, cy_int = int(cx), int(cy)
|
|
if not (0 <= cx_int < out_w and 0 <= cy_int < out_h):
|
|
continue
|
|
|
|
radius = max(1, int(gaussian_radius(box_h, box_w, self.min_overlap)))
|
|
draw_gaussian(heatmap[int(label)], (cx_int, cy_int), radius)
|
|
|
|
wh[:, cy_int, cx_int] = (box_w, box_h)
|
|
offset[:, cy_int, cx_int] = (cx - cx_int, cy - cy_int)
|
|
mask[cy_int, cx_int] = 1.0
|
|
count += 1
|
|
|
|
return {
|
|
"heatmap": torch.from_numpy(heatmap),
|
|
"wh": torch.from_numpy(wh),
|
|
"offset": torch.from_numpy(offset),
|
|
"mask": torch.from_numpy(mask),
|
|
}
|