|
|
|
|
|
""" |
|
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 |
|
|
