(batch, hyperparams, world2local, do_augmentation=False)
| 96 | |
| 97 | #transforms the data to a local space, put it on cuda device and reshapes it the way we expect it to be |
| 98 | def prepare_gt_batch(batch, hyperparams, world2local, do_augmentation=False): |
| 99 | gt_dict = {} |
| 100 | |
| 101 | tbn=batch['full_strands']["tbn"].cuda() |
| 102 | positions=batch['full_strands']["positions"].cuda() |
| 103 | root_normal=batch['full_strands']["root_normal"].cuda() |
| 104 | |
| 105 | #get it on local space |
| 106 | gt_strand_positions, gt_root_normals = world2local(tbn, positions, root_normal) |
| 107 | |
| 108 | #reshape it to be nr_strands, nr_points, dim |
| 109 | gt_strand_positions=gt_strand_positions.reshape(-1,256,3) |
| 110 | |
| 111 | if do_augmentation: |
| 112 | nr_strands = gt_strand_positions.shape[0] |
| 113 | |
| 114 | #do some random horizontal flip |
| 115 | rand_strand_mask=torch.rand(nr_strands, device="cuda")>0.5 |
| 116 | gt_strand_positions[rand_strand_mask,:,0] = -gt_strand_positions[rand_strand_mask,:,0] |
| 117 | |
| 118 | #a bit of rotation do it through quaternions since they allows for linear interpolation which actually does a slerp. If they were rotation matrices I would need to implement slerp |
| 119 | rotations_q = random_quaternions(nr_strands) |
| 120 | identity_q = torch.tensor([1, 0, 0, 0], device="cuda").view(1,4).repeat(nr_strands,1) |
| 121 | #interpolate more towards an identity rotation |
| 122 | rot_amount=0.1 |
| 123 | rotations_q = rotations_q*rot_amount + identity_q*(1.0-rot_amount) |
| 124 | rotations = quaternion_to_matrix(rotations_q) |
| 125 | #rotate positional data [Nr_strands, 3, 3] x [Nr_strands, nr_points_per_strand, 3] |
| 126 | rotations = rotations.reshape(nr_strands, 1, 3, 3) |
| 127 | gt_strand_positions = gt_strand_positions.reshape(nr_strands, -1, 3, 1) |
| 128 | gt_strand_positions= torch.matmul(rotations, gt_strand_positions) |
| 129 | gt_strand_positions=gt_strand_positions.reshape(-1,256,3) |
| 130 | |
| 131 | |
| 132 | |
| 133 | #center the data to be drawn from unit gaussian |
| 134 | # gt_strand_positions_normalized=whiten_gt_data(gt_strand_positions, normalization_dict, normalization_type="xyz") |
| 135 | |
| 136 | gt_dirs=compute_dirs(gt_strand_positions, append_last_dir=False) #nr_strands,256-1,3 |
| 137 | gt_curv=compute_dirs(gt_dirs, append_last_dir=False) #nr_strands,256-2,3 |
| 138 | |
| 139 | |
| 140 | gt_dict["strand_positions"]=gt_strand_positions |
| 141 | gt_dict["strand_directions"]=gt_dirs |
| 142 | gt_dict["strand_curvatures"]=gt_curv |
| 143 | |
| 144 | return gt_dict |
| 145 | |
| 146 | |
| 147 | def train(args, hyperparams, loader_train, loader_test, experiment_name, output_training_path): |
no test coverage detected