spam.py

import imgviz
import tqdm

import corvus_segmentation_models


dataset = corvus_segmentation_models.datasets.Warehouse20200526Dataset(
    split="sim"
)

viewer = imgviz.io.PygletThreadedImageViewer(height=480, width=640 * 2)
for i in tqdm.trange(len(dataset)):
    images = dataset[i]
    viz = imgviz.tile(
        [
            images["gray"],
            imgviz.label2rgb(
                images["class_label"],
                label_names=dataset.class_names,
                loc="rb",
            ),
        ],
        shape=(1, 2),
        border=(255, 255, 255),
    )
    viewer.imshow(viz)

# max_depth = []
# for i in tqdm.trange(len(dataset)):
#     images = dataset[i]
#     max_depth.append(np.nanmax(images["depth"]))
# print(np.mean(max_depth))