(batch)
| 38 | |
| 39 | #transforms the data to a local space, put it on cuda device and reshapes it the way we expect it to be |
| 40 | def prepare_gt_batch(batch): |
| 41 | gt_dict = {} |
| 42 | |
| 43 | tbn=batch['full_strands']["tbn"].cuda() |
| 44 | positions=batch['full_strands']["positions"].cuda() |
| 45 | root_normal=batch['full_strands']["root_normal"].cuda() |
| 46 | |
| 47 | #get it on local space |
| 48 | gt_strand_positions, gt_root_normals = world_to_tbn_space(tbn, |
| 49 | positions, |
| 50 | root_normal) |
| 51 | gt_strand_positions=gt_strand_positions.cuda() |
| 52 | |
| 53 | #reshape it to be nr_strands, nr_points, dim |
| 54 | gt_strand_positions=gt_strand_positions.reshape(-1,256,3) |
| 55 | |
| 56 | gt_dirs=compute_dirs(gt_strand_positions, append_last_dir=False) #nr_strands,256-1,3 |
| 57 | |
| 58 | |
| 59 | gt_dict["strand_positions"]=gt_strand_positions |
| 60 | gt_dict["strand_directions"]=gt_dirs |
| 61 | |
| 62 | |
| 63 | return gt_dict |
| 64 | |
| 65 | |
| 66 | def generate_scalp_textures(batch, model, normalization_dict, tex_size, output_scalp_texture_path): |
no test coverage detected