Upload demo.py

src/pixel3dmm/preprocessing/MICA/demo.py

@@ -0,0 +1,188 @@
+# -*- 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: [email protected]
+
+
+import argparse
+import os
+import random
+import traceback
+from glob import glob
+from pathlib import Path
+from PIL import Image
+
+import cv2
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+import trimesh
+from insightface.app.common import Face
+from insightface.utils import face_align
+from loguru import logger
+from skimage.io import imread
+from tqdm import tqdm
+#from retinaface.pre_trained_models import get_model
+#from retinaface.utils import vis_annotations
+#from matplotlib import pyplot as plt
+
+
+from pixel3dmm.preprocessing.MICA.configs.config import get_cfg_defaults
+from pixel3dmm.preprocessing.MICA.datasets.creation.util import get_arcface_input, get_center, draw_on
+from pixel3dmm.preprocessing.MICA.utils import util
+from pixel3dmm.preprocessing.MICA.utils.landmark_detector import LandmarksDetector, detectors
+from pixel3dmm import env_paths
+
+
+#model = get_model("resnet50_2020-07-20", max_size=512)
+#model.eval()
+
+
+def deterministic(rank):
+    torch.manual_seed(rank)
+    torch.cuda.manual_seed(rank)
+    np.random.seed(rank)
+    random.seed(rank)
+
+    cudnn.deterministic = True
+    cudnn.benchmark = False
+
+
+def process(args, app, image_size=224, draw_bbox=False):
+    dst = Path(args.a)
+    dst.mkdir(parents=True, exist_ok=True)
+    processes = []
+    image_paths = sorted(glob(args.i + '/*.*'))#[:1]
+    image_paths = image_paths[::max(1, len(image_paths)//10)]
+    for image_path in tqdm(image_paths):
+        name = Path(image_path).stem
+        img = cv2.imread(image_path)
+
+
+        # FOR pytorch retinaface use this: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        # I had issues with onnxruntime!
+        bboxes, kpss = app.detect(img)
+
+        #annotation = model.predict_jsons(img)
+        #Image.fromarray(vis_annotations(img, annotation)).show()
+
+        #bboxes = np.stack([np.array( annotation[0]['bbox'] + [annotation[0]['score']] ) for i in range(len(annotation))], axis=0)
+        #kpss = np.stack([np.array( annotation[0]['landmarks'] ) for i in range(len(annotation))], axis=0)
+        if bboxes.shape[0] == 0:
+            logger.error(f'[ERROR] Face not detected for {image_path}')
+            continue
+        i = get_center(bboxes, img)
+        bbox = bboxes[i, 0:4]
+        det_score = bboxes[i, 4]
+        kps = None
+        if kpss is not None:
+            kps = kpss[i]
+
+        ##for ikp in range(kps.shape[0]):
+        #    img[int(kps[ikp][1]), int(kps[ikp][0]), 0] = 255
+        #    img[int(kpss_[0][ikp][1]), int(kpss_[0][ikp][0]), 1] = 255
+        #Image.fromarray(img).show()
+        face = Face(bbox=bbox, kps=kps, det_score=det_score)
+        blob, aimg = get_arcface_input(face, img)
+        file = str(Path(dst, name))
+        np.save(file, blob)
+        processes.append(file + '.npy')
+        cv2.imwrite(file + '.jpg', face_align.norm_crop(img, landmark=face.kps, image_size=image_size))
+        if draw_bbox:
+            dimg = draw_on(img, [face])
+            cv2.imwrite(file + '_bbox.jpg', dimg)
+
+    return processes
+
+
+def to_batch(path):
+    src = path.replace('npy', 'jpg')
+    if not os.path.exists(src):
+        src = path.replace('npy', 'png')
+
+    image = imread(src)[:, :, :3]
+    image = image / 255.
+    image = cv2.resize(image, (224, 224)).transpose(2, 0, 1)
+    image = torch.tensor(image).cuda()[None]
+
+    arcface = np.load(path)
+    arcface = torch.tensor(arcface).cuda()[None]
+
+    return image, arcface
+
+
+def load_checkpoint(args, mica):
+    checkpoint = torch.load(args.m, weights_only=False)
+    if 'arcface' in checkpoint:
+        mica.arcface.load_state_dict(checkpoint['arcface'])
+    if 'flameModel' in checkpoint:
+        mica.flameModel.load_state_dict(checkpoint['flameModel'])
+
+
+def main(cfg, args):
+    device = 'cuda:0'
+    cfg.model.testing = True
+    mica = util.find_model_using_name(model_dir='micalib.models', model_name=cfg.model.name)(cfg, device)
+    load_checkpoint(args, mica)
+    mica.eval()
+
+    faces = mica.flameModel.generator.faces_tensor.cpu()
+    Path(args.o).mkdir(exist_ok=True, parents=True)
+
+    app = LandmarksDetector(model=detectors.RETINAFACE)
+
+    with torch.no_grad():
+        logger.info(f'Processing has started...')
+        paths = process(args, app, draw_bbox=False)
+        for path in tqdm(paths):
+            name = Path(path).stem
+            images, arcface = to_batch(path)
+            codedict = mica.encode(images, arcface)
+            opdict = mica.decode(codedict)
+            meshes = opdict['pred_canonical_shape_vertices']
+            code = opdict['pred_shape_code']
+            lmk = mica.flameModel.generator.compute_landmarks(meshes)
+
+            mesh = meshes[0]
+            landmark_51 = lmk[0, 17:]
+            landmark_7 = landmark_51[[19, 22, 25, 28, 16, 31, 37]]
+
+            dst = Path(args.o, name)
+            dst.mkdir(parents=True, exist_ok=True)
+            trimesh.Trimesh(vertices=mesh.cpu() * 1000.0, faces=faces, process=False).export(f'{dst}/mesh.ply')  # save in millimeters
+            trimesh.Trimesh(vertices=mesh.cpu() * 1000.0, faces=faces, process=False).export(f'{dst}/mesh.obj')
+            np.save(f'{dst}/identity', code[0].cpu().numpy())
+            np.save(f'{dst}/kpt7', landmark_7.cpu().numpy() * 1000.0)
+            np.save(f'{dst}/kpt68', lmk.cpu().numpy() * 1000.0)
+
+        logger.info(f'Processing finished. Results has been saved in {args.o}')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='MICA - Towards Metrical Reconstruction of Human Faces')
+    parser.add_argument('-video_name', required=True, type=str)
+    parser.add_argument('-a', default='demo/arcface', type=str, help='Processed images for MICA input')
+    parser.add_argument('-m', default='data/pretrained/mica.tar', type=str, help='Pretrained model path')
+
+    args = parser.parse_args()
+    cfg = get_cfg_defaults()
+    args.i = f'{env_paths.PREPROCESSED_DATA}/{args.video_name}/cropped/'
+    args.o = f'{env_paths.PREPROCESSED_DATA}/{args.video_name}/mica/'
+    if os.path.exists(f'{env_paths.PREPROCESSED_DATA}/{args.video_name}/mica/'):
+        if len(os.listdir(f'{env_paths.PREPROCESSED_DATA}/{args.video_name}/mica/')) >= 10:
+            print(f'''
+                            <<<<<<<< ALREADY COMPLETE MICA PREDICTION FOR {args.video_name}, SKIPPING >>>>>>>>
+                            ''')
+            exit()
+    main(cfg, args)