| 65 | |
| 66 | |
| 67 | def create_dataloaders(dataset_path): |
| 68 | difflocks_dataset_train = DiffLocksDataset(dataset_path, |
| 69 | processed_difflocks_path=None, |
| 70 | train=True, load_rgb_imgs=False, load_full_strands=True,\ |
| 71 | load_guide_strands=False, load_interpolated_strands=False, |
| 72 | load_cam=False, |
| 73 | compute_tbn_full_strands=True, |
| 74 | nr_full_strands_per_hairstyle=20, |
| 75 | check_validity=True, |
| 76 | overfit=False, |
| 77 | train_ratio=0.9) |
| 78 | difflocks_dataset_test = DiffLocksDataset(dataset_path, |
| 79 | processed_difflocks_path=None, |
| 80 | train=False, load_rgb_imgs=False, load_full_strands=True,\ |
| 81 | load_guide_strands=False, load_interpolated_strands=False, |
| 82 | load_cam=False, |
| 83 | compute_tbn_full_strands=True, |
| 84 | nr_full_strands_per_hairstyle=20, |
| 85 | check_validity=True, |
| 86 | overfit=False, |
| 87 | train_ratio=0.9) |
| 88 | |
| 89 | loader_train = DataLoader(difflocks_dataset_train, batch_size=10, num_workers=8, shuffle=True, pin_memory=True, persistent_workers=True, |
| 90 | prefetch_factor=3) |
| 91 | loader_test = DataLoader(difflocks_dataset_test, batch_size=10, num_workers=8, shuffle=False, pin_memory=True, persistent_workers=True, |
| 92 | prefetch_factor=3) |
| 93 | |
| 94 | return loader_train, loader_test |
| 95 | |
| 96 | |
| 97 | #transforms the data to a local space, put it on cuda device and reshapes it the way we expect it to be |