nodes/WeChatCV-Stand-In_Preprocessor_ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

change antelopv2 to yolov8

Browse Source
This commit is contained in:
KBRASK
2025-08-13 17:40:54 +08:00
parent d1e9f7141a
commit 3e80a86cbd
3 changed files with 120 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Stand-In_preprocessor.py
+119 -66
View File
@@ -5,28 +5,83 @@ import torch
import numpy as np
import PIL.Image
import PIL.ImageOps
from insightface.app import FaceAnalysis
# The ultralytics library is required for YOLOv8. Install with: pip install ultralytics
from ultralytics import YOLO
from facexlib.parsing import init_parsing_model
from torchvision.transforms.functional import normalize
from typing import Union, Optional
import folder_paths
from huggingface_hub import snapshot_download
# This code now requires the 'ultralytics' package.
# Please install it via pip:
# pip install ultralytics
LOADED_PROCESSORS = {}
def set_extra_config_model_path(extra_config_models_dir_key, models_dir_name: str):
"""Helper function to set up model directories within ComfyUI."""
models_dir_default = os.path.join(folder_paths.models_dir, models_dir_name)
if not os.path.exists(models_dir_default):
os.makedirs(models_dir_default, exist_ok=True)
if extra_config_models_dir_key not in folder_paths.folder_names_and_paths:
folder_paths.folder_names_and_paths[extra_config_models_dir_key] = (
[os.path.join(folder_paths.models_dir, models_dir_name)],
[models_dir_default],
folder_paths.supported_pt_extensions,
)
else:
if not os.path.exists(models_dir_default):
os.makedirs(models_dir_default, exist_ok=True)
folder_paths.add_model_folder_path(extra_config_models_dir_key, models_dir_default, is_default=True)
set_extra_config_model_path("insightface", "insightface")
# Set up paths for YOLO and face parsing models
set_extra_config_model_path("yolo", "yolo")
set_extra_config_model_path("face_parsing", "face_parsing")
def download_yolo_model(model_name="model.pt"):
"""
Checks if a YOLO model exists and downloads it from a public source if not.
"""
# Publicly accessible model URL
model_url = f"https://huggingface.co/arnabdhar/YOLOv8-Face-Detection/resolve/main/{model_name}"
yolo_dir = os.path.join(folder_paths.get_folder_paths("yolo")[0])
model_path = os.path.join(yolo_dir, model_name)
# Ensure the target directory exists
os.makedirs(yolo_dir, exist_ok=True)
if not os.path.exists(model_path):
print(f"Model '{model_name}' not found locally. Starting download from public source...")
try:
with requests.get(model_url, stream=True, timeout=60) as r:
if r.status_code == 401:
print(f"\nDownload failed: Received 401 Unauthorized. The link '{model_url}' may require authentication.")
raise requests.exceptions.HTTPError(f"401 Client Error: Unauthorized for url: {model_url}")
r.raise_for_status()
total_size = int(r.headers.get('content-length', 0))
with open(model_path, 'wb') as f:
downloaded = 0
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
downloaded += len(chunk)
# Simple progress bar
done = int(50 * downloaded / total_size) if total_size > 0 else 0
progress_mb = downloaded / (1024*1024)
total_mb = total_size / (1024*1024)
print(f"\r[{'=' * done}{' ' * (50-done)}] {progress_mb:.2f}MB / {total_mb:.2f}MB", end='')
print(f"\nModel successfully downloaded to: {model_path}")
except Exception as e:
print(f"\nFailed to download model: {e}")
print("Please check your internet connection or try downloading the model manually from:")
print(model_url)
print(f"And place it in the '{yolo_dir}' directory.")
# Clean up incomplete file if download failed
if os.path.exists(model_path):
os.remove(model_path)
raise e
return model_path
# --- Utility Functions ---
def _img2tensor(img: np.ndarray, bgr2rgb: bool = True) -> torch.Tensor:
if bgr2rgb:
@@ -39,17 +94,12 @@ def _pad_to_square(img: np.ndarray, pad_color: int = 255) -> np.ndarray:
h, w, _ = img.shape
if h == w:
return img
if h > w:
pad_size = (h - w) // 2
padded_img = cv2.copyMakeBorder(
img, 0, 0, pad_size, h - w - pad_size, cv2.BORDER_CONSTANT, value=[pad_color] * 3,
)
padded_img = cv2.copyMakeBorder(img, 0, 0, pad_size, h - w - pad_size, cv2.BORDER_CONSTANT, value=[pad_color] * 3)
else:
pad_size = (w - h) // 2
padded_img = cv2.copyMakeBorder(
img, pad_size, w - h - pad_size, 0, 0, cv2.BORDER_CONSTANT, value=[pad_color] * 3,
)
padded_img = cv2.copyMakeBorder(img, pad_size, w - h - pad_size, 0, 0, cv2.BORDER_CONSTANT, value=[pad_color] * 3)
return padded_img
def tensor_to_cv2_img(tensor: torch.Tensor) -> np.ndarray:
@@ -60,67 +110,61 @@ def tensor_to_cv2_img(tensor: torch.Tensor) -> np.ndarray:
def cv2_img_to_tensor(img: np.ndarray) -> torch.Tensor:
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_rgb = img_rgb.astype(np.float32) / 255.0
return torch.from_numpy(img_rgb).unsqueeze(0)
img_tensor = torch.from_numpy(img_rgb.astype(np.float32) / 255.0).unsqueeze(0)
return img_tensor
class FaceProcessorLoader:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"antelopev2_path": ("STRING", {"default": "models/insightface"})
}
}
return { "required": { "yolo_model_name": ("STRING", {"default": "model.pt"}) } }
RETURN_TYPES = ("FACE_PROCESSOR",)
RETURN_NAMES = ("face_processor",)
FUNCTION = "load_processor"
CATEGORY = "Stand-In"
def load_processor(self, antelopev2_path="models/insightface"):
if "face_processor" in LOADED_PROCESSORS:
print("Reusing cached FaceProcessor models.")
return (LOADED_PROCESSORS["face_processor"],)
antelopev2_path = folder_paths.get_folder_paths("insightface")[0]
print("Initializing FaceProcessor models...")
def load_processor(self, yolo_model_name="model.pt"):
if os.path.isdir(yolo_model_name):
print(f"Warning: You provided a directory ('{yolo_model_name}') instead of a filename. Using default model 'model.pt'.")
yolo_model_name = "model.pt"
processor_key = f"face_processor_{yolo_model_name}"
if processor_key in LOADED_PROCESSORS:
print("Reusing cached face processor model.")
return (LOADED_PROCESSORS[processor_key],)
print("Initializing face processor models...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
providers = ["CUDAExecutionProvider"] if device.type == "cuda" else ["CPUExecutionProvider"]
if not os.path.exists(antelopev2_path):
os.makedirs(antelopev2_path)
print(f"Created insightface model directory at: {antelopev2_path}")
try:
print("Attempting to load FaceAnalysis model from local files...")
app = FaceAnalysis(name="antelopev2", root=antelopev2_path, providers=providers)
if not yolo_model_name.endswith(('.pt', '.onnx')):
raise ValueError(f"Invalid model name: '{yolo_model_name}'. Model filename must end with '.pt' or '.onnx'.")
# Call the download function to ensure the model file exists
model_path = download_yolo_model(yolo_model_name)
print(f"Loading YOLO model from local path: {model_path}")
detection_model = YOLO(model_path)
detection_model.to(device)
print("YOLO model loaded successfully.")
except Exception as e:
print(f"Failed to load local FaceAnalysis model: {e}")
if snapshot_download is None:
raise ImportError("huggingface_hub is not installed. Please install it with 'pip install huggingface_hub' to enable automatic model downloading.") from e
print(f"Error loading YOLO model: {e}")
raise e
print("Attempting to download model from Hugging Face...")
download_target_path = os.path.join(antelopev2_path, "models", "antelopev2")
snapshot_download(
repo_id="DIAMONIK7777/antelopev2",
local_dir=download_target_path,
)
print("Model downloaded. Retrying to load FaceAnalysis model...")
app = FaceAnalysis(name="antelopev2", root=antelopev2_path, providers=providers)
app.prepare(ctx_id=0, det_size=(640, 640))
parsing_model = init_parsing_model(model_name="bisenet", device=device)
parsing_model_path = os.path.join(folder_paths.get_folder_paths("face_parsing")[0])
parsing_model = init_parsing_model(model_name="bisenet", half=False, model_rootpath=parsing_model_path, device=device)
parsing_model.eval()
print("Face parsing model loaded successfully.")
processor_tuple = (app, parsing_model, device)
LOADED_PROCESSORS["face_processor"] = processor_tuple
processor_tuple = (detection_model, parsing_model, device)
LOADED_PROCESSORS[processor_key] = processor_tuple
print("FaceProcessor models initialized and cached successfully.")
print("Face processor (YOLO detection + BiSeNet parsing) initialized and cached successfully.")
return (processor_tuple,)
class ApplyFaceProcessor:
@classmethod
def INPUT_TYPES(s):
@@ -131,6 +175,7 @@ class ApplyFaceProcessor:
"resize_to": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
"border_thresh": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}),
"face_crop_scale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 3.0, "step": 0.1}),
"confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}),
}
}
@@ -138,27 +183,32 @@ class ApplyFaceProcessor:
FUNCTION = "apply_processing"
CATEGORY = "FaceUtils"
def apply_processing(self, face_processor, image, resize_to, border_thresh, face_crop_scale):
app, parsing_model, device = face_processor
def apply_processing(self, face_processor, image, resize_to, border_thresh, face_crop_scale, confidence_threshold):
detection_model, parsing_model, device = face_processor
frame = tensor_to_cv2_img(image)
faces = app.get(frame)
h, w, _ = frame.shape
image_to_process = None
if not faces:
print("[Warning] No face detected. Using the whole image, padded to square.")
results = detection_model(frame, verbose=False)
boxes = results[0].boxes.xyxy
conf = results[0].boxes.conf
confident_boxes = boxes[conf > confidence_threshold]
if confident_boxes.shape[0] == 0:
print("[Warning] No confident face detected. Using the whole image padded to a square.")
image_to_process = _pad_to_square(frame, pad_color=255)
else:
largest_face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
x1, y1, x2, y2 = map(int, largest_face.bbox)
areas = (confident_boxes[:, 2] - confident_boxes[:, 0]) * (confident_boxes[:, 3] - confident_boxes[:, 1])
largest_face_idx = torch.argmax(areas)
x1, y1, x2, y2 = map(int, confident_boxes[largest_face_idx])
is_close_to_border = (x1 <= border_thresh and y1 <= border_thresh and
x2 >= w - border_thresh and y2 >= h - border_thresh)
is_close_to_border = (x1 <= border_thresh or y1 <= border_thresh or
x2 >= w - border_thresh or y2 >= h - border_thresh)
if is_close_to_border:
print("[Info] Face is close to border, padding original image to square.")
print("[Info] Face is close to the border. Padding the original image to a square.")
image_to_process = _pad_to_square(frame, pad_color=255)
else:
cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
@@ -174,15 +224,18 @@ class ApplyFaceProcessor:
image_resized = cv2.resize(image_to_process, (resize_to, resize_to), interpolation=cv2.INTER_AREA)
face_tensor = _img2tensor(image_resized, bgr2rgb=True).unsqueeze(0).to(device)
with torch.no_grad():
normalized_face = normalize(face_tensor, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
parsing_out = parsing_model(normalized_face)[0]
parsing_mask = parsing_out.argmax(dim=1, keepdim=True)
mask = (parsing_mask.squeeze().cpu().numpy() == 0).astype(np.uint8)
mask = (parsing_mask.squeeze().cpu().numpy() != 0).astype(np.uint8)
white_background = np.ones_like(image_resized, dtype=np.uint8) * 255
mask_3channel = cv2.cvtColor(mask * 255, cv2.COLOR_GRAY2BGR)
result_img_bgr = np.where(mask_3channel == 255, white_background, image_resized)
result_img_bgr = np.where(mask_3channel != 0, image_resized, white_background)
result_tensor = cv2_img_to_tensor(result_img_bgr)
@@ -197,4 +250,4 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceProcessorLoader": "Stand-In Processor Loader",
"ApplyFaceProcessor": "Apply Stand-In Processor"
}
}