656 lines
21 KiB
Python
Executable File
656 lines
21 KiB
Python
Executable File
import json
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from train_mono3d import resolve_data_yaml_for_roi
|
|
from ultralytics.data.dataset import Ground3DCalibrationError, YOLOGround3DDataset
|
|
from ultralytics.data.ground3d_augment import read_calib_from_path
|
|
from ultralytics.utils import YAML
|
|
|
|
|
|
def write_dataset_yaml(path: Path) -> None:
|
|
YAML.save(
|
|
file=path,
|
|
data={
|
|
"path": "/tmp/dataset",
|
|
"train": "train.txt",
|
|
"val": "val.txt",
|
|
"class_map": {"car": 0},
|
|
"default_roi": "roi0",
|
|
"roi_configs": {
|
|
"roi0": {
|
|
"roi": [1920, 880],
|
|
"virtual_fx": 537,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
"roi1": {
|
|
"roi": [768, 352],
|
|
"virtual_fx": 537,
|
|
"virtual_camera_prob": 0.5,
|
|
"virtual_camera_val_zoom": True,
|
|
"crop_center_mode": "vxvy",
|
|
},
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
def write_clip_level_camera4(calib_dir: Path, image_size: tuple[int, int], focal_u: float = 50.0) -> Path:
|
|
camera4_file = calib_dir / "L2_calib" / "camera4.json"
|
|
camera4_file.parent.mkdir(parents=True, exist_ok=True)
|
|
camera4_file.write_text(
|
|
json.dumps(
|
|
{
|
|
"focal_u": focal_u,
|
|
"focal_v": focal_u,
|
|
"cu": image_size[0] / 2,
|
|
"cv": image_size[1] / 2,
|
|
"pitch": 0.0,
|
|
"distort_coeffs": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
return camera4_file
|
|
|
|
|
|
def create_ground3d_dataset(
|
|
tmp_path: Path,
|
|
image_sizes: list[tuple[int, int]],
|
|
imgsz: tuple[int, int] = (64, 32),
|
|
roi: tuple[int, int] | None = None,
|
|
ori_img_size: tuple[int, int] | None = None,
|
|
) -> tuple[YOLOGround3DDataset, list[str]]:
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_labels = []
|
|
image_files = []
|
|
roi = roi or imgsz
|
|
ori_img_size = ori_img_size or imgsz
|
|
|
|
for idx, image_size in enumerate(image_sizes, start=1):
|
|
rel_label = Path(f"labels/seq{idx}/frame_{idx:04d}.txt")
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / f"seq{idx}" / f"frame_{idx:04d}.png"
|
|
clip_calib_dir = gt_root / "calib" / f"seq{idx}"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text("car 0.5 0.5 0.25 0.25 0\n", encoding="utf-8")
|
|
Image.new("RGB", image_size, color=(32, 64, 96)).save(image_file)
|
|
write_clip_level_camera4(clip_calib_dir, image_size)
|
|
|
|
rel_labels.append(rel_label.as_posix())
|
|
image_files.append(str(image_file.resolve()))
|
|
|
|
(gt_root / "train.txt").write_text("\n".join(rel_labels) + "\n", encoding="utf-8")
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=list(imgsz),
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"roi": list(roi),
|
|
"ori_img_size": list(ori_img_size),
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
return dataset, image_files
|
|
|
|
|
|
def test_resolve_data_yaml_for_roi_uses_default_roi(tmp_path):
|
|
data_yaml = tmp_path / "mono3d_ground.yaml"
|
|
write_dataset_yaml(data_yaml)
|
|
|
|
resolved_path, selected_roi = resolve_data_yaml_for_roi(str(data_yaml), None)
|
|
|
|
assert selected_roi == "roi0"
|
|
assert resolved_path != str(data_yaml)
|
|
|
|
resolved_cfg = YAML.load(resolved_path)
|
|
assert resolved_cfg["roi"] == [1920, 880]
|
|
assert resolved_cfg["virtual_camera_prob"] == -1.0
|
|
assert resolved_cfg["crop_center_mode"] == "cxvy"
|
|
assert "default_roi" not in resolved_cfg
|
|
assert "roi_configs" not in resolved_cfg
|
|
|
|
|
|
def test_resolve_data_yaml_for_roi_supports_explicit_override(tmp_path):
|
|
data_yaml = tmp_path / "mono3d_ground.yaml"
|
|
write_dataset_yaml(data_yaml)
|
|
|
|
resolved_path, selected_roi = resolve_data_yaml_for_roi(str(data_yaml), "roi1")
|
|
|
|
assert selected_roi == "roi1"
|
|
|
|
resolved_cfg = YAML.load(resolved_path)
|
|
assert resolved_cfg["roi"] == [768, 352]
|
|
assert resolved_cfg["virtual_camera_prob"] == 0.5
|
|
assert resolved_cfg["virtual_camera_val_zoom"] is True
|
|
assert resolved_cfg["crop_center_mode"] == "vxvy"
|
|
|
|
|
|
def test_resolve_data_yaml_for_roi_uses_unique_temp_paths(tmp_path):
|
|
data_yaml = tmp_path / "mono3d_ground.yaml"
|
|
write_dataset_yaml(data_yaml)
|
|
|
|
resolved_path_a, selected_roi_a = resolve_data_yaml_for_roi(str(data_yaml), "roi1")
|
|
resolved_path_b, selected_roi_b = resolve_data_yaml_for_roi(str(data_yaml), "roi1")
|
|
|
|
assert selected_roi_a == "roi1"
|
|
assert selected_roi_b == "roi1"
|
|
assert resolved_path_a != resolved_path_b
|
|
|
|
|
|
def test_resolve_data_yaml_for_roi_rejects_unknown_preset(tmp_path):
|
|
data_yaml = tmp_path / "mono3d_ground.yaml"
|
|
write_dataset_yaml(data_yaml)
|
|
|
|
with pytest.raises(ValueError, match="Available presets: roi0, roi1"):
|
|
resolve_data_yaml_for_roi(str(data_yaml), "roi2")
|
|
|
|
|
|
def test_resolve_data_yaml_for_roi_rejects_missing_required_ground3d_fields(tmp_path):
|
|
data_yaml = tmp_path / "mono3d_ground.yaml"
|
|
write_dataset_yaml(data_yaml)
|
|
data_cfg = YAML.load(data_yaml)
|
|
del data_cfg["roi_configs"]["roi1"]["crop_center_mode"]
|
|
YAML.save(data_yaml, data_cfg)
|
|
|
|
with pytest.raises(ValueError, match="crop_center_mode"):
|
|
resolve_data_yaml_for_roi(str(data_yaml), "roi1")
|
|
|
|
|
|
def test_ground3d_dataset_resolves_gt_list_to_image_and_calib(tmp_path):
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_label = Path("labels/seq0/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / "seq0" / "frame_0001.png"
|
|
clip_calib_dir = gt_root / "calib" / "seq0"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text("car 0.5 0.5 0.25 0.25 0\n", encoding="utf-8")
|
|
Image.new("RGB", (64, 32), color=(32, 64, 96)).save(image_file)
|
|
write_clip_level_camera4(clip_calib_dir, (64, 32))
|
|
(gt_root / "train.txt").write_text(f"{rel_label.as_posix()}\n", encoding="utf-8")
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
assert len(dataset.labels) == 1
|
|
assert dataset.labels[0] == (str(gt_root.resolve()), rel_label.as_posix())
|
|
|
|
raw_calib = read_calib_from_path(
|
|
str(image_file.resolve()),
|
|
image_root=image_root,
|
|
extra_calib_candidates=[str((gt_root / "calib" / "seq0" / "frame_0001.json").resolve())],
|
|
)
|
|
assert raw_calib["focal_u"] == 50.0
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
assert sample["im_file"] == str(image_file.resolve())
|
|
assert sample["img"].shape[:2] == (32, 64)
|
|
assert sample["calib"]["fx"] == pytest.approx(50.0)
|
|
|
|
|
|
def test_ground3d_dataset_prefers_label_root_calibration_over_image_root(tmp_path):
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_label = Path("labels/seq0/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / "seq0" / "frame_0001.png"
|
|
label_calib_dir = gt_root / "calib" / "seq0"
|
|
image_calib_file = image_root / "calib" / "seq0" / "frame_0001.json"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_calib_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text("car 0.5 0.5 0.25 0.25 0\n", encoding="utf-8")
|
|
Image.new("RGB", (64, 32), color=(32, 64, 96)).save(image_file)
|
|
write_clip_level_camera4(label_calib_dir, (64, 32), focal_u=80.0)
|
|
image_calib_file.write_text(
|
|
json.dumps({"focal_u": 50.0, "focal_v": 50.0, "cu": 32.0, "cv": 16.0, "pitch": 0.0, "distort_coeffs": []}),
|
|
encoding="utf-8",
|
|
)
|
|
(gt_root / "train.txt").write_text(f"{rel_label.as_posix()}\n", encoding="utf-8")
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
|
|
assert sample["calib"]["fx"] == pytest.approx(80.0)
|
|
|
|
|
|
def test_ground3d_dataset_reads_clip_level_camera4_from_label_root(tmp_path):
|
|
gt_root = tmp_path / "gt_20260320"
|
|
image_root = tmp_path / "dataset_20260202"
|
|
rel_label = Path("seq0/clip0/labels/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "seq0" / "clip0" / "images" / "frame_0001.png"
|
|
clip_calib_file = gt_root / "seq0" / "clip0" / "calib" / "L2_calib" / "camera4.json"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
clip_calib_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text(
|
|
"car 0.5 0.5 0.25 0.25 1 2 3 4 5 6 0.1 0.2 0.3 9 10 11 12 0\n",
|
|
encoding="utf-8",
|
|
)
|
|
Image.new("RGB", (1920, 1080), color=(32, 64, 96)).save(image_file)
|
|
clip_calib_file.write_text(
|
|
json.dumps(
|
|
{
|
|
"focal_u": 1450.9230324555967,
|
|
"focal_v": 1458.0023697476843,
|
|
"cu": 949.5149041625389,
|
|
"cv": 569.9146363123367,
|
|
"distort_coeffs": [-0.6, 0.7, -0.5, 0.2],
|
|
"pitch": 0.214,
|
|
"roll": 1.077,
|
|
"yaw": -0.643,
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
(gt_root / "train.txt").write_text(f"./{rel_label.as_posix()}\n", encoding="utf-8")
|
|
|
|
raw_calib = read_calib_from_path(str(image_file.resolve()), image_root=image_root, extra_calib_candidates=[
|
|
str((gt_root / "seq0" / "clip0" / "calib" / "frame_0001.json").resolve())
|
|
])
|
|
assert raw_calib is not None
|
|
assert raw_calib["focal_u"] == pytest.approx(1450.9230324555967)
|
|
assert raw_calib["pitch"] == pytest.approx(np.deg2rad(0.214))
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[768, 352],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"complete_3d_classes": [0],
|
|
"roi": [768, 352],
|
|
"ori_img_size": [1920, 1080],
|
|
"virtual_fx": 537,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
|
|
assert sample["im_file"] == str(image_file.resolve())
|
|
assert sample["calib"]["fx"] > 0
|
|
|
|
|
|
def test_ground3d_dataset_applies_class_filter_and_rect_lazily(tmp_path):
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_label = Path("labels/seq0/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / "seq0" / "frame_0001.png"
|
|
clip_calib_dir = gt_root / "calib" / "seq0"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text(
|
|
"car 0.5 0.5 0.25 0.25 0\ntruck 0.4 0.4 0.2 0.2 0\n",
|
|
encoding="utf-8",
|
|
)
|
|
Image.new("RGB", (64, 32), color=(32, 64, 96)).save(image_file)
|
|
write_clip_level_camera4(clip_calib_dir, (64, 32))
|
|
(gt_root / "train.txt").write_text(f"{rel_label.as_posix()}\n", encoding="utf-8")
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=True,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
classes=[1],
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0, "truck": 1},
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
assert dataset.batch.shape == (1,)
|
|
assert dataset.batch_shapes.shape == (1, 2)
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
assert sample["cls"].reshape(-1).tolist() == [1.0]
|
|
|
|
|
|
def test_ground3d_dataset_keeps_missing_3d_targets_as_nan_for_2d_only_labels(tmp_path):
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_label = Path("labels/seq0/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / "seq0" / "frame_0001.png"
|
|
clip_calib_dir = gt_root / "calib" / "seq0"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text("car 0.5 0.5 0.25 0.25 1 0\n", encoding="utf-8")
|
|
Image.new("RGB", (64, 32), color=(32, 64, 96)).save(image_file)
|
|
write_clip_level_camera4(clip_calib_dir, (64, 32))
|
|
(gt_root / "train.txt").write_text(f"{rel_label.as_posix()}\n", encoding="utf-8")
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
raw_sample = dataset.get_image_and_label(0)
|
|
assert raw_sample["labels_3d"].shape == (1, 42)
|
|
assert np.isnan(raw_sample["labels_3d"]).all()
|
|
|
|
sample = dataset[0]
|
|
|
|
assert sample["labels_3d"].shape == (1, 42)
|
|
assert sample["labels_3d"].isnan().all().item()
|
|
|
|
|
|
def test_ground3d_dataset_falls_back_to_jpg_when_png_is_missing(tmp_path):
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_label = Path("labels/seq0/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / "seq0" / "frame_0001.jpg"
|
|
clip_calib_dir = gt_root / "calib" / "seq0"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
label_file.write_text("car 0.5 0.5 0.25 0.25 0\n", encoding="utf-8")
|
|
Image.new("RGB", (64, 32), color=(32, 64, 96)).save(image_file)
|
|
write_clip_level_camera4(clip_calib_dir, (64, 32))
|
|
(gt_root / "train.txt").write_text(f"{rel_label.as_posix()}\n", encoding="utf-8")
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
|
|
assert sample["im_file"] == str(image_file.resolve())
|
|
assert sample["img"].shape[:2] == (32, 64)
|
|
|
|
|
|
def test_ground3d_dataset_skips_to_next_image_when_imread_fails(tmp_path, monkeypatch):
|
|
dataset, image_files = create_ground3d_dataset(tmp_path, [(64, 32), (64, 32)])
|
|
original_imread = cv2.imread
|
|
|
|
def fake_imread(path, flags):
|
|
if str(path) == image_files[0]:
|
|
return None
|
|
return original_imread(path, flags)
|
|
|
|
monkeypatch.setattr(cv2, "imread", fake_imread)
|
|
|
|
sample = dataset[0]
|
|
|
|
assert sample["im_file"] == image_files[1]
|
|
assert dataset._bad_image_mask[0]
|
|
|
|
|
|
def test_ground3d_dataset_allows_missing_calibration_for_2d_only_samples(tmp_path):
|
|
dataset, image_files = create_ground3d_dataset(
|
|
tmp_path,
|
|
[(128, 64), (128, 64)],
|
|
imgsz=(64, 32),
|
|
roi=(64, 32),
|
|
ori_img_size=(128, 64),
|
|
)
|
|
first_calib = tmp_path / "gt" / "calib" / "seq1" / "L2_calib" / "camera4.json"
|
|
first_calib.unlink()
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
|
|
assert sample["im_file"] == image_files[0]
|
|
assert sample["ori_shape"] == (32, 64)
|
|
assert sample["img"].shape[:2] == (32, 64)
|
|
assert sample["camera_mode"] == "roi"
|
|
|
|
|
|
def test_ground3d_dataset_fails_on_missing_calibration_for_3d_samples(tmp_path):
|
|
gt_root = tmp_path / "gt"
|
|
image_root = tmp_path / "dataset"
|
|
rel_label = Path("labels/seq0/frame_0001.txt")
|
|
|
|
label_file = gt_root / rel_label
|
|
image_file = image_root / "images" / "seq0" / "frame_0001.png"
|
|
clip_calib_file = gt_root / "calib" / "seq0" / "L2_calib" / "camera4.json"
|
|
|
|
label_file.parent.mkdir(parents=True, exist_ok=True)
|
|
image_file.parent.mkdir(parents=True, exist_ok=True)
|
|
clip_calib_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 19-col complete_3d label: class + 18 numeric fields.
|
|
label_file.write_text("car 0.5 0.5 0.25 0.25 1 2 3 4 5 6 0.1 0.2 0.3 9 10 11 12 0\n", encoding="utf-8")
|
|
Image.new("RGB", (64, 32), color=(32, 64, 96)).save(image_file)
|
|
clip_calib_file.write_text(
|
|
json.dumps({"focal_u": 50.0, "focal_v": 50.0, "cu": 32.0, "cv": 16.0, "pitch": 0.0, "distort_coeffs": []}),
|
|
encoding="utf-8",
|
|
)
|
|
(gt_root / "train.txt").write_text(f"{rel_label.as_posix()}\n", encoding="utf-8")
|
|
clip_calib_file.unlink()
|
|
|
|
dataset = YOLOGround3DDataset(
|
|
img_path=str(gt_root / "train.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(image_root),
|
|
"class_map": {"car": 0},
|
|
"complete_3d_classes": [0],
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
"virtual_camera_prob": -1.0,
|
|
"crop_center_mode": "cxvy",
|
|
},
|
|
)
|
|
|
|
with pytest.raises(Ground3DCalibrationError, match="calibration file not found"):
|
|
dataset[0]
|
|
|
|
|
|
def test_ground3d_dataset_rejects_missing_required_ground3d_fields(tmp_path):
|
|
with pytest.raises(ValueError, match="virtual_camera_prob, crop_center_mode"):
|
|
YOLOGround3DDataset(
|
|
img_path=str(tmp_path / "unused.txt"),
|
|
imgsz=[64, 32],
|
|
batch_size=1,
|
|
augment=False,
|
|
rect=False,
|
|
stride=32,
|
|
pad=0.5,
|
|
prefix="test: ",
|
|
task="detect",
|
|
data={
|
|
"path": str(tmp_path / "dataset"),
|
|
"class_map": {"car": 0},
|
|
"roi": [64, 32],
|
|
"ori_img_size": [64, 32],
|
|
"virtual_fx": 50,
|
|
},
|
|
)
|
|
|
|
|
|
def test_ground3d_dataset_skips_images_with_invalid_decoded_shape(tmp_path):
|
|
dataset, image_files = create_ground3d_dataset(tmp_path, [(32, 16), (64, 32)])
|
|
|
|
sample = dataset[0]
|
|
|
|
assert sample["im_file"] == image_files[1]
|
|
assert dataset._bad_image_mask[0]
|
|
|
|
|
|
def test_ground3d_dataset_resizes_in_half_steps_for_quarter_scale(tmp_path, monkeypatch):
|
|
dataset, _ = create_ground3d_dataset(
|
|
tmp_path,
|
|
[(256, 128)],
|
|
imgsz=(64, 32),
|
|
roi=(256, 128),
|
|
ori_img_size=(256, 128),
|
|
)
|
|
original_resize = cv2.resize
|
|
resize_calls = []
|
|
|
|
def tracked_resize(img, dsize, *args, **kwargs):
|
|
resize_calls.append(dsize)
|
|
return original_resize(img, dsize, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(cv2, "resize", tracked_resize)
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
|
|
assert resize_calls == [(128, 64), (64, 32)]
|
|
assert sample["img"].shape[:2] == (32, 64)
|
|
assert sample["calib"]["fx"] == pytest.approx(12.5)
|
|
|
|
|
|
def test_ground3d_dataset_resizes_in_half_steps_then_remainder(tmp_path, monkeypatch):
|
|
dataset, _ = create_ground3d_dataset(
|
|
tmp_path,
|
|
[(160, 80)],
|
|
imgsz=(64, 32),
|
|
roi=(160, 80),
|
|
ori_img_size=(160, 80),
|
|
)
|
|
original_resize = cv2.resize
|
|
resize_calls = []
|
|
|
|
def tracked_resize(img, dsize, *args, **kwargs):
|
|
resize_calls.append(dsize)
|
|
return original_resize(img, dsize, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(cv2, "resize", tracked_resize)
|
|
|
|
sample = dataset.get_image_and_label(0)
|
|
|
|
assert resize_calls == [(80, 40), (64, 32)]
|
|
assert sample["img"].shape[:2] == (32, 64)
|
|
assert sample["calib"]["fx"] == pytest.approx(20.0)
|