(self)
| 66 | trans_scale=0) |
| 67 | |
| 68 | def _create_model(self): |
| 69 | self.model_dict = {} |
| 70 | # Build all image encoders |
| 71 | # Hand encoder only works for right hand, for left hand, flip inputs and flip the results back |
| 72 | self.Encoder = {} |
| 73 | for key in self.cfg.network.encoder.keys(): |
| 74 | if self.cfg.network.encoder.get(key).type == 'resnet50': |
| 75 | self.Encoder[key] = ResnetEncoder() |
| 76 | elif self.cfg.network.encoder.get(key).type == 'hrnet': |
| 77 | self.Encoder[key] = HRNEncoder() |
| 78 | self.model_dict[f'Encoder_{key}'] = self.Encoder[key].state_dict() |
| 79 | |
| 80 | # Build the parameter regressors |
| 81 | self.Regressor = {} |
| 82 | for key in self.cfg.network.regressor.keys(): |
| 83 | n_output = sum(self.param_list_dict[f'{key}_list'].values()) |
| 84 | channels = [2048] + \ |
| 85 | self.cfg.network.regressor.get(key).channels + [n_output] |
| 86 | if self.cfg.network.regressor.get(key).type == 'mlp': |
| 87 | self.Regressor[key] = MLP(channels=channels) |
| 88 | self.model_dict[f'Regressor_{key}'] = self.Regressor[key].state_dict( |
| 89 | ) |
| 90 | |
| 91 | # Build the extractors |
| 92 | # to extract separate head/left hand/right hand feature from body feature |
| 93 | self.Extractor = {} |
| 94 | for key in self.cfg.network.extractor.keys(): |
| 95 | channels = [2048] + \ |
| 96 | self.cfg.network.extractor.get(key).channels + [2048] |
| 97 | if self.cfg.network.extractor.get(key).type == 'mlp': |
| 98 | self.Extractor[key] = MLP(channels=channels) |
| 99 | self.model_dict[f'Extractor_{key}'] = self.Extractor[key].state_dict( |
| 100 | ) |
| 101 | |
| 102 | # Build the moderators |
| 103 | self.Moderator = {} |
| 104 | for key in self.cfg.network.moderator.keys(): |
| 105 | detach_inputs = self.cfg.network.moderator.get(key).detach_inputs |
| 106 | detach_feature = self.cfg.network.moderator.get(key).detach_feature |
| 107 | channels = [2048*2] + \ |
| 108 | self.cfg.network.moderator.get(key).channels + [2] |
| 109 | self.Moderator[key] = TempSoftmaxFusion( |
| 110 | detach_inputs=detach_inputs, detach_feature=detach_feature, |
| 111 | channels=channels) |
| 112 | self.model_dict[f'Moderator_{key}'] = self.Moderator[key].state_dict( |
| 113 | ) |
| 114 | |
| 115 | # Build the SMPL-X body model, which we also use to represent faces and |
| 116 | # hands, using the relevant parts only |
| 117 | self.smplx = SMPLX('assets/SMPLX', n_shape=200, n_exp=50) |
| 118 | self.part_indices = self.smplx.part_indices |
| 119 | |
| 120 | # -- resume model |
| 121 | model_path = self.cfg.pretrained_modelpath |
| 122 | if os.path.exists(model_path): |
| 123 | checkpoint = torch.load(model_path) |
| 124 | for key in self.model_dict.keys(): |
| 125 | util.copy_state_dict(self.model_dict[key], checkpoint[key]) |
no test coverage detected