File size: 3,207 Bytes
0cfefd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """图像与运动学的数据增广。"""
from __future__ import annotations
import numpy as np
import torch
# DINOv3 的 ImageNet 标准化参数
DINOV3_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
DINOV3_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
def crop_top_half(image: torch.Tensor) -> torch.Tensor:
"""裁去图像上半部分（主要是天空）。
输入 ``[3, H, W]`` 或 ``[T, 3, H, W]``；返回相同维度但 H 减半。
"""
if image.dim() == 4:
h = image.shape[2]
return image[:, :, h // 2 :, :]
elif image.dim() == 3:
h = image.shape[1]
return image[:, h // 2 :, :]
raise ValueError(f"unsupported image dim: {image.dim()}")
def normalize_image(image: torch.Tensor, mean: torch.Tensor = DINOV3_MEAN, std: torch.Tensor = DINOV3_STD) -> torch.Tensor:
"""对 [0, 1] 范围的图像做标准化。支持 ``[3,H,W]``/``[T,3,H,W]``/``[B,T,3,H,W]``。"""
while mean.dim() < image.dim():
mean = mean.unsqueeze(0)
std = std.unsqueeze(0)
return (image - mean.to(image.device, image.dtype)) / std.to(image.device, image.dtype)
def add_gaussian_noise(image: torch.Tensor, std: float = 0.01) -> torch.Tensor:
"""高斯噪声增广。``image`` 应已归一化（mean=0,std=1 之后）。"""
if std <= 0:
return image
return image + torch.randn_like(image) * std
def perturb_kinematics(
ego_6d: torch.Tensor, # [T, 6]
intr_vec: torch.Tensor, # [14]
extr_6d: torch.Tensor, # [6]
translation_std_m: float,
rotation_std_deg: float,
intrinsic_std: float,
extrinsic_std: float,
rng: np.random.Generator,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""在 Stage1 中期对运动学和内外参添加微小扰动，作为校准训练增广。
返回扰动后值与扰动量（GT 残差 = -扰动量，因为校准网络要把扰动反推回去）。
返回
----
perturbed_ego, perturbed_intr, perturbed_extr,
gt_residual_concat (在 symlog 空间作为 calibration 监督，可选；
本文件仅返回扰动后的真实空间值，校准 GT 由 trainer 构造)
"""
rot_std_rad = np.deg2rad(rotation_std_deg)
# ego 8x6
delta_ego = np.zeros_like(ego_6d.numpy())
delta_ego[:, :3] = rng.normal(0.0, translation_std_m, size=(ego_6d.shape[0], 3))
delta_ego[:, 3:] = rng.normal(0.0, rot_std_rad, size=(ego_6d.shape[0], 3))
perturbed_ego = ego_6d + torch.from_numpy(delta_ego).to(ego_6d)
# intrinsic 14
delta_intr = rng.normal(0.0, intrinsic_std, size=(intr_vec.shape[0],))
perturbed_intr = intr_vec + torch.from_numpy(delta_intr).to(intr_vec)
# extrinsic 6
delta_extr = np.zeros_like(extr_6d.numpy())
delta_extr[:3] = rng.normal(0.0, extrinsic_std, size=(3,))
delta_extr[3:] = rng.normal(0.0, rot_std_rad, size=(3,))
perturbed_extr = extr_6d + torch.from_numpy(delta_extr).to(extr_6d)
return (
perturbed_ego,
perturbed_intr,
perturbed_extr,
torch.from_numpy(np.concatenate([delta_ego.flatten(), delta_intr, delta_extr])).to(ego_6d.dtype),
)
|