| 97 | |
| 98 | # |
| 99 | def compute_loss(phase, gt_dict, pred_dict, hyperparams): |
| 100 | |
| 101 | gt_material=gt_dict["material"] |
| 102 | pred_material=pred_dict["material"] |
| 103 | nr_batches=gt_material.shape[0] |
| 104 | #pred material is usually in the range 0,1 but the first two values are slightly different so we rescale those |
| 105 | |
| 106 | pred_material[:,0]*=30 |
| 107 | pred_material[:,1]*=360 |
| 108 | |
| 109 | |
| 110 | loss_per_elem = ((gt_material-pred_material)**2) |
| 111 | |
| 112 | |
| 113 | |
| 114 | |
| 115 | gt_melanin=gt_material[:,3] |
| 116 | root_darkness_strength=gt_material[:,-1] |
| 117 | |
| 118 | |
| 119 | #root_darkenss should be downweighted in loss if the melanin is high, so if the hair is dark, it doesn't matter if we predict the correct root_darkness |
| 120 | root_darkness_weight = 1.0-gt_melanin |
| 121 | |
| 122 | loss_per_elem[:,0]*=0.0 #material_wave_scale |
| 123 | loss_per_elem[:,1]*=0.0 #material_wave_phase_offset |
| 124 | loss_per_elem[:,2]*=0.0 #material_wave_strength |
| 125 | loss_per_elem[:,3]*=1.0 #material_melanin_amount |
| 126 | loss_per_elem[:,4]*=1.0 #bsdf_melanin_redness |
| 127 | loss_per_elem[:,5]*=0.0 #bsdf_roughness |
| 128 | loss_per_elem[:,6]*=0.0 #bsdf_radial_roughness |
| 129 | loss_per_elem[:,7]*=0.0 #bsdf_coat |
| 130 | loss_per_elem[:,8]*=root_darkness_strength*root_darkness_weight #root_darkness_start |
| 131 | loss_per_elem[:,9]*=root_darkness_strength*root_darkness_weight #root_darkness_end |
| 132 | loss_per_elem[:,10]*=1.0*root_darkness_weight #root_darkness_strength |
| 133 | |
| 134 | |
| 135 | loss = loss_per_elem.mean() |
| 136 | |
| 137 | loss_dict={} |
| 138 | loss_dict["loss"]=loss |
| 139 | |
| 140 | return loss_dict |
| 141 | |
| 142 | |
| 143 | def prepare_gt_batch(batch, hyperparams, do_augmentation=False): |