From 22a5eb2b1b38bf98a30a1ecc201a98bfe26eced6 Mon Sep 17 00:00:00 2001
From: LiaoFJ <50604988+LiaoFJ@users.noreply.github.com>
Date: Tue, 23 Apr 2024 16:19:37 +0800
Subject: [PATCH] add function "adjust_contrast" in the function callback for
 better image displaying.

---
 app/service/generate_image/service.py         |  3 +-
 .../generate_image/utils/adjust_contrast.py   | 31 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 app/service/generate_image/utils/adjust_contrast.py

diff --git a/app/service/generate_image/service.py b/app/service/generate_image/service.py
index d0bda9a..45346bf 100644
--- a/app/service/generate_image/service.py
+++ b/app/service/generate_image/service.py
@@ -23,7 +23,7 @@ from tritonclient.utils import np_to_triton_dtype
 from app.core.config import *
 from app.schemas.generate_image import GenerateImageModel
 from app.service.generate_image.utils.upload_sd_image import upload_png_sd
-
+from app.service.generate_image.utils.adjust_contrast import adjust_contrast
 logger = logging.getLogger()
 
 
@@ -71,6 +71,7 @@ class GenerateImage:
             self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
         else:
             image_result = result.as_numpy("generated_image")[0]
+            image_result = adjust_contrast(image_result)
             image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")
             self.generate_data['status'] = "SUCCESS"
             self.generate_data['message'] = "success"
diff --git a/app/service/generate_image/utils/adjust_contrast.py b/app/service/generate_image/utils/adjust_contrast.py
new file mode 100644
index 0000000..4ed110e
--- /dev/null
+++ b/app/service/generate_image/utils/adjust_contrast.py
@@ -0,0 +1,31 @@
+import cv2
+
+def adjust_contrast(image, alpha=1.5, beta=-60):
+    """
+    调整图像的对比度和亮度。
+
+    参数:
+    image_path (numpy): 图像的路径。
+    alpha (float): 控制对比度的系数。alpha > 1 增加对比度，alpha < 1 减少对比度。
+    beta (int): 用于调整亮度的值，可以是正或负。
+
+    返回:
+    adjusted_image (ndarray): 调整对比度后的图像。
+    """
+
+    adjusted_image = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
+    return adjusted_image
+
+
+# 使用示例
+if __name__ == "__main__":
+    image = cv2.imread('output_6.png') # 替换为你的图片路径
+    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    alpha = 1.5 # 对比度系数，大于1增加对比度
+    beta = -60  # 亮度调整，这里设置为0，不改变亮度
+
+    # 调整图像对比度
+    result_image = adjust_contrast(image, alpha, beta)
+        # 可以选择保存调整后的图像
+    cv2.imwrite('adjusted_image.jpg', result_image)  # 保存调整后的图片
\ No newline at end of file