From 350223c5adb663c601524a37241785f83223b9f7 Mon Sep 17 00:00:00 2001 From: YULI <1575108428@qq.com> Date: Wed, 12 Aug 2026 17:27:17 +0800 Subject: [PATCH] tijiao --- .gitignore | 0 README.md | 89 ------ extract_head_mask_single.py | 537 ------------------------------------ requirements.txt | 9 - 4 files changed, 635 deletions(-) delete mode 100644 .gitignore delete mode 100644 README.md delete mode 100644 extract_head_mask_single.py delete mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore deleted file mode 100644 index e69de29..0000000 diff --git a/README.md b/README.md deleted file mode 100644 index 521f0c3..0000000 --- a/README.md +++ /dev/null @@ -1,89 +0,0 @@ -# 人脸裁剪与背景去除 - -`extract_faces.py` 会批量检测正脸及左右侧脸，按 MediaPipe 三维人脸网格生成面部轮廓，再用 GrabCut 细化边缘。结果保存为带透明通道的 PNG，只保留面部区域，不包含头发、耳朵和衣服。 - -## 运行 - -```powershell -python -m pip install -r requirements.txt -python extract_faces.py . --output faces_no_bg -``` - -也可以处理单张图片： - -```powershell -python extract_faces.py yolo_white_left.jpg --output faces_no_bg -``` - -## 常用参数 - -- `--padding 0.04`：人脸周围透明边距比例。 -- `--expand 0.004`：向外扩展面部蒙版，避免切掉边缘皮肤。 -- `--feather 0.002`：透明边缘羽化比例。 -- `--edge-band 0.05`：覆盖自动姿态估计，手动指定轮廓细化带宽度。 -- `--max-refine-size 960`：轮廓细化使用的最大局部图尺寸。 -- `--no-edge-refine`：关闭 GrabCut，仅使用人脸网格轮廓。 -- `--recursive`：递归处理输入目录中的图片。 -- `--rotate-left`：检测前将每张输入图逆时针旋转 90 度，不修改原图。 - -原图不会被覆盖，输出文件名格式为 `<原文件名>_face.png`。 - -## 仅保留面部皮肤 - -`extract_face_skin.py` 使用语义分割保留皮肤、鼻子、眼睛、眉毛和嘴唇，排除头发、耳朵、颈部与背景。第一次运行会下载通用预训练模型。 - -```powershell -python extract_face_skin.py "输入目录" --output "输出目录" --recursive --rotate-left -``` - -默认按模型训练分辨率 `512×512` 做语义分割。 - -## 保留头发并输出 Mask - -`extract_head_mask.py` 先根据人脸位置和朝向自适应扩展发顶、两侧与后脑区域，再分割脸、耳朵和头发。每张输入图会得到： - -- `<原名>_mask.png`：与旋转后原图等大的灰度 Alpha Mask。 -- `<原名>_head.png`：使用平滑 Mask 截取的黑色背景 RGB PNG。 - -```powershell -python extract_head_mask.py "输入目录" --output "输出目录" --recursive --rotate-left -``` - -输入根目录含有其它结果文件夹时，可以重复使用 `--exclude-dir` 排除： - -```powershell -python extract_head_mask.py "输入目录" --output "输出目录" --recursive --rotate-left --exclude-dir faces_no_bg --exclude-dir output -``` - -如需只分发一个业务脚本，可使用功能相同且不导入其它本地脚本的单文件版： - -```powershell -python -m pip install -r requirements-head-mask.txt -python extract_head_mask_single.py "输入目录" --output "输出目录" --recursive --rotate-left -``` - -最小分发内容为 `extract_head_mask_single.py`、`requirements-head-mask.txt` 和本说明文档。建议使用 Windows 64 位及 Python 3.9；GPU 可选，语义分割默认自动选择 CUDA 或 CPU，人脸检测固定使用 CPU。 - -## 模型下载与离线运行 - -首次运行会下载以下两个模型，模型文件不会嵌入 Python 脚本： - -- 语义分割模型：[`jonathandinu/face-parsing`](https://huggingface.co/jonathandinu/face-parsing)，默认缓存到 `%USERPROFILE%\.cache\huggingface\hub\models--jonathandinu--face-parsing`。 -- 人脸检测模型：[`buffalo_l.zip`](https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip)，默认解压到 `%USERPROFILE%\.insightface\models\buffalo_l`，实际检测文件是 `det_10g.onnx`。 - -离线分发时，在联网电脑上先运行一次脚本，再准备以下文件： - -1. 将 Hugging Face 缓存中同一 `snapshots\<版本>` 目录内的 `config.json`、`preprocessor_config.json` 和 `model.safetensors` 复制到分发目录的 `models\face-parsing`。 -2. 将整个 `buffalo_l` 目录复制到离线电脑的 `%USERPROFILE%\.insightface\models\buffalo_l`，确保 `det_10g.onnx` 直接位于该目录内。 -3. 使用本地语义模型运行： - -```powershell -python extract_head_mask_single.py "输入目录" --output "输出目录" --recursive --rotate-left --model ".\models\face-parsing" -``` - -Python 第三方依赖也需要预先安装。可在相同 Windows/Python 环境的联网电脑下载离线安装包： - -```powershell -python -m pip download -r requirements-head-mask.txt --dest wheels -python -m pip install --no-index --find-links .\wheels -r requirements-head-mask.txt -``` diff --git a/extract_head_mask_single.py b/extract_head_mask_single.py deleted file mode 100644 index a4551e1..0000000 --- a/extract_head_mask_single.py +++ /dev/null @@ -1,537 +0,0 @@ -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()) diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 4c16f91..0000000 --- a/requirements.txt +++ /dev/null @@ -1,9 +0,0 @@ -mediapipe==0.10.9 -numpy>=1.24,<2.1 -opencv-python>=4.8 -Pillow>=10.0 -torch>=2.0 -transformers>=4.40,<5 -safetensors>=0.4 -insightface>=0.7 -onnxruntime>=1.16