()
| 18 | |
| 19 | |
| 20 | def main(): |
| 21 | |
| 22 | #argparse |
| 23 | parser = argparse.ArgumentParser(description='Create latents') |
| 24 | parser.add_argument('--dataset_path', required=True, help='Path to the hair_synth dataset') |
| 25 | args = parser.parse_args() |
| 26 | |
| 27 | |
| 28 | difflocks_dataset = DiffLocksDataset(args.dataset_path, |
| 29 | check_validity=False, |
| 30 | load_full_strands=True |
| 31 | ) |
| 32 | loader = DataLoader(difflocks_dataset, batch_size=1, num_workers=8, shuffle=False, pin_memory=True, persistent_workers=True) |
| 33 | |
| 34 | |
| 35 | |
| 36 | progress_bar = tqdm(range(0, len(difflocks_dataset)), desc="Training progress") |
| 37 | |
| 38 | for batch in loader: |
| 39 | progress_bar.update() |
| 40 | |
| 41 | # print("batch",batch) |
| 42 | |
| 43 | path_hairstyle=batch["path"][0] |
| 44 | print("path", path_hairstyle) |
| 45 | |
| 46 | positions=batch["full_strands"]["positions"].squeeze(0).cpu().numpy() |
| 47 | root_uv=batch["full_strands"]["root_uv"].squeeze(0).cpu().numpy() |
| 48 | root_normal=batch["full_strands"]["root_normal"].squeeze(0).cpu().numpy() |
| 49 | # print("positions", positions.shape) |
| 50 | # print("root_uv", root_uv.shape) |
| 51 | # print("root_normal", root_normal.shape) |
| 52 | |
| 53 | |
| 54 | select_nr_random_strands=10000 |
| 55 | nr_strands_per_chunk_list=[1000,100] |
| 56 | |
| 57 | #select random strands because the original 100K is too much |
| 58 | if select_nr_random_strands: |
| 59 | nr_strands_left = positions.shape[0] |
| 60 | per_curve_keep_random = np.random.choice(nr_strands_left, select_nr_random_strands, replace=False) |
| 61 | positions=positions[per_curve_keep_random,:,:] |
| 62 | root_uv=root_uv[per_curve_keep_random,:] |
| 63 | root_normal=root_normal[per_curve_keep_random,:] |
| 64 | |
| 65 | |
| 66 | #break the strands into chunks if needed and write those too |
| 67 | nr_strands_total = positions.shape[0] |
| 68 | if nr_strands_per_chunk_list: |
| 69 | for nr_strands_cur_chunk in nr_strands_per_chunk_list: |
| 70 | nr_chunks = math.ceil(nr_strands_total/nr_strands_cur_chunk) |
| 71 | |
| 72 | positions_chunked = np.array_split(positions, nr_chunks,axis=0) |
| 73 | root_uv_chunked = np.array_split(root_uv, nr_chunks,axis=0) |
| 74 | root_normal_chunked = np.array_split(root_normal, nr_chunks,axis=0) |
| 75 | |
| 76 | #make path for this chunked data |
| 77 | chunked_data_path=os.path.join(path_hairstyle,"full_strands_chunked","nr_strands_"+str(nr_strands_cur_chunk)) |
no test coverage detected