Merge branch 'develop'

app/api/api_generate_image.py

@@ -10,6 +10,7 @@ logger = logging.getLogger()
 @router.post("/generate_image")
 def generate_image(request_item: GenerateImageModel, background_tasks: BackgroundTasks):
     try:
+        logger.info(f"request data ### : {request_item}")
         service = GenerateImage(request_item)
         background_tasks.add_task(service.get_result)
         code = 200

app/core/config.py

@@ -23,9 +23,11 @@ DEBUG = False
 if DEBUG:
     LOGS_PATH = "logs/"
     CATEGORY_PATH = "service/attribute/config/descriptor/category/category_dis.csv"
+    FACE_CLASSIFIER = "service/generate_image/utils/haarcascade_frontalface_alt.xml"
 else:
     LOGS_PATH = "app/logs/"
     CATEGORY_PATH = "app/service/attribute/config/descriptor/category/category_dis.csv"
+    FACE_CLASSIFIER = 'app/service/generate_image/utils/haarcascade_frontalface_alt.xml'
 
 RABBITMQ_ENV = ""  # 生产环境
 # RABBITMQ_ENV = "-dev"  # 开发环境
@@ -52,6 +54,13 @@ RABBITMQ_PARAMS = {
     "virtual_host": "/"
 }
 
+# milvus 配置
+MILVUS_DB_HOST = "10.1.1.240"
+MILVUS_ALIAS = "default"
+MILVUS_PORT = "19530"
+MILVUS_TABLE_KEYPOINT = "keypoint_cache"
+MILVUS_TABLE_SEG = "seg_cache"
+
 # attribute service config
 ATT_TRITON_URL = "10.1.1.240:8020"
 
@@ -62,15 +71,50 @@ SR_MINIO_BUCKET = "aida-users"
 SR_RABBITMQ_QUEUES = os.getenv("SR_RABBITMQ_QUEUES", f"SuperResolution{RABBITMQ_ENV}")
 
 # GenerateImage service config
-GI_MODEL_NAME = 'stable_diffusion_xl_lcm'
+GI_MODEL_NAME = 'stable_diffusion_xl'
-GI_MODEL_URL = '10.1.1.150:8001'
+GI_MODEL_URL = '10.1.1.240:10041'
 GI_MINIO_BUCKET = "aida-users"
 GI_RABBITMQ_QUEUES = os.getenv("GI_RABBITMQ_QUEUES", f"GenerateImage{RABBITMQ_ENV}")
+GI_SYS_IMAGE_URL = "aida-sys-image/generate_image/white_image.jpg"
 
 # SEG service config
 SEG_MODEL_URL = '10.1.1.240:10000'
 SEGMENTATION = {
+    "new_model_name": "seg_knet",
     "name": "seg_ocrnet_hr18",
     "input": "seg_input__0",
     "output": "seg_output__0",
 }
+
+# DESIGN config
+DESIGN_MODEL_URL = '10.1.1.240:9000'
+
+AIDA_CLOTHING = "aida-clothing"
+
+# 优先级
+PRIORITY_DICT = {
+    'earring_front': 99,
+    'bag_front': 98,
+    'hairstyle_front': 97,
+    'outwear_front': 20,
+    'tops_front': 19,
+    'dress_front': 18,
+    'blouse_front': 17,
+    'skirt_front': 16,
+    'trousers_front': 15,
+    'bottoms_front': 14,
+    'shoes_right': 1,
+    'shoes_left': 1,
+    'body': 0,
+    'bottoms_back': -14,
+    'trousers_back': -15,
+    'skirt_back': -16,
+    'blouse_back': -17,
+    'dress_back': -18,
+    'tops_back': -19,
+    'outwear_back': -20,
+    'hairstyle_back': -97,
+    'bag_back': -98,
+    'earring_back': -99,
+}
+

app/schemas/generate_image.py

@@ -7,3 +7,4 @@ class GenerateImageModel(BaseModel):
     image_url: str
     mode: str
     category: str
+    gender: str

app/service/attribute/config/descriptor/category/category_dis.csv

@@ -1,7 +1,7 @@
 labelName,join_attr,taskName,taskId
-top,attr_top,category,1
+Blouse,attr_top,category,1
-pants,attr_pants,category,1
+Trousers,attr_pants,category,1
-skirt,attr_skirt,category,1
+Skirt,attr_skirt,category,1
-dress,attr_dress,category,1
+Dress,attr_dress,category,1
-outwear,attr_outwear,category,1
+Outwear,attr_outwear,category,1
-jumpsuit,attr_jumpsuit,category,1
+Jumpsuit,attr_jumpsuit,category,1

app/service/generate_image/service.py

@@ -22,19 +22,25 @@ 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
+from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust, face_detect_pic
+from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_stain_png_sd
 
 logger = logging.getLogger()
 
 
 class GenerateImage:
     def __init__(self, request_data):
+        if DEBUG is False:
+            self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+            self.channel = self.connection.channel()
+        # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+        # self.channel = self.connection.channel()
         self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
         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.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
-        self.channel = self.connection.channel()
         if request_data.mode == "img2img":
+            # cv2 读图片是BGR PIL读图片是RGB
             self.image = self.get_image(request_data.image_url)
             self.prompt = request_data.prompt
         else:
@@ -47,35 +53,69 @@ class GenerateImage:
         self.batch_size = 1
         self.category = request_data.category
         self.index = 0
-        self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'data': ''}
+        self.gender = request_data.gender
+        self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': '', 'category': ''}
         self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
         self.redis_client.expire(self.tasks_id, 600)
 
     def get_image(self, image_url):
         # Get data of an object.
         # Read data from response.
+        # read image use cv2
         try:
             response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
             image_file = BytesIO(response.data)
             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.resize(image_rbg, (1024, 1024))
         except minio.error.S3Error:
-            image_cv2 = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
+            image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
-        return image_cv2
+        return image
 
     def callback(self, result, error):
         if error:
             self.generate_data['status'] = "FAILURE"
             self.generate_data['message'] = str(error)
