()
| 179 | |
| 180 | |
| 181 | def main(): |
| 182 | |
| 183 | #argparse |
| 184 | parser = argparse.ArgumentParser(description='Create scalp textures') |
| 185 | parser.add_argument('--dataset_path', required=True, help='Path to the hair_synth dataset') |
| 186 | parser.add_argument('--path_strand_vae_model', required=True, help='Path to .pt of the strandvae') |
| 187 | parser.add_argument('--out_path', required=True, type=str, help='Where to output the processed hair_synth dataset') |
| 188 | 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') |
| 189 | args = parser.parse_args() |
| 190 | |
| 191 | |
| 192 | tex_size=256 |
| 193 | |
| 194 | print("args.check_validity",args.check_validity) |
| 195 | |
| 196 | difflocks_dataset = DiffLocksDataset(args.dataset_path, |
| 197 | check_validity=args.check_validity, |
| 198 | load_full_strands=True, |
| 199 | compute_tbn_full_strands=True, |
| 200 | # restrict_to_single_hairstyle_name="base_31_idx_9378" #for running local on output v5 |
| 201 | ) |
| 202 | loader = DataLoader(difflocks_dataset, batch_size=1, num_workers=8, shuffle=False, pin_memory=True, persistent_workers=True) |
| 203 | normalization_dict=difflocks_dataset.get_normalization_data() |
| 204 | |
| 205 | model = StrandCodec(do_vae=False, |
| 206 | decode_type="dir", |
| 207 | scale_init=30.0, |
| 208 | nr_verts_per_strand=256, nr_values_to_decode=255, |
| 209 | dim_per_value_decoded=3).cuda() |
| 210 | model.load_state_dict(torch.load(args.path_strand_vae_model)) |
| 211 | model = torch.compile(model) |
| 212 | |
| 213 | scalp_mesh, scalp_mesh_data = difflocks_dataset.get_scalp() |
| 214 | |
| 215 | progress_bar = tqdm(range(0, len(difflocks_dataset)), desc="Training progress") |
| 216 | |
| 217 | # for data in hair_synth_dataset: |
| 218 | for batch in loader: |
| 219 | progress_bar.update() |
| 220 | |
| 221 | #make the output path |
| 222 | output_scalp_texture_path=os.path.join(args.out_path, "processed_hairstyles", batch["file"][0], "scalp_textures") |
| 223 | os.makedirs(output_scalp_texture_path, exist_ok=True) |
| 224 | if not os.path.isfile( os.path.join(output_scalp_texture_path,"x_done.txt")): |
| 225 | #if it doesn't exist or we can't load it we create it |
| 226 | generate_scalp_textures(batch, model, normalization_dict, tex_size, output_scalp_texture_path) |
| 227 | |
| 228 | |
| 229 | #generate also a flipped texture, the reason being that just flipping the scalp texture does not result in a flipped hairstyle so we have to horizontally flip the data in the batch then encode a new flipped texture |
| 230 | #make the output path |
| 231 | output_scalp_texture_path=os.path.join(args.out_path, "processed_hairstyles", batch["file"][0], "scalp_textures_flip") |
| 232 | os.makedirs(output_scalp_texture_path, exist_ok=True) |
| 233 | if not os.path.isfile( os.path.join(output_scalp_texture_path,"x_done.txt")): |
| 234 | batch=horizontally_flip(batch, scalp_mesh_data) |
| 235 | #if it doesn't exist or we can't load it we create it |
| 236 | generate_scalp_textures(batch, model, normalization_dict, tex_size, output_scalp_texture_path) |
| 237 | |
| 238 |
no test coverage detected