From 3e80a86cbd760ba3fa89a661d0ab82062fc5042a Mon Sep 17 00:00:00 2001 From: KBRASK Date: Wed, 13 Aug 2025 17:40:54 +0800 Subject: [PATCH] change antelopv2 to yolov8 --- Stand-In.png | Bin Stand-In_preprocessor.py | 185 +++++++++++++++++++++++++-------------- requirements.txt | 2 +- 3 files changed, 120 insertions(+), 67 deletions(-) mode change 100755 => 100644 Stand-In.png diff --git a/Stand-In.png b/Stand-In.png old mode 100755 new mode 100644 diff --git a/Stand-In_preprocessor.py b/Stand-In_preprocessor.py index c329c37..2cba3cd 100644 --- a/Stand-In_preprocessor.py +++ b/Stand-In_preprocessor.py @@ -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" -} \ No newline at end of file +} diff --git a/requirements.txt b/requirements.txt index 5c7088d..902bccc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -insightface==0.7.3 +ultralytics huggingface_hub==0.31.1 facexlib==0.3.0 onnxruntime-gpu==1.22.0 \ No newline at end of file