Upload landmark_detection.py

landmark_detection.py

@@ -1,5 +1,7 @@
+import os
 import torch
 import torchvision
+import huggingface_hub
 from torchvision.transforms import InterpolationMode
 from network.models.facexformer import FaceXFormer
 from dataclasses import dataclass
@@ -12,8 +14,8 @@ import numpy as np
 # device = "cuda:0"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float32
-# weights_path = "ckpts/model.pt"
-weights_path = "ckpts/pytorch_model.bin"
+weights_path = "ckpts/model.pt"
+# weights_path = "ckpts/pytorch_model.bin"
 # face_model_path = "ckpts/blaze_face_short_range.tflite"
 
 # import mediapipe as mp
@@ -46,6 +48,15 @@ transforms_image = torchvision.transforms.Compose(
 
 def load_model(weights_path):
     model = FaceXFormer().to(device)
+    if not os.path.exists(weights_path):
+        name = os.path.basename(weights_path)
+        path = os.path.dirname(weights_path)
+        huggingface_hub.hf_hub_download(
+            "kartiknarayan/facexformer",
+            "ckpts/model.pt",
+            repo_type="model",
+            local_dir=path,
+        )
     checkpoint = torch.load(weights_path, map_location=device)
     model.load_state_dict(checkpoint)
     # model.load_state_dict(checkpoint["state_dict_backbone"])