fix path / load face model

app.py

@@ -11,7 +11,7 @@ import retinaface
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 from diffusers import UniPCMultistepScheduler
 
-from spiga.inference.config import ModelConfig
+from modelconfig import ModelConfig
 from spiga.inference.framework import SPIGAFramework
 import spiga.demo.analyze.track.retinasort.config as cfg
 
@@ -26,7 +26,7 @@ face_detector = retinaface.RetinaFaceDetector(model=config['retina']['model_name
                                               extra_features=config['retina']['extra_features'],
                                               cfg_postreat=config['retina']['postreat'])
 # Landmark extraction
-spiga_extractor = SPIGAFramework(ModelConfig("300wpublic"))
+spiga_extractor = SPIGAFramework(ModelConfig("300wpublic", False))
 
 uncanny_controlnet = ControlNetModel.from_pretrained(
     "multimodalart/uncannyfaces_25K", torch_dtype=torch.float16

modelconfig.py

@@ -0,0 +1,62 @@
+from collections import OrderedDict
+
+from spiga.data.loaders.dl_config import DatabaseStruct
+
+MODELS_URL = {
+    "wflw": "https://drive.google.com/uc?export=download&confirm=yes&id=1h0qA5ysKorpeDNRXe9oYkVcVe8UYyzP7",
+    "300wpublic": "https://drive.google.com/uc?export=download&confirm=yes&id=1YrbScfMzrAAWMJQYgxdLZ9l57nmTdpQC",
+    "300wprivate": "https://drive.google.com/uc?export=download&confirm=yes&id=1fYv-Ie7n14eTD0ROxJYcn6SXZY5QU9SM",
+    "merlrav": "https://drive.google.com/uc?export=download&confirm=yes&id=1GKS1x0tpsTVivPZUk_yrSiMhwEAcAkg6",
+    "cofw68": "https://drive.google.com/uc?export=download&confirm=yes&id=1fYv-Ie7n14eTD0ROxJYcn6SXZY5QU9SM",
+}
+
+
+class ModelConfig(object):
+
+    def __init__(self, dataset_name=None, load_model_url=True):
+        # Model configuration
+        self.model_weights = None
+        self.model_weights_path = "./"
+        self.load_model_url = load_model_url
+        self.model_weights_url = None
+        # Pretreatment
+        self.focal_ratio = 1.5  # Camera matrix focal length ratio.
+        self.target_dist = 1.6  # Target distance zoom in/out around face.
+        self.image_size = (256, 256)
+        # Outputs
+        self.ftmap_size = (64, 64)
+        # Dataset
+        self.dataset = None
+
+        if dataset_name is not None:
+            self.update_with_dataset(dataset_name)
+
+    def update_with_dataset(self, dataset_name):
+
+        config_dict = {
+            "dataset": DatabaseStruct(dataset_name),
+            "model_weights": "spiga_%s.pt" % dataset_name,
+        }
+
+        if dataset_name == "cofw68":  # Test only
+            config_dict["model_weights"] = "spiga_300wprivate.pt"
+
+        if self.load_model_url:
+            config_dict["model_weights_url"] = MODELS_URL[dataset_name]
+
+        self.update(config_dict)
+
+    def update(self, params_dict):
+        state_dict = self.state_dict()
+        for k, v in params_dict.items():
+            if k in state_dict or hasattr(self, k):
+                setattr(self, k, v)
+            else:
+                raise Warning("Unknown option: {}: {}".format(k, v))
+
+    def state_dict(self):
+        state_dict = OrderedDict()
+        for k in self.__dict__.keys():
+            if not k.startswith("_"):
+                state_dict[k] = getattr(self, k)
+        return state_dict

spiga_300wpublic.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98f014611ac25d549e89083992d9e6ade15da133c634a3883473abf2953cee2d
+size 254397265