From d9c97178f8e72bcdb8f38c9fa85ec567be3cdcf9 Mon Sep 17 00:00:00 2001 From: magicse Date: Wed, 11 Sep 2024 23:49:16 +0300 Subject: [PATCH] Update inference.py Face detection for CPU and GPU --- inference.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/inference.py b/inference.py index fe65937d8..5e1522d25 100644 --- 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")