Update inference.py

Face detection for CPU and GPU
justinjohn0306 · Sep 11, 2024 · d9c9717 · d9c9717
1 parent b725d4f
commit d9c9717
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/inference.py b/inference.py
@@ -294,9 +294,13 @@ def do_load(checkpoint_path):
     model = load_model(checkpoint_path)
 
     # SFDDetector.load_model(device)
-    detector = RetinaFace(gpu_id=0, model_path="checkpoints/mobilenet.pth", network="mobilenet")
+    # detector = RetinaFace(gpu_id=0, model_path="checkpoints/mobilenet.pth", network="mobilenet")
     # detector = RetinaFace(gpu_id=0, model_path="checkpoints/resnet50.pth", network="resnet50")
-
+    if torch.cuda.is_available():
+        detector = RetinaFace(gpu_id=0, model_path="checkpoints/mobilenet.pth", network="mobilenet")
+    else:
+        detector = RetinaFace( model_path="checkpoints/mobilenet.pth", network="mobilenet")
+
     detector_model = detector.model
 
     print("Models loaded")