Update utils.py

utils.py

@@ -8,6 +8,7 @@ import cv2
 from models.mtcnn import MTCNN
 from models.inception_resnet_v1 import InceptionResnetV1
 
+
 device = 'cpu'
 
 SINGLE_FACE_PATH = "data/test_images"
@@ -17,11 +18,11 @@ THRESHOLD = 0.8
 mtcnn = MTCNN(
     image_size=160, margin=0, min_face_size=20,
     thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
-    device='cpu',
+    device=device,
     keep_all=True
 )
 
-resnet = InceptionResnetV1(pretrained='vggface2').eval().to('cpu')
+resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)
 
 def extract_face(box, img, margin=20):
     img_size = (len(img), len(img[0]))
@@ -52,9 +53,9 @@ def trans(img):
 
 def find_and_compare(model, face):
     face = trans(face)
-    face_embedded = model(face.unsqueeze(dim=0).to('cpu'))
+    face_embedded = model(face.unsqueeze(dim=0).to(device))
     for usr in os.listdir(SINGLE_FACE_PATH):
-        usr_embedding = torch.load(SINGLE_FACE_PATH + "/" + usr + "/" + usr + ".pt").to('cpu')
+        usr_embedding = torch.load(SINGLE_FACE_PATH + "/" + usr + "/" + usr + ".pt", map_location=torch.device(device)).to(device)
         if (face_embedded - usr_embedding).norm().item() < THRESHOLD:
             print("Found " + usr + " in the image!")
             return usr