Meehai/dronescapes

1#!/usr/bin/env python32import sys3from pathlib import Path4sys.path.append(Path(__file__).parents[1].__str__())5from dronescapes_reader import MultiTaskDataset6from dronescapes_reader.dronescapes_representations import dronescapes_task_types7from pprint import pprint8from torch.utils.data import DataLoader9import random10import numpy as np11 12def main():13    assert len(sys.argv) == 2, f"Usage ./dronescapes_viewer.py /path/to/dataset"14    reader = MultiTaskDataset(sys.argv[1], task_names=list(dronescapes_task_types.keys()),15                              task_types=dronescapes_task_types, handle_missing_data="fill_nan",16                              normalization="min_max", cache_task_stats=True)17    print(reader)18 19    print("== Shapes ==")20    pprint(reader.data_shape)21 22    print("== Random loaded item ==")23    rand_ix = random.randint(0, len(reader) - 1)24    data, name, repr_names = reader[rand_ix] # get a random item25    pprint({k: v for k, v in data.items()})26 27    img_data = {}28    for k, v in data.items():29        img_data[k] = reader.name_to_task[k].plot_fn(v) if v is not None else np.zeros((*reader.data_shape[k][0:2], 3))30 31    print("== Random loaded batch ==")32    batch_data, name, repr_names = reader[rand_ix: min(len(reader), rand_ix + 5)] # get a random batch33    pprint({k: v for k, v in batch_data.items()}) # Nones are converted to 0s automagically34 35    print("== Random loaded batch using torch DataLoader ==")36    loader = DataLoader(reader, collate_fn=reader.collate_fn, batch_size=5, shuffle=True)37    batch_data, name, repr_names = next(iter(loader))38    pprint({k: v for k, v in batch_data.items()}) # Nones are converted to 0s automagically39 40if __name__ == "__main__":41    main()42