(batch)
| 28 | |
| 29 | #transforms the data to a local space, put it on cuda device and reshapes it the way we expect it to be |
| 30 | def prepare_gt_batch(batch): |
| 31 | gt_dict = {} |
| 32 | |
| 33 | tbn=batch['full_strands']["tbn"].cuda() |
| 34 | positions=batch['full_strands']["positions"].cuda() |
| 35 | root_normal=batch['full_strands']["root_normal"].cuda() |
| 36 | |
| 37 | #get it on local space |
| 38 | gt_strand_positions, gt_root_normals = world_to_tbn_space(tbn, |
| 39 | positions, |
| 40 | root_normal) |
| 41 | gt_strand_positions=gt_strand_positions.cuda() |
| 42 | |
| 43 | #reshape it to be nr_strands, nr_points, dim |
| 44 | gt_strand_positions=gt_strand_positions.reshape(-1,256,3) |
| 45 | |
| 46 | gt_dirs=compute_dirs(gt_strand_positions, append_last_dir=False) #nr_strands,256-1,3 |
| 47 | |
| 48 | |
| 49 | gt_dict["strand_positions"]=gt_strand_positions |
| 50 | gt_dict["strand_directions"]=gt_dirs |
| 51 | |
| 52 | |
| 53 | return gt_dict |
| 54 | |
| 55 | |
| 56 | class HyperParamsStrandVAE: |
nothing calls this directly
no test coverage detected