diff --git a/README.md b/README.md index 449e17f..cd1a379 100644 --- a/README.md +++ b/README.md @@ -52,13 +52,13 @@ - 在线体验： [![SwanHub Demo](https://img.shields.io/static/v1?label=Demo&message=SwanHub%20Demo&color=blue)](https://swanhub.co/ZeYiLin/HivisionIDPhotos/demo)、[![Spaces](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue)](https://huggingface.co/spaces/TheEeeeLin/HivisionIDPhotos) +- 2024.09.10: 增加新的**人脸检测模型** Retinaface-resnet50，以稍弱于mtcnn的速度换取更高的检测精度，推荐使用 - 2024.09.09: 增加新的**抠图模型** [BiRefNet-v1-lite](https://github.com/ZhengPeng7/BiRefNet) | Gradio增加**高级参数设置**和**水印**选项卡 - 2024.09.08: 增加新的**抠图模型** [RMBG-1.4](https://huggingface.co/briaai/RMBG-1.4) | **ComfyUI工作流** - [HivisionIDPhotos-ComfyUI](https://github.com/AIFSH/HivisionIDPhotos-ComfyUI) 贡献 by [AIFSH](https://github.com/AIFSH/HivisionIDPhotos-ComfyUI) - 2024.09.07: 增加**人脸检测API选项** [Face++](docs/face++_CN.md)，实现更高精度的人脸检测 - 2024.09.06: 增加新的抠图模型 [modnet_photographic_portrait_matting.onnx](https://github.com/ZHKKKe/MODNet) - 2024.09.05: 更新 [Restful API 文档](docs/api_CN.md) - 2024.09.02: 更新**调整照片 KB 大小**，[DockerHub](https://hub.docker.com/r/linzeyi/hivision_idphotos/tags) -- 2023.12.01: 更新**API 部署（基于 fastapi）**
@@ -142,7 +142,8 @@ python scripts/download_model.py --models all | 拓展人脸检测模型 | 介绍 | 使用文档 | | -- | -- | -- | -| MTCNN | **离线**人脸检测模型，高性能CPU推理，为默认模型，检测精度较低 | Clone此项目后直接使用 | +| MTCNN | **离线**人脸检测模型，高性能CPU推理（毫秒级），为默认模型，检测精度较低 | Clone此项目后直接使用 | +| RetinaFace | **离线**人脸检测模型，CPU推理速度中等（秒级），精度较高| [下载](https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/retinaface-resnet50.onnx)后放到`hivision/creator/retinaface/weights`目录下 | | Face++ | 旷视推出的在线人脸检测API，检测精度较高，[官方文档](https://console.faceplusplus.com.cn/documents/4888373) | [使用文档](docs/face++_CN.md)| ## 5. GPU推理加速（可选） diff --git a/README_EN.md b/README_EN.md index 8f60727..5f943f7 100644 --- a/README_EN.md +++ b/README_EN.md @@ -50,13 +50,13 @@ English / [中文](README.md) / [日本語](README_JP.md) / [한국어](README_K - Online Experience: [![SwanHub Demo](https://img.shields.io/static/v1?label=Demo&message=SwanHub%20Demo&color=blue)](https://swanhub.co/ZeYiLin/HivisionIDPhotos/demo)、[![Spaces](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue)](https://huggingface.co/spaces/TheEeeeLin/HivisionIDPhotos) +- 2024.09.10: Added a new **face detection model** Retinaface-resnet50, which offers higher detection accuracy at a slightly slower speed compared to mtcnn. Recommended for use. - 2024.09.09: Added a new **Background Removal Model** [BiRefNet-v1-lite](https://github.com/ZhengPeng7/BiRefNet) | Gradio added **Advanced Parameter Settings** and **Watermark** tabs - 2024.09.08: Added new **Matting Model** [RMBG-1.4](https://huggingface.co/briaai/RMBG-1.4) | **ComfyUI Workflow** - [HivisionIDPhotos-ComfyUI](https://github.com/AIFSH/HivisionIDPhotos-ComfyUI) contributed by [AIFSH](https://github.com/AIFSH/HivisionIDPhotos-ComfyUI) - 2024.09.07: Added **Face Detection API Option** [Face++](docs/face++_EN.md), achieving higher precision in face detection - 2024.09.06: Added new matting model [modnet_photographic_portrait_matting.onnx](https://github.com/ZHKKKe/MODNet) - 2024.09.05: Updated [Restful API Documentation](docs/api_EN.md) - 2024.09.02: Updated **Adjust Photo KB Size**, [DockerHub](https://hub.docker.com/r/linzeyi/hivision_idphotos/tags) -- 2023.12.01: Updated **API Deployment (based on fastapi)**
@@ -140,6 +140,7 @@ Store in the project's `hivision/creator/weights` directory: | Extended Face Detection Model | Description | Documentation | | -- | -- | -- | | MTCNN | **Offline** face detection model, high-performance CPU inference, default model, lower detection accuracy | Use it directly after cloning this project | +| RetinaFace | **Offline** face detection model, moderate CPU inference speed (in seconds), and high accuracy | [Download](https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/retinaface-resnet50.onnx) and place it in the `hivision/creator/retinaface/weights` directory | | Face++ | Online face detection API launched by Megvii, higher detection accuracy, [official documentation](https://console.faceplusplus.com.cn/documents/4888373) | [Usage Documentation](docs/face++_EN.md)| ## 5. GPU Inference Acceleration (Optional) diff --git a/app.py b/app.py index ecd5bbc..a318e52 100644 --- a/app.py +++ b/app.py @@ -16,6 +16,18 @@ HUMAN_MATTING_MODELS = [ model for model in HUMAN_MATTING_MODELS if model in HUMAN_MATTING_MODELS_EXIST ] +FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn"] +FACE_DETECT_MODELS_EXPAND = ( + ["retinaface-resnet50"] + if os.path.exists( + os.path.join( + root_dir, "hivision/creator/retinaface/weights/retinaface-resnet50.onnx" + ) + ) + else [] +) +FACE_DETECT_MODELS += FACE_DETECT_MODELS_EXPAND + if __name__ == "__main__": argparser = argparse.ArgumentParser() argparser.add_argument( @@ -34,7 +46,9 @@ if __name__ == "__main__": processor = IDPhotoProcessor() - demo = create_ui(processor, root_dir, HUMAN_MATTING_MODELS_EXIST) + demo = create_ui( + processor, root_dir, HUMAN_MATTING_MODELS_EXIST, FACE_DETECT_MODELS + ) demo.launch( server_name=args.host, server_port=args.port, diff --git a/demo/ui.py b/demo/ui.py index adcb89b..9c3462f 100644 --- a/demo/ui.py +++ b/demo/ui.py @@ -11,7 +11,9 @@ def load_description(fp): return content -def create_ui(processor, root_dir, human_matting_models: list): +def create_ui( + processor, root_dir, human_matting_models: list, face_detect_models: list +): DEFAULT_LANG = "zh" DEFAULT_HUMAN_MATTING_MODEL = "modnet_photographic_portrait_matting" DEFAULT_FACE_DETECT_MODEL = "mtcnn" @@ -61,7 +63,7 @@ def create_ui(processor, root_dir, human_matting_models: list): ) face_detect_model_options = gr.Dropdown( - choices=FACE_DETECT_MODELS, + choices=face_detect_models, label=LOCALES["face_model"][DEFAULT_LANG]["label"], value=DEFAULT_FACE_DETECT_MODEL, ) diff --git a/hivision/creator/choose_handler.py b/hivision/creator/choose_handler.py index 81b1ddb..d0885a0 100644 --- a/hivision/creator/choose_handler.py +++ b/hivision/creator/choose_handler.py @@ -9,7 +9,7 @@ HUMAN_MATTING_MODELS = [ "rmbg-1.4", ] -FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn"] +FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn", "retinaface-resnet50"] def choose_handler(creator, matting_model_option=None, face_detect_option=None): @@ -29,5 +29,7 @@ def choose_handler(creator, matting_model_option=None, face_detect_option=None): or face_detect_option == "face++ (联网Online API)" ): creator.detection_handler = detect_face_face_plusplus + elif face_detect_option == "retinaface-resnet50": + creator.detection_handler = detect_face_retinaface else: creator.detection_handler = detect_face_mtcnn diff --git a/hivision/creator/face_detector.py b/hivision/creator/face_detector.py index 4f18c13..2843521 100644 --- a/hivision/creator/face_detector.py +++ b/hivision/creator/face_detector.py @@ -16,12 +16,15 @@ except ImportError: from .context import Context from hivision.error import FaceError, APIError from hivision.utils import resize_image_to_kb_base64 +from hivision.creator.retinaface import retinaface_detect_faces import requests import cv2 import os mtcnn = None +base_dir = os.path.dirname(os.path.abspath(__file__)) +RETINAFCE_SESS = None def detect_face_mtcnn(ctx: Context, scale: int = 2): @@ -129,3 +132,45 @@ def detect_face_face_plusplus(ctx: Context): f"Face++ Status code {status_code} Request entity too large: The image exceeds the 2MB limit.", status_code, ) + + +def detect_face_retinaface(ctx: Context): + """ + 基于RetinaFace模型的人脸检测处理器，只进行人脸数量的检测 + :param ctx: 上下文，此时已获取到原始图和抠图结果，但是我们只需要原始图 + :raise FaceError: 人脸检测错误，多个人脸或者没有人脸 + """ + from time import time + + global RETINAFCE_SESS + + if RETINAFCE_SESS is None: + print("首次加载RetinaFace模型...") + # 计算用时 + tic = time() + faces_dets, sess = retinaface_detect_faces( + ctx.origin_image, + os.path.join(base_dir, "retinaface/weights/retinaface-resnet50.onnx"), + sess=None, + ) + RETINAFCE_SESS = sess + print("首次RetinaFace模型推理用时: {:.4f}s".format(time() - tic)) + else: + tic = time() + faces_dets, _ = retinaface_detect_faces( + ctx.origin_image, + os.path.join(base_dir, "retinaface/weights/retinaface-resnet50.onnx"), + sess=RETINAFCE_SESS, + ) + print("二次RetinaFace模型推理用时: {:.4f}s".format(time() - tic)) + + faces_num = len(faces_dets) + if faces_num != 1: + raise FaceError("Expected 1 face, but got {}".format(faces_num), faces_num) + face_det = faces_dets[0] + ctx.face = ( + face_det[0], + face_det[1], + face_det[2] - face_det[0] + 1, + face_det[3] - face_det[1] + 1, + ) diff --git a/hivision/creator/retinaface/__init__.py b/hivision/creator/retinaface/__init__.py new file mode 100644 index 0000000..9b5e94d --- /dev/null +++ b/hivision/creator/retinaface/__init__.py @@ -0,0 +1 @@ +from .inference import retinaface_detect_faces diff --git a/hivision/creator/retinaface/box_utils.py b/hivision/creator/retinaface/box_utils.py new file mode 100644 index 0000000..f1c8c21 --- /dev/null +++ b/hivision/creator/retinaface/box_utils.py @@ -0,0 +1,57 @@ +import numpy as np + + +def decode(loc, priors, variances): + """Decode locations from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + loc (tensor): location predictions for loc layers, + Shape: [num_priors,4] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded bounding box predictions + """ + + boxes = None + + boxes = np.concatenate( + ( + priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], + priors[:, 2:] * np.exp(loc[:, 2:] * variances[1]), + ), + axis=1, + ) + + boxes[:, :2] -= boxes[:, 2:] / 2 + boxes[:, 2:] += boxes[:, :2] + return boxes + + +def decode_landm(pre, priors, variances): + """Decode landm from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + pre (tensor): landm predictions for loc layers, + Shape: [num_priors,10] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded landm predictions + """ + landms = None + + landms = np.concatenate( + ( + priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:], + ), + axis=1, + ) + + return landms diff --git a/hivision/creator/retinaface/inference.py b/hivision/creator/retinaface/inference.py new file mode 100644 index 0000000..0622662 --- /dev/null +++ b/hivision/creator/retinaface/inference.py @@ -0,0 +1,190 @@ +import numpy as np +import cv2 +import onnxruntime as ort +from hivision.creator.retinaface.box_utils import decode, decode_landm +from hivision.creator.retinaface.prior_box import PriorBox +import argparse + + +def py_cpu_nms(dets, thresh): + """Pure Python NMS baseline.""" + x1 = dets[:, 0] + y1 = dets[:, 1] + x2 = dets[:, 2] + y2 = dets[:, 3] + scores = dets[:, 4] + + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + inds = np.where(ovr <= thresh)[0] + order = order[inds + 1] + + return keep + + +parser = argparse.ArgumentParser(description="Retinaface") + +parser.add_argument( + "--network", default="resnet50", help="Backbone network mobile0.25 or resnet50" +) +parser.add_argument( + "--cpu", action="store_true", default=False, help="Use cpu inference" +) +parser.add_argument( + "--confidence_threshold", default=0.8, type=float, help="confidence_threshold" +) +parser.add_argument("--top_k", default=5000, type=int, help="top_k") +parser.add_argument("--nms_threshold", default=0.2, type=float, help="nms_threshold") +parser.add_argument("--keep_top_k", default=750, type=int, help="keep_top_k") +parser.add_argument( + "-s", + "--save_image", + action="store_true", + default=True, + help="show detection results", +) +parser.add_argument( + "--vis_thres", default=0.6, type=float, help="visualization_threshold" +) +args = parser.parse_args() + + +def load_model_ort(model_path): + ort_session = ort.InferenceSession(model_path) + return ort_session + + +def retinaface_detect_faces(image, model_path: str, sess=None): + cfg = { + "name": "Resnet50", + "min_sizes": [[16, 32], [64, 128], [256, 512]], + "steps": [8, 16, 32], + "variance": [0.1, 0.2], + "clip": False, + "loc_weight": 2.0, + "gpu_train": True, + "batch_size": 24, + "ngpu": 4, + "epoch": 100, + "decay1": 70, + "decay2": 90, + "image_size": 840, + "pretrain": True, + "return_layers": {"layer2": 1, "layer3": 2, "layer4": 3}, + "in_channel": 256, + "out_channel": 256, + } + + # Load ONNX model + if sess is None: + retinaface = load_model_ort(model_path) + else: + retinaface = sess + + resize = 1 + + # Read and preprocess the image + img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + img = np.float32(img_rgb) + + im_height, im_width, _ = img.shape + scale = np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]) + img -= (104, 117, 123) + img = img.transpose(2, 0, 1) + img = np.expand_dims(img, axis=0) + + # Run the model + inputs = {"input": img} + loc, conf, landms = retinaface.run(None, inputs) + + # tic = time.time() + priorbox = PriorBox(cfg, image_size=(im_height, im_width)) + priors = priorbox.forward() + + prior_data = priors + + boxes = decode(np.squeeze(loc, axis=0), prior_data, cfg["variance"]) + boxes = boxes * scale / resize + scores = np.squeeze(conf, axis=0)[:, 1] + + landms = decode_landm(np.squeeze(landms.data, axis=0), prior_data, cfg["variance"]) + + scale1 = np.array( + [ + img.shape[3], + img.shape[2], + img.shape[3], + img.shape[2], + img.shape[3], + img.shape[2], + img.shape[3], + img.shape[2], + img.shape[3], + img.shape[2], + ] + ) + landms = landms * scale1 / resize + + # ignore low scores + inds = np.where(scores > args.confidence_threshold)[0] + boxes = boxes[inds] + landms = landms[inds] + scores = scores[inds] + + # keep top-K before NMS + order = scores.argsort()[::-1][: args.top_k] + boxes = boxes[order] + landms = landms[order] + scores = scores[order] + + # do NMS + dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) + keep = py_cpu_nms(dets, args.nms_threshold) + # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu) + dets = dets[keep, :] + landms = landms[keep] + + # keep top-K faster NMS + dets = dets[: args.keep_top_k, :] + landms = landms[: args.keep_top_k, :] + + dets = np.concatenate((dets, landms), axis=1) + # print("post processing time: {:.4f}s".format(time.time() - tic)) + + return dets, retinaface + + +if __name__ == "__main__": + import gradio as gr + + # Create Gradio interface + iface = gr.Interface( + fn=retinaface_detect_faces, + inputs=[ + gr.Image( + type="numpy", label="上传图片", height=400 + ), # Set the height to 400 + gr.Textbox(value="./FaceDetector.onnx", label="ONNX模型路径"), + ], + outputs=gr.Number(label="检测到的人脸数量"), + title="人脸检测", + description="上传图片并提供ONNX模型路径以检测人脸数量。", + ) + + # Launch the Gradio app + iface.launch() diff --git a/hivision/creator/retinaface/prior_box.py b/hivision/creator/retinaface/prior_box.py new file mode 100644 index 0000000..341a35c --- /dev/null +++ b/hivision/creator/retinaface/prior_box.py @@ -0,0 +1,41 @@ +from itertools import product as product +import numpy as np +from math import ceil + + +class PriorBox(object): + def __init__(self, cfg, image_size=None): + super(PriorBox, self).__init__() + self.min_sizes = cfg["min_sizes"] + self.steps = cfg["steps"] + self.clip = cfg["clip"] + self.image_size = image_size + self.feature_maps = [ + [ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)] + for step in self.steps + ] + self.name = "s" + + def forward(self): + anchors = [] + for k, f in enumerate(self.feature_maps): + min_sizes = self.min_sizes[k] + for i, j in product(range(f[0]), range(f[1])): + for min_size in min_sizes: + s_kx = min_size / self.image_size[1] + s_ky = min_size / self.image_size[0] + dense_cx = [ + x * self.steps[k] / self.image_size[1] for x in [j + 0.5] + ] + dense_cy = [ + y * self.steps[k] / self.image_size[0] for y in [i + 0.5] + ] + for cy, cx in product(dense_cy, dense_cx): + anchors += [cx, cy, s_kx, s_ky] + + output = np.array(anchors).reshape(-1, 4) + + if self.clip: + output = np.clip(output, 0, 1) + + return output diff --git a/hivision/creator/retinaface/weights/.gitkeep b/hivision/creator/retinaface/weights/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/scripts/download_model.py b/scripts/download_model.py index 2a4311d..d27d3cd 100644 --- a/scripts/download_model.py +++ b/scripts/download_model.py @@ -32,14 +32,14 @@ def download_file(url, save_path): def download_models(model_urls): - # 指定下载保存的目录 - save_dir = "hivision/creator/weights" - - # 创建目录（如果不存在的话） - os.makedirs(os.path.join(base_path, save_dir), exist_ok=True) - # 下载每个模型 for model_name, model_info in model_urls.items(): + # 指定下载保存的目录 + save_dir = model_info["location"] + + # 创建目录（如果不存在的话） + os.makedirs(os.path.join(base_path, save_dir), exist_ok=True) + url = model_info["url"] file_format = model_info["format"] @@ -63,10 +63,12 @@ def main(models_to_download): "hivision_modnet": { "url": "https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/hivision_modnet.onnx", "format": "onnx", + "location": "hivision/creator/weights", }, "modnet_photographic_portrait_matting": { "url": "https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/modnet_photographic_portrait_matting.onnx", "format": "onnx", + "location": "hivision/creator/weights", }, # "mnn_hivision_modnet": { # "url": "https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/mnn_hivision_modnet.mnn", @@ -75,10 +77,17 @@ def main(models_to_download): "rmbg-1.4": { "url": "https://huggingface.co/briaai/RMBG-1.4/resolve/main/onnx/model.onnx?download=true", "format": "onnx", + "location": "hivision/creator/weights", }, "birefnet-v1-lite": { "url": "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx", "format": "onnx", + "location": "hivision/creator/weights", + }, + "retinaface-resnet50": { + "url": "https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/retinaface-resnet50.onnx", + "format": "onnx", + "location": "hivision/creator/retinaface/weights", }, }