Increased GPU usage

facelib/detection/retinaface/retinaface.py

@@ -11,7 +11,13 @@ from facelib.detection.retinaface.retinaface_net import FPN, SSH, MobileNetV1, m
 from facelib.detection.retinaface.retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,
                                                  py_cpu_nms)
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif torch.backends.mps.is_available():
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
 
 
 def generate_config(network_name):
@@ -367,4 +373,4 @@ class RetinaFace(nn.Module):
         # self.total_frame += len(frames)
         # print(self.batch_time / self.total_frame)
 
-        return final_bounding_boxes, final_landmarks
+        return final_bounding_boxes, final_landmarks

facelib/detection/yolov5face/face_detector.py

@@ -139,4 +139,4 @@ class YoloDetector:
             return None
 
     def __call__(self, *args):
-        return self.predict(*args)
+        return self.predict(*args)

facelib/detection/yolov5face/utils/extract_ckpt.py

@@ -1,5 +1,21 @@
 import torch
 import sys
-sys.path.insert(0,'./facelib/detection/yolov5face')
-model = torch.load('facelib/detection/yolov5face/yolov5n-face.pt', map_location='cpu')['model']
-torch.save(model.state_dict(),'weights/facelib/yolov5n-face.pth')
+import os
+
+# Setup dynamic device selection
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif torch.backends.mps.is_available():
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
+
+sys.path.insert(0, './facelib/detection/yolov5face')
+
+# Load the model to the selected device
+ckpt = torch.load('facelib/detection/yolov5face/yolov5n-face.pt', map_location=device)
+model = ckpt['model'].to(device)
+
+# Save only the weights
+os.makedirs('weights/facelib', exist_ok=True)
+torch.save(model.state_dict(), 'weights/facelib/yolov5n-face.pth')