()
| 76 | |
| 77 | |
| 78 | def main(): |
| 79 | |
| 80 | #argparse |
| 81 | parser = argparse.ArgumentParser(description='Create latents') |
| 82 | parser.add_argument('--dataset_path', required=True, help='Path to the hair_synth dataset') |
| 83 | parser.add_argument('--out_path', required=True, type=str, help='Where to output the processed hair_synth dataset') |
| 84 | parser.add_argument('--subsample_factor', default=1, type=int, help='Subsample factor for the RGB img') |
| 85 | parser.add_argument('--skip_validity_check', dest='check_validity', action='store_false', help='Wether to check for the validity of each hairstyle we read from the dataset. Some older dataset versions might need this turned to false') |
| 86 | args = parser.parse_args() |
| 87 | |
| 88 | |
| 89 | #v2 from torch |
| 90 | image_size = int(768/(2**(args.subsample_factor-1))) |
| 91 | print("Selected dino with img size", image_size) |
| 92 | #going to the nearest multiple of 14 because 14 is the patch size |
| 93 | if image_size==768: |
| 94 | image_size=770 |
| 95 | else: |
| 96 | print("I haven't implemented the other ones yet") |
| 97 | latents_preprocessor = T.Compose([ |
| 98 | T.Resize(image_size, interpolation=T.InterpolationMode.BICUBIC), |
| 99 | T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), |
| 100 | ]) |
| 101 | latents_model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg') |
| 102 | latents_model.cuda() |
| 103 | |
| 104 | latents_model.eval() |
| 105 | |
| 106 | |
| 107 | |
| 108 | |
| 109 | print("args.check_validity",args.check_validity) |
| 110 | |
| 111 | difflocks_dataset = DiffLocksDataset(args.dataset_path, |
| 112 | check_validity=args.check_validity, |
| 113 | load_rgb_imgs=True, |
| 114 | processed_difflocks_path = args.out_path, |
| 115 | subsample_factor=args.subsample_factor, |
| 116 | ) |
| 117 | loader = DataLoader(difflocks_dataset, batch_size=1, num_workers=8, shuffle=False, pin_memory=True, persistent_workers=True) |
| 118 | |
| 119 | |
| 120 | |
| 121 | progress_bar = tqdm(range(0, len(difflocks_dataset)), desc="Training progress") |
| 122 | |
| 123 | for batch in loader: |
| 124 | progress_bar.update() |
| 125 | |
| 126 | #make the output path |
| 127 | output_latents_path=os.path.join(args.out_path, "processed_hairstyles", batch["file"][0], "latents_"+"dinov2"+"_subsample_"+str(args.subsample_factor)) |
| 128 | os.makedirs(output_latents_path, exist_ok=True) |
| 129 | #check if we already created this one |
| 130 | if not os.path.isfile( os.path.join(output_latents_path,"x_done.txt")): |
| 131 | # if True: |
| 132 | #if it doesn't exist or we can't load it we create it |
| 133 | generate_latents_dinov2(args, batch, latents_preprocessor, latents_model, output_latents_path) |
| 134 | |
| 135 |
no test coverage detected