Spaces:
Sleeping
Sleeping
File size: 3,598 Bytes
5e014de 7bf182f 5e014de 7bf182f 5e014de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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)
|