From 24a22f72d845b4404c1c352f8906bd051960f918 Mon Sep 17 00:00:00 2001 From: Ze-Yi LIN <58305964+Zeyi-Lin@users.noreply.github.com> Date: Sat, 14 Sep 2024 18:07:29 +0800 Subject: [PATCH] feat: image dpi (#120) --- app.py | 3 + demo/locals.py | 22 +- demo/processor.py | 572 ++++++++++++++++++++++++++++------------------ demo/ui.py | 72 ++++-- hivision/utils.py | 34 ++- 5 files changed, 448 insertions(+), 255 deletions(-) diff --git a/app.py b/app.py index fc2880bf..92013ed7 100644 --- a/app.py +++ b/app.py @@ -6,6 +6,9 @@ root_dir = os.path.dirname(os.path.abspath(__file__)) +# 获取存在的人像分割模型列表 +# 通过检查 hivision/creator/weights 目录下的 .onnx 和 .mnn 文件 +# 只保留文件名（不包括扩展名） HUMAN_MATTING_MODELS_EXIST = [ os.path.splitext(file)[0] for file in os.listdir(os.path.join(root_dir, "hivision/creator/weights")) diff --git a/demo/locals.py b/demo/locals.py index 03c2334c..e1188f47 100644 --- a/demo/locals.py +++ b/demo/locals.py @@ -138,6 +138,24 @@ "label": "KB 大小", }, }, + "image_dpi": { + "en": { + "label": "Set DPI", + "choices": ["Not Set", "Custom"], + }, + "zh": { + "label": "设置 DPI 大小", + "choices": ["不设置", "自定义"], + }, + }, + "image_dpi_size": { + "en": { + "label": "DPI size", + }, + "zh": { + "label": "DPI 大小", + }, + }, "render_mode": { "en": { "label": "Render mode", @@ -278,10 +296,10 @@ }, "download": { "en": { - "label": "Download the photo after adjusting the KB size", + "label": "Download the photo after adjusting the DPI or KB size", }, "zh": { - "label": "下载调整 KB 大小后的照片", + "label": "下载调整 DPI 或 KB 大小后的照片", }, }, "matting_image": { diff --git a/demo/processor.py b/demo/processor.py index 28b979e1..a9ad8d30 100644 --- a/demo/processor.py +++ b/demo/processor.py @@ -1,7 +1,12 @@ import numpy as np from hivision import IDCreator from hivision.error import FaceError, APIError -from hivision.utils import add_background, resize_image_to_kb, add_watermark +from hivision.utils import ( + add_background, + resize_image_to_kb, + add_watermark, + save_image_dpi_to_bytes, +) from hivision.creator.layout_calculator import ( generate_layout_photo, generate_layout_image, @@ -42,265 +47,382 @@ def process( head_measure_ratio=0.2, top_distance_max=0.12, whitening_strength=0, + image_dpi_option=False, + custom_image_dpi=None, ): + # 初始化参数 top_distance_min = top_distance_max - 0.02 + idphoto_json = self._initialize_idphoto_json( + mode_option, color_option, render_option, image_kb_options + ) - idphoto_json = { + # 处理尺寸模式 + size_result = self._process_size_mode( + idphoto_json, + language, + size_list_option, + custom_size_height, + custom_size_width, + ) + if isinstance(size_result, list): + return size_result # 返回错误信息 + + # 处理颜色模式 + self._process_color_mode( + idphoto_json, + language, + color_option, + custom_color_R, + custom_color_G, + custom_color_B, + ) + + # 如果设置了自定义KB大小 + if ( + idphoto_json["image_kb_mode"] + == LOCALES["image_kb"][language]["choices"][-1] + ): + idphoto_json["custom_image_kb"] = custom_image_kb + + # 如果设置了自定义DPI大小 + if image_dpi_option == LOCALES["image_dpi"][language]["choices"][-1]: + idphoto_json["custom_image_dpi"] = custom_image_dpi + + # 创建IDCreator实例并设置处理器 + creator = IDCreator() + choose_handler(creator, matting_model_option, face_detect_option) + + # 生成证件照 + try: + result = self._generate_id_photo( + creator, + input_image, + idphoto_json, + language, + head_measure_ratio, + top_distance_max, + top_distance_min, + whitening_strength, + ) + except (FaceError, APIError): + return self._handle_photo_generation_error(language) + + # 后处理生成的照片 + return self._process_generated_photo( + result, + idphoto_json, + language, + watermark_option, + watermark_text, + watermark_text_size, + watermark_text_opacity, + watermark_text_angle, + watermark_text_space, + watermark_text_color, + ) + + def _initialize_idphoto_json( + self, + mode_option, + color_option, + render_option, + image_kb_options, + ): + """初始化idphoto_json字典""" + return { "size_mode": mode_option, "color_mode": color_option, "render_mode": render_option, "image_kb_mode": image_kb_options, "custom_image_kb": None, + "custom_image_dpi": None, } - # 如果尺寸模式选择的是尺寸列表 + def _process_size_mode( + self, + idphoto_json, + language, + size_list_option, + custom_size_height, + custom_size_width, + ): + """处理尺寸模式""" if idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][0]: idphoto_json["size"] = LOCALES["size_list"][language]["develop"][ size_list_option ] - # 如果尺寸模式选择的是自定义尺寸 elif idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][2]: - id_height = int(custom_size_height) - id_width = int(custom_size_width) + id_height, id_width = int(custom_size_height), int(custom_size_width) if ( id_height < id_width or min(id_height, id_width) < 100 or max(id_height, id_width) > 1800 ): - return [ - gr.update(value=None), # img_output_standard - gr.update(value=None), # img_output_standard_hd - gr.update(value=None), # img_output_standard_png - gr.update(value=None), # img_output_standard_hd_png - None, # img_output_layout (assuming it should be None or not updated) - gr.update( # notification - value=LOCALES["size_mode"][language]["custom_size_eror"], - visible=True, - ), - None, # file_download (assuming it should be None or not updated) - ] - + return self._create_error_response(language) idphoto_json["size"] = (id_height, id_width) else: idphoto_json["size"] = (None, None) - # 如果颜色模式选择的是自定义底色 + def _process_color_mode( + self, + idphoto_json, + language, + color_option, + custom_color_R, + custom_color_G, + custom_color_B, + ): + """处理颜色模式""" if idphoto_json["color_mode"] == LOCALES["bg_color"][language]["choices"][-1]: - idphoto_json["color_bgr"] = ( - range_check(custom_color_R), - range_check(custom_color_G), - range_check(custom_color_B), + idphoto_json["color_bgr"] = tuple( + map(range_check, [custom_color_R, custom_color_G, custom_color_B]) ) else: - hex_color = idphoto_json["color_bgr"] = LOCALES["bg_color"][language][ - "develop" - ][color_option] - # 转为 RGB + hex_color = LOCALES["bg_color"][language]["develop"][color_option] idphoto_json["color_bgr"] = tuple( int(hex_color[i : i + 2], 16) for i in (0, 2, 4) ) - # 如果输出 KB 大小选择的是自定义 - if ( - idphoto_json["image_kb_mode"] - == LOCALES["image_kb"][language]["choices"][-1] - ): - idphoto_json["custom_image_kb"] = custom_image_kb - - creator = IDCreator() - choose_handler(creator, matting_model_option, face_detect_option) - - # 是否只换底 + def _generate_id_photo( + self, + creator, + input_image, + idphoto_json, + language, + head_measure_ratio, + top_distance_max, + top_distance_min, + whitening_strength, + ): + """生成证件照""" change_bg_only = ( idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1] ) + return creator( + input_image, + change_bg_only=change_bg_only, + size=idphoto_json["size"], + head_measure_ratio=head_measure_ratio, + head_top_range=(top_distance_max, top_distance_min), + whitening_strength=whitening_strength, + ) - # 生成证件照 - try: - result = creator( - input_image, - change_bg_only=change_bg_only, - size=idphoto_json["size"], - head_measure_ratio=head_measure_ratio, - head_top_range=(top_distance_max, top_distance_min), - whitening_strength=whitening_strength, + def _handle_photo_generation_error(self, language): + """处理照片生成错误""" + return [gr.update(value=None) for _ in range(4)] + [ + gr.update(visible=False), + gr.update( + value=LOCALES["notification"][language]["face_error"], visible=True + ), + None, + ] + + def _process_generated_photo( + self, + result, + idphoto_json, + language, + watermark_option, + watermark_text, + watermark_text_size, + watermark_text_opacity, + watermark_text_angle, + watermark_text_space, + watermark_text_color, + ): + """处理生成的照片""" + result_image_standard, result_image_hd, _, _, _, _ = result + result_image_standard_png = np.uint8(result_image_standard) + result_image_hd_png = np.uint8(result_image_hd) + + # 渲染背景 + result_image_standard, result_image_hd = self._render_background( + result_image_standard, result_image_hd, idphoto_json + ) + + # 生成排版照片 + result_layout_image = self._generate_layout_image( + idphoto_json, + result_image_standard, + language, + watermark_option, + watermark_text, + watermark_text_size, + watermark_text_opacity, + watermark_text_angle, + watermark_text_space, + watermark_text_color, + ) + + # 添加水印 + if watermark_option == LOCALES["watermark_switch"][language]["choices"][1]: + result_image_standard, result_image_hd = self._add_watermark( + result_image_standard, + result_image_hd, + watermark_text, + watermark_text_size, + watermark_text_opacity, + watermark_text_angle, + watermark_text_space, + watermark_text_color, ) - # 如果检测到人脸数量不等于1 - except FaceError: - return [ - gr.update(value=None), # img_output_standard - gr.update(value=None), # img_output_standard_hd - gr.update(value=None), # img_output_standard_png - gr.update(value=None), # img_output_standard_hd_png - gr.update(visible=False), # img_output_layout - gr.update( # notification - value=LOCALES["notification"][language]["face_error"], - visible=True, - ), - None, # file_download (assuming it should be None or have no update) - ] - # 如果 API 错误 - except APIError as e: - return [ - gr.update(value=None), # img_output_standard - gr.update(value=None), # img_output_standard_hd - gr.update(value=None), # img_output_standard_png - gr.update(value=None), # img_output_standard_hd_png - gr.update(visible=False), # img_output_layout - gr.update( # notification - value=LOCALES["notification"][language]["face_error"], - visible=True, - ), - None, # file_download (assuming it should be None or have no update) - ] - # 证件照生成正常 - else: - (result_image_standard, result_image_hd, _, _, _, _) = result - result_image_standard_png = np.uint8(result_image_standard) - result_image_hd_png = np.uint8(result_image_hd) + # 调整图片大小 + output_image_path = self._resize_image_if_needed( + result_image_standard, idphoto_json + ) + + return self._create_response( + result_image_standard, + result_image_hd, + result_image_standard_png, + result_image_hd_png, + result_layout_image, + output_image_path, + ) - # 纯色渲染 - if ( - idphoto_json["render_mode"] - == LOCALES["render_mode"][language]["choices"][0] - ): - result_image_standard = np.uint8( - add_background(result_image_standard, bgr=idphoto_json["color_bgr"]) - ) - result_image_hd = np.uint8( - add_background(result_image_hd, bgr=idphoto_json["color_bgr"]) - ) - # 上下渐变渲染 - elif ( - idphoto_json["render_mode"] - == LOCALES["render_mode"][language]["choices"][1] - ): - result_image_standard = np.uint8( - add_background( - result_image_standard, - bgr=idphoto_json["color_bgr"], - mode="updown_gradient", - ) - ) - result_image_hd = np.uint8( - add_background( - result_image_hd, - bgr=idphoto_json["color_bgr"], - mode="updown_gradient", - ) - ) - # 中心渐变渲染 - else: - result_image_standard = np.uint8( - add_background( - result_image_standard, - bgr=idphoto_json["color_bgr"], - mode="center_gradient", - ) - ) - result_image_hd = np.uint8( - add_background( - result_image_hd, - bgr=idphoto_json["color_bgr"], - mode="center_gradient", - ) - ) - - # 如果只换底，就不生成排版照 - if change_bg_only: - result_layout_image = gr.update(visible=False) - else: - typography_arr, typography_rotate = generate_layout_photo( - input_height=idphoto_json["size"][0], - input_width=idphoto_json["size"][1], - ) - - if ( - watermark_option - == LOCALES["watermark_switch"][language]["choices"][1] - ): - result_layout_image = gr.update( - value=generate_layout_image( - add_watermark( - image=result_image_standard, - text=watermark_text, - size=watermark_text_size, - opacity=watermark_text_opacity, - angle=watermark_text_angle, - space=watermark_text_space, - color=watermark_text_color, - ), - typography_arr, - typography_rotate, - height=idphoto_json["size"][0], - width=idphoto_json["size"][1], - ), - visible=True, - ) - else: - result_layout_image = gr.update( - value=generate_layout_image( - result_image_standard, - typography_arr, - typography_rotate, - height=idphoto_json["size"][0], - width=idphoto_json["size"][1], - ), - visible=True, - ) - - # 如果添加水印 - if watermark_option == LOCALES["watermark_switch"][language]["choices"][1]: - result_image_standard = add_watermark( - image=result_image_standard, - text=watermark_text, - size=watermark_text_size, - opacity=watermark_text_opacity, - angle=watermark_text_angle, - space=watermark_text_space, - color=watermark_text_color, - ) - result_image_hd = add_watermark( - image=result_image_hd, - text=watermark_text, - size=watermark_text_size, - opacity=watermark_text_opacity, - angle=watermark_text_angle, - space=watermark_text_space, - color=watermark_text_color, - ) - - # 如果输出 KB 大小选择的是自定义 - if idphoto_json["custom_image_kb"]: - print("调整 kb 大小到", idphoto_json["custom_image_kb"], "kb") - output_image_path = f"{os.path.join(os.path.dirname(os.path.dirname(__file__)), 'demo/kb_output')}/{int(time.time())}.jpg" - resize_image_to_kb( - result_image_standard, - output_image_path, - idphoto_json["custom_image_kb"], - ) - else: - output_image_path = None - - # 返回结果 - if output_image_path: - return [ - result_image_standard, # img_output_standard - result_image_hd, # img_output_standard_hd - result_image_standard_png, # img_output_standard_png - result_image_hd_png, # img_output_standard_hd_png - result_layout_image, # img_output_layout - gr.update(visible=False), # notification - gr.update(visible=True, value=output_image_path), # file_download - ] - else: - return [ - result_image_standard, # img_output_standard - result_image_hd, # img_output_standard_hd - result_image_standard_png, # img_output_standard_png - result_image_hd_png, # img_output_standard_hd_png - result_layout_image, # img_output_layout - gr.update(visible=False), # notification - gr.update(visible=False), # file_download - ] + def _render_background(self, result_image_standard, result_image_hd, idphoto_json): + """渲染背景""" + render_modes = {0: "pure_color", 1: "updown_gradient", 2: "center_gradient"} + render_mode = render_modes.get(idphoto_json["render_mode"], "pure_color") + + result_image_standard = np.uint8( + add_background( + result_image_standard, bgr=idphoto_json["color_bgr"], mode=render_mode + ) + ) + result_image_hd = np.uint8( + add_background( + result_image_hd, bgr=idphoto_json["color_bgr"], mode=render_mode + ) + ) + + return result_image_standard, result_image_hd + + def _generate_layout_image( + self, + idphoto_json, + result_image_standard, + language, + watermark_option, + watermark_text, + watermark_text_size, + watermark_text_opacity, + watermark_text_angle, + watermark_text_space, + watermark_text_color, + ): + """生成排版照片""" + if idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]: + return gr.update(visible=False) + + typography_arr, typography_rotate = generate_layout_photo( + input_height=idphoto_json["size"][0], + input_width=idphoto_json["size"][1], + ) + + image = result_image_standard + if watermark_option == LOCALES["watermark_switch"][language]["choices"][1]: + image = add_watermark( + image=image, + text=watermark_text, + size=watermark_text_size, + opacity=watermark_text_opacity, + angle=watermark_text_angle, + space=watermark_text_space, + color=watermark_text_color, + ) + + return gr.update( + value=generate_layout_image( + image, + typography_arr, + typography_rotate, + height=idphoto_json["size"][0], + width=idphoto_json["size"][1], + ), + visible=True, + ) + + def _add_watermark( + self, + result_image_standard, + result_image_hd, + watermark_text, + watermark_text_size, + watermark_text_opacity, + watermark_text_angle, + watermark_text_space, + watermark_text_color, + ): + """添加水印""" + watermark_params = { + "text": watermark_text, + "size": watermark_text_size, + "opacity": watermark_text_opacity, + "angle": watermark_text_angle, + "space": watermark_text_space, + "color": watermark_text_color, + } + result_image_standard = add_watermark( + image=result_image_standard, **watermark_params + ) + result_image_hd = add_watermark(image=result_image_hd, **watermark_params) + return result_image_standard, result_image_hd + + def _resize_image_if_needed(self, result_image_standard, idphoto_json): + """如果需要，调整图片大小""" + output_image_path = f"{os.path.join(os.path.dirname(os.path.dirname(__file__)), 'demo/kb_output')}/{int(time.time())}.jpg" + # 如果设置了自定义KB大小 + if idphoto_json["custom_image_kb"]: + resize_image_to_kb( + result_image_standard, + output_image_path, + idphoto_json["custom_image_kb"], + dpi=idphoto_json["custom_image_dpi"], + ) + return output_image_path + # 如果只设置了dpi + elif idphoto_json["custom_image_dpi"]: + save_image_dpi_to_bytes( + result_image_standard, + output_image_path, + dpi=idphoto_json["custom_image_dpi"], + ) + return output_image_path + + return None + + def _create_response( + self, + result_image_standard, + result_image_hd, + result_image_standard_png, + result_image_hd_png, + result_layout_image, + output_image_path, + ): + """创建响应""" + response = [ + result_image_standard, + result_image_hd, + result_image_standard_png, + result_image_hd_png, + result_layout_image, + gr.update(visible=False), + ] + if output_image_path: + response.append(gr.update(visible=True, value=output_image_path)) + else: + response.append(gr.update(visible=False)) + return response + + def _create_error_response(self, language): + """创建错误响应""" + return [gr.update(value=None) for _ in range(4)] + [ + None, + gr.update( + value=LOCALES["size_mode"][language]["custom_size_eror"], visible=True + ), + None, + ] diff --git a/demo/ui.py b/demo/ui.py index b9b4887b..9294d33a 100644 --- a/demo/ui.py +++ b/demo/ui.py @@ -133,14 +133,28 @@ def create_ui( value=LOCALES["image_kb"][DEFAULT_LANG]["choices"][0], ) - with gr.Row(visible=False) as custom_image_kb: - custom_image_kb_size = gr.Slider( - minimum=10, - maximum=1000, - value=50, - label=LOCALES["image_kb_size"][DEFAULT_LANG]["label"], - interactive=True, - ) + custom_image_kb_size = gr.Slider( + minimum=10, + maximum=1000, + value=50, + label=LOCALES["image_kb_size"][DEFAULT_LANG]["label"], + interactive=True, + visible=False, + ) + + image_dpi_options = gr.Radio( + choices=LOCALES["image_dpi"][DEFAULT_LANG]["choices"], + label=LOCALES["image_dpi"][DEFAULT_LANG]["label"], + value=LOCALES["image_dpi"][DEFAULT_LANG]["choices"][0], + ) + custom_image_dpi_size = gr.Slider( + minimum=72, + maximum=600, + value=300, + label=LOCALES["image_dpi_size"][DEFAULT_LANG]["label"], + interactive=True, + visible=False, + ) # TAB3 - 美颜 with gr.Tab( @@ -409,26 +423,30 @@ def change_language(language): whitening_option: gr.update( label=LOCALES["whitening_strength"][language]["label"] ), + image_dpi_options: gr.update( + label=LOCALES["image_dpi"][language]["label"], + choices=LOCALES["image_dpi"][language]["choices"], + value=LOCALES["image_dpi"][language]["choices"][0], + ), + custom_image_dpi_size: gr.update( + label=LOCALES["image_dpi"][language]["label"] + ), } def change_color(colors): - if colors == "自定义底色" or colors == "Custom Color": + if colors == LOCALES["bg_color"][DEFAULT_LANG]["choices"][1]: return {custom_color: gr.update(visible=True)} else: return {custom_color: gr.update(visible=False)} def change_size_mode(size_option_item): - if ( - size_option_item == "自定义尺寸" - or size_option_item == "Custom Size" - ): + if size_option_item == LOCALES["size_mode"][DEFAULT_LANG]["choices"][2]: return { custom_size: gr.update(visible=True), size_list_row: gr.update(visible=False), } elif ( - size_option_item == "只换底" - or size_option_item == "Only Change Background" + size_option_item == LOCALES["size_mode"][DEFAULT_LANG]["choices"][1] ): return { custom_size: gr.update(visible=False), @@ -441,10 +459,16 @@ def change_size_mode(size_option_item): } def change_image_kb(image_kb_option): - if image_kb_option == "自定义" or image_kb_option == "Custom": - return {custom_image_kb: gr.update(visible=True)} + if image_kb_option == LOCALES["image_kb"][DEFAULT_LANG]["choices"][1]: + return {custom_image_kb_size: gr.update(visible=True)} else: - return {custom_image_kb: gr.update(visible=False)} + return {custom_image_kb_size: gr.update(visible=False)} + + def change_image_dpi(image_dpi_option): + if image_dpi_option == LOCALES["image_dpi"][DEFAULT_LANG]["choices"][1]: + return {custom_image_dpi_size: gr.update(visible=True)} + else: + return {custom_image_dpi_size: gr.update(visible=False)} # ---------------- 绑定事件 ---------------- # 语言切换 @@ -497,7 +521,15 @@ def change_image_kb(image_kb_option): ) image_kb_options.input( - change_image_kb, inputs=[image_kb_options], outputs=[custom_image_kb] + change_image_kb, + inputs=[image_kb_options], + outputs=[custom_image_kb_size], + ) + + image_dpi_options.input( + change_image_dpi, + inputs=[image_dpi_options], + outputs=[custom_image_dpi_size], ) img_but.click( @@ -528,6 +560,8 @@ def change_image_kb(image_kb_option): head_measure_ratio_option, top_distance_option, whitening_option, + image_dpi_options, + custom_image_dpi_size, ], outputs=[ img_output_standard, diff --git a/hivision/utils.py b/hivision/utils.py index da8452ee..c087b388 100644 --- a/hivision/utils.py +++ b/hivision/utils.py @@ -1,12 +1,5 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -r""" -@DATE: 2024/9/5 21:52 -@File: utils.py -@IDE: pycharm -@Description: - hivision提供的工具函数 -""" from PIL import Image import io import numpy as np @@ -15,7 +8,30 @@ from hivision.plugin.watermark import Watermarker, WatermarkerStyles -def resize_image_to_kb(input_image, output_image_path, target_size_kb): +def save_image_dpi_to_bytes(image, output_image_path, dpi=300): + """ + 设置图像的DPI（每英寸点数）并返回字节流 + + :param image: numpy.ndarray, 输入的图像数组 + :param output_image_path: Path to save the resized image. 保存调整大小后的图像的路径。 + :param dpi: int, 要设置的DPI值，默认为300 + """ + image = Image.fromarray(image) + # 创建一个字节流对象 + byte_stream = io.BytesIO() + # 将图像保存到字节流 + image.save(byte_stream, format="PNG", dpi=(dpi, dpi)) + # 获取字节流的内容 + image_bytes = byte_stream.getvalue() + + # Save the image to the output path + with open(output_image_path, "wb") as f: + f.write(image_bytes) + + return image_bytes + + +def resize_image_to_kb(input_image, output_image_path, target_size_kb, dpi=300): """ Resize an image to a target size in KB. 将图像调整大小至目标文件大小（KB）。 @@ -47,7 +63,7 @@ def resize_image_to_kb(input_image, output_image_path, target_size_kb): img_byte_arr = io.BytesIO() # Save the image to the BytesIO object with the current quality - img.save(img_byte_arr, format="JPEG", quality=quality) + img.save(img_byte_arr, format="JPEG", quality=quality, dpi=(dpi, dpi)) # Get the size of the image in KB img_size_kb = len(img_byte_arr.getvalue()) / 1024