Brings back the similarity of the face to the Opal version part 3.

Alucard24 · Jun 5, 2024 · d5af208 · d5af208
1 parent f5a306d
commit d5af208
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/rope/GUI.py b/rope/GUI.py
@@ -1359,8 +1359,6 @@ def load_input_faces(self):
                     img = cv2.imread(file)
 
                     if img is not None:
-                        # convert to RGB format
-                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                         img = torch.from_numpy(img.astype('uint8')).to('cuda')
 
                         pad_scale = 0.2
@@ -1385,13 +1383,13 @@ def load_input_faces(self):
                             print('Image cropped too close:', file)
                         else:
                             face_emb, cropped_image = self.models.run_recognize(img, kpss)
-                            # PIL follows RGB color convention and cropped image come in RGB format
-                            crop = torchvision.transforms.functional.to_pil_image(cropped_image)
-                            crop = crop.resize((85, 85))
+                            crop = cv2.cvtColor(cropped_image.cpu().numpy(), cv2.COLOR_BGR2RGB)
+                            crop = cv2.resize(crop, (85, 85))
+
                             new_source_face = self.source_face.copy()
                             self.source_faces.append(new_source_face)
 
-                            self.source_faces[-1]["Image"] = ImageTk.PhotoImage(image=crop)
+                            self.source_faces[-1]["Image"] = ImageTk.PhotoImage(image=Image.fromarray(crop))
                             self.source_faces[-1]["Embedding"] = face_emb
                             self.source_faces[-1]["TKButton"] = tk.Button(self.source_faces_canvas, style.media_button_off_3, image=self.source_faces[-1]["Image"], height=90, width=90)
                             self.source_faces[-1]["ButtonState"] = False
@@ -1451,8 +1449,7 @@ def find_faces(self):
 
                     # If we dont find any existing simularities, it means that this is a new face and should be added to our found faces
                     if not found:
-                        crop = torchvision.transforms.functional.to_pil_image(face[2])
-                        crop = crop.resize((82, 82))
+                        crop = cv2.resize(face[2].cpu().numpy(), (82, 82))
 
                         new_target_face = self.target_face.copy()
                         self.target_faces.append(new_target_face)
@@ -1461,7 +1458,7 @@ def find_faces(self):
                         self.target_faces[last_index]["TKButton"] = tk.Button(self.found_faces_canvas, style.media_button_off_3, height = 86, width = 86)
                         self.target_faces[last_index]["TKButton"].bind("<MouseWheel>", self.target_faces_mouse_wheel)
                         self.target_faces[last_index]["ButtonState"] = False
-                        self.target_faces[last_index]["Image"] = ImageTk.PhotoImage(image=crop)
+                        self.target_faces[last_index]["Image"] = ImageTk.PhotoImage(image=Image.fromarray(crop))
                         self.target_faces[last_index]["Embedding"] = face[1]
                         self.target_faces[last_index]["EmbeddingNumber"] = 1
 

diff --git a/rope/Models.py b/rope/Models.py
@@ -1848,9 +1848,9 @@ def recognize(self, img, face_kps):
         img = v2.functional.affine(img, tform.rotation*57.2958, (tform.translation[0], tform.translation[1]) , tform.scale, 0, center = (0,0) ) 
         img = v2.functional.crop(img, 0,0, 112, 112)
 
-        cropped_image = img
         # Switch to BGR and normalize
         img = img.permute(1,2,0) #112,112,3
+        cropped_image = img
         img = img[:, :, [2,1,0]]
         img = torch.sub(img, 127.5)
         img = torch.div(img, 127.5)