def get_dataset()

in pt/vmz/datasets/data.py [0:0]

• 33 lines of code
• 12 McCabe index (conditional complexity)

def get_dataset(args, transform, split="train"):
    metadata = None
    if split == "val" or split == "validataion":
        if args.val_file and os.path.isfile(args.val_file):
            metadata = torch.load(args.val_file)
        root = args.valdir
        train = False

    elif split == "train":
        if args.train_file and os.path.isfile(args.train_file):
            metadata = torch.load(args.train_file)
        root = args.traindir
        train = True

    if args.dataset == "kinetics400":
        _dataset = Kinetics(
            root, args.num_frames, transform=transform, _precomputed_metadata=metadata
        )
    elif args.dataset == "ucf101":
        _dataset = UCF(
            root,
            args.annotation_path,
            frames_per_clip=args.num_frames,
            train=train,
            transform=transform,
            fold=args.fold,
            _precomputed_metadata=metadata,
        )

    _dataset.video_clips.compute_clips(args.num_frames, 1)
    if args.train_file is None or not os.path.isfile(args.train_file):
        utils.save_on_master(
            _dataset.metadata,
            "{}_{}_{}fms.pth".format(args.dataset, split, args.num_frames),
        )
    return _dataset