Class PixieEncoder

src/modules/pixie/pixie_encoder.py:35–641

class PixieEncoder(nn.Module):
    def __init__(self, config=None, ckpt_fp=''):
        super().__init__()
        if config is None:
            self.cfg = cfg
        else:
            self.cfg = config

        if os.path.exists(ckpt_fp):
            self.cfg.pretrained_modelpath = ckpt_fp

        # parameters setting
        self.param_list_dict = {}
        for lst in self.cfg.params.keys():
            param_list = cfg.params.get(lst)
            self.param_list_dict[lst] = {
                i: cfg.model.get('n_'+i) for i in param_list}

        # Build the models
        self._create_model()
        # Set up the cropping modules used to generate face/hand crops from the body predictions
        self._setup_cropper()

    def _setup_cropper(self):
        self.Cropper = {}
        for crop_part in ['head', 'hand']:
            data_cfg = self.cfg.dataset[crop_part]
            scale_size = (data_cfg.scale_min + data_cfg.scale_max)*0.5
            self.Cropper[crop_part] = tensor_cropper.Cropper(
                crop_size=data_cfg.image_size,
                scale=[scale_size, scale_size],
                trans_scale=0)

    def _create_model(self):
        self.model_dict = {}
        # Build all image encoders
        # Hand encoder only works for right hand, for left hand, flip inputs and flip the results back
        self.Encoder = {}
        for key in self.cfg.network.encoder.keys():
            if self.cfg.network.encoder.get(key).type == 'resnet50':
                self.Encoder[key] = ResnetEncoder()
            elif self.cfg.network.encoder.get(key).type == 'hrnet':
                self.Encoder[key] = HRNEncoder()
            self.model_dict[f'Encoder_{key}'] = self.Encoder[key].state_dict()

        # Build the parameter regressors
        self.Regressor = {}
        for key in self.cfg.network.regressor.keys():
            n_output = sum(self.param_list_dict[f'{key}_list'].values())
            channels = [2048] + \
                self.cfg.network.regressor.get(key).channels + [n_output]
            if self.cfg.network.regressor.get(key).type == 'mlp':
                self.Regressor[key] = MLP(channels=channels)
            self.model_dict[f'Regressor_{key}'] = self.Regressor[key].state_dict(
            )

        # Build the extractors
        # to extract separate head/left hand/right hand feature from body feature