| 39 | |
| 40 | |
| 41 | class PIXIE(object): |
| 42 | def __init__(self, config=None, device='cuda:0'): |
| 43 | if config is None: |
| 44 | self.cfg = cfg |
| 45 | else: |
| 46 | self.cfg = config |
| 47 | |
| 48 | self.device = device |
| 49 | # parameters setting |
| 50 | self.param_list_dict = {} |
| 51 | for lst in self.cfg.params.keys(): |
| 52 | param_list = cfg.params.get(lst) |
| 53 | self.param_list_dict[lst] = { |
| 54 | i: cfg.model.get('n_'+i) for i in param_list} |
| 55 | |
| 56 | # Build the models |
| 57 | self._create_model() |
| 58 | # Set up the cropping modules used to generate face/hand crops from the body predictions |
| 59 | self._setup_cropper() |
| 60 | |
| 61 | def _setup_cropper(self): |
| 62 | self.Cropper = {} |
| 63 | for crop_part in ['head', 'hand']: |
| 64 | data_cfg = self.cfg.dataset[crop_part] |
| 65 | scale_size = (data_cfg.scale_min + data_cfg.scale_max)*0.5 |
| 66 | self.Cropper[crop_part] = tensor_cropper.Cropper( |
| 67 | crop_size=data_cfg.image_size, |
| 68 | scale=[scale_size, scale_size], |
| 69 | trans_scale=0) |
| 70 | |
| 71 | def _create_model(self): |
| 72 | self.model_dict = {} |
| 73 | # Build all image encoders |
| 74 | # Hand encoder only works for right hand, for left hand, flip inputs and flip the results back |
| 75 | self.Encoder = {} |
| 76 | for key in self.cfg.network.encoder.keys(): |
| 77 | if self.cfg.network.encoder.get(key).type == 'resnet50': |
| 78 | self.Encoder[key] = ResnetEncoder().to(self.device) |
| 79 | elif self.cfg.network.encoder.get(key).type == 'hrnet': |
| 80 | self.Encoder[key] = HRNEncoder().to(self.device) |
| 81 | self.model_dict[f'Encoder_{key}'] = self.Encoder[key].state_dict() |
| 82 | |
| 83 | # Build the parameter regressors |
| 84 | self.Regressor = {} |
| 85 | for key in self.cfg.network.regressor.keys(): |
| 86 | n_output = sum(self.param_list_dict[f'{key}_list'].values()) |
| 87 | channels = [2048] + \ |
| 88 | self.cfg.network.regressor.get(key).channels + [n_output] |
| 89 | if self.cfg.network.regressor.get(key).type == 'mlp': |
| 90 | self.Regressor[key] = MLP(channels=channels).to(self.device) |
| 91 | self.model_dict[f'Regressor_{key}'] = self.Regressor[key].state_dict( |
| 92 | ) |
| 93 | |
| 94 | # Build the extractors |
| 95 | # to extract separate head/left hand/right hand feature from body feature |
| 96 | self.Extractor = {} |
| 97 | for key in self.cfg.network.extractor.keys(): |
| 98 | channels = [2048] + \ |
nothing calls this directly
no outgoing calls
no test coverage detected