| |
|
|
| """ |
| Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image |
| """ |
|
|
| from torch import nn |
| import torch |
|
|
| from .convnextv2 import convnextv2_tiny |
| from .util import filter_state_dict |
|
|
| model_dict = { |
| 'convnextv2_tiny': convnextv2_tiny, |
| } |
|
|
|
|
| class MotionExtractor(nn.Module): |
| def __init__(self, **kwargs): |
| super(MotionExtractor, self).__init__() |
|
|
| |
| backbone = kwargs.get('backbone', 'convnextv2_tiny') |
| self.detector = model_dict.get(backbone)(**kwargs) |
|
|
| def load_pretrained(self, init_path: str): |
| if init_path not in (None, ''): |
| state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model'] |
| state_dict = filter_state_dict(state_dict, remove_name='head') |
| ret = self.detector.load_state_dict(state_dict, strict=False) |
| print(f'Load pretrained model from {init_path}, ret: {ret}') |
|
|
| def forward(self, x): |
| out = self.detector(x) |
| return out |
|
|
