Update stf/stf-api-alternative/src/stf_alternative/preprocess_dir/utils/face_finder.py

stf/stf-api-alternative/src/stf_alternative/preprocess_dir/utils/face_finder.py

@@ -30,7 +30,7 @@ def init_face_finder(device="cuda:0"):
     if g_mtcnn is None and g_recognizer is None:
         g_mtcnn = MTCNN(image_size=166, device=device)
         print("load MTCNN ", "success ^ ^" if g_mtcnn is not None else "fail ㅠㅠ")
-        g_recognizer = InceptionResnetV1(pretrained="vggface2").eval().to(device)
+        g_recognizer = InceptionResnetV1(pretrained="vggface2").eval().cuda(0) # // .to(device)
         print(
             "load g_recognizer ",
             "success ^ ^" if g_recognizer is not None else "fail ㅠㅠ",
@@ -70,11 +70,11 @@ def find_face(img):
         img = imageio.imread(img)
     frame = np.array(img)
     df_non_face = pd.DataFrame({"box": [np.nan], "ebd": [np.nan]})
-    with torch.no_grad():
-        boxes = g_mtcnn.detect(frame)
-        if boxes[0] is None:
-            return df_non_face, None
-        boxes = boxes[0].round().astype(np.int32)
+    #with torch.no_grad():
+    boxes = g_mtcnn.detect(frame)
+    if boxes[0] is None:
+        return df_non_face, None
+    boxes = boxes[0].round().astype(np.int32)
 
     org = np.array(frame)
 
@@ -84,8 +84,8 @@ def find_face(img):
         sz = g_mtcnn.image_size
         resized = cv2.resize(crop, (sz, sz), cv2.INTER_AREA)
         x = torchvision.transforms.functional.to_tensor(resized)
-        with torch.no_grad():
-            ebd = g_recognizer(x.unsqueeze(0).to(g_device))
+        #with torch.no_grad():
+        ebd = g_recognizer(x.unsqueeze(0).cuda(0)) # to(g_device))
         return ebd[0].cpu(), crop
 
     def check_box(x1, y1, x2, y2):
@@ -103,7 +103,7 @@ def find_face(img):
 class FaceFinder:
     def __init__(self, device="cuda:0"):
         self.mtcnn = MTCNN(image_size=166, device=device)
-        self.recognizer = InceptionResnetV1(pretrained="vggface2").eval().to(device)
+        self.recognizer = InceptionResnetV1(pretrained="vggface2").eval().cuda(0) #to(device)
         self.device = device
         self.tracker = cv2.TrackerCSRT_create()  # cv2.legacy.TrackerMOSSE_create()
         self.tracker_initialized = False
@@ -116,8 +116,8 @@ class FaceFinder:
         sz = self.mtcnn.image_size
         resized = cv2.resize(crop, (sz, sz), cv2.INTER_AREA)
         x = torchvision.transforms.functional.to_tensor(resized)
-        with torch.no_grad():
-            ebd = self.recognizer(x.unsqueeze(0).to(g_device))
+        #with torch.no_grad():
+        ebd = self.recognizer(x.unsqueeze(0).cuda(0)) #to(g_device))
         return ebd[0].cpu()
 
     def find_face(self, frame):