-            self.generate_data['data'] = str(error)
+            # self.generate_data['data'] = str(error)
             self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
         else:
-            image_result = result.as_numpy("generated_image")[0]
+            # pil图像转成numpy数组
-            image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")
+            image = result.as_numpy("generated_image")
-            self.generate_data['status'] = "SUCCESS"
+            image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
-            self.generate_data['message'] = "success"
+            is_smudge = True
-            self.generate_data['data'] = str(image_url)
+            if self.category == "sketch":
-            self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
+                # 色阶调整
+                cutoff = 1
+                levels_img = autoLevels(image_result, cutoff)
+                # 亮度调整
+                luminance = luminance_adjust(0.3, levels_img)
+                # 去背景
+                remove_bg_image = remove_background(luminance)
+                # 人脸检测
+                if face_detect_pic(remove_bg_image, self.user_id, self.category, self.tasks_id) > 0:
+                    is_smudge = False
+                else:
+                    # 污点/
+                    is_smudge, not_smudge_image = stain_detection(remove_bg_image, self.user_id, self.category, self.tasks_id)
+                    # 类型识别
+                    category, scores, not_smudge_image = generate_category_recognition(image=remove_bg_image, gender=self.gender)
+                    self.generate_data['category'] = str(category)
+                    image_result = not_smudge_image
+            if is_smudge:  # 无污点
+                # 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")
+                # logger.info(f"upload image SUCCESS ： {image_url}")
+                self.generate_data['status'] = "SUCCESS"
+                self.generate_data['message'] = "success"
+                self.generate_data['image_url'] = str(image_url)
+                self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
+            else:  # 有污点 保存图片到本地 测试用
+                self.generate_data['status'] = "SUCCESS"
+                self.generate_data['message'] = "success"
+                self.generate_data['image_url'] = str(GI_SYS_IMAGE_URL)
+                self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
+            # logger.info(f"stain_detection result : {self.generate_data}")
 
     def read_tasks_status(self):
         status_data = self.redis_client.get(self.tasks_id)
@@ -108,10 +148,11 @@ class GenerateImage:
 
             inputs = [input_text, input_image, input_mode]
             ctx = self.infer(inputs)
-            time_out = 60
+            time_out = 600
             generate_data = None
             while time_out > 0:
                 generate_data, _ = self.read_tasks_status()
+                # logger.info(generate_data)
                 if generate_data['status'] in ["REVOKED", "FAILURE"]:
                     ctx.cancel()
                     break
@@ -119,16 +160,18 @@ class GenerateImage:
                     break
                 time_out -= 1
                 time.sleep(0.1)
+                # logger.info(time_out, generate_data)
             return generate_data
         except Exception as e:
             self.generate_data['status'] = "FAILURE"
-            self.generate_data['message'] = "failure"
+            self.generate_data['message'] = str(e)
-            self.generate_data['data'] = str(e)
             self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
             raise Exception(str(e))
         finally:
             dict_generate_data, str_generate_data = self.read_tasks_status()
-            self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=str_generate_data)
+            if DEBUG is False:
+                self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=str_generate_data)
+            # self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=str_generate_data)
             logger.info(f" [x] Sent {json.dumps(dict_generate_data, indent=4)}")
 
 

app/service/generate_image/test.py

