from __future__ import annotations import argparse import sys from pathlib import Path import cv2 import numpy as np import torch import torch.nn.functional as torch_functional from insightface.app import FaceAnalysis from PIL import Image, ImageOps from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} DEFAULT_MODEL = "jonathandinu/face-parsing" INSIGHTFACE_MODEL = "buffalo_l" HEAD_LABELS = { "skin", "nose", "eye_g", "l_eye", "r_eye", "l_brow", "r_brow", "mouth", "u_lip", "l_lip", "l_ear", "r_ear", "hair", } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Extract complete heads to black-background PNG files and grayscale masks." ) parser.add_argument("input", nargs="?", default=".", help="Input image or directory.") parser.add_argument("-o", "--output", default="heads_no_bg", help="Output directory.") parser.add_argument("--recursive", action="store_true", help="Search directories recursively.") parser.add_argument( "--exclude-dir", action="append", default=[], help="Directory name to skip while searching; may be supplied more than once.", ) parser.add_argument( "--rotate-left", action="store_true", help="Rotate each image 90 degrees counterclockwise before processing.", ) parser.add_argument( "--save-mask", action=argparse.BooleanOptionalAction, default=True, help="Save the full-resolution grayscale mask (default: enabled).", ) parser.add_argument( "--padding", type=float, default=0.04, help="Black crop padding as a fraction of head size (default: 0.04).", ) parser.add_argument( "--expand", type=float, default=0.002, help="Mask expansion as a fraction of head size (default: 0.002).", ) parser.add_argument( "--feather", type=float, default=0.002, help="Edge feathering as a fraction of head size (default: 0.002).", ) parser.add_argument("--model", default=DEFAULT_MODEL, help="Hugging Face face-parsing model ID.") parser.add_argument( "--inference-size", type=int, default=512, help="Square semantic inference size (default: 512).", ) parser.add_argument( "--device", choices=("auto", "cpu", "cuda"), default="auto", help="Inference device (default: auto).", ) return parser.parse_args() def validate_args(args: argparse.Namespace) -> None: for name in ("padding", "expand", "feather"): if getattr(args, name) < 0: raise ValueError(f"--{name} must be non-negative") if args.device == "cuda" and not torch.cuda.is_available(): raise ValueError("CUDA was requested but is not available") if args.inference_size < 256: raise ValueError("--inference-size must be at least 256") def find_images( input_path: Path, output_dir: Path, recursive: bool, excluded_dirs: list[str], ) -> list[Path]: if input_path.is_file(): if input_path.suffix.lower() not in SUPPORTED_EXTENSIONS: raise ValueError(f"Unsupported image format: {input_path.suffix}") return [input_path] if not input_path.is_dir(): raise FileNotFoundError(f"Input does not exist: {input_path}") iterator = input_path.rglob("*") if recursive else input_path.glob("*") output_resolved = output_dir.resolve() excluded = {name.casefold() for name in excluded_dirs} images = [] for path in iterator: if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS: continue relative_parts = path.relative_to(input_path).parts[:-1] if any(part.casefold() in excluded for part in relative_parts): continue try: path.resolve().relative_to(output_resolved) continue except ValueError: images.append(path) return sorted(images, key=lambda path: str(path).lower()) def load_rgb(path: Path, rotate_left: bool) -> np.ndarray: with Image.open(path) as image: rgb = np.asarray(ImageOps.exif_transpose(image).convert("RGB")) if rotate_left: rgb = np.rot90(rgb) return np.ascontiguousarray(rgb) def create_detector() -> FaceAnalysis: detector = FaceAnalysis( name=INSIGHTFACE_MODEL, allowed_modules=["detection"], providers=["CPUExecutionProvider"], ) detector.prepare(ctx_id=-1, det_size=(640, 640), det_thresh=0.15) return detector def largest_head_roi( detector: FaceAnalysis, rgb: np.ndarray, ) -> tuple[tuple[int, int, int, int], tuple[int, int, int, int], float]: height, width = rgb.shape[:2] faces = detector.get(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) if not faces: fallback_face = ( round(width * 0.2), round(height * 0.15), round(width * 0.8), round(height * 0.9), ) return (0, 0, width, height), fallback_face, float(max(width, height) * 0.7) face = max( faces, key=lambda item: float((item.bbox[2] - item.bbox[0]) * (item.bbox[3] - item.bbox[1])), ) left, top, right, bottom = (float(value) for value in face.bbox) face_width = right - left face_height = bottom - top face_size = max(face_width, face_height) face_box = ( max(0, int(np.floor(left))), max(0, int(np.floor(top))), min(width, int(np.ceil(right))), min(height, int(np.ceil(bottom))), ) center_x = (left + right) * 0.5 nose_x = center_x if getattr(face, "kps", None) is not None and len(face.kps) >= 3: nose_x = float(face.kps[2, 0]) yaw = float(np.clip((nose_x - center_x) / max(face_width * 0.5, 1.0), -1.0, 1.0)) side_pose = abs(yaw) front_margin = 0.2 + 0.05 * side_pose back_margin = 0.3 + 0.2 * side_pose if yaw > 0: left_margin, right_margin = back_margin, front_margin elif yaw < 0: left_margin, right_margin = front_margin, back_margin else: left_margin = right_margin = 0.3 roi = ( max(0, int(np.floor(left - face_width * left_margin))), max(0, int(np.floor(top - face_height * 0.5))), min(width, int(np.ceil(right + face_width * right_margin))), min(height, int(np.ceil(bottom + face_height * 0.1))), ) return roi, face_box, face_size def semantic_label_ids(model: AutoModelForSemanticSegmentation) -> list[int]: ids = [ int(label_id) for label_id, label_name in model.config.id2label.items() if str(label_name).lower() in HEAD_LABELS ] if not ids: raise ValueError("The segmentation model does not expose supported head labels") return ids def segment_head( rgb: np.ndarray, detector: FaceAnalysis, processor: AutoImageProcessor, model: AutoModelForSemanticSegmentation, label_ids: list[int], device: torch.device, inference_size: int, ) -> np.ndarray: (left, top, right, bottom), face_box, face_size = largest_head_roi(detector, rgb) roi = rgb[top:bottom, left:right] inputs = processor( images=Image.fromarray(roi), return_tensors="pt", size={"height": inference_size, "width": inference_size}, ) inputs = {name: tensor.to(device) for name, tensor in inputs.items()} with torch.inference_mode(): logits = model(**inputs).logits logits = torch_functional.interpolate( logits, size=roi.shape[:2], mode="bilinear", align_corners=False, ) labels = logits.argmax(dim=1)[0].cpu().numpy() roi_mask = np.isin(labels, label_ids).astype(np.uint8) * 255 ear_label_ids = [ int(label_id) for label_id, label_name in model.config.id2label.items() if str(label_name).lower() in {"l_ear", "r_ear"} ] local_face_box = ( face_box[0] - left, face_box[1] - top, face_box[2] - left, face_box[3] - top, ) roi_mask = remove_posterior_ear_artifacts( roi_mask, labels, ear_label_ids, local_face_box, face_size, ) mask = np.zeros(rgb.shape[:2], dtype=np.uint8) mask[top:bottom, left:right] = roi_mask return clean_head_mask(mask, face_box, face_size) def remove_posterior_ear_artifacts( mask: np.ndarray, labels: np.ndarray, ear_label_ids: list[int], face_box: tuple[int, int, int, int], face_size: float, ) -> np.ndarray: if not ear_label_ids: return mask ear_count, _, ear_stats, _ = cv2.connectedComponentsWithStats( np.isin(labels, ear_label_ids).astype(np.uint8), connectivity=8, ) if ear_count <= 1: return mask face_left, _, face_right, _ = face_box face_width = max(1, face_right - face_left) face_center = (face_left + face_right) * 0.5 minimum_ear_area = max(16, round(face_size * face_size * 0.0005)) maximum_artifact_area = max(64, round(face_size * face_size * 0.04)) maximum_artifact_width = max(8, round(face_size * 0.22)) side_margin = round(face_size * 0.02) cleaned = mask.copy() for ear_label in range(1, ear_count): ear_area = int(ear_stats[ear_label, cv2.CC_STAT_AREA]) if ear_area < minimum_ear_area: continue ear_left = int(ear_stats[ear_label, cv2.CC_STAT_LEFT]) ear_top = int(ear_stats[ear_label, cv2.CC_STAT_TOP]) ear_width = int(ear_stats[ear_label, cv2.CC_STAT_WIDTH]) ear_height = int(ear_stats[ear_label, cv2.CC_STAT_HEIGHT]) ear_right = ear_left + ear_width ear_bottom = ear_top + ear_height ear_center = (ear_left + ear_right) * 0.5 if abs(ear_center - face_center) < face_width * 0.35 or ear_bottom >= mask.shape[0]: continue lower_mask = np.where(cleaned[ear_bottom:] > 0, 1, 0).astype(np.uint8) component_count, components, stats, _ = cv2.connectedComponentsWithStats( lower_mask, connectivity=8, ) for component in range(1, component_count): component_left = int(stats[component, cv2.CC_STAT_LEFT]) component_width = int(stats[component, cv2.CC_STAT_WIDTH]) component_right = component_left + component_width component_area = int(stats[component, cv2.CC_STAT_AREA]) if component_width > maximum_artifact_width or component_area > maximum_artifact_area: continue if ear_center > face_center: is_posterior = component_left >= max(face_right, ear_left) - side_margin else: is_posterior = component_right <= min(face_left, ear_right) + side_margin if is_posterior: cleaned[ear_bottom:][components == component] = 0 return cleaned def clean_head_mask( mask: np.ndarray, face_box: tuple[int, int, int, int], face_size: float, ) -> np.ndarray: bridge = max(1, round(face_size * 0.008)) kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (bridge * 2 + 1, bridge * 2 + 1)) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) component_count, labels, stats, _ = cv2.connectedComponentsWithStats( np.where(mask > 0, 1, 0).astype(np.uint8), connectivity=8, ) if component_count <= 1: raise ValueError("The semantic model did not find a head") face_left, face_top, face_right, face_bottom = face_box face_labels = labels[face_top:face_bottom, face_left:face_right] best_label = 0 best_score = -1.0 for label in range(1, component_count): overlap = int(np.count_nonzero(face_labels == label)) area = int(stats[label, cv2.CC_STAT_AREA]) score = overlap * 10.0 + area if score > best_score: best_label = label best_score = score cleaned = np.where(labels == best_label, 255, 0).astype(np.uint8) # Preserve semantic gaps between the ear, hair and face. Filling the full # outer contour would also keep equipment enclosed by the head outline. hole_limit = max(16, round(face_size * face_size * 0.00015)) inverted = cv2.bitwise_not(cleaned) hole_count, hole_labels, hole_stats, _ = cv2.connectedComponentsWithStats( np.where(inverted > 0, 1, 0).astype(np.uint8), connectivity=8, ) for label in range(1, hole_count): x = int(hole_stats[label, cv2.CC_STAT_LEFT]) y = int(hole_stats[label, cv2.CC_STAT_TOP]) width = int(hole_stats[label, cv2.CC_STAT_WIDTH]) height = int(hole_stats[label, cv2.CC_STAT_HEIGHT]) area = int(hole_stats[label, cv2.CC_STAT_AREA]) touches_border = ( x == 0 or y == 0 or x + width == cleaned.shape[1] or y + height == cleaned.shape[0] ) if not touches_border and area <= hole_limit: cleaned[hole_labels == label] = 255 return cleaned def finish_mask(mask: np.ndarray, expand_ratio: float, feather_ratio: float) -> np.ndarray: rows, columns = np.nonzero(mask) if not len(columns): raise ValueError("Generated head mask is empty") head_size = max(int(columns.max() - columns.min() + 1), int(rows.max() - rows.min() + 1)) expand_pixels = round(head_size * expand_ratio) if expand_pixels > 0: size = expand_pixels * 2 + 1 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) mask = cv2.dilate(mask, kernel) smoothing = max(1, round(head_size * 0.0015)) size = smoothing * 2 + 1 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) feather_pixels = head_size * feather_ratio if feather_pixels >= 0.5: mask = cv2.GaussianBlur(mask, (0, 0), feather_pixels, feather_pixels) return mask def crop_rgba(rgb: np.ndarray, alpha: np.ndarray, padding_ratio: float) -> np.ndarray: rows, columns = np.nonzero(alpha > 1) if not len(columns): raise ValueError("Generated head mask is empty") left = int(columns.min()) right = int(columns.max()) + 1 top = int(rows.min()) bottom = int(rows.max()) + 1 head_size = max(right - left, bottom - top) padding = round(head_size * padding_ratio) desired_left = left - padding desired_top = top - padding desired_right = right + padding desired_bottom = bottom + padding source_left = max(0, desired_left) source_top = max(0, desired_top) source_right = min(rgb.shape[1], desired_right) source_bottom = min(rgb.shape[0], desired_bottom) output = np.zeros( (desired_bottom - desired_top, desired_right - desired_left, 4), dtype=np.uint8, ) offset_x = source_left - desired_left offset_y = source_top - desired_top copy_height = source_bottom - source_top copy_width = source_right - source_left output[offset_y : offset_y + copy_height, offset_x : offset_x + copy_width, :3] = rgb[ source_top:source_bottom, source_left:source_right ] output[offset_y : offset_y + copy_height, offset_x : offset_x + copy_width, 3] = alpha[ source_top:source_bottom, source_left:source_right ] output[output[:, :, 3] == 0, :3] = 0 return output def composite_on_black(rgba: np.ndarray) -> np.ndarray: alpha = rgba[:, :, 3:4].astype(np.uint16) return ((rgba[:, :, :3].astype(np.uint16) * alpha + 127) // 255).astype(np.uint8) def output_paths_for( image_path: Path, input_path: Path, output_dir: Path, ) -> tuple[Path, Path]: if input_path.is_dir(): try: relative_parent = image_path.parent.relative_to(input_path) except ValueError: relative_parent = Path() else: relative_parent = Path() parent = output_dir / relative_parent return parent / f"{image_path.stem}_head.png", parent / f"{image_path.stem}_mask.png" def resolve_device(name: str) -> torch.device: if name == "cuda" or (name == "auto" and torch.cuda.is_available()): return torch.device("cuda") return torch.device("cpu") def process_images(args: argparse.Namespace) -> int: input_path = Path(args.input).expanduser() output_dir = Path(args.output).expanduser() images = find_images(input_path, output_dir, args.recursive, args.exclude_dir) if not images: print("No supported images found.", file=sys.stderr) return 1 output_dir.mkdir(parents=True, exist_ok=True) device = resolve_device(args.device) print(f"Loading semantic head parser '{args.model}' on {device}...") processor = AutoImageProcessor.from_pretrained(args.model, use_fast=False) model = AutoModelForSemanticSegmentation.from_pretrained(args.model).to(device).eval() label_ids = semantic_label_ids(model) print(f"Loading InsightFace detector '{INSIGHTFACE_MODEL}' on CPU...") detector = create_detector() failures = 0 for image_path in images: try: rgb = load_rgb(image_path, args.rotate_left) mask = segment_head( rgb, detector, processor, model, label_ids, device, args.inference_size, ) alpha = finish_mask(mask, args.expand, args.feather) output_image = composite_on_black(crop_rgba(rgb, alpha, args.padding)) destination, mask_destination = output_paths_for(image_path, input_path, output_dir) destination.parent.mkdir(parents=True, exist_ok=True) if args.save_mask: Image.fromarray(alpha).save(mask_destination, optimize=True) Image.fromarray(output_image).save(destination, optimize=True) mask_message = f", mask: {mask_destination}" if args.save_mask else "" print( f"OK {image_path} -> {destination} " f"({output_image.shape[1]}x{output_image.shape[0]}{mask_message})" ) except Exception as error: failures += 1 print(f"FAIL {image_path}: {error}", file=sys.stderr) succeeded = len(images) - failures print(f"Processed {succeeded}/{len(images)} image(s). Output: {output_dir.resolve()}") return 0 if failures == 0 else 2 def main() -> int: args = parse_args() try: validate_args(args) return process_images(args) except (FileNotFoundError, ValueError) as error: print(f"Error: {error}", file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())