mirror of
https://github.com/Zeyi-Lin/HivisionIDPhotos.git
synced 2026-09-01 15:07:25 +08:00
KAAANG
120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
import requests
|
||
import base64
|
||
import argparse
|
||
import os
|
||
|
||
|
||
def base64_save(_base64_image_data, save_path):
|
||
# 解码 Base64 数据并保存为 PNG 文件
|
||
img_data = base64.b64decode(_base64_image_data)
|
||
with open(save_path, "wb") as file:
|
||
file.write(img_data)
|
||
|
||
|
||
# 读取本地图像文件并转换为Base64编码
|
||
def file_2_base64(file_path):
|
||
with open(file_path, "rb") as file:
|
||
encoded_string = base64.b64encode(file.read()).decode("utf-8")
|
||
return encoded_string
|
||
|
||
|
||
# 发送请求到 /idphoto 接口
|
||
def request_idphoto(file_path, height, width):
|
||
files = {"input_image": open(file_path, "rb")}
|
||
data = {"height": int(height), "width": int(width)}
|
||
response = requests.post(url, files=files, data=data)
|
||
return response.json()
|
||
|
||
|
||
# 发送请求到 /add_background 接口
|
||
def request_add_background(file_path, color, kb=None):
|
||
files = {"input_image": open(file_path, "rb")}
|
||
data = {"color": str(color), "kb": kb}
|
||
response = requests.post(url, files=files, data=data)
|
||
return response.json()
|
||
|
||
|
||
# 发送请求到 /generate_layout_photos 接口
|
||
def request_generate_layout_photos(file_path, height, width, kb=None):
|
||
files = {"input_image": open(file_path, "rb")}
|
||
data = {"height": height, "width": width, "kb": kb}
|
||
response = requests.post(url, files=files, data=data)
|
||
return response.json()
|
||
|
||
|
||
# 示例调用
|
||
if __name__ == "__main__":
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="HivisionIDPhotos 证件照制作推理程序。"
|
||
)
|
||
parser.add_argument(
|
||
"-u", "--url", help="API 服务的 URL", default="http://localhost:8080"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"-t",
|
||
"--type",
|
||
help="请求 API 的种类，有 idphoto、add_background 和 generate_layout_photos 可选",
|
||
default="idphoto",
|
||
)
|
||
parser.add_argument("-i", "--input_image_dir", help="输入图像路径", required=True)
|
||
parser.add_argument("-o", "--output_image_dir", help="保存图像路径", required=True)
|
||
parser.add_argument("--height", help="证件照尺寸-高", default=413)
|
||
parser.add_argument("--width", help="证件照尺寸-宽", default=295)
|
||
parser.add_argument("-c", "--color", help="证件照背景色", default="638cce")
|
||
parser.add_argument(
|
||
"-k", "--kb", help="输出照片的 KB 值，仅对换底和制作排版照生效", default=None
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
url = f"{args.url}/{args.type}" # 替换为实际的接口 URL
|
||
# color = hex_to_rgb(args.color)
|
||
# color = (color[2], color[1], color[0])
|
||
|
||
if args.type == "idphoto":
|
||
# 调用 /idphoto 接口
|
||
idphoto_response = request_idphoto(
|
||
args.input_image_dir, int(args.height), int(args.width)
|
||
)
|
||
|
||
if idphoto_response["status"]:
|
||
# 解码 Base64 数据并保存为 PNG 文件
|
||
base64_image_data_standard = idphoto_response["image_base64_standard"]
|
||
base64_image_data_standard_hd = idphoto_response["image_base64_hd"]
|
||
|
||
file_name, file_extension = os.path.splitext(args.output_image_dir)
|
||
# 定义新的文件路径（在原有的文件名后添加"_hd"）
|
||
new_file_name = file_name + "_hd" + file_extension
|
||
|
||
# 解码 Base64 数据并保存为 PNG 文件
|
||
base64_save(base64_image_data_standard, args.output_image_dir)
|
||
base64_save(base64_image_data_standard_hd, new_file_name)
|
||
|
||
print(f"请求{args.type}接口成功，已保存图像。")
|
||
else:
|
||
print("人脸数量不等于 1，请上传单张人脸的图像。")
|
||
|
||
elif args.type == "add_background":
|
||
# 调用 /add_background 接口
|
||
add_background_response = request_add_background(
|
||
args.input_image_dir, args.color, kb=args.kb
|
||
)
|
||
base64_image_data = add_background_response["image_base64"]
|
||
base64_save(base64_image_data, args.output_image_dir)
|
||
|
||
print(f"请求{args.type}接口成功，已保存图像。")
|
||
|
||
elif args.type == "generate_layout_photos":
|
||
# 调用 /generate_layout_photos 接口
|
||
generate_layout_response = request_generate_layout_photos(
|
||
args.input_image_dir, int(args.height), int(args.width), args.kb
|
||
)
|
||
base64_image_data = generate_layout_response["image_base64"]
|
||
base64_save(base64_image_data, args.output_image_dir)
|
||
|
||
print(f"请求{args.type}接口成功，已保存图像。")
|
||
|
||
else:
|
||
print("不支持的 API 类型，请检查输入参数。")
|