@@ -9,224 +9,169 @@
 """
 import json
 import logging
-
-import minio
-import numpy as np
-import random
-import redis
-import tritonclient
-import tritonclient.grpc as grpc_client
-from minio import Minio
-import cv2
-from PIL import Image
 import time
+from io import BytesIO
+
+import cv2
+import minio
+import redis
+import tritonclient.grpc as grpcclient
+import numpy as np
+from minio import Minio
+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.remove_background import remove_background
+from app.service.generate_image.utils.adjust_contrast import adjust_contrast
+from app.service.generate_image.utils.image_processing import remove_background, stain_detection
 from app.service.generate_image.utils.upload_sd_image import upload_png_sd
-from app.service.utils.decorator import RunTime
-from app.service.utils.generate_uuid import generate_uuid
 
 logger = logging.getLogger()
 
 
 class GenerateImage:
     def __init__(self, request_data):
-        self.tasks_id = request_data.tasks_id
+        if DEBUG is False:
-        self.model = request_data.model
+            self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
-        self.request_count = request_data.request_count
+            self.channel = self.connection.channel()
-        self.prompt = request_data.prompt
+        # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
-        self.image = request_data.image
+        # self.channel = self.connection.channel()
-        self.mode = request_data.mode
-        self.batch_size = request_data.batch_size
-
-        self.image_url = request_data.image_url
-        self.user_id = request_data.user_id
-        self.content = request_data.content
-        self.category = request_data.category
-        self.model_name = f"{self.category}{GI_MODEL_NAME}"
-        self.mode = request_data.mode
-        self.version = request_data.version
-        self.triton_client = grpc_client.InferenceServerClient(url="1")
-        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
-        self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
-        self.channel = self.connection.channel()
         self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
-        self.samples = 4  # no.of images to generate
+        self.grpc_client = grpcclient.InferenceServerClient(url=GI_MODEL_URL)
-        self.steps = 24
+        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
-        self.guidance_scale = 7
+        if request_data.mode == "img2img":
-        self.seed = random.randint(0, 2000000000)
+            self.image = self.get_image(request_data.image_url)
-        self.batch_size = 1
+            self.prompt = request_data.prompt
-        self.generate_data = json.dumps({'status': 'PENDING', 'message': "pending", 'data': ''})
-        self.redis_client.set(self.tasks_id, self.generate_data)
-
-    def get_result(self):
-
-        pass
-
-    @staticmethod
-    def image_grid(imgs, rows, cols):
-        assert len(imgs) == rows * cols
-
-        w, h = imgs[0].size
-        grid = Image.new('RGB', size=(cols * w, rows * h))
-
-        for i, img in enumerate(imgs):
-            grid.paste(img, box=(i % cols * w, i // cols * h))
-        return grid
-
-    @staticmethod
-    def preprocess_image(image, category):
-        height, width, _ = image.shape
-
-        if category == "print" or category == "moodboard":
-            square_size = min(height, width)
-            start_x = (width - square_size) // 2
-            start_y = (height - square_size) // 2
-            cropped = image[start_y: start_y + square_size, start_x: start_x + square_size]
-            resized_image = cv2.resize(cropped, (512, 512))
-
-        elif category == "sketch":
-            # below is the way that get "bigger" square image.
-            max_dimension = max(height, width)
-            square_image = np.ones((max_dimension, max_dimension, 3), dtype=np.uint8) * 255
-            start_h = (max_dimension - height) // 2
-            start_w = (max_dimension - width) // 2
-            square_image[start_h:start_h + height, start_w:start_w + width] = image
-            resized_image = cv2.resize(square_image, (512, 512))
-
         else:
-            raise ValueError(f"wrong category {category}, only in moodboard, print and sketch!")
+            self.image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
+            self.prompt = request_data.prompt
 
-        return resized_image
+        self.tasks_id = request_data.tasks_id
+        self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
+        self.mode = request_data.mode
+        self.batch_size = 1
+        self.category = request_data.category
+        self.index = 0
+        self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'data': ''}
+        self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
+        self.redis_client.expire(self.tasks_id, 600)
 
-    def get_image(self):
+    def get_image(self, image_url):
         # Get data of an object.
         # Read data from response.
         try:
-            response = self.minio_client.get_object(self.image_url.split('/')[0], self.image_url[self.image_url.find('/') + 1:])
+            response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
-            img = np.frombuffer(response.data, np.uint8)  # 转成8位无符号整型
+            image_file = BytesIO(response.data)
-            img = cv2.imdecode(img, cv2.IMREAD_COLOR)  # 解码
+            image_array = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
-            img = self.preprocess_image(img, self.category)
+            image_cv2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            image = cv2.resize(image_cv2, (1024, 1024))
         except minio.error.S3Error:
-            img = np.random.randn(512, 512, 3)
+            image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
-        return img
+        return image
 
     def callback(self, result, error):
         if error:
-            generate_data = json.dumps({'status': 'FAILURE', 'message': f"{error}", 'data': f"{error}"})
+            self.generate_data['status'] = "FAILURE"
-            self.redis_client.set(self.tasks_id, generate_data)
+            self.generate_data['message'] = str(error)
+            self.generate_data['data'] = str(error)
+            self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
         else:
-            images = result.as_numpy("IMAGES")
+            image_result = result.as_numpy("generated_image")[0]
-            if images.ndim == 3:
+            is_smudge = True
-                images = images[None, ...]
+            if self.category == "sketch":
-            images = (images * 255).round().astype("uint8")
+                # 去背景
-            pil_images = [Image.fromarray(image) for image in images]
+                remove_bg_image = remove_background(np.asarray(image_result))
-
+                # 污点检测
-            # for i in range(len(pil_images)):
+                is_smudge, not_smudge_image = stain_detection(remove_bg_image)
-            #     pil = pil_images[i]
+                image_result = not_smudge_image
-            #     pil.save(f'./temp_i2_{i}.png')
+            if is_smudge:  # 无污点
-            # self.image_grid(pil_images, rows, cols)
+                image_result = adjust_contrast(image_result)
-            url_list = []
+                image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")
-            for i, image in enumerate(pil_images):
+                # logger.info(f"upload image SUCCESS ： {image_url}")
-
+                self.generate_data['status'] = "SUCCESS"
-                if self.category == "sketch":
+                self.generate_data['message'] = "success"
-                    image = remove_background(np.asarray(image))
+                self.generate_data['data'] = str(image_url)
-                image_url = upload_png_sd(image, user_id=self.user_id, category=f"{self.category}", object_name=f"{generate_uuid()}_{i}.png", )
+                self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
-                url_list.append(image_url)
+            else:  # 有污点
-            generate_data = json.dumps({'status': 'SUCCESS', 'message': 'success', 'data': f'{url_list}'})
+                self.generate_data['status'] = "SUCCESS"
-            self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=generate_data)
+                self.generate_data['message'] = "success"
-            logger.info(f" [x] Sent {generate_data}")
+                self.generate_data['data'] = str(GI_SYS_IMAGE_URL)
-            self.redis_client.set(self.tasks_id, generate_data)
+                self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
+            # logger.info(f"stain_detection result : {self.generate_data}")
 
     def read_tasks_status(self):
-        status_data = json.loads(self.redis_client.get(self.tasks_id))
+        status_data = self.redis_client.get(self.tasks_id)
-        logging.info(f"{self.tasks_id} ===> {status_data}")
+        return json.loads(status_data), status_data
-        return status_data
+
+    def infer(self, inputs):
+        return self.grpc_client.infer(
+            model_name=GI_MODEL_NAME,
+            inputs=inputs,
+            # callback=self.callback
+        )
 
-    # @RunTime
     def get_result(self):
-        self.triton_client.get_model_metadata(model_name=self.model_name, model_version=self.version)
+        try:
-        self.triton_client.get_model_config(model_name=self.model_name, model_version=self.version)
+            prompts = [self.prompt] * self.batch_size
+            modes = [self.mode] * self.batch_size
+            images = [self.image.astype(np.float16)] * self.batch_size
 
-        image = self.get_image()
+            text_obj = np.array(prompts, dtype="object").reshape((-1, 1))
+            mode_obj = np.array(modes, dtype="object").reshape((-1, 1))
+            image_obj = np.array(images, dtype=np.float16).reshape((-1, 1024, 1024, 3))
 
-        # Input placeholder
+            input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
-        prompt_in = tritonclient.grpc.InferInput(name="PROMPT", shape=(self.batch_size,), datatype="BYTES")
+            input_image = grpcclient.InferInput("input_image", image_obj.shape, "FP16")
-        samples_in = tritonclient.grpc.InferInput("SAMPLES", (self.batch_size,), "INT32")
+            input_mode = grpcclient.InferInput("mode", mode_obj.shape, np_to_triton_dtype(text_obj.dtype))
-        steps_in = tritonclient.grpc.InferInput("STEPS", (self.batch_size,), "INT32")
-        guidance_scale_in = tritonclient.grpc.InferInput("GUIDANCE_SCALE", (self.batch_size,), "FP32")
-        seed_in = tritonclient.grpc.InferInput("SEED", (self.batch_size,), "INT64")
-        input_images_in = tritonclient.grpc.InferInput("INPUT_IMAGES", image.shape, "FP16")
-        images = tritonclient.grpc.InferRequestedOutput(name="IMAGES",
-                                                        # binary_data=False
-                                                        )
-        mode_in = tritonclient.grpc.InferInput("MODE", (self.batch_size,), "INT32")
 
-        # Setting inputs
+            input_text.set_data_from_numpy(text_obj)
-        prompt_in.set_data_from_numpy(np.asarray([self.content] * self.batch_size, dtype=object))
+            input_image.set_data_from_numpy(image_obj)
-        samples_in.set_data_from_numpy(np.asarray([self.samples], dtype=np.int32))
+            input_mode.set_data_from_numpy(mode_obj)
-        steps_in.set_data_from_numpy(np.asarray([self.steps], dtype=np.int32))
-        guidance_scale_in.set_data_from_numpy(np.asarray([self.guidance_scale], dtype=np.float32))
-        seed_in.set_data_from_numpy(np.asarray([self.seed], dtype=np.int64))
-        input_images_in.set_data_from_numpy(image.astype(np.float16))
-        mode_in.set_data_from_numpy(np.asarray([self.mode], dtype=np.int32))
 
-        # inference
+            inputs = [input_text, input_image, input_mode]
-        # @RunTime
+            ctx = self.infer(inputs)
-        def infer():
+            time_out = 600
-            return self.triton_client.async_infer(
+            generate_data = None
-                model_name=self.model_name,
+            while time_out > 0:
-                model_version=self.version,
+                generate_data, _ = self.read_tasks_status()
-                inputs=[prompt_in, samples_in, steps_in, guidance_scale_in, seed_in, input_images_in, mode_in],
+                # logger.info(generate_data)
-                outputs=[images],
+                if generate_data['status'] in ["REVOKED", "FAILURE"]:
-                callback=self.callback
+                    ctx.cancel()
-            )
+                    break
-
+                elif generate_data['status'] == "SUCCESS":
-        ctx = infer()
+                    break
-        time_out = 60
+                time_out -= 1
-        while time_out > 0:
+                time.sleep(0.1)
-            generate_data = self.read_tasks_status()
+                # logger.info(time_out, generate_data)
-            if generate_data['status'] in ["REVOKED", "FAILURE"]:
+            return generate_data
-                ctx.cancel()
+        except Exception as e:
-                self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=json.dumps(generate_data))
+        #     self.generate_data['status'] = "FAILURE"
-                logger.info(f" [x] Sent {generate_data}")
+        #     self.generate_data['message'] = "failure"
-                break
+        #     self.generate_data['data'] = str(e)
-            elif generate_data['status'] == "SUCCESS":
+        #     self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
-                break
+            raise Exception(str(e))
-            time_out -= 1
+        # finally:
-            time.sleep(1)
+        #     dict_generate_data, str_generate_data = self.read_tasks_status()
-        return self.read_tasks_status()
+        #     if DEBUG is False:
+        #         self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=str_generate_data)
+        #     logger.info(f" [x] Sent {json.dumps(dict_generate_data, indent=4)}")
 
 
 def infer_cancel(tasks_id):
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
-    data = {'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
+    data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
-    generate_data = json.dumps({'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'})
+    generate_data = json.dumps(data)
     redis_client.set(tasks_id, generate_data)
     return data
 
 
 if __name__ == '__main__':
-    # request_data = {
-    #     "user_id": 78,
-    #     "image_url": "123_123.png",
-    #     "category": "print",
-    #     "mode": 1,
-    #     "str": "a simple print",
-    #     "version": "1"
-    # }
     rd = GenerateImageModel(
-        mode=1,
+        tasks_id="123-89",
-        content='a blouse',
+        prompt='skeleton sitting by the side of a river looking soulful, concert poster, 4k, artistic',
-        gender='',
+        image_url="",
-        user_id=89,
+        mode='txt2img',
-        image_url='test/微信图片_20231206133428.jpg',
+        category="test"
-        category='sketch',
-        version='1',
-        tasks_id='123456'
     )
     server = GenerateImage(rd)
-    server.get_result()
+    print(server.get_result())
-    # print(infer_cancel(123456))

app/service/generate_image/utils/adjust_contrast.py

@@ -0,0 +1,30 @@
+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)  # 保存调整后的图片

app/service/generate_image/utils/image_processing.py

@@ -0,0 +1,395 @@
+import logging
+import time
+
+import mmcv
+import numpy as np
+import torch
+import tritonclient.http as httpclient
+import torch.nn.functional as F
+from app.core.config import *
+import cv2
+
+from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_stain_png_sd, upload_face_png_sd
+
+logger = logging.getLogger()
+
+
+def seg_preprocess(img_path):
+    img = mmcv.imread(img_path)
+    ori_shape = img.shape[:2]
+    img_scale = ori_shape
+    scale_factor = []
+    img, x, y = mmcv.imresize(img, img_scale, return_scale=True)
+    scale_factor.append(x)
+    scale_factor.append(y)
+    img = mmcv.imnormalize(img, mean=np.array([123.675, 116.28, 103.53]), std=np.array([58.395, 57.12, 57.375]), to_rgb=True)
+    preprocessed_img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
+    return preprocessed_img, ori_shape
+
+
+def get_mask(image_obj):
+    pre_mask = None
+    if len(image_obj.shape) == 2:
+        image_obj = cv2.cvtColor(image_obj, cv2.COLOR_GRAY2RGB)
+    if image_obj.shape[2] == 4:  # 如果是四通道 mask
+        pre_mask = image_obj[:, :, 3]
+        image_obj = image_obj[:, :, :3]
+
+    Contour = get_contours(image_obj)
+    Mask = np.zeros(image_obj.shape[:2], np.uint8)
+    if len(Contour):
+        Max_contour = Contour[0]
+        Epsilon = 0.001 * cv2.arcLength(Max_contour, True)
+        Approx = cv2.approxPolyDP(Max_contour, Epsilon, True)
+        cv2.drawContours(Mask, [Approx], -1, 255, -1)
+    else:
+        Mask = np.ones(image_obj.shape[:2], np.uint8) * 255
+
+    if pre_mask is None:
+        mask = Mask
+    else:
+        mask = cv2.bitwise_and(Mask, pre_mask)
+    return image_obj, mask
+
+
+def get_contours(image):
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    Edge = cv2.Canny(gray, 10, 150)
+    kernel = np.ones((5, 5), np.uint8)
+    Edge = cv2.dilate(Edge, kernel=kernel, iterations=1)
+    Edge = cv2.erode(Edge, kernel=kernel, iterations=1)
+    Contour, _ = cv2.findContours(Edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    Contour = sorted(Contour, key=cv2.contourArea, reverse=True)
+    return Contour
+
+
+# def seg_infer_image(image_obj):
+#     image, ori_shape = seg_preprocess(image_obj)
+#     client = httpclient.InferenceServerClient(url=f"{SEG_MODEL_URL}")
+#     transformed_img = image.astype(np.float32)
+#     # 输入集
+#     inputs = [
+#         httpclient.InferInput(SEGMENTATION['input'], transformed_img.shape, datatype="FP32")
+#     ]
+#     inputs[0].set_data_from_numpy(transformed_img, binary_data=True)
+#     # 输出集
+#     outputs = [
+#         httpclient.InferRequestedOutput(SEGMENTATION['output'], binary_data=True),
+#     ]
+#     results = client.infer(model_name=SEGMENTATION['name'], inputs=inputs, outputs=outputs)
+#     # 推理
+#     # 取结果
+#     inference_output1 = torch.from_numpy(results.as_numpy(SEGMENTATION['output']))
+#     seg_result = seg_postprocess(inference_output1, ori_shape)
+#     return seg_result
+
+def seg_infer_image(image_obj):
+    image, ori_shape = seg_preprocess(image_obj)
+    client = httpclient.InferenceServerClient(url=f"{SEG_MODEL_URL}")
+    transformed_img = image.astype(np.float32)
+    # 输入集
+    inputs = [
+        httpclient.InferInput(SEGMENTATION['input'], transformed_img.shape, datatype="FP32")
+    ]
+    inputs[0].set_data_from_numpy(transformed_img, binary_data=True)
+    # 输出集
+    outputs = [
+        httpclient.InferRequestedOutput(SEGMENTATION['output'], binary_data=True),
+    ]
+    start_time = time.time()
+    results = client.infer(model_name=SEGMENTATION['new_model_name'], inputs=inputs, outputs=outputs)
+    print(f"KNet infer time is :{time.time() - start_time}")
+    # 推理
+    # 取结果
+    inference_output1 = results.as_numpy(SEGMENTATION['output'])
+    seg_result = seg_postprocess(inference_output1, ori_shape)
+    return seg_result
+
+# def seg_postprocess(output, ori_shape):
+#     seg_logit = F.interpolate(output, size=ori_shape, scale_factor=None, mode='bilinear', align_corners=False)
+#     seg_logit = F.softmax(seg_logit, dim=1)
+#     seg_pred = seg_logit.argmax(dim=1)
+#     seg_pred = seg_pred.cpu().numpy()
+#     return seg_pred
+
+# KNet
+def seg_postprocess(output, ori_shape):
+    # seg_logit = F.interpolate(torch.tensor(output).float(), size=ori_shape, scale_factor=None, mode='bilinear', align_corners=False)
+    # seg_logit = F.softmax(seg_logit, dim=1)
+    # seg_pred = seg_logit.argmax(dim=1)
+    # seg_pred = output.cpu().numpy()
+    return output[0]
+
+def remove_background(image):
+    image_obj, mask = get_mask(image)
+    seg_result = seg_infer_image(image_obj)
+
+    temp_front = seg_result == 1
+    front_mask = (mask * (temp_front + 0).astype(np.uint8))
+    temp_back = seg_result == 2
+    back_mask = (mask * (temp_back + 0).astype(np.uint8))
+
+    if len(front_mask.shape) > 2:
+        front_mask = front_mask[0]
+    else:
+        front_mask = front_mask
+
+    if len(back_mask.shape) > 2:
+        back_mask = back_mask[0]
+    else:
+        back_mask = back_mask
+
+    result_mask = front_mask + back_mask
+    white_background = np.ones_like(image_obj) * 255
+    remove_bg_image = np.where(result_mask[:, :, None].astype(bool), image_obj, white_background)
+    # cv2.imwrite("source_image", image)
+    # cv2.imwrite("remove_bg_image", remove_bg_image)
+
+    return remove_bg_image
+
+
+def bounding_box(image):
+    edges = cv2.Canny(image, 50, 150)
+    # 查找轮廓
+    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # 初始化包围所有外接矩形的大矩形的坐标
+    x_min, y_min, x_max, y_max = float('inf'), float('inf'), -1, -1
+    # 遍历所有外接矩形，更新大矩形的坐标
+    for contour in contours:
+        x, y, w, h = cv2.boundingRect(contour)
+        x_min = min(x_min, x)
+        y_min = min(y_min, y)
+        x_max = max(x_max, x + w)
+        y_max = max(y_max, y + h)
+
+    # 根据大矩形的坐标来裁剪原始图像
+    result_image = image[y_min:y_max, x_min:x_max]
+    # cv2.imshow("result_image", result_image)
+    # cv2.waitKey(0)
+    return result_image
+
+
+def stain_detection(image, user_id, category, tasks_id, spot_size=100):
+    height, width, _ = image.shape
+
+    corners = [
+        image[0:spot_size, 0:spot_size],  # top left
+        image[0:spot_size, width - spot_size:width],  # top right
+        # image[height - spot_size:height, 0:spot_size],  # bottom left
+        # image[height - spot_size:height, width - spot_size:width]  # bottom right
+    ]
+
+    for index, corner in enumerate(corners):
+        num_white_pixels = (corner == [255, 255, 255]).all(axis=2).sum()
+        if num_white_pixels != spot_size * spot_size:
+            logger.info(f"第{index + 1}发现了污点")
+            return False, None
+    # 中心区域检测
+    # 将图像转换为灰度图像
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    # 获取图像的中心点坐标
+    center_x, center_y = image.shape[1] // 2, image.shape[0] // 2
+    # 定义中心区域的大小
+    patch_size = 100
+    half_patch = patch_size // 2
+    # 提取中心区域
+    center_patch = gray[center_y - half_patch:center_y + half_patch, center_x - half_patch:center_x + half_patch]
+    # 设置阈值来检测纯白区域
+    _, thresh = cv2.threshold(center_patch, 254, 255, cv2.THRESH_BINARY)
+    # 寻找轮廓
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # 过滤非连续的纯白区域
+    filtered_contours = [cnt for cnt in contours if cv2.contourArea(cnt) >= 300]  # 根据面积进行过滤，这里假设面积大于30的为连续区域
+    # 如果有连续的纯白区域存在
+    if filtered_contours:
+        # 将纯白区域替换为灰色
+        if DEBUG:
+            for cnt in filtered_contours:
+                x, y, w, h = cv2.boundingRect(cnt)
+                # 在原始图像上进行替换
+                image[y + center_y - half_patch:y + center_y - half_patch + h, x + center_x - half_patch:x + center_x - half_patch + w][thresh[y:y + h, x:x + w] == 255] = (128, 128, 128)
+            # 显示图像
+            cv2.imshow('Marked Image', image)
+            cv2.waitKey(0)
+        logger.info("中心区域存在连续的纯白区域")
+        is_pure_white = True
+    else:
+        logger.info("中心区域不存在连续的纯白区域")
+        is_pure_white = False
+
+    if is_pure_white:
+        return False, None
+    if DEBUG:
+        for corner_coords in [
+            (0, 0),
+            # (0, width - spot_size),
+            (height - spot_size, 0),
+            # (height - spot_size, width - spot_size)
+            # 中心点
+        ]:
+            cv2.rectangle(image, corner_coords, (corner_coords[0] + spot_size, corner_coords[1] + spot_size), (0, 0, 255), 2)
+        cv2.rectangle(image, (center_x - spot_size // 2, center_y - spot_size // 2), (center_x + spot_size // 2, center_y + spot_size // 2), (0, 255, 0), 2)  # 在原始图像上绘制矩形框
+    dst = image.copy()
+    for corner_coords in [
+        (0, 0),
+        # (0, width - spot_size),
+        (height - spot_size, 0),
+        # (height - spot_size, width - spot_size)
+        # 中心点
+    ]:
+        cv2.rectangle(dst, corner_coords, (corner_coords[0] + spot_size, corner_coords[1] + spot_size), (0, 0, 255), 2)
+    cv2.rectangle(dst, (center_x - spot_size // 2, center_y - spot_size // 2), (center_x + spot_size // 2, center_y + spot_size // 2), (0, 255, 0), 2)  # 在原始图像上绘制矩形框
+    image_url = upload_stain_png_sd(dst, user_id=user_id, category=f"{category}", object_name=f"{tasks_id}.png")
+    return True, image
+
+
+def generate_category_recognition(image, gender):
+    def preprocess(img):
+        img = mmcv.imread(img)
+        # ori_shape = img.shape[:2]
+        img_scale = (224, 224)
+        scale_factor = []
+        img, x, y = mmcv.imresize(img, img_scale, return_scale=True)
+        scale_factor.append(x)
+        scale_factor.append(y)
+        img = mmcv.imnormalize(img, mean=np.array([123.675, 116.28, 103.53]), std=np.array([58.395, 57.12, 57.375]), to_rgb=True)
+        preprocessed_img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
+        return preprocessed_img
+
+    preprocessed_img = preprocess(image)
+    triton_client = httpclient.InferenceServerClient(url=ATT_TRITON_URL)
+
+    inputs = [
+        httpclient.InferInput("input__0", preprocessed_img.shape, datatype="FP32")
+    ]
+    inputs[0].set_data_from_numpy(preprocessed_img, binary_data=True)
+    results = triton_client.infer(model_name="attr_retrieve_category", inputs=inputs)
+    inference_output = torch.from_numpy(results.as_numpy(f'output__0'))
+
+    scores = inference_output.detach().numpy()
+    import pandas as pd
+
+    attr_type = pd.read_csv(CATEGORY_PATH)
+    colattr = list(attr_type['labelName'])
+
+    task = attr_type['taskName'][0]
+
+    maxsc = np.max(scores[0][:5])
+    indexs = np.argwhere(scores == maxsc)[:, 1]
+    category = colattr[indexs[0]]
+
+    if gender == "Male":
+        if category == 'Trousers' or category == 'Skirt':
+            category = 'Bottoms'
+        elif category == 'Blouse' or category == 'Dress':
+            category = 'Tops'
+        else:
+            category = 'Outwear'
+    return category, scores, image
+
+
+def autoLevels(img, cutoff=0.1):
+    channels = img.shape[2]  # h,w,ch
+    table = np.zeros((1, 256, 3), np.uint8)
+    for ch in range(channels):
+        # cutoff=0.1, 计算 0.1%, 99.9% 分位的灰度值
+        low = np.percentile(img[:, :, ch], q=cutoff)  # ch 通道, cutoff=0.1, 0.1 分位的灰度值
+        high = np.percentile(img[:, :, ch], q=100 - cutoff)  # 99.9 分位的灰度值, [0, high] 占比99.9%
+        # 输入动态线性拉伸
+        Sin = min(max(low, 0), high - 2)  # Sin, 黑场阈值, 0<=Sin<Hin
+        Hin = min(high, 255)  # Hin, 白场阈值, Sin<Hin<=255
+        difIn = Hin - Sin
+        V1 = np.array([(min(max(255 * (i - Sin) / difIn, 0), 255)) for i in range(256)])
+        # 灰场伽马调节
+        gradMed = np.median(img[:, :, ch])  # 拉伸前的中值
+        Mt = V1[int(gradMed)] / 128.  # 拉伸后的映射值
+        V2 = 255 * np.power(V1 / 255, 1 / Mt)  # 伽马调节
+        # 输出线性拉伸
+        Sout, Hout = 5, 250  # Sout 输出黑场阈值, Hout 输出白场阈值
+        difOut = Hout - Sout
+        table[0, :, ch] = np.array([(min(max(Sout + difOut * V2[i] / 255, 0), 255)) for i in range(256)])
+    return cv2.LUT(img, table)
+
+
+def luminance_adjust(alpha, img):
+    if alpha > 0:
+        img_out = img * (1 - alpha) + alpha * 255.0
+    else:
+        img_out = img * (1 + alpha)
+
+    return np.array(img_out, dtype='uint8')
+
+    # 14.14 Photoshop 自动色阶调整算法
+
+
+def face_detect_pic(image, user_id, category, tasks_id):
+    # 1、转灰度图
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    # cv2.imshow("gray", gray)
+
+    # 2、训练一组人脸
+    face_detector = cv2.CascadeClassifier(FACE_CLASSIFIER)
+
+    # 3、检测人脸（用灰度图检测，返回人脸矩形坐标(4个角)）
+    faces_rect = face_detector.detectMultiScale(gray, 1.05, 3)
+
+    if DEBUG:
+        dst = image.copy()
+        for x, y, w, h in faces_rect:
+            cv2.rectangle(dst, (x, y), (x + w, y + h), (0, 0, 255), 3)  # 画出矩形框
+        # cv2.imshow("", dst)
+        # cv2.waitKey(0)
+    # TODO 暂时保留
+    dst = image.copy()
+    for x, y, w, h in faces_rect:
+        cv2.rectangle(dst, (x, y), (x + w, y + h), (0, 0, 255), 3)  # 画出矩形框
+    image_url = upload_face_png_sd(dst, user_id=user_id, category=f"{category}", object_name=f"{tasks_id}.png")
+    return len(faces_rect)
+
+
+if __name__ == '__main__':
+    # Photoshop 自动色阶调整算法
+    img = cv2.imread("2.png", flags=1)  # 读取彩色图像
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转换为灰度图像
+    print("cutoff={}, minG={}, maxG={}".format(0.0, gray.min(), gray.min()))
+
+    # 色阶手动调整
+    # equManual = levelsAdjust(img, 63, 205, 0.8, 10, 245)  # 手动调节
+    # 色阶自动调整
+    cutoff = 1.0  # 截断比例, 建议范围 [0.0,1.0]
+    # cv2.imwrite("source.png", img)
+    equAuto = autoLevels(img, cutoff)
+    # cv2.imwrite("levels.png", equAuto)
+    luminance = luminance_adjust(0.3, equAuto)
+    # cv2.imwrite("luminance.png", luminance)
+    #
+    # # 将图像转换为灰度
+    # gray = cv2.cvtColor(luminance, cv2.COLOR_BGR2GRAY)
+    #
+    # # 使用Canny边缘检测算法检测图像的边缘
+    # edges = cv2.Canny(gray, 150, 200)
+    #
+    # # 对边缘进行膨胀操作，增强轮廓
+    # kernel = np.ones((1, 1), np.uint8)
+    # dilated_edges = cv2.dilate(edges, kernel, iterations=1)
+    #
+    # # 创建一个与原始图像相同大小的空白图像
+    # # result = np.zeros_like(luminance)
+    #
+    # # 将增强后的轮廓叠加到原始图像上
+    # luminance[dilated_edges != 0] = (255, 255, 255)
+
+    remove_bg_img = remove_background(luminance)
+    # cv2.imwrite("remove_bg_img.png", remove_bg_img)
+
+    print(1)
+    cv2.imshow("source", img)
+    cv2.imshow("levels", equAuto)
+    cv2.imshow("luminance", luminance)
+    # cv2.imshow("dilated_edges", luminance)
+    cv2.imshow("remove_bg_img", remove_bg_img)
+
+    cv2.waitKey(0)
+
+    image = cv2.imread("1.png")
+    remove_background(image)

app/service/generate_image/utils/remove_background.py

@@ -1,112 +0,0 @@
-import cv2
-import mmcv
-import numpy as np
-import torch
-from PIL import Image
-import tritonclient.http as httpclient
-import torch.nn.functional as F
-from app.core.config import *
-
-
-def seg_preprocess(img_path):
-    img = mmcv.imread(img_path)
-    ori_shape = img.shape[:2]
-    img_scale = (224, 224)
-    scale_factor = []
-    img, x, y = mmcv.imresize(img, img_scale, return_scale=True)
-    scale_factor.append(x)
-    scale_factor.append(y)
-    img = mmcv.imnormalize(img, mean=np.array([123.675, 116.28, 103.53]), std=np.array([58.395, 57.12, 57.375]), to_rgb=True)
-    preprocessed_img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
-    return preprocessed_img, ori_shape
-
-
-def get_mask(image_obj):
-    pre_mask = None
-    if len(image_obj.shape) == 2:
-        image_obj = cv2.cvtColor(image_obj, cv2.COLOR_GRAY2RGB)
-    if image_obj.shape[2] == 4:  # 如果是四通道 mask
-        pre_mask = image_obj[:, :, 3]
-        image_obj = image_obj[:, :, :3]
-
-    Contour = get_contours(image_obj)
-    Mask = np.zeros(image_obj.shape[:2], np.uint8)
-    if len(Contour):
-        Max_contour = Contour[0]
-        Epsilon = 0.001 * cv2.arcLength(Max_contour, True)
-        Approx = cv2.approxPolyDP(Max_contour, Epsilon, True)
-        cv2.drawContours(Mask, [Approx], -1, 255, -1)
-    else:
-        Mask = np.ones(image_obj.shape[:2], np.uint8) * 255
-
-    if pre_mask is None:
-        mask = Mask
-    else:
-        mask = cv2.bitwise_and(Mask, pre_mask)
-    return image_obj, mask
-
-
-def get_contours(image):
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    Edge = cv2.Canny(gray, 10, 150)
-    kernel = np.ones((5, 5), np.uint8)
-    Edge = cv2.dilate(Edge, kernel=kernel, iterations=1)
-    Edge = cv2.erode(Edge, kernel=kernel, iterations=1)
-    Contour, _ = cv2.findContours(Edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    Contour = sorted(Contour, key=cv2.contourArea, reverse=True)
-    return Contour
-
-
-def seg_infer_image(image_obj):
-    image, ori_shape = seg_preprocess(image_obj)
-    client = httpclient.InferenceServerClient(url=f"{SEG_MODEL_URL}")
-    transformed_img = image.astype(np.float32)
-    # 输入集
-    inputs = [
-        httpclient.InferInput(SEGMENTATION['input'], transformed_img.shape, datatype="FP32")
-    ]
-    inputs[0].set_data_from_numpy(transformed_img, binary_data=True)
-    # 输出集
-    outputs = [
-        httpclient.InferRequestedOutput(SEGMENTATION['output'], binary_data=True),
-    ]
-    results = client.infer(model_name=SEGMENTATION['name'], inputs=inputs, outputs=outputs)
-    # 推理
-    # 取结果
-    inference_output1 = torch.from_numpy(results.as_numpy(SEGMENTATION['output']))
-    seg_result = seg_postprocess(inference_output1, ori_shape)
-    return seg_result
-
-
-def seg_postprocess(output, ori_shape):
-    seg_logit = F.interpolate(output, size=ori_shape, scale_factor=None, mode='bilinear', align_corners=False)
-    seg_logit = F.softmax(seg_logit, dim=1)
-    seg_pred = seg_logit.argmax(dim=1)
-    seg_pred = seg_pred.cpu().numpy()
-    return seg_pred
-
-
-def remove_background(image):
-    image_obj, mask = get_mask(image)
-    seg_result = seg_infer_image(image_obj)
-
-    temp_front = seg_result == 1
-    front_mask = (mask * (temp_front + 0).astype(np.uint8))
-    temp_back = seg_result == 2
-    back_mask = (mask * (temp_back + 0).astype(np.uint8))
-
-    if len(front_mask.shape) > 2:
-        front_mask = front_mask[0]
-    else:
-        front_mask = front_mask
-
-    if len(back_mask.shape) > 2:
-        back_mask = back_mask[0]
-    else:
-        back_mask = back_mask
-
-    result_mask = front_mask + back_mask
-    white_background = np.ones_like(image_obj) * 255
-    result_image = np.where(result_mask[:, :, None].astype(bool), image_obj, white_background)
-
-    return Image.fromarray(result_image)

app/service/generate_image/utils/upload_sd_image.py

@@ -10,6 +10,7 @@
 import io
 import logging
 
+import cv2
 from PIL import Image
 from minio import Minio
 
@@ -20,18 +21,47 @@ minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET
 
 def upload_png_sd(image, user_id, category, object_name):
     try:
-        image_file = io.BytesIO()
+        _, img_byte_array = cv2.imencode('.jpg', image)
-        image = Image.fromarray(image)
-        image.save(image_file, format='JPEG')
-        image_file.seek(0)
         minio_req = minio_client.put_object(
             GI_MINIO_BUCKET,
             f'{user_id}/{category}/{object_name}',
-            image_file,
+            io.BytesIO(img_byte_array),
-            len(image_file.getvalue()),
+            len(img_byte_array),
             content_type='image/jpeg'
         )
         image_url = f"aida-users/{minio_req.object_name}"
         return image_url
     except Exception as e:
         logging.warning(f"upload_png_mask runtime exception : {e}")
+
+
+def upload_stain_png_sd(image, user_id, category, object_name):
+    try:
+        _, img_byte_array = cv2.imencode('.jpg', image)
+        minio_req = minio_client.put_object(
+            "test",
+            f'generate_result/stain/{user_id}_{category}_{object_name}',
+            io.BytesIO(img_byte_array),
+            len(img_byte_array),
+            content_type='image/jpeg'
+        )
+        image_url = f"test/{minio_req.object_name}"
+        return image_url
+    except Exception as e:
+        logging.warning(f"upload_png_mask runtime exception : {e}")
+
+
+def upload_face_png_sd(image, user_id, category, object_name):
+    try:
+        _, img_byte_array = cv2.imencode('.jpg', image)
+        minio_req = minio_client.put_object(
+            "test",
+            f'generate_result/face/{user_id}_{category}_{object_name}',
+            io.BytesIO(img_byte_array),
+            len(img_byte_array),
+            content_type='image/jpeg'
+        )
+        image_url = f"test/{minio_req.object_name}"
+        return image_url
+    except Exception as e:
+        logging.warning(f"upload_png_mask runtime exception : {e}")