| 7 | |
| 8 | |
| 9 | class RGB2MaterialModel(nn.Module): |
| 10 | |
| 11 | def __init__(self, input_dim, out_dim, hidden_dim): |
| 12 | super().__init__() |
| 13 | |
| 14 | self.out_dim=out_dim |
| 15 | |
| 16 | |
| 17 | |
| 18 | #attempt 2 |
| 19 | self.dino2conf=nn.Sequential( |
| 20 | nn.Conv2d(in_channels=input_dim, out_channels=hidden_dim, kernel_size=1, padding=0, bias=True), |
| 21 | nn.SiLU(), |
| 22 | nn.Conv2d(in_channels=hidden_dim, out_channels=1, kernel_size=1, padding=0, bias=True), |
| 23 | nn.Sigmoid() |
| 24 | ) |
| 25 | self.dino2mat=nn.Sequential( |
| 26 | nn.Conv2d(in_channels=input_dim, out_channels=hidden_dim, kernel_size=1, padding=0, bias=True), |
| 27 | nn.SiLU(), |
| 28 | nn.Conv2d(in_channels=hidden_dim, out_channels=out_dim, kernel_size=1, padding=0, bias=True), |
| 29 | nn.Sigmoid() |
| 30 | ) |
| 31 | |
| 32 | |
| 33 | |
| 34 | |
| 35 | self.apply(lambda x: kaiming_init(x, False, nonlinearity="silu")) |
| 36 | |
| 37 | def save(self, root_folder, experiment_name, hyperparams, iter_nr, info=None): |
| 38 | name=str(iter_nr) |
| 39 | if info is not None: |
| 40 | name+="_"+info |
| 41 | models_path = os.path.join(root_folder, experiment_name, name, "models") |
| 42 | if not os.path.exists(models_path): |
| 43 | os.makedirs(models_path, exist_ok=True) |
| 44 | torch.save(self.state_dict(), os.path.join(models_path, "rgb2material.pt")) |
| 45 | |
| 46 | hyperparams_params_path=os.path.join(models_path, "hyperparams.json") |
| 47 | with open(hyperparams_params_path, 'w', encoding='utf-8') as f: |
| 48 | json.dump(vars(hyperparams), f, ensure_ascii=False, indent=4) |
| 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 |