Function encode_depth_codes

experiments/scripts/generate_depth_annotation.py:855–941

(
    bgr_images: list[np.ndarray],
    depth_model: DepthAnythingV2,
    vae: VQVAE,
    *,
    depth_input_size: int,
    vae_image_size: int,
    batch_size: int,
    device: torch.device,
)

@torch.no_grad()
def encode_depth_codes(
    bgr_images: list[np.ndarray],
    depth_model: DepthAnythingV2,
    vae: VQVAE,
    *,
    depth_input_size: int,
    vae_image_size: int,
    batch_size: int,
    device: torch.device,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    exact = batch_size == 1
    rgb_resized: list[np.ndarray] = []
    codes: list[np.ndarray] = []
    for start in range(0, len(bgr_images), batch_size):
        batch_images = bgr_images[start : start + batch_size]
        tensors: list[Tensor] = []
        original_sizes: list[tuple[int, int]] = []
        for bgr in batch_images:
            tensor, size = depth_image_to_tensor(bgr, depth_input_size, device)
            tensors.append(tensor)
            original_sizes.append(size)
            rgb_resized.append(cv2.resize(bgr, (vae_image_size, vae_image_size), interpolation=cv2.INTER_LINEAR))

        if exact:
            for idx, tensor in enumerate(tensors):
                depth = depth_model(tensor[None])[0]
                height, width = original_sizes[idx]
                resized_depth = nn_func.interpolate(
                    depth[None, None],
                    (height, width),
                    mode="bilinear",
                    align_corners=True,
                )[0, 0]
                raw = resized_depth.detach().cpu().numpy()
                dmin = raw.min()
                dmax = raw.max()
                depth_vis = ((raw - dmin) / (dmax - dmin + 1e-8) * 255.0).astype(np.uint8)
                depth_vis = cv2.resize(
                    depth_vis,
                    (vae_image_size, vae_image_size),
                    interpolation=cv2.INTER_NEAREST,
                )
                depth_tensor = torch.from_numpy(depth_vis).float() / 255.0
                depth_tensor = (depth_tensor - 0.5) / 0.5
                depth_tensor = depth_tensor.unsqueeze(0).unsqueeze(0).to(device)
                code_grid = vae(img=depth_tensor, return_indices=True)[0]
                codes.append(code_grid.detach().cpu().numpy().astype(np.int16))
        else:
            codes_by_index: list[np.ndarray | None] = [None] * len(tensors)
            shape_groups: dict[tuple[int, int, int, int], list[int]] = {}
            for idx, tensor in enumerate(tensors):
                input_height, input_width = tensor.shape[-2:]
                original_height, original_width = original_sizes[idx]
                shape_groups.setdefault((input_height, input_width, original_height, original_width), []).append(idx)
            for indices in shape_groups.values():
                input_height, input_width = tensors[indices[0]].shape[-2:]
                original_height, original_width = original_sizes[indices[0]]
                max_depth_batch = max(1, (2**31 - 1) // (64 * input_height * input_width))