Update app.py

app.py

@@ -43,10 +43,15 @@ def gen_sources(deepfake_img):
     #----------------DeepFake Face Segmentation-----------------
     ##----------------------Initialize:Face Segmentation----------------------------------
     segmenter = FaceSegmenter(threshold=0.5)
-    deepfake_seg = segmenter.segment_face(deepfake_img)
-        
-    config_path = "./models/config.yaml"
+    # Convert PIL Image to BGR numpy array for segmentation
+    img_np = np.array(deepfake_img.convert('RGB'))
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    # Segment the face
+    segmented_np = segmenter.segment_face(img_bgr)
+    # Convert segmented numpy array (BGR) back to PIL Image
+    deepfake_seg = Image.fromarray(cv2.cvtColor(segmented_np, cv2.COLOR_BGR2RGB))
     #------------Initialize:Decoder-F------------------------
+    config_path = "./models/config.yaml"
     checkpoint_path_f = "./models/model_vaq1_ff.pth"
     checkpoint_f = torch.load(checkpoint_path_f, map_location=device)
     model_vaq_f = Generator(config_path, device)