(batch, model, normalization_dict, tex_size, output_scalp_texture_path)
| 64 | |
| 65 | |
| 66 | def generate_scalp_textures(batch, model, normalization_dict, tex_size, output_scalp_texture_path): |
| 67 | #gt batch |
| 68 | gt_dict = prepare_gt_batch(batch) |
| 69 | nr_strands = gt_dict["strand_positions"].shape[0] |
| 70 | # print("nr strands", nr_strands) |
| 71 | |
| 72 | create_pca_pngs=False |
| 73 | |
| 74 | |
| 75 | #get the values by encoding the strands |
| 76 | gt_dict_normalized = normalize_gt_data(gt_dict, normalization_dict) |
| 77 | #do it in chunks because otherwise we run out of memory |
| 78 | nr_chunks=10 |
| 79 | gt_strands_pos = gt_dict_normalized["strand_positions"] |
| 80 | gt_strands_dirs = gt_dict_normalized["strand_directions"] |
| 81 | gt_strands_pos_chunked = torch.chunk(gt_strands_pos, nr_chunks) |
| 82 | gt_strands_dirs_chunked = torch.chunk(gt_strands_dirs, nr_chunks) |
| 83 | latents_chunked= [] |
| 84 | for i in range(len(gt_strands_pos_chunked)): |
| 85 | cur_strands_pos = gt_strands_pos_chunked[i] |
| 86 | cur_strands_dirs = gt_strands_dirs_chunked[i] |
| 87 | gt_dict_chunk = {} |
| 88 | gt_dict_chunk["strand_positions"]=cur_strands_pos |
| 89 | gt_dict_chunk["strand_directions"]=cur_strands_dirs |
| 90 | #run model and get latent for this chunk |
| 91 | encoded_dict = model.encoder(gt_dict_chunk) |
| 92 | latents=encoded_dict["z"] |
| 93 | # print("latent shape",latents.shape) |
| 94 | latents_chunked.append(latents) |
| 95 | latents = torch.cat(latents_chunked) |
| 96 | |
| 97 | |
| 98 | #assign directly to pixels |
| 99 | strand_uv = batch["full_strands"]["root_uv"] |
| 100 | strand_uv=strand_uv.cuda().squeeze(0) |
| 101 | strand_uv_orig_flip = strand_uv |
| 102 | strand_uv_orig_flip[:,1]=1.0-strand_uv_orig_flip[:,1] |
| 103 | strand_indices = (strand_uv_orig_flip*tex_size).floor().int() #TODO Do we need to add a +0.5? |
| 104 | scalp_texture = torch.zeros(1,64,tex_size,tex_size).cuda() |
| 105 | scalp_texture[:,:,strand_indices[:,1],strand_indices[:,0]] = latents.transpose(0,1).unsqueeze(0) #make latents 1,64,nr_strands |
| 106 | |
| 107 | #get mask |
| 108 | scalp_mask = torch.zeros(1,1,tex_size,tex_size).cuda() |
| 109 | homogeneous_coord = torch.ones(nr_strands,1).cuda() |
| 110 | scalp_mask[:,:,strand_indices[:,1],strand_indices[:,0]]=homogeneous_coord.transpose(0,1).unsqueeze(0) |
| 111 | torchvision.utils.save_image(scalp_mask.squeeze(0), os.path.join(output_scalp_texture_path,"scalp_mask_"+str(tex_size)+".png")) |
| 112 | |
| 113 | |
| 114 | #inpaint |
| 115 | scalp_texture_inpainted=push_pull_inpaint.push_pull_inpaint(scalp_mask.squeeze(1), scalp_texture) |
| 116 | if create_pca_pngs: |
| 117 | scalp_texture_inpainted_pca = img_2_pca(scalp_texture_inpainted) |
| 118 | torchvision.utils.save_image(scalp_texture_inpainted_pca.squeeze(0), os.path.join(output_scalp_texture_path,"scalp_texture_inpainted_pca_"+str(tex_size)+".png")) |
| 119 | torch.save(scalp_texture_inpainted, os.path.join(output_scalp_texture_path, "scalp_texture_inpainted_"+str(tex_size)+".pt")) |
| 120 | |
| 121 | #make also downsized versions |
| 122 | for i in range(4): |
| 123 | scale_factor = 1.0/ (np.power(2,i+1)) |
no test coverage detected