224 lines
7.7 KiB
Python
224 lines
7.7 KiB
Python
"""Preprocess face images by detecting, cropping, resizing, and augmenting them."""
|
||
|
||
import argparse
|
||
import random
|
||
import shutil
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from rembg import remove as rembg_remove
|
||
|
||
VALID_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
|
||
|
||
|
||
def _collect_class_dirs(root_dir: Path) -> List[Path]:
|
||
"""Return the class folders inside the source dataset."""
|
||
return sorted([p for p in root_dir.iterdir() if p.is_dir()])
|
||
|
||
|
||
def _collect_image_files(root_dir: Path) -> List[Path]:
|
||
"""Return all supported image files below one class folder."""
|
||
return sorted(p for p in root_dir.rglob("*") if p.is_file() and p.suffix.lower() in VALID_EXTENSIONS)
|
||
|
||
|
||
def load_face_cascade() -> cv2.CascadeClassifier:
|
||
"""Load the built-in Haar cascade for face detection."""
|
||
cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
|
||
face_cascade = cv2.CascadeClassifier(cascade_path)
|
||
if face_cascade.empty():
|
||
raise RuntimeError("Gagal memuat Haar Cascade untuk deteksi wajah.")
|
||
return face_cascade
|
||
|
||
|
||
def detect_largest_face(image_bgr: np.ndarray, face_cascade: cv2.CascadeClassifier) -> Optional[Tuple[int, int, int, int]]:
|
||
"""Find the biggest face box in one image."""
|
||
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
|
||
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(40, 40))
|
||
if len(faces) == 0:
|
||
return None
|
||
x, y, w, h = max(faces, key=lambda face: face[2] * face[3])
|
||
return int(x), int(y), int(w), int(h)
|
||
|
||
|
||
def remove_background(image_bgr: np.ndarray, bg_color: Tuple[int, int, int] = (255, 255, 255)) -> np.ndarray:
|
||
"""Remove the background from a BGR image using rembg and composite onto a solid colour.
|
||
|
||
Args:
|
||
image_bgr: Input image in BGR format (as returned by cv2.imread).
|
||
bg_color: Background fill colour in BGR order. Defaults to white.
|
||
|
||
Returns:
|
||
BGR image with background replaced by *bg_color*.
|
||
"""
|
||
# rembg expects PNG bytes; encode the BGR frame to PNG in-memory
|
||
success, encoded = cv2.imencode(".png", image_bgr)
|
||
if not success:
|
||
return image_bgr # fall back to original if encoding fails
|
||
|
||
png_bytes = encoded.tobytes()
|
||
result_bytes = rembg_remove(png_bytes) # returns RGBA PNG bytes
|
||
|
||
# Decode the RGBA result
|
||
result_array = np.frombuffer(result_bytes, dtype=np.uint8)
|
||
rgba = cv2.imdecode(result_array, cv2.IMREAD_UNCHANGED)
|
||
if rgba is None or rgba.shape[2] != 4:
|
||
return image_bgr # fall back if decoding fails
|
||
|
||
# cv2.imdecode returns BGR data even for RGBA PNGs (channels: B, G, R, A)
|
||
# So we treat the first 3 channels as BGR directly — no further conversion needed.
|
||
alpha = rgba[:, :, 3:4].astype(np.float32) / 255.0
|
||
bgr = rgba[:, :, :3].astype(np.float32)
|
||
background = np.full_like(bgr, fill_value=bg_color, dtype=np.float32) # already BGR
|
||
composited = (bgr * alpha + background * (1.0 - alpha)).astype(np.uint8)
|
||
return composited
|
||
|
||
|
||
def crop_and_resize(image_bgr: np.ndarray, face_box: Tuple[int, int, int, int], target_size: int) -> np.ndarray:
|
||
"""Crop the face area and resize it to the model input size."""
|
||
x, y, w, h = face_box
|
||
face_crop = image_bgr[y : y + h, x : x + w]
|
||
return cv2.resize(face_crop, (target_size, target_size), interpolation=cv2.INTER_AREA)
|
||
|
||
|
||
def random_augment(color_img: np.ndarray, rng: random.Random) -> np.ndarray:
|
||
"""Create a slightly changed copy of one face image."""
|
||
image = color_img.copy()
|
||
height, width = image.shape[:2]
|
||
|
||
rotation = rng.uniform(-18, 18)
|
||
scale = rng.uniform(0.95, 1.05)
|
||
rotation_matrix = cv2.getRotationMatrix2D((width // 2, height // 2), rotation, scale)
|
||
image = cv2.warpAffine(
|
||
image,
|
||
rotation_matrix,
|
||
(width, height),
|
||
flags=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_REFLECT_101,
|
||
)
|
||
|
||
shift_x = rng.randint(-10, 10)
|
||
shift_y = rng.randint(-10, 10)
|
||
translation_matrix = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
|
||
image = cv2.warpAffine(
|
||
image,
|
||
translation_matrix,
|
||
(width, height),
|
||
flags=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_REFLECT_101,
|
||
)
|
||
|
||
if rng.random() < 0.5:
|
||
image = cv2.flip(image, 1)
|
||
|
||
brightness = rng.uniform(0.85, 1.20)
|
||
contrast = rng.uniform(-20, 20)
|
||
image = cv2.convertScaleAbs(image, alpha=brightness, beta=contrast)
|
||
|
||
if rng.random() < 0.3:
|
||
kernel_size = rng.choice([3, 5])
|
||
image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
|
||
|
||
return image
|
||
|
||
|
||
def preprocess_dataset(
|
||
source_dir: Path,
|
||
output_dir: Path,
|
||
target_size: int,
|
||
min_images_per_class: int,
|
||
seed: int,
|
||
overwrite: bool = False,
|
||
) -> Dict[str, int]:
|
||
"""Process every class folder and save cleaned plus augmented images."""
|
||
if not source_dir.is_dir():
|
||
raise FileNotFoundError(f"Folder sumber tidak ditemukan: {source_dir}")
|
||
|
||
class_dirs = _collect_class_dirs(source_dir)
|
||
if not class_dirs:
|
||
raise RuntimeError("Dataset raw tidak memiliki subfolder kelas.")
|
||
|
||
# Always remove old preprocessed dataset to ensure clean data
|
||
# This prevents mixing old and new preprocessed images
|
||
if output_dir.exists():
|
||
shutil.rmtree(output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
face_cascade = load_face_cascade()
|
||
rng = random.Random(seed)
|
||
stats = {"class_count": len(class_dirs), "processed": 0, "skipped": 0, "generated": 0, "total_output": 0}
|
||
|
||
for class_dir in class_dirs:
|
||
image_files = _collect_image_files(class_dir)
|
||
if not image_files:
|
||
continue
|
||
|
||
class_output_dir = output_dir / class_dir.name
|
||
class_output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
clean_images: List[np.ndarray] = []
|
||
|
||
for index, src_path in enumerate(image_files, start=1):
|
||
image = cv2.imread(str(src_path), cv2.IMREAD_COLOR)
|
||
if image is None:
|
||
stats["skipped"] += 1
|
||
continue
|
||
|
||
# Step 1 – remove background
|
||
image = remove_background(image)
|
||
|
||
# Step 2 – detect the largest face on the clean image
|
||
face_box = detect_largest_face(image, face_cascade)
|
||
if face_box is None:
|
||
stats["skipped"] += 1
|
||
continue
|
||
|
||
# Step 3 – crop and resize
|
||
resized = crop_and_resize(image, face_box, target_size)
|
||
output_path = class_output_dir / f"orig_{index:04d}.jpg"
|
||
if cv2.imwrite(str(output_path), resized):
|
||
stats["processed"] += 1
|
||
clean_images.append(resized)
|
||
else:
|
||
stats["skipped"] += 1
|
||
|
||
if not clean_images:
|
||
continue
|
||
|
||
needed_images = max(0, min_images_per_class - len(clean_images))
|
||
for index in range(needed_images):
|
||
base_image = clean_images[index % len(clean_images)]
|
||
augmented_image = random_augment(base_image, rng)
|
||
augmented_path = class_output_dir / f"aug_{index + 1:04d}.jpg"
|
||
if cv2.imwrite(str(augmented_path), augmented_image):
|
||
stats["generated"] += 1
|
||
|
||
stats["total_output"] = stats["processed"] + stats["generated"]
|
||
return stats
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="Preprocess dataset: detect, crop, resize, augment.")
|
||
parser.add_argument("--source", type=str, default="dataset/Dataset_Raw")
|
||
parser.add_argument("--output", type=str, default="dataset/Dataset_Preprocessed")
|
||
parser.add_argument("--size", type=int, default=224)
|
||
parser.add_argument("--min_images", type=int, default=30)
|
||
parser.add_argument("--seed", type=int, default=42)
|
||
parser.add_argument("--overwrite", action="store_true")
|
||
args = parser.parse_args()
|
||
|
||
stats = preprocess_dataset(
|
||
source_dir=Path(args.source),
|
||
output_dir=Path(args.output),
|
||
target_size=args.size,
|
||
min_images_per_class=args.min_images,
|
||
seed=args.seed,
|
||
overwrite=args.overwrite,
|
||
)
|
||
print(stats)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|