from fastapi import FastAPI, UploadFile, Form, File from hivision import IDCreator from hivision.error import FaceError from hivision.creator.layout_calculator import ( generate_layout_array, generate_layout_image, ) from hivision.creator.choose_handler import choose_handler from hivision.utils import ( add_background, resize_image_to_kb, bytes_2_base64, base64_2_numpy, hex_to_rgb, add_watermark, save_image_dpi_to_bytes, ) import numpy as np import cv2 from starlette.middleware.cors import CORSMiddleware from starlette.formparsers import MultiPartParser # 设置Starlette表单字段大小限制 MultiPartParser.max_part_size = 10 * 1024 * 1024 # 10MB # 设置Starlette文件上传大小限制 MultiPartParser.max_file_size = 20 * 1024 * 1024 # 20MB app = FastAPI() creator = IDCreator() # 添加 CORS 中间件 解决跨域问题 app.add_middleware( CORSMiddleware, allow_origins=["*"], # 允许的请求来源 allow_credentials=True, # 允许携带 Cookie allow_methods=[ "*" ], # 允许的请求方法，例如：GET, POST 等，也可以指定 ["GET", "POST"] allow_headers=["*"], # 允许的请求头，也可以指定具体的头部 ) # 证件照智能制作接口 @app.post("/idphoto") async def idphoto_inference( input_image: UploadFile = File(None), input_image_base64: str = Form(None), height: int = Form(413), width: int = Form(295), human_matting_model: str = Form("modnet_photographic_portrait_matting"), face_detect_model: str = Form("mtcnn"), hd: bool = Form(True), dpi: int = Form(300), face_align: bool = Form(False), whitening_strength: int = Form(0), head_measure_ratio: float = Form(0.2), head_height_ratio: float = Form(0.45), top_distance_max: float = Form(0.12), top_distance_min: float = Form(0.10), brightness_strength: float = Form(0), contrast_strength: float = Form(0), sharpen_strength: float = Form(0), saturation_strength: float = Form(0), ): # 如果传入了base64，则直接使用base64解码 if input_image_base64: img = base64_2_numpy(input_image_base64) # 否则使用上传的图片 else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # 将BGR转换为RGB img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # ------------------- 选择抠图与人脸检测模型 ------------------- choose_handler(creator, human_matting_model, face_detect_model) # 将字符串转为元组 size = (int(height), int(width)) try: result = creator( img, size=size, head_measure_ratio=head_measure_ratio, head_height_ratio=head_height_ratio, head_top_range=(top_distance_max, top_distance_min), face_alignment=face_align, whitening_strength=whitening_strength, brightness_strength=brightness_strength, contrast_strength=contrast_strength, sharpen_strength=sharpen_strength, saturation_strength=saturation_strength, ) except FaceError: result_message = {"status": False} # 如果检测到人脸数量等于1, 则返回标准证和高清照结果（png 4通道图像） else: result_image_standard_bytes = save_image_dpi_to_bytes(result.standard, None, dpi) result_message = { "status": True, "image_base64_standard": bytes_2_base64(result_image_standard_bytes), } # 如果hd为True, 则增加高清照结果（png 4通道图像） if hd: result_image_hd_bytes = save_image_dpi_to_bytes(result.hd, None, dpi) result_message["image_base64_hd"] = bytes_2_base64(result_image_hd_bytes) return result_message # 人像抠图接口 @app.post("/human_matting") async def human_matting_inference( input_image: UploadFile = File(None), input_image_base64: str = Form(None), human_matting_model: str = Form("hivision_modnet"), dpi: int = Form(300), ): if input_image_base64: img = base64_2_numpy(input_image_base64) else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # ------------------- 选择抠图与人脸检测模型 ------------------- choose_handler(creator, human_matting_model, None) try: result = creator( img, change_bg_only=True, ) except FaceError: result_message = {"status": False} else: result_image_standard_bytes = save_image_dpi_to_bytes(cv2.cvtColor(result.standard, cv2.COLOR_RGBA2BGRA), None, dpi) result_message = { "status": True, "image_base64": bytes_2_base64(result_image_standard_bytes), } return result_message # 透明图像添加纯色背景接口 @app.post("/add_background") async def photo_add_background( input_image: UploadFile = File(None), input_image_base64: str = Form(None), color: str = Form("000000"), kb: int = Form(None), dpi: int = Form(300), render: int = Form(0), ): render_choice = ["pure_color", "updown_gradient", "center_gradient"] if input_image_base64: img = base64_2_numpy(input_image_base64) else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED) color = hex_to_rgb(color) color = (color[2], color[1], color[0]) result_image = add_background( img, bgr=color, mode=render_choice[render], ).astype(np.uint8) result_image = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR) if kb: result_image_bytes = resize_image_to_kb(result_image, None, int(kb), dpi=dpi) else: result_image_bytes = save_image_dpi_to_bytes(result_image, None, dpi=dpi) result_messgae = { "status": True, "image_base64": bytes_2_base64(result_image_bytes), } return result_messgae # 六寸排版照生成接口 @app.post("/generate_layout_photos") async def generate_layout_photos( input_image: UploadFile = File(None), input_image_base64: str = Form(None), height: int = Form(413), width: int = Form(295), kb: int = Form(None), dpi: int = Form(300), ): # try: if input_image_base64: img = base64_2_numpy(input_image_base64) else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) size = (int(height), int(width)) typography_arr, typography_rotate = generate_layout_array( input_height=size[0], input_width=size[1] ) result_layout_image = generate_layout_image( img, typography_arr, typography_rotate, height=size[0], width=size[1] ).astype(np.uint8) result_layout_image = cv2.cvtColor(result_layout_image, cv2.COLOR_RGB2BGR) if kb: result_layout_image_bytes = resize_image_to_kb( result_layout_image, None, int(kb), dpi=dpi ) else: result_layout_image_bytes = save_image_dpi_to_bytes(result_layout_image, None, dpi=dpi) result_layout_image_base64 = bytes_2_base64(result_layout_image_bytes) result_messgae = { "status": True, "image_base64": result_layout_image_base64, } return result_messgae # 透明图像添加水印接口 @app.post("/watermark") async def watermark( input_image: UploadFile = File(None), input_image_base64: str = Form(None), text: str = Form("Hello"), size: int = 20, opacity: float = 0.5, angle: int = 30, color: str = "#000000", space: int = 25, kb: int = Form(None), dpi: int = Form(300), ): if input_image_base64: img = base64_2_numpy(input_image_base64) else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) try: result_image = add_watermark(img, text, size, opacity, angle, color, space) result_image = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR) if kb: result_image_bytes = resize_image_to_kb(result_image, None, int(kb), dpi=dpi) else: result_image_bytes = save_image_dpi_to_bytes(result_image, None, dpi=dpi) result_image_base64 = bytes_2_base64(result_image_bytes) result_messgae = { "status": True, "image_base64": result_image_base64, } except Exception as e: result_messgae = { "status": False, "error": str(e), } return result_messgae # 设置照片KB值接口(RGB图) @app.post("/set_kb") async def set_kb( input_image: UploadFile = File(None), input_image_base64: str = Form(None), dpi: int = Form(300), kb: int = Form(50), ): if input_image_base64: img = base64_2_numpy(input_image_base64) else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) try: result_image = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) result_image_bytes = resize_image_to_kb(result_image, None, int(kb), dpi=dpi) result_image_base64 = bytes_2_base64(result_image_bytes) result_messgae = { "status": True, "image_base64": result_image_base64, } except Exception as e: result_messgae = { "status": False, "error": e, } return result_messgae # 证件照智能裁剪接口 @app.post("/idphoto_crop") async def idphoto_crop_inference( input_image: UploadFile = File(None), input_image_base64: str = Form(None), height: int = Form(413), width: int = Form(295), face_detect_model: str = Form("mtcnn"), hd: bool = Form(True), dpi: int = Form(300), head_measure_ratio: float = Form(0.2), head_height_ratio: float = Form(0.45), top_distance_max: float = Form(0.12), top_distance_min: float = Form(0.10), ): if input_image_base64: img = base64_2_numpy(input_image_base64) else: image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED) # 读取图像(4通道) # ------------------- 选择抠图与人脸检测模型 ------------------- choose_handler(creator, face_detect_option=face_detect_model) # 将字符串转为元组 size = (int(height), int(width)) try: result = creator( img, size=size, head_measure_ratio=head_measure_ratio, head_height_ratio=head_height_ratio, head_top_range=(top_distance_max, top_distance_min), crop_only=True, ) except FaceError: result_message = {"status": False} # 如果检测到人脸数量等于1, 则返回标准证和高清照结果（png 4通道图像） else: result_image_standard_bytes = save_image_dpi_to_bytes(cv2.cvtColor(result.standard, cv2.COLOR_RGBA2BGRA), None, dpi) result_message = { "status": True, "image_base64_standard": bytes_2_base64(result_image_standard_bytes), } # 如果hd为True, 则增加高清照结果（png 4通道图像） if hd: result_image_hd_bytes = save_image_dpi_to_bytes(cv2.cvtColor(result.hd, cv2.COLOR_RGBA2BGRA), None, dpi) result_message["image_base64_hd"] = bytes_2_base64(result_image_hd_bytes) return result_message if __name__ == "__main__": import uvicorn # 在8080端口运行推理服务 uvicorn.run(app, host="0.0.0.0", port=8080)