102 lines
4.2 KiB
Python
102 lines
4.2 KiB
Python
|
import cv2
|
||
|
|
import mediapipe as mp
|
||
|
|
import numpy as np
|
||
|
|
from minio import Minio
|
||
|
|
|
||
|
|
from app.core.config import MINIO_URL, MINIO_ACCESS, MINIO_SECRET, MINIO_SECURE
|
||
|
|
from app.schemas.mannequin_edit import MannequinModel
|
||
|
|
from app.service.utils.new_oss_client import oss_get_image, oss_upload_image
|
||
|
|
|
||
|
|
minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
|
||
|
|
|
||
|
|
|
||
|
|
class MannequinEditService():
|
||
|
|
def __init__(self, request_data):
|
||
|
|
self.scale = request_data.scale
|
||
|
|
self.image = oss_get_image(oss_client=minio_client, bucket=request_data.mannequins.split('/')[0], object_name=request_data.mannequins[request_data.mannequins.find('/') + 1:], data_type="cv2")
|
||
|
|
self.mannequin_name = request_data.mannequin_name
|
||
|
|
self.bucket_name = request_data.bucket_name
|
||
|
|
if self.image.shape[2] == 4:
|
||
|
|
self.bgr = self.image[:, :, :3]
|
||
|
|
self.alpha = self.image[:, :, 3]
|
||
|
|
self.bgr = cv2.bitwise_and(self.bgr, self.bgr, mask=cv2.normalize(self.alpha, None, 0, 1, cv2.NORM_MINMAX))
|
||
|
|
self.h, self.w, _ = self.bgr.shape
|
||
|
|
else:
|
||
|
|
self.bgr = self.image
|
||
|
|
self.h, self.w, _ = self.bgr.shape
|
||
|
|
self.alpha = None
|
||
|
|
|
||
|
|
def __call__(self, *args, **kwargs):
|
||
|
|
leg_top, leg_bottom = self.attitude_detection()
|
||
|
|
if leg_top and leg_bottom:
|
||
|
|
new_mannequin = self.resize_leg(leg_top, leg_bottom)
|
||
|
|
_, encoded_image = cv2.imencode('.png', new_mannequin)
|
||
|
|
image_bytes = encoded_image.tobytes()
|
||
|
|
req = oss_upload_image(oss_client=minio_client, bucket=self.bucket_name, object_name=f"{self.mannequin_name}.png", image_bytes=image_bytes)
|
||
|
|
return req.bucket_name + "/" + req.object_name
|
||
|
|
else:
|
||
|
|
return "No leg detected"
|
||
|
|
|
||
|
|
def attitude_detection(self):
|
||
|
|
mp_pose = mp.solutions.pose
|
||
|
|
pose = mp_pose.Pose()
|
||
|
|
|
||
|
|
# 将 BGR 图像转换为 RGB 格式
|
||
|
|
image_rgb = cv2.cvtColor(self.bgr, cv2.COLOR_BGR2RGB)
|
||
|
|
leg_top, leg_bottom = None, None
|
||
|
|
# 进行姿态检测
|
||
|
|
results = pose.process(image_rgb)
|
||
|
|
if results.pose_landmarks:
|
||
|
|
# 获取腿部关键点
|
||
|
|
landmarks = results.pose_landmarks.landmark
|
||
|
|
|
||
|
|
# 找到腿部上边界和下边界
|
||
|
|
leg_top = int(landmarks[mp_pose.PoseLandmark.LEFT_HIP].y * self.h)
|
||
|
|
leg_bottom = int(max(landmarks[mp_pose.PoseLandmark.LEFT_ANKLE].y,
|
||
|
|
landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE].y) * self.h)
|
||
|
|
|
||
|
|
return leg_top, leg_bottom
|
||
|
|
|
||
|
|
def resize_leg(self, leg_top, leg_bottom):
|
||
|
|
# 上半身
|
||
|
|
top_part_bgr = self.bgr[:leg_top, :]
|
||
|
|
top_part_bgr_alpha = self.alpha[:leg_top, :]
|
||
|
|
|
||
|
|
# 小腿
|
||
|
|
part_thigh = self.bgr[leg_top:leg_bottom, :]
|
||
|
|
part_thigh_alpha = self.alpha[leg_top:leg_bottom, :]
|
||
|
|
|
||
|
|
# 大腿
|
||
|
|
part_calf = self.bgr[leg_bottom:, :]
|
||
|
|
part_calf_alpha = self.alpha[leg_bottom:, :]
|
||
|
|
|
||
|
|
new_thigh_height = int((leg_bottom - leg_top) * self.scale[0])
|
||
|
|
new_calf_height = int((self.h - leg_bottom) * self.scale[1])
|
||
|
|
|
||
|
|
resized_thigh = cv2.resize(part_thigh, (self.w, new_thigh_height), interpolation=cv2.INTER_LINEAR)
|
||
|
|
resized_thigh_alpha = cv2.resize(part_thigh_alpha, (self.w, new_thigh_height), interpolation=cv2.INTER_LINEAR)
|
||
|
|
resized_calf = cv2.resize(part_calf, (self.w, new_calf_height), interpolation=cv2.INTER_LINEAR)
|
||
|
|
resized_calf_alpha = cv2.resize(part_calf_alpha, (self.w, new_calf_height), interpolation=cv2.INTER_LINEAR)
|
||
|
|
|
||
|
|
new_bgr = np.vstack((top_part_bgr, resized_thigh, resized_calf))
|
||
|
|
new_bgr_alpha = np.vstack((top_part_bgr_alpha, resized_thigh_alpha, resized_calf_alpha))
|
||
|
|
|
||
|
|
if self.alpha is not None:
|
||
|
|
# 拼接 alpha 通道
|
||
|
|
# 合并 BGR 通道和 alpha 通道
|
||
|
|
new_image = np.dstack((new_bgr, new_bgr_alpha))
|
||
|
|
else:
|
||
|
|
new_image = new_bgr
|
||
|
|
return new_image
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
request_data = MannequinModel(
|
||
|
|
mannequins="aida-sys-image/models/male/dc36ce58-46c3-4b6f-8787-5ca7d6fc26e6.png",
|
||
|
|
scale=[0.75, 0.75],
|
||
|
|
bucket_name="test",
|
||
|
|
mannequin_name="mannequin_name"
|
||
|
|
)
|
||
|
|
service = MannequinEditService(request_data)
|
||
|
|
print(service())
|