<gh_stars>0
import argparse

import torch
import torchvision
import tqdm
from torch.utils.data.dataloader import DataLoader

from carbrain.datasets import NuScenes
from carbrain.datasets.util import AnnotationType


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('root')
    parser.add_argument('version', choices=['v1.0-mini', 'v1.0-trainval', 'v1.0-test', 'v1.0-private-test'])
    args = parser.parse_args()

    nuscenes = NuScenes(
        root=args.root,
        split=args.version,
        annotation_type=AnnotationType.NONE,
        transforms=torchvision.transforms.ToTensor(),
    )
    loader = DataLoader(
        dataset=nuscenes,
        batch_size=100,
        shuffle=False,
        num_workers=8,
    )

    print('Computing mean')
    mean = torch.zeros(3, dtype=torch.float64)
    for imgs, _ in tqdm.tqdm(loader):
        img_mean = imgs.mean(dim=1)
        mean += img_mean.sum(dim=0)
    mean /= len(nuscenes)

    print('Computing std')
    std = torch.zeros(3, dtype=torch.float64)
    for imgs, _ in tqdm.tqdm(loader):
        img_mean = imgs.mean(dim=1)
        diff = img_mean - mean
        std += diff.sum(dim=0)
    std /= len(nuscenes)

    # output
    print('NuScenes {}: mean={}, std={}'.format(args.version, str(mean), str(std)))


if __name__ == '__main__':
    main()
import argparse

import torch
import torchvision
import tqdm
from torch.utils.data.dataloader import DataLoader

from carbrain.datasets import NuScenes
from carbrain.datasets.util import AnnotationType


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('root')
    parser.add_argument('version', choices=['v1.0-mini', 'v1.0-trainval', 'v1.0-test', 'v1.0-private-test'])
    args = parser.parse_args()

    nuscenes = NuScenes(
        root=args.root,
        split=args.version,
        annotation_type=AnnotationType.NONE,
        transforms=torchvision.transforms.ToTensor(),
    )
    loader = DataLoader(
        dataset=nuscenes,
        batch_size=100,
        shuffle=False,
        num_workers=8,
    )

    print('Computing mean')
    mean = torch.zeros(3, dtype=torch.float64)
    for imgs, _ in tqdm.tqdm(loader):
        img_mean = imgs.mean(dim=1)
        mean += img_mean.sum(dim=0)
    mean /= len(nuscenes)

    print('Computing std')
    std = torch.zeros(3, dtype=torch.float64)
    for imgs, _ in tqdm.tqdm(loader):
        img_mean = imgs.mean(dim=1)
        diff = img_mean - mean
        std += diff.sum(dim=0)
    std /= len(nuscenes)

    # output
    print('NuScenes {}: mean={}, std={}'.format(args.version, str(mean), str(std)))


if __name__ == '__main__':
    main()