mirror of
https://github.com/Zeyi-Lin/HivisionIDPhotos.git
synced 2026-08-28 19:45:12 +08:00
汉图-汪钧剑
374 lines
12 KiB
Python
374 lines
12 KiB
Python
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)
|