feat flux 取消污点检测 增加类别判断

fix

app/service/generate_image/service_generate_image.py

@@ -35,7 +35,12 @@ class GenerateImage:
         # 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.version = request_data.version
+        if request_data.version == "fast":
+            self.grpc_client = grpcclient.InferenceServerClient(url=FAST_GI_MODEL_URL)
+        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)
         if request_data.mode == "img2img":
             # cv2 读图片是BGR PIL读图片是RGB
@@ -87,23 +92,28 @@ class GenerateImage:
             image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
             is_smudge = True
             if self.category == "sketch":
-                # 色阶调整
-                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 self.version == "fast":
+                    # 色阶调整
+                    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
+                else:
+                    category, scores, not_smudge_image = generate_category_recognition(image=image_result, 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}", file_name=f"{self.tasks_id}.png")
@@ -134,15 +144,19 @@ class GenerateImage:
             image_obj = np.array(images, dtype=np.float16).reshape((-1, 1024, 1024, 3))
 
             input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
-            input_image = grpcclient.InferInput("input_image", image_obj.shape, "FP16")
-            input_mode = grpcclient.InferInput("mode", mode_obj.shape, np_to_triton_dtype(text_obj.dtype))
+            input_image = grpcclient.InferInput("input_image", image_obj.shape, np_to_triton_dtype(image_obj.dtype))
+            input_mode = grpcclient.InferInput("mode", mode_obj.shape, np_to_triton_dtype(mode_obj.dtype))
 
             input_text.set_data_from_numpy(text_obj)
             input_image.set_data_from_numpy(image_obj)
             input_mode.set_data_from_numpy(mode_obj)
 
             inputs = [input_text, input_image, input_mode]
-            ctx = self.grpc_client.async_infer(model_name=GI_MODEL_NAME, inputs=inputs, callback=self.callback)
+            if self.version == "fast":
+                ctx = self.grpc_client.async_infer(model_name=FAST_GI_MODEL_NAME, inputs=inputs, callback=self.callback)
+            else:
+                ctx = self.grpc_client.async_infer(model_name=GI_MODEL_NAME, inputs=inputs, callback=self.callback)
+
             time_out = 600
             generate_data = None
             while time_out > 0:
@@ -181,11 +195,12 @@ def infer_cancel(tasks_id):
 if __name__ == '__main__':
     rd = GenerateImageModel(
         tasks_id="123-89",
-        prompt='skeleton sitting by the side of a river looking soulful, concert poster, 4k, artistic',
+        prompt='a single item of sketch of Wabi-sabi, skirt, tiered, 4k, white background',
         image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg",
         mode='txt2img',
         category="test",
-        gender="male"
+        gender="male",
+        version="high"
     )
     server = GenerateImage(rd)
     print(server.get_result())