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:49:27 +08:00
parent 6a0c581e5d
commit 87d3848088
1 changed files with 23 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -92,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" and 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
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")
@@ -190,7 +195,7 @@ def infer_cancel(tasks_id):
if __name__ == '__main__':
rd = GenerateImageModel(
tasks_id="123-89",
prompt='a fabric print, flower, yellow, 4k, hud',
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",