pre-trained models adjusted

src/pixel3dmm/env_paths.py

@@ -1,7 +1,7 @@
 import json
 from pathlib import Path
 from environs import Env
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 
 env = Env(expand_vars=True)
 env_file_path = Path(f"{Path.home()}/.config/pixel3dmm/.env")
@@ -14,36 +14,18 @@ with env.prefixed("PIXEL3DMM_"):
     PREPROCESSED_DATA = env("PREPROCESSED_DATA")
     TRACKING_OUTPUT = env("TRACKING_OUTPUT")
 
-FLAME_ASSETS = hf_hub_download(
-            repo_id="alexnasa/pixel3dmm",
-            filename="generic_model.pkl",
-            repo_type="model", 
+base = snapshot_download(
+            repo_id="alexnasa/pixel3dmm",     # your model repo
+            repo_type="model",                # model vs dataset
         )
 
-MICA_TAR_ASSET = hf_hub_download(
-            repo_id="alexnasa/pixel3dmm",
-            filename="mica.tar",
-            repo_type="model", 
-        )
-
-
-PIPNET_LOCAL_ASSET = hf_hub_download(
-            repo_id="alexnasa/pixel3dmm",
-            filename="epoch59.pth",
-            repo_type="model", 
-        )
-
-CKPT_N_PRED = hf_hub_download(
-            repo_id="alexnasa/pixel3dmm",
-            filename="normals.ckpt",
-            repo_type="model", 
-        )
-
-CKPT_UV_PRED = hf_hub_download(
-            repo_id="alexnasa/pixel3dmm",
-            filename="uv.ckpt",
-            repo_type="model", 
-        )
+FLAME_ASSET       = os.path.join(base, "generic_model.pkl")
+MICA_TAR_ASSET    = os.path.join(base, "mica.tar")
+PIPNET_LOCAL_ASSET= os.path.join(base, "epoch59.pth")
+CKPT_N_PRED       = os.path.join(base, "normals.ckpt")
+CKPT_UV_PRED      = os.path.join(base, "uv.ckpt")
+ANT_DIR           = os.path.join(base_dir, "antelopev2")
+BUFFALO_DIR       = os.path.join(base_dir, "buffalo_l")
 
 head_template = f'{CODE_BASE}/assets/head_template.obj'
 head_template_color = f'{CODE_BASE}/assets/head_template_color.obj'

src/pixel3dmm/preprocessing/MICA/datasets/creation/generator.py

@@ -24,6 +24,7 @@ from typing import List
 import cv2
 import numpy as np
 from insightface.app import FaceAnalysis
+from pixel3dmm import env_paths
 from insightface.app.common import Face
 from insightface.utils import face_align
 from loguru import logger
@@ -66,7 +67,9 @@ class Generator:
             instance.preprocess()
 
     def arcface(self):
-        app = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider'])
+        app = FaceAnalysis(name='antelopev2', 
+                           root=f'{env_paths.ANT_DIR}', 
+                           providers=['CUDAExecutionProvider'])
         app.prepare(ctx_id=0, det_size=(224, 224))
 
         logger.info('Start arcface...')

src/pixel3dmm/preprocessing/MICA/micalib/tester.py

@@ -23,6 +23,8 @@ import numpy as np
 import torch
 import torch.distributed as dist
 from insightface.app import FaceAnalysis
+from pixel3dmm import env_paths
+ 
 from insightface.app.common import Face
 from insightface.utils import face_align
 from loguru import logger
@@ -183,7 +185,9 @@ class Tester(object):
         else:
             cache = {}
 
-        app = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider'])
+        app = FaceAnalysis(name='antelopev2', 
+                            root=f'{env_paths.ANT_DIR}',
+                            providers=['CUDAExecutionProvider'])
         app.prepare(ctx_id=0, det_size=(224, 224), det_thresh=0.4)
 
         for actor in tqdm(sorted(os.listdir(NOW_PICTURES))):
@@ -202,7 +206,9 @@ class Tester(object):
         else:
             cache = {}
 
-        app = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider'])
+        app = FaceAnalysis(name='antelopev2', 
+                            root=f'{env_paths.ANT_DIR}',
+                            providers=['CUDAExecutionProvider'])
         app.prepare(ctx_id=0, det_size=(224, 224), det_thresh=0.1)
 
         cache['HQ'] = {}

src/pixel3dmm/preprocessing/MICA/utils/landmark_detector.py

@@ -17,6 +17,7 @@
 import face_alignment
 import numpy as np
 from insightface.app import FaceAnalysis
+from pixel3dmm import env_paths
 from loguru import logger
 
 from datasets.creation.util import get_bbox
@@ -36,7 +37,9 @@ class LandmarksDetector:
         model = model.upper()
         self.predictor = model
         if model == detectors.RETINAFACE:
-            self._face_detector = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider'])
+            self._face_detector = FaceAnalysis(name='antelopev2', 
+                            root=f'{env_paths.ANT_DIR}',
+                            providers=['CUDAExecutionProvider'])
             self._face_detector.prepare(ctx_id=0, det_size=(224, 224))
         elif model == detectors.FAN:
             self._face_detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device=device)