From f66677f5b80fa11575c06630dbd3c0589edd510e Mon Sep 17 00:00:00 2001
From: Ze-Yi LIN <58305964+Zeyi-Lin@users.noreply.github.com>
Date: Thu, 12 Sep 2024 00:59:17 +0800
Subject: [PATCH] feat: api set kb (#106)

---
 deploy_api.py     | 27 +++++++++++++++++++++++++++
 docs/api_CN.md    | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 hivision/utils.py | 15 ++++++++-------
 3 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/deploy_api.py b/deploy_api.py
index 1b512023..43904f29 100644
--- a/deploy_api.py
+++ b/deploy_api.py
@@ -221,6 +221,33 @@ async def watermark(
     return result_messgae
 
 
+# 设置照片KB值接口(RGB图)
+@app.post("/set_kb")
+async def set_kb(
+    input_image: UploadFile,
+    kb: int = Form(50),
+):
+    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_base64 = resize_image_to_kb_base64(result_image, int(kb))
+
+        result_messgae = {
+            "status": True,
+            "image_base64": result_image_base64,
+        }
+    except Exception as e:
+        result_messgae = {
+            "status": False,
+            "error": e,
+        }
+
+    return result_messgae
+
+
 if __name__ == "__main__":
     import uvicorn
 
diff --git a/docs/api_CN.md b/docs/api_CN.md
index 0230e01d..8d318ed3 100644
--- a/docs/api_CN.md
+++ b/docs/api_CN.md
@@ -56,6 +56,19 @@ python deploy_api.py
 
 `人像抠图`接口的逻辑是发送一张 RGB 图像，输出一张标准抠图人像照和高清抠图人像照（无任何背景填充）。
 
+### 5.图像加水印
+
+接口名：`watermark`
+
+`图像加水印`接口的功能是接收一个水印文本，然后在原图上添加指定的水印。用户可以指定水印的位置、透明度和大小等属性，以便将水印无缝地融合到原图中。
+
+
+### 6.设置图像KB大小
+
+接口名：`set_kb`
+
+`设置图像KB大小`接口的功能是接收一张图像和目标文件大小（以KB为单位），如果设置的KB值小于原文件，则调整压缩率；如果设置的KB值大于源文件，则通过给文件头添加信息的方式调大KB值，目标是让图像的最终大小与设置的KB值一致。
+
 <br>
 
 ## cURL 请求示例
@@ -111,6 +124,16 @@ curl -X 'POST' \
   -F 'text=Hello'
 ```
 
+### 6. 设置图像KB大小
+```bash
+curl -X 'POST' \
+  'http://127.0.0.1:8080/set_kb' \
+  -H 'accept: application/json' \
+  -H 'Content-Type: multipart/form-data' \
+  -F 'input_image=@demo/images/test0.jpg;type=image/jpeg' \
+  -F 'kb=50'
+```
+
 
 <br>
 
@@ -211,6 +234,30 @@ else:
     print(f"Request failed with status code {response.status_code}: {response.text}")
 ```
 
+### 6. 设置图像KB大小
+```bash
+import requests
+
+# 设置请求的 URL
+url = "http://127.0.0.1:8080/set_kb"
+
+# 设置文件和其他表单数据
+input_image_path = "demo/images/test0.jpg"
+files = {"input_image": open(input_image_path, "rb")}
+data = {"kb": 50}
+
+# 发送 POST 请求
+response = requests.post(url, files=files, data=data)
+
+# 检查响应
+if response.ok:
+    # 输出响应内容
+    print(response.json())
+else:
+    # 输出错误信息
+    print(f"Request failed with status code {response.status_code}: {response.text}")
+```
+
 <br>
 
 ## Java 请求示例
diff --git a/hivision/utils.py b/hivision/utils.py
index 6bf0465e..afef56c0 100644
--- a/hivision/utils.py
+++ b/hivision/utils.py
@@ -75,12 +75,6 @@ def resize_image_to_kb(input_image, output_image_path, target_size_kb):
             quality = 1
 
 
-import numpy as np
-from PIL import Image
-import io
-import base64
-
-
 def resize_image_to_kb_base64(input_image, target_size_kb, mode="exact"):
     """
     Resize an image to a target size in KB and return it as a base64 encoded string.
@@ -154,13 +148,20 @@ def resize_image_to_kb_base64(input_image, target_size_kb, mode="exact"):
     return img_base64
 
 
-def numpy_2_base64(img: np.ndarray):
+def numpy_2_base64(img: np.ndarray) -> str:
     _, buffer = cv2.imencode(".png", img)
     base64_image = base64.b64encode(buffer).decode("utf-8")
 
     return base64_image
 
 
+def base64_2_numpy(base64_image: str) -> np.ndarray:
+    img = base64.b64decode(base64_image)
+    img = np.frombuffer(img, np.uint8)
+
+    return img
+
+
 def save_numpy_image(numpy_img, file_path):
     # 检查数组的形状
     if numpy_img.shape[2] == 4: