From bb46282987d1c788a21c4ca4d71cb8dd1494f11c Mon Sep 17 00:00:00 2001 From: zhouchengrong Date: Mon, 13 May 2024 11:35:27 +0800 Subject: [PATCH] =?UTF-8?q?feat=20generate=20=E6=B1=A1=E7=82=B9=E5=9B=BE?= =?UTF-8?q?=E4=BF=9D=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/service/generate_image/service.py | 7 ++-- .../generate_image/utils/image_processing.py | 26 ++++++++++++--- .../generate_image/utils/upload_sd_image.py | 32 +++++++++++++++++++ 3 files changed, 57 insertions(+), 8 deletions(-) diff --git a/app/service/generate_image/service.py b/app/service/generate_image/service.py index 51d1c4d..0c8bca3 100644 --- a/app/service/generate_image/service.py +++ b/app/service/generate_image/service.py @@ -24,7 +24,7 @@ from app.core.config import * from app.schemas.generate_image import GenerateImageModel 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 +from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_stain_png_sd logger = logging.getLogger() @@ -93,11 +93,11 @@ class GenerateImage: # 去背景 remove_bg_image = remove_background(luminance) # 人脸检测 - if face_detect_pic(remove_bg_image) > 0: + 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) + 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) @@ -111,7 +111,6 @@ class GenerateImage: self.generate_data['image_url'] = str(image_url) self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) else: # 有污点 保存图片到本地 测试用 - cv2.imwrite(f"{self.tasks_id}.png", image_result) self.generate_data['status'] = "SUCCESS" self.generate_data['message'] = "success" self.generate_data['image_url'] = str(GI_SYS_IMAGE_URL) diff --git a/app/service/generate_image/utils/image_processing.py b/app/service/generate_image/utils/image_processing.py index d4f5974..2883129 100644 --- a/app/service/generate_image/utils/image_processing.py +++ b/app/service/generate_image/utils/image_processing.py @@ -7,6 +7,8 @@ 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() @@ -137,7 +139,7 @@ def bounding_box(image): return result_image -def stain_detection(image, spot_size=100): +def stain_detection(image, user_id, category, tasks_id, spot_size=100): height, width, _ = image.shape corners = [ @@ -197,6 +199,17 @@ def stain_detection(image, spot_size=100): ]: 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 @@ -279,7 +292,7 @@ def luminance_adjust(alpha, img): # 14.14 Photoshop 自动色阶调整算法 -def face_detect_pic(image): +def face_detect_pic(image, user_id, category, tasks_id): # 1、转灰度图 gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) # cv2.imshow("gray", gray) @@ -294,8 +307,13 @@ def face_detect_pic(image): 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) + # 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) diff --git a/app/service/generate_image/utils/upload_sd_image.py b/app/service/generate_image/utils/upload_sd_image.py index 0c3baec..0e8e542 100644 --- a/app/service/generate_image/utils/upload_sd_image.py +++ b/app/service/generate_image/utils/upload_sd_image.py @@ -33,3 +33,35 @@ def upload_png_sd(image, user_id, category, 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}")