update kb resize

.gitignore

@@ -0,0 +1,6 @@
+*.pyc
+**/__pycache__
+.idea
+.DS_Store
+hivision_modnet.onnx
+output/*.jpg

app.py

@@ -5,10 +5,20 @@ from hivisionai.hycv.vision import add_background
 from src.layoutCreate import generate_layout_photo, generate_layout_image
 import pathlib
 import numpy as np
-
-size_list_dict = {"一寸": (413, 295), "二寸": (626, 413),
-                  "教师资格证": (413, 295), "国家公务员考试": (413, 295), "初级会计考试": (413, 295)}
-color_list_dict = {"蓝色": (86, 140, 212), "白色": (255, 255, 255), "红色": (233, 51, 35)}
+from image_utils import resize_image_to_kb
+
+size_list_dict = {
+    "一寸": (413, 295),
+    "二寸": (626, 413),
+    "教师资格证": (413, 295),
+    "国家公务员考试": (413, 295),
+    "初级会计考试": (413, 295),
+}
+color_list_dict = {
+    "蓝色": (86, 140, 212),
+    "白色": (255, 255, 255),
+    "红色": (233, 51, 35),
+}
 
 
 # 设置Gradio examples
@@ -26,25 +36,30 @@ def range_check(value, min_value=0, max_value=255):
     return value
 
 
-def idphoto_inference(input_image,
-                      mode_option,
-                      size_list_option,
-                      color_option,
-                      render_option,
-                      custom_color_R,
-                      custom_color_G,
-                      custom_color_B,
-                      custom_size_height,
-                      custom_size_width,
-                      head_measure_ratio=0.2,
-                      head_height_ratio=0.45,
-                      top_distance_max=0.12,
-                      top_distance_min=0.10):
+def idphoto_inference(
+    input_image,
+    mode_option,
+    size_list_option,
+    color_option,
+    render_option,
+    image_kb_options,
+    custom_color_R,
+    custom_color_G,
+    custom_color_B,
+    custom_size_height,
+    custom_size_width,
+    custom_image_kb,
+    head_measure_ratio=0.2,
+    head_height_ratio=0.45,
+    top_distance_max=0.12,
+    top_distance_min=0.10,
+):
 
     idphoto_json = {
         "size_mode": mode_option,
         "color_mode": color_option,
         "render_mode": render_option,
+        "image_kb_mode": image_kb_options,
     }
 
     # 如果尺寸模式选择的是尺寸列表
@@ -54,97 +69,182 @@ def idphoto_inference(input_image,
     elif idphoto_json["size_mode"] == "自定义尺寸":
         id_height = int(custom_size_height)
         id_width = int(custom_size_width)
-        if id_height < id_width or min(id_height, id_width) < 100 or max(id_height, id_width) > 1800:
+        if (
+            id_height < id_width
+            or min(id_height, id_width) < 100
+            or max(id_height, id_width) > 1800
+        ):
             return {
                 img_output_standard: gr.update(value=None),
                 img_output_standard_hd: gr.update(value=None),
-                notification: gr.update(value="宽度应不大于长度；长宽不应小于100，大于1800", visible=True)}
+                notification: gr.update(
+                    value="宽度应不大于长度；长宽不应小于100，大于1800", visible=True
+                ),
+            }
         idphoto_json["size"] = (id_height, id_width)
     else:
         idphoto_json["size"] = (None, None)
 
     # 如果颜色模式选择的是自定义底色
     if idphoto_json["color_mode"] == "自定义底色":
-        idphoto_json["color_bgr"] = (range_check(custom_color_R),
-                                     range_check(custom_color_G),
-                                     range_check(custom_color_B))
+        idphoto_json["color_bgr"] = (
+            range_check(custom_color_R),
+            range_check(custom_color_G),
+            range_check(custom_color_B),
+        )
     else:
         idphoto_json["color_bgr"] = color_list_dict[color_option]
 
-    result_image_hd, result_image_standard, typography_arr, typography_rotate, \
-    _, _, _, _, status = IDphotos_create(input_image,
-                                         mode=idphoto_json["size_mode"],
-                                         size=idphoto_json["size"],
-                                         head_measure_ratio=head_measure_ratio,
-                                         head_height_ratio=head_height_ratio,
-                                         align=False,
-                                         beauty=False,
-                                         fd68=None,
-                                         human_sess=sess,
-                                         IS_DEBUG=False,
-                                         top_distance_max=top_distance_max,
-                                         top_distance_min=top_distance_min)
+    # 如果输出KB大小选择的是自定义
+    if idphoto_json["image_kb_mode"] == "自定义":
+        idphoto_json["custom_image_kb"] = custom_image_kb
+    else:
+        idphoto_json["custom_image_kb"] = None
+
+    # 生成证件照
+    (
+        result_image_hd,
+        result_image_standard,
+        typography_arr,
+        typography_rotate,
+        _,
+        _,
+        _,
+        _,
+        status,
+    ) = IDphotos_create(
+        input_image,
+        mode=idphoto_json["size_mode"],
+        size=idphoto_json["size"],
+        head_measure_ratio=head_measure_ratio,
+        head_height_ratio=head_height_ratio,
+        align=False,
+        beauty=False,
+        fd68=None,
+        human_sess=sess,
+        IS_DEBUG=False,
+        top_distance_max=top_distance_max,
+        top_distance_min=top_distance_min,
+    )
 
     # 如果检测到人脸数量不等于1
     if status == 0:
         result_messgae = {
             img_output_standard: gr.update(value=None),
             img_output_standard_hd: gr.update(value=None),
-            notification: gr.update(value="人脸数量不等于1", visible=True)
+            notification: gr.update(value="人脸数量不等于1", visible=True),
         }
 
     # 如果检测到人脸数量等于1
     else:
         if idphoto_json["render_mode"] == "纯色":
             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"]))
+                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"] == "上下渐变(白)":
             result_image_standard = np.uint8(
-                add_background(result_image_standard, bgr=idphoto_json["color_bgr"], mode="updown_gradient"))
+                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"))
+                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"))
+                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"))
+                add_background(
+                    result_image_hd,
+                    bgr=idphoto_json["color_bgr"],
+                    mode="center_gradient",
+                )
+            )
 
         if idphoto_json["size_mode"] == "只换底":
             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])
-
-            result_layout_image = generate_layout_image(result_image_standard, typography_arr,
-                                                        typography_rotate,
-                                                        height=idphoto_json["size"][0],
-                                                        width=idphoto_json["size"][1])
-
-        result_messgae = {
-            img_output_standard: result_image_standard,
-            img_output_standard_hd: result_image_hd,
-            img_output_layout: result_layout_image,
-            notification: gr.update(visible=False)}
+            typography_arr, typography_rotate = generate_layout_photo(
+                input_height=idphoto_json["size"][0],
+                input_width=idphoto_json["size"][1],
+            )
+
+            result_layout_image = generate_layout_image(
+                result_image_standard,
+                typography_arr,
+                typography_rotate,
+                height=idphoto_json["size"][0],
+                width=idphoto_json["size"][1],
+            )
+
+        # 如果输出KB大小选择的是自定义
+        if idphoto_json["custom_image_kb"]:
+            # 将标准照大小调整至目标大小
+            print("调整kb大小到", idphoto_json["custom_image_kb"], "kb")
+            # 输出路径为一个根据时间戳+哈希值生成的随机文件名
+            import time
+
+            output_image_path = f"./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:
+            result_messgae = {
+                img_output_standard: result_image_standard,
+                img_output_standard_hd: result_image_hd,
+                img_output_layout: result_layout_image,
+                notification: gr.update(visible=False),
+                file_download: gr.update(visible=True, value=output_image_path),
+            }
+        else:
+            result_messgae = {
+                img_output_standard: result_image_standard,
+                img_output_standard_hd: result_image_hd,
+                img_output_layout: result_layout_image,
+                notification: gr.update(visible=False),
+                file_download: gr.update(visible=False),
+            }
 
     return result_messgae
 
 
 if __name__ == "__main__":
+    # 预加载ONNX模型
     HY_HUMAN_MATTING_WEIGHTS_PATH = "./hivision_modnet.onnx"
     sess = onnxruntime.InferenceSession(HY_HUMAN_MATTING_WEIGHTS_PATH)
+
     size_mode = ["尺寸列表", "只换底", "自定义尺寸"]
     size_list = ["一寸", "二寸", "教师资格证", "国家公务员考试", "初级会计考试"]
     colors = ["蓝色", "白色", "红色", "自定义底色"]
     render = ["纯色", "上下渐变(白)", "中心渐变(白)"]
+    image_kb = ["不设置", "自定义"]
 
     title = "<h1 id='title'>HivisionIDPhotos</h1>"
-    description = "<h3>😎6.20更新：新增尺寸选择列表</h3>"
-    css = '''
+    description = "<h3>😎9.2更新：新增照片大小KB调整</h3>"
+    css = """
     h1#title, h3 {
       text-align: center;
     }
-    '''
+    """
 
     demo = gr.Blocks(css=css)
 
@@ -152,39 +252,90 @@ if __name__ == "__main__":
         gr.Markdown(title)
         gr.Markdown(description)
         with gr.Row():
+            # ------------ 左半边UI ----------------
             with gr.Column():
                 img_input = gr.Image().style(height=350)
-                mode_options = gr.Radio(choices=size_mode, label="证件照尺寸选项", value="尺寸列表", elem_id="size")
+                mode_options = gr.Radio(
+                    choices=size_mode,
+                    label="证件照尺寸选项",
+                    value="尺寸列表",
+                    elem_id="size",
+                )
                 # 预设尺寸下拉菜单
                 with gr.Row(visible=True) as size_list_row:
-                    size_list_options = gr.Dropdown(choices=size_list, label="预设尺寸", value="一寸", elem_id="size_list")
+                    size_list_options = gr.Dropdown(
+                        choices=size_list,
+                        label="预设尺寸",
+                        value="一寸",
+                        elem_id="size_list",
+                    )
 
                 with gr.Row(visible=False) as custom_size:
-                    custom_size_height = gr.Number(value=413, label="height", interactive=True)
-                    custom_size_wdith = gr.Number(value=295, label="width", interactive=True)
-
-                color_options = gr.Radio(choices=colors, label="背景色", value="蓝色", elem_id="color")
+                    custom_size_height = gr.Number(
+                        value=413, label="height", interactive=True
+                    )
+                    custom_size_wdith = gr.Number(
+                        value=295, label="width", interactive=True
+                    )
+
+                # 左：背景色选项
+                color_options = gr.Radio(
+                    choices=colors, label="背景色", value="蓝色", elem_id="color"
+                )
+
+                # 左：如果选择「自定义底色」，显示RGB输入框
                 with gr.Row(visible=False) as custom_color:
                     custom_color_R = gr.Number(value=0, label="R", interactive=True)
                     custom_color_G = gr.Number(value=0, label="G", interactive=True)
                     custom_color_B = gr.Number(value=0, label="B", interactive=True)
 
-                render_options = gr.Radio(choices=render, label="渲染方式", value="纯色", elem_id="render")
+                # 左：渲染方式选项
+                render_options = gr.Radio(
+                    choices=render,
+                    label="渲染方式",
+                    value="纯色",
+                    elem_id="render",
+                )
+
+                # 左：输出KB大小选项
+                image_kb_options = gr.Radio(
+                    choices=image_kb,
+                    label="设置KB大小（结果在右边最底的组件下载）",
+                    value="不设置",
+                    elem_id="image_kb",
+                )
+
+                # 自定义KB大小, 滑动条，最小10KB，最大200KB
+                with gr.Row(visible=False) as custom_image_kb:
+                    custom_image_kb_size = gr.Slider(
+                        minimum=10,
+                        maximum=1000,
+                        value=50,
+                        label="KB大小",
+                        interactive=True,
+                    )
+
+                img_but = gr.Button("开始制作")
 
-                img_but = gr.Button('开始制作')
                 # 案例图片
-                example_images = gr.Dataset(components=[img_input],
-                                            samples=[[path.as_posix()]
-                                                     for path in sorted(pathlib.Path('images').rglob('*.jpg'))])
-
+                example_images = gr.Dataset(
+                    components=[img_input],
+                    samples=[
+                        [path.as_posix()]
+                        for path in sorted(pathlib.Path("images").rglob("*.jpg"))
+                    ],
+                )
+
+            # ---------------- 右半边UI ----------------
             with gr.Column():
                 notification = gr.Text(label="状态", visible=False)
                 with gr.Row():
                     img_output_standard = gr.Image(label="标准照").style(height=350)
                     img_output_standard_hd = gr.Image(label="高清照").style(height=350)
                 img_output_layout = gr.Image(label="六寸排版照").style(height=350)
+                file_download = gr.File(label="下载调整KB大小后的照片", visible=False)
 
-
+            # ---------------- 设置隐藏/显示组件 ----------------
             def change_color(colors):
                 if colors == "自定义底色":
                     return {custom_color: gr.update(visible=True)}
@@ -193,23 +344,69 @@ if __name__ == "__main__":
 
             def change_size_mode(size_option_item):
                 if size_option_item == "自定义尺寸":
-                    return {custom_size: gr.update(visible=True),
-                            size_list_row: gr.update(visible=False)}
+                    return {
+                        custom_size: gr.update(visible=True),
+                        size_list_row: gr.update(visible=False),
+                    }
                 elif size_option_item == "只换底":
-                    return {custom_size: gr.update(visible=False),
-                            size_list_row: gr.update(visible=False)}
+                    return {
+                        custom_size: gr.update(visible=False),
+                        size_list_row: gr.update(visible=False),
+                    }
+                else:
+                    return {
+                        custom_size: gr.update(visible=False),
+                        size_list_row: gr.update(visible=True),
+                    }
+
+            def change_image_kb(image_kb_option):
+                if image_kb_option == "自定义":
+                    return {custom_image_kb: gr.update(visible=True)}
                 else:
-                    return {custom_size: gr.update(visible=False),
-                            size_list_row: gr.update(visible=True)}
-
-        color_options.input(change_color, inputs=[color_options], outputs=[custom_color])
-        mode_options.input(change_size_mode, inputs=[mode_options], outputs=[custom_size, size_list_row])
-
-        img_but.click(idphoto_inference,
-                      inputs=[img_input, mode_options, size_list_options, color_options, render_options,
-                              custom_color_R, custom_color_G, custom_color_B,
-                              custom_size_height, custom_size_wdith],
-                      outputs=[img_output_standard, img_output_standard_hd, img_output_layout, notification])
-        example_images.click(fn=set_example_image, inputs=[example_images], outputs=[img_input])
+                    return {custom_image_kb: gr.update(visible=False)}
+
+        # ---------------- 绑定事件 ----------------
+        color_options.input(
+            change_color, inputs=[color_options], outputs=[custom_color]
+        )
+
+        mode_options.input(
+            change_size_mode,
+            inputs=[mode_options],
+            outputs=[custom_size, size_list_row],
+        )
+
+        image_kb_options.input(
+            change_image_kb, inputs=[image_kb_options], outputs=[custom_image_kb]
+        )
+
+        img_but.click(
+            idphoto_inference,
+            inputs=[
+                img_input,
+                mode_options,
+                size_list_options,
+                color_options,
+                render_options,
+                image_kb_options,
+                custom_color_R,
+                custom_color_G,
+                custom_color_B,
+                custom_size_height,
+                custom_size_wdith,
+                custom_image_kb_size,
+            ],
+            outputs=[
+                img_output_standard,
+                img_output_standard_hd,
+                img_output_layout,
+                notification,
+                file_download,
+            ],
+        )
+
+        example_images.click(
+            fn=set_example_image, inputs=[example_images], outputs=[img_input]
+        )
 
     demo.launch()

image_utils.py

@@ -0,0 +1,58 @@
+from PIL import Image
+import io
+import numpy as np
+
+
+def resize_image_to_kb(input_image: np.ndarray, output_image_path, target_size_kb):
+    """
+    Resize an image to a target size in KB.
+    将图像调整大小至目标文件大小（KB）。
+
+    :param input_image_path: Path to the input image. 输入图像的路径。
+    :param output_image_path: Path to save the resized image. 保存调整大小后的图像的路径。
+    :param target_size_kb: Target size in KB. 目标文件大小（KB）。
+
+    Example:
+    resize_image_to_kb('input_image.jpg', 'output_image.jpg', 50)
+    """
+
+    # Open an image file
+    with Image.fromarray(input_image) as img:
+        # Convert image to RGB mode if it's not
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+
+        # Initial quality value
+        quality = 95
+
+        while True:
+            # Create a BytesIO object to hold the image data in memory
+            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)
+
+            # Get the size of the image in KB
+            img_size_kb = len(img_byte_arr.getvalue()) / 1024
+
+            # Check if the image size is within the target size
+            if img_size_kb <= target_size_kb:
+                # If the image is smaller than the target size, add padding
+                if img_size_kb < target_size_kb:
+                    padding_size = int(
+                        (target_size_kb * 1024) - len(img_byte_arr.getvalue())
+                    )
+                    padding = b"\x00" * padding_size
+                    img_byte_arr.write(padding)
+
+                # Save the image to the output path
+                with open(output_image_path, "wb") as f:
+                    f.write(img_byte_arr.getvalue())
+                break
+
+            # Reduce the quality if the image is still too large
+            quality -= 5
+
+            # Ensure quality does not go below 1
+            if quality < 1:
+                quality = 1