103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
# "Design-focused node suite for ComfyUI."
|
|
# Copyright 2026 Augment Studio
|
|
# Augmentstudio.app
|
|
|
|
import json
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
import os
|
|
|
|
class FaceCrop:
|
|
"""Fast face crop using OpenCV. Image in, cropped image + coordinates out."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"padding": ("INT", {"default": 20, "min": 0, "max": 512, "step": 5}),
|
|
"square": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "INT", "AUGMENT_JSON", "TRIGGER")
|
|
RETURN_NAMES = ("cropped", "x", "y", "width", "height", "json", "trigger")
|
|
OUTPUT_NODE = False
|
|
FUNCTION = "detect_and_crop"
|
|
CATEGORY = "Augment/Image"
|
|
|
|
def detect_and_crop(self, image, padding=20, square=True):
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
img_h, img_w = img_np.shape[:2]
|
|
|
|
gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
|
|
|
|
# Try multiple paths to find the Haar cascade file
|
|
cascade_file = None
|
|
|
|
# Try cv2.data path first (if available)
|
|
if hasattr(cv2, 'data'):
|
|
try:
|
|
cascade_file = cv2.data.haarcascades + "haarcascade_frontalface_alt2.xml"
|
|
if os.path.exists(cascade_file):
|
|
pass # Found it
|
|
else:
|
|
cascade_file = None
|
|
except:
|
|
cascade_file = None
|
|
|
|
# Try other common paths
|
|
if not cascade_file:
|
|
cascade_paths = [
|
|
os.path.join(cv2.__path__[0], "data", "haarcascade_frontalface_alt2.xml"),
|
|
"/usr/share/opencv4/haarcascades/haarcascade_frontalface_alt2.xml",
|
|
"/usr/local/share/opencv4/haarcascades/haarcascade_frontalface_alt2.xml",
|
|
"haarcascade_frontalface_alt2.xml", # System path fallback
|
|
]
|
|
|
|
for path in cascade_paths:
|
|
if os.path.exists(path):
|
|
cascade_file = path
|
|
break
|
|
|
|
if not cascade_file:
|
|
cascade_file = "haarcascade_frontalface_alt2.xml"
|
|
|
|
cascade = cv2.CascadeClassifier(cascade_file)
|
|
faces = cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=6, minSize=(40, 40))
|
|
|
|
face_count = len(faces) if isinstance(faces, np.ndarray) else 0
|
|
|
|
if face_count == 0:
|
|
print("[augment] FaceCrop: no faces detected")
|
|
result_json = json.dumps({"node": "FaceCrop", "x": 0, "y": 0, "width": img_w, "height": img_h})
|
|
return (image, 0, 0, img_w, img_h, result_json, "done")
|
|
|
|
faces = sorted(faces, key=lambda f: f[2] * f[3], reverse=True)
|
|
fx, fy, fw, fh = [int(v) for v in faces[0]]
|
|
|
|
cx, cy = fx + fw // 2, fy + fh // 2
|
|
fw += padding * 2
|
|
fh += padding * 2
|
|
|
|
if square:
|
|
side = max(fw, fh)
|
|
fw, fh = side, side
|
|
|
|
x0 = max(0, cx - fw // 2)
|
|
y0 = max(0, cy - fh // 2)
|
|
x1 = min(img_w, x0 + fw)
|
|
y1 = min(img_h, y0 + fh)
|
|
|
|
out_x, out_y = int(x0), int(y0)
|
|
out_w, out_h = int(x1 - x0), int(y1 - y0)
|
|
|
|
print(f"[augment] FaceCrop: {face_count} face(s), crop x={out_x} y={out_y} w={out_w} h={out_h}")
|
|
|
|
cropped = image[:, out_y:out_y + out_h, out_x:out_x + out_w, :]
|
|
result_json = json.dumps({"node": "FaceCrop", "x": out_x, "y": out_y, "width": out_w, "height": out_h})
|
|
return (cropped, out_x, out_y, out_w, out_h, result_json, "done")
|