feat : 代码梳理 移除所有敏感密钥 通过环境变量方式配置

app/service/generate_image/service_generate_image.py

@@ -16,16 +16,18 @@ import minio
 import numpy as np
 import redis
 import tritonclient.grpc as grpcclient
+from minio import Minio
 from tritonclient.utils import np_to_triton_dtype
 
-from app.core.config import *
+from app.core.config import settings, FAST_GI_MODEL_URL, GI_MODEL_URL, FAST_GI_MODEL_NAME, GI_MODEL_NAME, GI_RABBITMQ_QUEUES
 from app.schemas.generate_image import GenerateImageModel
 from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust
 from app.service.generate_image.utils.mq import publish_status
 from app.service.generate_image.utils.upload_sd_image import upload_png_sd
-from app.service.utils.oss_client import oss_get_image
+from app.service.utils.new_oss_client import oss_get_image
 
 logger = logging.getLogger()
+minio_client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ACCESS, secret_key=settings.MINIO_SECRET, secure=settings.MINIO_SECURE)
 
 
 class GenerateImage:
@@ -36,7 +38,7 @@ class GenerateImage:
         else:
             self.grpc_client = grpcclient.InferenceServerClient(url=GI_MODEL_URL)
 
-        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+        self.redis_client = redis.StrictRedis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=settings.REDIS_DB, decode_responses=True)
         if request_data.mode == "img2img":
             # cv2 读图片是BGR PIL读图片是RGB
             self.image = self.get_image(request_data.image_url)
@@ -67,8 +69,7 @@ class GenerateImage:
             # image_array = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
             # image_cv2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
             # image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
-
-            image_cv2 = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="cv2")
+            image_cv2 = oss_get_image(oss_client=minio_client, bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="cv2")
             image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
             image = cv2.resize(image_rbg, (1024, 1024))
         except minio.error.S3Error:
@@ -120,7 +121,7 @@ class GenerateImage:
             else:  # 有污点 保存图片到本地 测试用
                 self.generate_data['status'] = "SUCCESS"
                 self.generate_data['message'] = "success"
-                self.generate_data['image_url'] = str(GI_SYS_IMAGE_URL)
+                self.generate_data['image_url'] = "aida-sys-image/generate_image/white_image.jpg"
                 self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
             # logger.info(f"stain_detection result : {self.generate_data}")
 
@@ -171,12 +172,12 @@ class GenerateImage:
             raise Exception(str(e))
         finally:
             dict_generate_data, str_generate_data = self.read_tasks_status()
-            if not DEBUG:
+            if not settings.DEBUG:
                 publish_status(str_generate_data, GI_RABBITMQ_QUEUES)
 
 
 def infer_cancel(tasks_id):
-    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+    redis_client = redis.StrictRedis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=settings.REDIS_DB, decode_responses=True)
     data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
     generate_data = json.dumps(data)
     redis_client.set(tasks_id, generate_data)
@@ -186,12 +187,12 @@ def infer_cancel(tasks_id):
 if __name__ == '__main__':
     rd = GenerateImageModel(
         tasks_id="123-89",
-        prompt="Women's clothing ,dress,technical drawing style, clean line art, no shading, no texture, flat sketch, no human body, no face, centered composition, pure white background, single garmentsingle garment only, front flat view",
-        image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg",
-        mode='txt2img',
-        category="test",
-        gender="male",
-        version="high"
+        prompt="a single item of sketch of dress, 4k, white background",
+        image_url="aida-collection-element/89/Sketchboard/95f20cdc-e059-435c-b8b1-d04cc9e80c3d.png",
+        mode='img2img',
+        category="sketch",
+        gender="Female",
+        version="fast"
     )
     server = GenerateImage(rd)
     print(server.get_result())