RabbitRUI/ruispace

1"""This script defines the face reconstruction model for Deep3DFaceRecon_pytorch2"""3 4import numpy as np5import torch6from src.face3d.models.base_model import BaseModel7from src.face3d.models import networks8from src.face3d.models.bfm import ParametricFaceModel9from src.face3d.models.losses import perceptual_loss, photo_loss, reg_loss, reflectance_loss, landmark_loss10from src.face3d.util import util 11from src.face3d.util.nvdiffrast import MeshRenderer12# from src.face3d.util.preprocess import estimate_norm_torch13 14import trimesh15from scipy.io import savemat16 17class FaceReconModel(BaseModel):18 19    @staticmethod20    def modify_commandline_options(parser, is_train=False):21        """  Configures options specific for CUT model22        """23        # net structure and parameters24        parser.add_argument('--net_recon', type=str, default='resnet50', choices=['resnet18', 'resnet34', 'resnet50'], help='network structure')25        parser.add_argument('--init_path', type=str, default='./checkpoints/init_model/resnet50-0676ba61.pth')26        parser.add_argument('--use_last_fc', type=util.str2bool, nargs='?', const=True, default=False, help='zero initialize the last fc')27        parser.add_argument('--bfm_folder', type=str, default='./checkpoints/BFM_Fitting/')28        parser.add_argument('--bfm_model', type=str, default='BFM_model_front.mat', help='bfm model')29 30        # renderer parameters31        parser.add_argument('--focal', type=float, default=1015.)32        parser.add_argument('--center', type=float, default=112.)33        parser.add_argument('--camera_d', type=float, default=10.)34        parser.add_argument('--z_near', type=float, default=5.)35        parser.add_argument('--z_far', type=float, default=15.)36 37        if is_train:38            # training parameters39            parser.add_argument('--net_recog', type=str, default='r50', choices=['r18', 'r43', 'r50'], help='face recog network structure')40            parser.add_argument('--net_recog_path', type=str, default='checkpoints/recog_model/ms1mv3_arcface_r50_fp16/backbone.pth')41            parser.add_argument('--use_crop_face', type=util.str2bool, nargs='?', const=True, default=False, help='use crop mask for photo loss')42            parser.add_argument('--use_predef_M', type=util.str2bool, nargs='?', const=True, default=False, help='use predefined M for predicted face')43 44            45            # augmentation parameters46            parser.add_argument('--shift_pixs', type=float, default=10., help='shift pixels')47            parser.add_argument('--scale_delta', type=float, default=0.1, help='delta scale factor')48            parser.add_argument('--rot_angle', type=float, default=10., help='rot angles, degree')49 50            # loss weights51            parser.add_argument('--w_feat', type=float, default=0.2, help='weight for feat loss')52            parser.add_argument('--w_color', type=float, default=1.92, help='weight for loss loss')53            parser.add_argument('--w_reg', type=float, default=3.0e-4, help='weight for reg loss')54            parser.add_argument('--w_id', type=float, default=1.0, help='weight for id_reg loss')55            parser.add_argument('--w_exp', type=float, default=0.8, help='weight for exp_reg loss')56            parser.add_argument('--w_tex', type=float, default=1.7e-2, help='weight for tex_reg loss')57            parser.add_argument('--w_gamma', type=float, default=10.0, help='weight for gamma loss')58            parser.add_argument('--w_lm', type=float, default=1.6e-3, help='weight for lm loss')59            parser.add_argument('--w_reflc', type=float, default=5.0, help='weight for reflc loss')60 61        opt, _ = parser.parse_known_args()62        parser.set_defaults(63                focal=1015., center=112., camera_d=10., use_last_fc=False, z_near=5., z_far=15.64            )65        if is_train:66            parser.set_defaults(67                use_crop_face=True, use_predef_M=False68            )69        return parser70 71    def __init__(self, opt):72        """Initialize this model class.73 74        Parameters:75            opt -- training/test options76 77        A few things can be done here.78        - (required) call the initialization function of BaseModel79        - define loss function, visualization images, model names, and optimizers80        """81        BaseModel.__init__(self, opt)  # call the initialization method of BaseModel82        83        self.visual_names = ['output_vis']84        self.model_names = ['net_recon']85        self.parallel_names = self.model_names + ['renderer']86 87        self.facemodel = ParametricFaceModel(88            bfm_folder=opt.bfm_folder, camera_distance=opt.camera_d, focal=opt.focal, center=opt.center,89            is_train=self.isTrain, default_name=opt.bfm_model90        )91        92        fov = 2 * np.arctan(opt.center / opt.focal) * 180 / np.pi93        self.renderer = MeshRenderer(94            rasterize_fov=fov, znear=opt.z_near, zfar=opt.z_far, rasterize_size=int(2 * opt.center)95        )96 97        if self.isTrain:98            self.loss_names = ['all', 'feat', 'color', 'lm', 'reg', 'gamma', 'reflc']99 100            self.net_recog = networks.define_net_recog(101                net_recog=opt.net_recog, pretrained_path=opt.net_recog_path102                )103            # loss func name: (compute_%s_loss) % loss_name104            self.compute_feat_loss = perceptual_loss105            self.comupte_color_loss = photo_loss106            self.compute_lm_loss = landmark_loss107            self.compute_reg_loss = reg_loss108            self.compute_reflc_loss = reflectance_loss109 110            self.optimizer = torch.optim.Adam(self.net_recon.parameters(), lr=opt.lr)111            self.optimizers = [self.optimizer]112            self.parallel_names += ['net_recog']113        # Our program will automatically call <model.setup> to define schedulers, load networks, and print networks114 115    def set_input(self, input):116        """Unpack input data from the dataloader and perform necessary pre-processing steps.117 118        Parameters:119            input: a dictionary that contains the data itself and its metadata information.120        """121        self.input_img = input['imgs'].to(self.device) 122        self.atten_mask = input['msks'].to(self.device) if 'msks' in input else None123        self.gt_lm = input['lms'].to(self.device)  if 'lms' in input else None124        self.trans_m = input['M'].to(self.device) if 'M' in input else None125        self.image_paths = input['im_paths'] if 'im_paths' in input else None126 127    def forward(self, output_coeff, device):128        self.facemodel.to(device)129        self.pred_vertex, self.pred_tex, self.pred_color, self.pred_lm = \130            self.facemodel.compute_for_render(output_coeff)131        self.pred_mask, _, self.pred_face = self.renderer(132            self.pred_vertex, self.facemodel.face_buf, feat=self.pred_color)133        134        self.pred_coeffs_dict = self.facemodel.split_coeff(output_coeff)135 136 137    def compute_losses(self):138        """Calculate losses, gradients, and update network weights; called in every training iteration"""139 140        assert self.net_recog.training == False141        trans_m = self.trans_m142        if not self.opt.use_predef_M:143            trans_m = estimate_norm_torch(self.pred_lm, self.input_img.shape[-2])144 145        pred_feat = self.net_recog(self.pred_face, trans_m)146        gt_feat = self.net_recog(self.input_img, self.trans_m)147        self.loss_feat = self.opt.w_feat * self.compute_feat_loss(pred_feat, gt_feat)148 149        face_mask = self.pred_mask150        if self.opt.use_crop_face:151            face_mask, _, _ = self.renderer(self.pred_vertex, self.facemodel.front_face_buf)152        153        face_mask = face_mask.detach()154        self.loss_color = self.opt.w_color * self.comupte_color_loss(155            self.pred_face, self.input_img, self.atten_mask * face_mask)156        157        loss_reg, loss_gamma = self.compute_reg_loss(self.pred_coeffs_dict, self.opt)158        self.loss_reg = self.opt.w_reg * loss_reg159        self.loss_gamma = self.opt.w_gamma * loss_gamma160 161        self.loss_lm = self.opt.w_lm * self.compute_lm_loss(self.pred_lm, self.gt_lm)162 163        self.loss_reflc = self.opt.w_reflc * self.compute_reflc_loss(self.pred_tex, self.facemodel.skin_mask)164 165        self.loss_all = self.loss_feat + self.loss_color + self.loss_reg + self.loss_gamma \166                        + self.loss_lm + self.loss_reflc167            168 169    def optimize_parameters(self, isTrain=True):170        self.forward()               171        self.compute_losses()172        """Update network weights; it will be called in every training iteration."""173        if isTrain:174            self.optimizer.zero_grad()  175            self.loss_all.backward()         176            self.optimizer.step()        177 178    def compute_visuals(self):179        with torch.no_grad():180            input_img_numpy = 255. * self.input_img.detach().cpu().permute(0, 2, 3, 1).numpy()181            output_vis = self.pred_face * self.pred_mask + (1 - self.pred_mask) * self.input_img182            output_vis_numpy_raw = 255. * output_vis.detach().cpu().permute(0, 2, 3, 1).numpy()183            184            if self.gt_lm is not None:185                gt_lm_numpy = self.gt_lm.cpu().numpy()186                pred_lm_numpy = self.pred_lm.detach().cpu().numpy()187                output_vis_numpy = util.draw_landmarks(output_vis_numpy_raw, gt_lm_numpy, 'b')188                output_vis_numpy = util.draw_landmarks(output_vis_numpy, pred_lm_numpy, 'r')189            190                output_vis_numpy = np.concatenate((input_img_numpy, 191                                    output_vis_numpy_raw, output_vis_numpy), axis=-2)192            else:193                output_vis_numpy = np.concatenate((input_img_numpy, 194                                    output_vis_numpy_raw), axis=-2)195 196            self.output_vis = torch.tensor(197                    output_vis_numpy / 255., dtype=torch.float32198                ).permute(0, 3, 1, 2).to(self.device)199 200    def save_mesh(self, name):201 202        recon_shape = self.pred_vertex  # get reconstructed shape203        recon_shape[..., -1] = 10 - recon_shape[..., -1] # from camera space to world space204        recon_shape = recon_shape.cpu().numpy()[0]205        recon_color = self.pred_color206        recon_color = recon_color.cpu().numpy()[0]207        tri = self.facemodel.face_buf.cpu().numpy()208        mesh = trimesh.Trimesh(vertices=recon_shape, faces=tri, vertex_colors=np.clip(255. * recon_color, 0, 255).astype(np.uint8))209        mesh.export(name)210 211    def save_coeff(self,name):212 213        pred_coeffs = {key:self.pred_coeffs_dict[key].cpu().numpy() for key in self.pred_coeffs_dict}214        pred_lm = self.pred_lm.cpu().numpy()215        pred_lm = np.stack([pred_lm[:,:,0],self.input_img.shape[2]-1-pred_lm[:,:,1]],axis=2) # transfer to image coordinate216        pred_coeffs['lm68'] = pred_lm217        savemat(name,pred_coeffs)218 219 220 221