import warnings from typing import Optional import cv2 import numpy as np from PIL import Image def _load_mtcnn(device='cpu'): """Attempt to import and return MTCNN detector. Returns None on failure.""" try: from facenet_pytorch import MTCNN return MTCNN(keep_all=True, device=device, post_process=False) except ImportError: warnings.warn( "facenet-pytorch not installed — falling back to Haar cascade. " "Install with: pip install facenet-pytorch", RuntimeWarning, stacklevel=3 ) return None def _pil_to_bgr(image: Image.Image) -> np.ndarray: rgb = np.array(image.convert('RGB')) return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) def _bgr_to_pil(frame: np.ndarray) -> Image.Image: rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) return Image.fromarray(rgb) def _crop_with_margin(image: Image.Image, x1: int, y1: int, x2: int, y2: int, margin: int) -> tuple: """Apply margin, clamp to image bounds, return (cropped_pil, bbox_xyxy).""" w, h = image.size x1 = max(0, x1 - margin) y1 = max(0, y1 - margin) x2 = min(w, x2 + margin) y2 = min(h, y2 + margin) return image.crop((x1, y1, x2, y2)), (x1, y1, x2 - x1, y2 - y1) # bbox as (x,y,w,h) def detect_and_crop_faces( image, method: str = 'mtcnn', margin: int = 20, device: str = 'cpu' ) -> list[tuple]: """ Detect faces and return crops. Args: image: PIL Image, numpy array (H,W,3 BGR or RGB), or file path str. method: 'mtcnn' (with auto-fallback) or 'haar'. margin: pixel margin to add around each detected bbox. device: torch device string used by MTCNN. Returns: List of (face_crop_PIL, bbox_xywh) tuples. bbox_xywh is (x, y, w, h) in original image coords, or None if no face detected. If no face is detected the full image (resized to 224) is returned with bbox=None. """ # Normalise to PIL RGB if isinstance(image, str): pil_img = Image.open(image).convert('RGB') elif isinstance(image, np.ndarray): if image.ndim == 3 and image.shape[2] == 3: # Assume BGR (OpenCV convention) pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) else: pil_img = Image.fromarray(image).convert('RGB') elif isinstance(image, Image.Image): pil_img = image.convert('RGB') else: raise TypeError(f"Unsupported image type: {type(image)}") faces = [] if method == 'mtcnn': mtcnn = _load_mtcnn(device=device) if mtcnn is not None: faces = _detect_mtcnn(mtcnn, pil_img, margin) else: method = 'haar' if method == 'haar' or (method == 'mtcnn' and not faces): faces = _detect_haar(pil_img, margin) if not faces: warnings.warn( "No face detected — running inference on full image.", RuntimeWarning, stacklevel=2 ) fallback = pil_img.resize((224, 224)) return [(fallback, None)] return faces def _detect_mtcnn(mtcnn, pil_img: Image.Image, margin: int) -> list[tuple]: import torch boxes, probs = mtcnn.detect(pil_img) if boxes is None or len(boxes) == 0: return [] results = [] for box in boxes: x1, y1, x2, y2 = (int(v) for v in box) crop, bbox = _crop_with_margin(pil_img, x1, y1, x2, y2, margin) results.append((crop, bbox)) return results def _detect_haar(pil_img: Image.Image, margin: int) -> list[tuple]: bgr = _pil_to_bgr(pil_img) gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml' cascade = cv2.CascadeClassifier(cascade_path) detected = cascade.detectMultiScale( gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30) ) if len(detected) == 0: return [] results = [] for (x, y, w, h) in detected: crop, bbox = _crop_with_margin(pil_img, x, y, x + w, y + h, margin) results.append((crop, bbox)) return results