Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,295 Bytes
cf92dec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2023 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: mica@tue.mpg.de
import os
import sys
sys.path.append("./nfclib")
import torch
import torch.nn.functional as F
from models.arcface import Arcface
from models.generator import Generator
from micalib.base_model import BaseModel
from loguru import logger
class MICA(BaseModel):
def __init__(self, config=None, device=None, tag='MICA'):
super(MICA, self).__init__(config, device, tag)
self.initialize()
def create_model(self, model_cfg):
mapping_layers = model_cfg.mapping_layers
pretrained_path = None
if not model_cfg.use_pretrained:
pretrained_path = model_cfg.arcface_pretrained_model
self.arcface = Arcface(pretrained_path=pretrained_path).to(self.device)
self.flameModel = Generator(512, 300, self.cfg.model.n_shape, mapping_layers, model_cfg, self.device)
def load_model(self):
model_path = os.path.join(self.cfg.output_dir, 'model.tar')
if os.path.exists(self.cfg.pretrained_model_path) and self.cfg.model.use_pretrained:
model_path = self.cfg.pretrained_model_path
if os.path.exists(model_path):
logger.info(f'[{self.tag}] Trained model found. Path: {model_path} | GPU: {self.device}')
checkpoint = torch.load(model_path, weights_only=False)
if 'arcface' in checkpoint:
self.arcface.load_state_dict(checkpoint['arcface'])
if 'flameModel' in checkpoint:
self.flameModel.load_state_dict(checkpoint['flameModel'])
else:
logger.info(f'[{self.tag}] Checkpoint not available starting from scratch!')
def model_dict(self):
return {
'flameModel': self.flameModel.state_dict(),
'arcface': self.arcface.state_dict()
}
def parameters_to_optimize(self):
return [
{'params': self.flameModel.parameters(), 'lr': self.cfg.train.lr},
{'params': self.arcface.parameters(), 'lr': self.cfg.train.arcface_lr},
]
def encode(self, images, arcface_imgs):
codedict = {}
codedict['arcface'] = F.normalize(self.arcface(arcface_imgs))
codedict['images'] = images
return codedict
def decode(self, codedict, epoch=0):
self.epoch = epoch
flame_verts_shape = None
shapecode = None
if not self.testing:
flame = codedict['flame']
shapecode = flame['shape_params'].view(-1, flame['shape_params'].shape[2])
shapecode = shapecode.to(self.device)[:, :self.cfg.model.n_shape]
with torch.no_grad():
flame_verts_shape, _, _ = self.flame(shape_params=shapecode)
identity_code = codedict['arcface']
pred_canonical_vertices, pred_shape_code = self.flameModel(identity_code)
output = {
'flame_verts_shape': flame_verts_shape,
'flame_shape_code': shapecode,
'pred_canonical_shape_vertices': pred_canonical_vertices,
'pred_shape_code': pred_shape_code,
'faceid': codedict['arcface']
}
return output
def compute_losses(self, input, encoder_output, decoder_output):
losses = {}
pred_verts = decoder_output['pred_canonical_shape_vertices']
gt_verts = decoder_output['flame_verts_shape'].detach()
pred_verts_shape_canonical_diff = (pred_verts - gt_verts).abs()
if self.use_mask:
pred_verts_shape_canonical_diff *= self.vertices_mask
losses['pred_verts_shape_canonical_diff'] = torch.mean(pred_verts_shape_canonical_diff) * 1000.0
return losses
|