spam.py

#!/usr/bin/env python

import numpy as np
import torch

import corvus_segmentation_models

from train import TransformDataset


def transform(i, images):
    K = np.array(dataset.meta_root["intrinsic_matrix"])
    images["hha"] = corvus_segmentation_models.geometry.depth2hha(images["depth_filled"], K)

    out_file = f"out/{i:08d}/images.npz"
    np.savez_compressed(out_file, **images)
    return images


dataset = corvus_segmentation_models.datasets.NyuDepthV2Dataset("all")
dataset = TransformDataset(dataset, transform)
loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=12)

import tqdm
for _ in tqdm.tqdm(loader):
    pass