deepfake_ai / preprocess.py
noorzz's picture
Update preprocess.py
0f973c9 verified
Raw
History Blame Contribute Delete
1.65 kB
# preprocess.py
import cv2
import os
import numpy as np
from tqdm import tqdm
from mtcnn import MTCNN
print("Starting preprocessing with face detection...")
detector = MTCNN()
def process_images(input_dir, output_dir, label_name):
os.makedirs(output_dir, exist_ok=True)
files = os.listdir(input_dir)
for file in tqdm(files, desc=f"Processing {label_name}"):
if not file.endswith(('.jpg', '.png')):
continue
img_path = os.path.join(input_dir, file)
img = cv2.imread(img_path)
if img is None:
continue
# Detect faces
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
faces = detector.detect_faces(rgb)
if len(faces) > 0:
# Get largest face
largest = max(faces, key=lambda x: x['box'][2] * x['box'][3])
x, y, w, h = largest['box']
# Add padding
x = max(0, x - 20)
y = max(0, y - 20)
w = min(img.shape[1] - x, w + 40)
h = min(img.shape[0] - y, h + 40)
# Crop face
face = img[y:y+h, x:x+w]
face_resized = cv2.resize(face, (224, 224))
else:
# No face - use center crop
h, w = img.shape[:2]
face_resized = cv2.resize(img[h//4:3*h//4, w//4:3*w//4], (224, 224))
cv2.imwrite(os.path.join(output_dir, file), face_resized)
# Process real and fake images
process_images("data/real", "processed/real", "REAL")
process_images("data/fake", "processed/fake", "FAKE")
print("✅ Preprocessing complete!")