(self, batch_dict)
| 49 | |
| 50 | |
| 51 | def forward(self, batch_dict): |
| 52 | |
| 53 | x=batch_dict["dinov2_latents"] #BCHW (1,1024,55,55) |
| 54 | |
| 55 | |
| 56 | #attempt 2, each patch predicts a confidence and a material, then we average all the materials across all patches, weighted by confidence |
| 57 | conf=self.dino2conf(x) |
| 58 | mat=self.dino2mat(x) |
| 59 | |
| 60 | |
| 61 | #average the mat across the pixels |
| 62 | avg_mat = (mat*conf).sum((2,3)) / (conf.sum((2,3)) +1e-6) #sum across all H and W dimensions |
| 63 | x=avg_mat |
| 64 | |
| 65 | |
| 66 | #split the material in parameters, at least the ones that are meaningfull and actually have a loss applied to them |
| 67 | melanin=x[:,3] |
| 68 | redness=x[:,4] |
| 69 | root_darkness_start=x[:,8] |
| 70 | root_darkness_end=x[:,9] |
| 71 | root_darkness_strength=x[:,10] |
| 72 | |
| 73 | |
| 74 | |
| 75 | pred_dict={} |
| 76 | pred_dict["material"]=x |
| 77 | pred_dict["melanin"]=melanin |
| 78 | pred_dict["redness"]=redness |
| 79 | pred_dict["root_darkness_start"]=root_darkness_start |
| 80 | pred_dict["root_darkness_end"]=root_darkness_end |
| 81 | pred_dict["root_darkness_strength"]=root_darkness_strength |
| 82 | |
| 83 | |
| 84 | return pred_dict |
nothing calls this directly
no outgoing calls
no test coverage detected