From 5ccb0a5e687cd6a1df951422d8f367c2129b83be Mon Sep 17 00:00:00 2001 From: toto Date: Fri, 24 May 2024 16:36:32 +0800 Subject: [PATCH] Auto update --- facechain/common/model_processor.py | 4 +++- facechain/nodes.py | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/facechain/common/model_processor.py b/facechain/common/model_processor.py index 9a69f6f..a552642 100644 --- a/facechain/common/model_processor.py +++ b/facechain/common/model_processor.py @@ -9,7 +9,7 @@ from facechain.utils.convert_utils import * def debug(*args): print(f"==== face chain debug ====", *args) -def facechain_detect_crop(source_image_pil, face_index=0, crop_ratio=1, mode='normal'): +def facechain_detect_crop(source_image_pil, face_index=0, crop_ratio=1, mode='normal', error_when_no_face=False): det_result = get_face_detection()(source_image_pil) mask = np.zeros_like(source_image_pil) bboxes = det_result['boxes'] @@ -24,6 +24,8 @@ def facechain_detect_crop(source_image_pil, face_index=0, crop_ratio=1, mode='no # idx = i try: bbox = bboxes[face_index] + if error_when_no_face: + return source_image_pil, None, None, None except: raise RuntimeError('No face detected or face index error/没有检测到人脸或者人脸的索引错误') diff --git a/facechain/nodes.py b/facechain/nodes.py index 755f479..09cf844 100644 --- a/facechain/nodes.py +++ b/facechain/nodes.py @@ -33,6 +33,7 @@ class FaceDetectCrop: "face_index": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1}), "crop_ratio": ("FLOAT", {"default": 1.0, "min": 0, "max": 10, "step": 0.1}), "mode": (["normal", "square 512 width height"],), + "error_when_no_face": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), } } @@ -40,9 +41,11 @@ class FaceDetectCrop: FUNCTION = "face_detection" CATEGORY = "facechain/model" - def face_detection(self, source_image, face_index, crop_ratio, mode): + def face_detection(self, source_image, face_index, crop_ratio, mode, error_when_no_face): pil_image = tensor_to_img(source_image) - corp_img_pil, mask, bbox, points_array = facechain_detect_crop(pil_image, face_index, crop_ratio, mode) + corp_img_pil, mask, bbox, points_array = facechain_detect_crop(pil_image, face_index, crop_ratio, mode, error_when_no_face) + if mask is None and bbox is None and points_array is None: + return (source_image, source_image, None, None) return (image_to_tensor(corp_img_pil), mask_np3_to_mask_tensor(mask), bbox, points_array,) class FCFaceSegment: