538 lines
19 KiB
Python
538 lines
19 KiB
Python
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())
|