Upload face_cropper.py

face_cropper.py

@@ -0,0 +1,103 @@
+import cv2
+import mediapipe as mp
+import os
+from gradio_client import Client
+# from test_image_fusion import Test
+# from test_image_fusion import Test
+from test_image import Test
+import numpy as np
+
+
+
+from PIL import Image
+import numpy as np
+import cv2
+
+# client = Client("https://tbvl-real-and-fake-face-detection.hf.space/--replicas/40d41jxhhx/")
+
+data = 'faceswap'
+dct = 'fft'
+
+
+# testet = Test(model_paths = [f"weights/{data}-hh-best_model.pth",
+#                             f"weights/{data}-fft-best_model.pth"],
+#                             multi_modal = ['hh', 'fft'])
+
+testet = Test(model_path =f"weights/{data}-hh-best_model.pth",
+                            multi_modal ='hh')
+
+# Initialize MediaPipe Face Detection
+mp_face_detection = mp.solutions.face_detection
+mp_drawing = mp.solutions.drawing_utils
+face_detection = mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.35)
+
+# Create a directory to save the cropped face images if it does not exist
+save_dir = "cropped_faces"
+os.makedirs(save_dir, exist_ok=True)
+
+# def detect_and_label_faces(image_path):
+
+
+# Function to crop faces from a video and save them as images
+# def crop_faces_from_video(video_path):
+#     # Read the video
+#     cap = cv2.VideoCapture(video_path)
+#     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+#     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+#     fps = int(cap.get(cv2.CAP_PROP_FPS))
+#     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+#     # Define the codec and create VideoWriter object
+#     out = cv2.VideoWriter(f'output_{real}_{data}_fusion.avi', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (frame_width, frame_height))
+
+#     if not cap.isOpened():
+#         print("Error: Could not open video.")
+#         return
+# Convert PIL Image to NumPy array for OpenCV
+def pil_to_opencv(pil_image):
+    open_cv_image = np.array(pil_image)
+    # Convert RGB to BGR for OpenCV
+    open_cv_image = open_cv_image[:, :, ::-1].copy()
+    return open_cv_image
+
+# Convert OpenCV NumPy array to PIL Image
+def opencv_to_pil(opencv_image):
+    # Convert BGR to RGB
+    pil_image = Image.fromarray(opencv_image[:, :, ::-1])
+    return pil_image
+
+
+
+
+def detect_and_label_faces(frame):
+    frame = pil_to_opencv(frame)
+
+
+    print(type(frame))
+    # Convert the frame to RGB
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    # Perform face detection
+    results = face_detection.process(frame_rgb)
+
+    # If faces are detected, crop and save each face as an image
+    if results.detections:
+        for face_count,detection in enumerate(results.detections):
+            bboxC = detection.location_data.relative_bounding_box
+            ih, iw, _ = frame.shape
+            x, y, w, h = int(bboxC.xmin * iw), int(bboxC.ymin * ih), int(bboxC.width * iw), int(bboxC.height * ih)
+            # Crop the face region and make sure the bounding box is within the frame dimensions
+            crop_img = frame[max(0, y):min(ih, y+h), max(0, x):min(iw, x+w)]
+            if crop_img.size > 0:
+                face_filename = os.path.join(save_dir, f'face_{face_count}.jpg')
+                cv2.imwrite(face_filename, crop_img)
+
+                label = testet.testimage(face_filename)
+
+                if os.path.exists(face_filename):
+                    os.remove(face_filename)
+
+                color = (0, 0, 255) if label == 'fake' else (0, 255, 0)
+                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
+                cv2.putText(frame, label, (x, y + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+    return opencv_to_pil(frame)
+