aidlab/AiDA_Python
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
fix
This commit is contained in:
zhouchengrong
2024-12-01 19:52:33 +08:00
parent 13e3f8ac3d
commit 6e621038f6
1 changed files with 38 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
app/service/generate_image/service_generate_image.py
View File

@@ -35,7 +35,12 @@ class GenerateImage:
# self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS)) # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
# self.channel = self.connection.channel() # self.channel = self.connection.channel()
# self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE) # 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) self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
if request_data.mode == "img2img": if request_data.mode == "img2img":
# cv2 读图片是BGR PIL读图片是RGB # cv2 读图片是BGR PIL读图片是RGB
@@ -87,23 +92,28 @@ class GenerateImage:
image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR) image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
is_smudge = True is_smudge = True
if self.category == "sketch": if self.category == "sketch":
# 色阶调整 if self.version == "fast":
cutoff = 1 # 色阶调整
levels_img = autoLevels(image_result, cutoff) cutoff = 1
# 亮度调整 levels_img = autoLevels(image_result, cutoff)
luminance = luminance_adjust(0.3, levels_img) # 亮度调整
# 去背景 luminance = luminance_adjust(0.3, levels_img)
remove_bg_image = remove_background(luminance) # 去背景
# 人脸检测 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 # if face_detect_pic(remove_bg_image, self.user_id, self.category, self.tasks_id) > 0:
# else: # is_smudge = False
# 污点/ # else:
is_smudge, not_smudge_image = stain_detection(remove_bg_image, self.user_id, self.category, self.tasks_id) # 污点/
# 类型识别 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) category, scores, not_smudge_image = generate_category_recognition(image=remove_bg_image, gender=self.gender)
image_result = not_smudge_image 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: # 无污点 if is_smudge: # 无污点
# image_result = adjust_contrast(image_result) # 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") 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)) 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_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_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(text_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_text.set_data_from_numpy(text_obj)
input_image.set_data_from_numpy(image_obj) input_image.set_data_from_numpy(image_obj)
input_mode.set_data_from_numpy(mode_obj) input_mode.set_data_from_numpy(mode_obj)
inputs = [input_text, input_image, input_mode] 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 time_out = 600
generate_data = None generate_data = None
while time_out > 0: while time_out > 0:
@@ -181,11 +195,12 @@ def infer_cancel(tasks_id):
if __name__ == '__main__': if __name__ == '__main__':
rd = GenerateImageModel( rd = GenerateImageModel(
tasks_id="123-89", 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", image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg",
mode='txt2img', mode='txt2img',
category="test", category="test",
gender="male" gender="male",
version="high"
) )
server = GenerateImage(rd) server = GenerateImage(rd)
print(server.get_result()) print(server.get_result())