114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
import io
|
||
import logging
|
||
import os
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
from app.service.utils.new_oss_client import oss_get_image
|
||
|
||
logger = logging.getLogger()
|
||
|
||
|
||
class LoadBodyImage:
|
||
name = "LoadBodyImage"
|
||
|
||
def __init__(self, minio_client):
|
||
self.minio_client = minio_client
|
||
|
||
@classmethod
|
||
def get_name(cls):
|
||
return cls.name
|
||
|
||
def __call__(self, result):
|
||
result["name"] = "mannequin"
|
||
result['body_image'] = oss_get_image(oss_client=self.minio_client, bucket=result['body_path'].split("/", 1)[0], object_name=result['body_path'].split("/", 1)[1], data_type="PIL")
|
||
return result
|
||
|
||
|
||
class LoadImage:
|
||
name = "LoadImage"
|
||
|
||
def __init__(self, minio_client):
|
||
self.minio_client = minio_client
|
||
|
||
@classmethod
|
||
def get_name(cls):
|
||
return cls.name
|
||
|
||
def __call__(self, result):
|
||
result['image'], result['pre_mask'] = self.read_image(result['path'])
|
||
# if 'extract_lines' in result.keys():
|
||
# if result['extract_lines']:
|
||
# result['gray'] = self.get_lines(cv2.cvtColor(result['image'], cv2.COLOR_BGR2GRAY), result['path'])
|
||
# else:
|
||
# result['gray'] = cv2.cvtColor(result['image'], cv2.COLOR_BGR2GRAY)
|
||
# else:
|
||
# result['gray'] = cv2.cvtColor(result['image'], cv2.COLOR_BGR2GRAY)
|
||
|
||
result['gray'] = self.get_lines(cv2.cvtColor(result['image'], cv2.COLOR_BGR2GRAY), result['path'])
|
||
result['keypoint'] = self.get_keypoint(result['name'])
|
||
result['img_shape'] = result['image'].shape
|
||
result['ori_shape'] = result['image'].shape
|
||
return result
|
||
|
||
def get_lines(self, img, path):
|
||
binary = cv2.adaptiveThreshold(img, 255,
|
||
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
||
cv2.THRESH_BINARY_INV,
|
||
25, 10)
|
||
|
||
# 步骤2：细化边缘（可选，让线条更干净）
|
||
# kernel = np.ones((1, 1), np.uint8)
|
||
# clean = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
|
||
|
||
thinned = cv2.ximgproc.thinning(binary, thinningType=cv2.ximgproc.THINNING_ZHANGSUEN) # thinning算法细化线条
|
||
mask = thinned > 0
|
||
result = np.ones_like(img) * 255
|
||
result[mask] = img[mask]
|
||
|
||
# 步骤3：反转回 白底黑线
|
||
# lines = cv2.bitwise_not(thinned)
|
||
# cv2.imwrite(os.path.join('/home/user/PycharmProjects/trinity_client_aida/test/lines_original_result_5', f"Original_{path.replace('/', '-')}.png"), img)
|
||
# cv2.imwrite(os.path.join('/home/user/PycharmProjects/trinity_client_aida/test/lines_original_result_5', f"Line_{path.replace('/', '-')}.png"), result)
|
||
|
||
return result
|
||
|
||
def read_image(self, image_path):
|
||
image_mask = None
|
||
image = oss_get_image(oss_client=self.minio_client, bucket=image_path.split("/", 1)[0], object_name=image_path.split("/", 1)[1], data_type="cv2")
|
||
if len(image.shape) == 2:
|
||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
||
if image.shape[2] == 4: # 如果是四通道 mask
|
||
image_mask = image[:, :, 3]
|
||
image = image[:, :, :3]
|
||
|
||
if image.shape[:2] <= (50, 50):
|
||
# 计算新尺寸
|
||
new_size = (image.shape[1] * 2, image.shape[0] * 2)
|
||
# 调整大小
|
||
image = cv2.resize(image, new_size, interpolation=cv2.INTER_LINEAR)
|
||
return image, image_mask
|
||
|
||
@staticmethod
|
||
def get_keypoint(name):
|
||
if name == 'blouse' or name == 'outwear' or name == 'dress' or name == 'tops':
|
||
keypoint = 'shoulder'
|
||
elif name == 'trousers' or name == 'skirt' or name == 'bottoms':
|
||
keypoint = 'waistband'
|
||
elif name == 'bag':
|
||
keypoint = 'hand_point'
|
||
elif name == 'shoes':
|
||
keypoint = 'toe'
|
||
elif name == 'hairstyle':
|
||
keypoint = 'head_point'
|
||
elif name == 'earring':
|
||
keypoint = 'ear_point'
|
||
elif name == 'accessories':
|
||
keypoint = "accessories"
|
||
else:
|
||
raise KeyError(f"{name} does not belong to item category list: blouse, outwear, dress, trousers, skirt, "
|
||
f"bag, shoes, hairstyle, earring.")
|
||
return keypoint
|