File size: 7,049 Bytes
f71ac1d | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """Postprocess cache exporter (test-time).
This evaluator is used with vis4d's EvaluatorCallback to export per-image caches
needed for depth-based 3D box post-processing, without changing the normal
evaluation flow.
Cache layout:
{cache_root}/{dataset_name}/{image_id}.npz
We intentionally store the full metric depth map (aligned to original_hw) to
avoid coordinate-system bugs from cropping.
"""
from __future__ import annotations
import os
from typing import Any
import numpy as np
import torch
from vis4d.common.array import array_to_numpy
from vis4d.common.typing import GenericFunc, MetricLogs, NDArrayNumber
from vis4d.eval.base import Evaluator
class PostprocessCacheExporter(Evaluator):
"""Exports model outputs needed for post-processing into .npz cache files."""
def __init__(
self,
cache_root: str,
compress: bool = True,
overwrite: bool = False,
depth_dtype: str = "float32",
) -> None:
super().__init__()
self.cache_root = cache_root
self.compress = compress
self.overwrite = overwrite
if depth_dtype not in {"float16", "float32"}:
raise ValueError(f"Unsupported depth_dtype: {depth_dtype}")
self.depth_dtype = depth_dtype
self._num_written = 0
self._num_skipped = 0
@property
def metrics(self) -> list[str]:
# Not a real evaluator; we only export.
return []
def reset(self) -> None: # pragma: no cover
self._num_written = 0
self._num_skipped = 0
def gather(self, gather_func: GenericFunc) -> None: # pragma: no cover
# Nothing to gather; each rank writes its own files (safe because image_id is unique).
return
def process_batch(
self,
coco_image_id: list[int],
dataset_names: list[str],
pred_boxes: list[NDArrayNumber],
pred_scores: list[NDArrayNumber],
pred_classes: list[NDArrayNumber],
pred_boxes3d: list[NDArrayNumber] | None = None,
pred_categories: list[list[str]] | None = None,
depth_maps: list[torch.Tensor] | None = None,
intrinsics: list[NDArrayNumber] | NDArrayNumber | None = None,
original_hw: list[tuple[int, int]] | None = None,
) -> None:
"""Write one .npz per image."""
if pred_boxes3d is None:
# No 3D boxes -> nothing to export for depth alignment.
print("[PostprocessCacheExporter] Skipping: pred_boxes3d is None")
return
if depth_maps is None:
# Depth backend disabled -> nothing to export.
print("[PostprocessCacheExporter] Skipping: depth_maps is None")
return
if intrinsics is None:
print("[PostprocessCacheExporter] Skipping: intrinsics is None")
return
if original_hw is None:
print("[PostprocessCacheExporter] Skipping: original_hw is None")
return
print(f"[PostprocessCacheExporter] Processing batch: {len(coco_image_id)} images")
# Normalize intrinsics to per-sample list
if torch.is_tensor(intrinsics):
# intrinsics: Tensor [B, 3, 3] (may be on GPU)
intrinsics_np = intrinsics.detach().cpu().numpy()
intrinsics_list = [intrinsics_np[j] for j in range(intrinsics_np.shape[0])]
elif isinstance(intrinsics, np.ndarray):
# intrinsics: ndarray [3,3] or [B,3,3]
if intrinsics.ndim == 2:
intrinsics_list = [intrinsics for _ in range(len(coco_image_id))]
else:
intrinsics_list = [intrinsics[j] for j in range(intrinsics.shape[0])]
else:
# intrinsics: sequence of arrays/tensors
intrinsics_list = list(intrinsics)
for i, image_id in enumerate(coco_image_id):
dataset_name = dataset_names[i]
out_dir = os.path.join(self.cache_root, str(dataset_name))
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, f"{int(image_id)}.npz")
if (not self.overwrite) and os.path.exists(out_path):
self._num_skipped += 1
continue
boxes2d = array_to_numpy(
pred_boxes[i].to(torch.float32) if hasattr(pred_boxes[i], "to") else pred_boxes[i],
n_dims=None,
dtype=np.float32,
)
scores = array_to_numpy(
pred_scores[i].to(torch.float32) if hasattr(pred_scores[i], "to") else pred_scores[i],
n_dims=None,
dtype=np.float32,
)
class_ids = array_to_numpy(
pred_classes[i].to(torch.int64) if hasattr(pred_classes[i], "to") else pred_classes[i],
n_dims=None,
dtype=np.int64,
)
boxes3d = array_to_numpy(
pred_boxes3d[i].to(torch.float32) if hasattr(pred_boxes3d[i], "to") else pred_boxes3d[i],
n_dims=None,
dtype=np.float32,
)
# depth_maps is list[Tensor] where each Tensor is [H, W] or [1, H, W]
depth = depth_maps[i]
if depth.ndim == 3 and depth.shape[0] == 1:
depth = depth[0]
depth_np = depth.detach().cpu().numpy()
depth_np = depth_np.astype(np.float16 if self.depth_dtype == "float16" else np.float32)
Ki = intrinsics_list[i]
if torch.is_tensor(Ki):
K = Ki.detach().cpu().numpy().astype(np.float32)
else:
K = np.asarray(Ki, dtype=np.float32)
hw = original_hw[i]
meta: dict[str, Any] = {
"dataset_name": str(dataset_name),
"image_id": int(image_id),
"original_hw": np.asarray(hw, dtype=np.int32),
}
# Categories are variable-length strings; store as object array.
if pred_categories is not None and i < len(pred_categories) and pred_categories[i] is not None:
cats = np.asarray(pred_categories[i], dtype=object)
else:
cats = np.asarray([], dtype=object)
save_fn = np.savez_compressed if self.compress else np.savez
save_fn(
out_path,
boxes2d=boxes2d,
scores=scores,
class_ids=class_ids,
boxes3d_raw=boxes3d,
categories=cats,
depth_map=depth_np,
intrinsics=K,
meta=np.asarray(meta, dtype=object),
)
self._num_written += 1
def evaluate(self, metric: str) -> tuple[MetricLogs, str]:
# No evaluation; return empty.
return {}, f"PostprocessCacheExporter: wrote={self._num_written}, skipped={self._num_skipped}"
def save(self, metric: str, output_dir: str, prefix: str | None = None) -> None: # pragma: no cover
# Nothing to save beyond the cache files.
return
|