from vdataset import LabelMap, VDataset

from torch.utils.data import DataLoader

from torchvideotransforms.volume_transforms import ClipToTensor # https://github.com/hassony2/torch_videovision
from torchvideotransforms import video_transforms, volume_transforms # https://github.com/hassony2/torch_videovision

# Create Label Map
label_map = LabelMap(labels_csv="/path-to-csv/csv_file.csv", labels_col_name="label") 

print(label_map)
label_map.print() # printing the labels on label-map

# Use Video Transformers
video_transform_list = [video_transforms.RandomRotation(30),
            video_transforms.Resize((100, 100)),
            volume_transforms.ClipToTensor()]
video_transforms = video_transforms.Compose(video_transform_list)

# Create Vdataset (No frame limitation)
full_dataset = VDataset(csv_file='/path-to-csv/csv_file.csv', root_dir='/path-to-root/', video_transforms=video_transforms, label_map=label_map)

# Create Vdataset (Manual frames limitation, remove first 5 frames and last 5 frames)
frames_limited_dataset = VDataset(csv_file='/path-to-csv/csv_file.csv', root_dir='/path-to-root/', video_transforms=video_transforms, frames_limit_mode="manual",  frames_limit={"start": 5, "end": -5} label_map=label_map)

full_dataloader = DataLoader(full_dataset, batch_size=64, shuffle=True, num_workers=2, pin_memory=True)
print(full_dataloader)

frames_limited_dataloader = DataLoader(frames_limited_dataset, batch_size=64, shuffle=True, num_workers=2, pin_memory=True)
print(frames_limited_dataloader)

for image, label in full_dataloader: # Do what do you want in dataset
    print(image, label)
    print(image.size())
    break

for image, label in frames_limited_dataloader: # Do what do you want in dataset
    print(image, label)
    print(image.size())
    break