| 45 | |
| 46 | |
| 47 | def create_dataloaders(dataset_path, dataset_processed_path): |
| 48 | difflocks_dataset_train = DiffLocksDataset(dataset_path, |
| 49 | processed_difflocks_path=dataset_processed_path, |
| 50 | train=True, load_rgb_imgs=False, load_full_strands=False,\ |
| 51 | load_guide_strands=False, load_interpolated_strands=False, |
| 52 | load_cam=False, |
| 53 | subsample_factor=1, #needed in order to load the dino latents from subsample1 |
| 54 | load_material=True, |
| 55 | compute_tbn_full_strands=False, |
| 56 | load_latents=True, |
| 57 | latents_type_list=["dinov2"], |
| 58 | load_latents_layers=[ |
| 59 | ["final_latent"] |
| 60 | ], |
| 61 | #not really needd but we want to filter to only those samples that have a scalp texture because on local we don't have all of them downloaded |
| 62 | # load_scalp_texture=True, |
| 63 | # scalp_texture_resolution=64, |
| 64 | |
| 65 | check_validity=True, |
| 66 | do_pedantic_checks=False, |
| 67 | overfit=False, |
| 68 | train_ratio=0.9) |
| 69 | difflocks_dataset_test = DiffLocksDataset(dataset_path, |
| 70 | processed_difflocks_path=dataset_processed_path, |
| 71 | train=False, load_rgb_imgs=False, load_full_strands=False,\ |
| 72 | load_guide_strands=False, load_interpolated_strands=False, |
| 73 | load_cam=False, |
| 74 | subsample_factor=1, #needed in order to load the dino latents from subsample1 |
| 75 | load_material=True, |
| 76 | compute_tbn_full_strands=False, |
| 77 | load_latents=True, |
| 78 | latents_type_list=["dinov2"], |
| 79 | load_latents_layers=[ |
| 80 | ["final_latent"] |
| 81 | ], |
| 82 | #not really needd but we want to filter to only those samples that have a scalp texture because on local we don't have all of them downloaded |
| 83 | # load_scalp_texture=True, |
| 84 | # scalp_texture_resolution=64, |
| 85 | |
| 86 | check_validity=True, |
| 87 | do_pedantic_checks=False, |
| 88 | overfit=False, |
| 89 | train_ratio=0.9) |
| 90 | |
| 91 | loader_train = DataLoader(difflocks_dataset_train, batch_size=8, num_workers=8, shuffle=True, pin_memory=True, persistent_workers=True, |
| 92 | prefetch_factor=3) |
| 93 | loader_test = DataLoader(difflocks_dataset_test, batch_size=8, num_workers=8, shuffle=False, pin_memory=True, persistent_workers=True, |
| 94 | prefetch_factor=3) |
| 95 | |
| 96 | return loader_train, loader_test |
| 97 | |
| 98 | # |
| 99 | def compute_loss(phase, gt_dict, pred_dict, hyperparams): |