import logging import time import cv2 import mmcv import numpy as np import torch import tritonclient.http as httpclient from app.core.config import settings, DESIGN_MODEL_URL, DESIGN_MODEL_NAME from app.service.generate_image.utils.upload_sd_image import upload_stain_png_sd, upload_face_png_sd logger = logging.getLogger() def seg_preprocess(img_path): img = mmcv.imread(img_path) ori_shape = img.shape[:2] img_scale = ori_shape img = cv2.resize(img, img_scale) img = mmcv.imnormalize(img, mean=np.array([123.675, 116.28, 103.53]), std=np.array([58.395, 57.12, 57.375]), to_rgb=True) preprocessed_img = np.expand_dims(img.transpose(2, 0, 1), axis=0) return preprocessed_img, ori_shape def get_mask(image_obj): pre_mask = None if len(image_obj.shape) == 2: image_obj = cv2.cvtColor(image_obj, cv2.COLOR_GRAY2RGB) if image_obj.shape[2] == 4: # 如果是四通道 mask pre_mask = image_obj[:, :, 3] image_obj = image_obj[:, :, :3] Contour = get_contours(image_obj) Mask = np.zeros(image_obj.shape[:2], np.uint8) if len(Contour): Max_contour = Contour[0] Epsilon = 0.001 * cv2.arcLength(Max_contour, True) Approx = cv2.approxPolyDP(Max_contour, Epsilon, True) cv2.drawContours(Mask, [Approx], -1, 255, -1) else: Mask = np.ones(image_obj.shape[:2], np.uint8) * 255 if pre_mask is None: mask = Mask else: mask = cv2.bitwise_and(Mask, pre_mask) return image_obj, mask def get_contours(image): gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) Edge = cv2.Canny(gray, 10, 150) kernel = np.ones((5, 5), np.uint8) Edge = cv2.dilate(Edge, kernel=kernel, iterations=1) Edge = cv2.erode(Edge, kernel=kernel, iterations=1) Contour, _ = cv2.findContours(Edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) Contour = sorted(Contour, key=cv2.contourArea, reverse=True) return Contour # def seg_infer_image(image_obj): # image, ori_shape = seg_preprocess(image_obj) # client = httpclient.InferenceServerClient(url=f"{SEG_MODEL_URL}") # transformed_img = image.astype(np.float32) # # 输入集 # inputs = [ # httpclient.InferInput(DESIGN_MODEL_NAME, transformed_img.shape, datatype="FP32") # ] # inputs[0].set_data_from_numpy(transformed_img, binary_data=True) # # 输出集 # outputs = [ # httpclient.InferRequestedOutput("seg_input__0", binary_data=True), # ] # results = client.infer(model_name=SEGMENTATION['name'], inputs=inputs, outputs=outputs) # # 推理 # # 取结果 # inference_output1 = torch.from_numpy(results.as_numpy("seg_input__0")) # seg_result = seg_postprocess(inference_output1, ori_shape) # return seg_result def seg_infer_image(image_obj): image, ori_shape = seg_preprocess(image_obj) client = httpclient.InferenceServerClient(url=DESIGN_MODEL_URL) transformed_img = image.astype(np.float32) # 输入集 inputs = [ httpclient.InferInput("seg_input__0", transformed_img.shape, datatype="FP32") ] inputs[0].set_data_from_numpy(transformed_img, binary_data=True) # 输出集 outputs = [ httpclient.InferRequestedOutput("seg_output__0", binary_data=True), ] start_time = time.time() results = client.infer(model_name=DESIGN_MODEL_NAME, inputs=inputs, outputs=outputs) print(f"KNet infer time is :{time.time() - start_time}") # 推理 # 取结果 inference_output1 = results.as_numpy("seg_output__0") seg_result = seg_postprocess(inference_output1) return seg_result # def seg_postprocess(output, ori_shape): # seg_logit = F.interpolate(output, size=ori_shape, scale_factor=None, mode='bilinear', align_corners=False) # seg_logit = F.softmax(seg_logit, dim=1) # seg_pred = seg_logit.argmax(dim=1) # seg_pred = seg_pred.cpu().numpy() # return seg_pred # KNet def seg_postprocess(output): # seg_logit = F.interpolate(torch.tensor(output).float(), size=ori_shape, scale_factor=None, mode='bilinear', align_corners=False) # seg_logit = F.softmax(seg_logit, dim=1) # seg_pred = seg_logit.argmax(dim=1) # seg_pred = output.cpu().numpy() return output[0] def remove_background(image): image_obj, mask = get_mask(image) seg_result = seg_infer_image(image_obj) temp_front = seg_result == 1 front_mask = (mask * (temp_front + 0).astype(np.uint8)) temp_back = seg_result == 2 back_mask = (mask * (temp_back + 0).astype(np.uint8)) if len(front_mask.shape) > 2: front_mask = front_mask[0] else: front_mask = front_mask if len(back_mask.shape) > 2: back_mask = back_mask[0] else: back_mask = back_mask result_mask = front_mask + back_mask white_background = np.ones_like(image_obj) * 255 remove_bg_image = np.where(result_mask[:, :, None].astype(bool), image_obj, white_background) # cv2.imwrite("source_image", image) # cv2.imwrite("remove_bg_image", remove_bg_image) return remove_bg_image def bounding_box(image): edges = cv2.Canny(image, 50, 150) # 查找轮廓 contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # 初始化包围所有外接矩形的大矩形的坐标 x_min, y_min, x_max, y_max = float('inf'), float('inf'), -1, -1 # 遍历所有外接矩形，更新大矩形的坐标 for contour in contours: x, y, w, h = cv2.boundingRect(contour) x_min = min(x_min, x) y_min = min(y_min, y) x_max = max(x_max, x + w) y_max = max(y_max, y + h) # 根据大矩形的坐标来裁剪原始图像 result_image = image[y_min:y_max, x_min:x_max] # cv2.imshow("result_image", result_image) # cv2.waitKey(0) return result_image def stain_detection(image, user_id, category, tasks_id, spot_size=100): height, width, _ = image.shape corners = [ image[0:spot_size, 0:spot_size], # top left image[0:spot_size, width - spot_size:width], # top right # image[height - spot_size:height, 0:spot_size], # bottom left # image[height - spot_size:height, width - spot_size:width] # bottom right ] for index, corner in enumerate(corners): num_white_pixels = (corner == [255, 255, 255]).all(axis=2).sum() if num_white_pixels != spot_size * spot_size: logger.info(f"第{index + 1}发现了污点") return False, None # 中心区域检测 # 将图像转换为灰度图像 gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # 获取图像的中心点坐标 center_x, center_y = image.shape[1] // 2, image.shape[0] // 2 # 定义中心区域的大小 patch_size = 100 half_patch = patch_size // 2 # 提取中心区域 center_patch = gray[center_y - half_patch:center_y + half_patch, center_x - half_patch:center_x + half_patch] # 设置阈值来检测纯白区域 _, thresh = cv2.threshold(center_patch, 254, 255, cv2.THRESH_BINARY) # 寻找轮廓 contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # 过滤非连续的纯白区域 filtered_contours = [cnt for cnt in contours if cv2.contourArea(cnt) >= 300] # 根据面积进行过滤，这里假设面积大于30的为连续区域 # 如果有连续的纯白区域存在 if filtered_contours: # 将纯白区域替换为灰色 if settings.DEBUG: for cnt in filtered_contours: x, y, w, h = cv2.boundingRect(cnt) # 在原始图像上进行替换 image[y + center_y - half_patch:y + center_y - half_patch + h, x + center_x - half_patch:x + center_x - half_patch + w][thresh[y:y + h, x:x + w] == 255] = (128, 128, 128) # 显示图像 cv2.imshow('Marked Image', image) cv2.waitKey(0) logger.info("中心区域存在连续的纯白区域") is_pure_white = True else: is_pure_white = False if is_pure_white: return False, None if settings.DEBUG: for corner_coords in [ (0, 0), # (0, width - spot_size), (height - spot_size, 0), # (height - spot_size, width - spot_size) # 中心点 ]: 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) # 在原始图像上绘制矩形框 upload_stain_png_sd(dst, user_id=user_id, category=f"{category}", object_name=f"{tasks_id}.png") return True, image def generate_category_recognition(image, gender): def preprocess(img): img = mmcv.imread(img) img_scale = (224, 224) img = cv2.resize(img, img_scale) img = mmcv.imnormalize(img, mean=np.array([123.675, 116.28, 103.53]), std=np.array([58.395, 57.12, 57.375]), to_rgb=True) preprocessed_img = np.expand_dims(img.transpose(2, 0, 1), axis=0) return preprocessed_img preprocessed_img = preprocess(image) triton_client = httpclient.InferenceServerClient(url=DESIGN_MODEL_URL) inputs = [ httpclient.InferInput("input__0", preprocessed_img.shape, datatype="FP32") ] inputs[0].set_data_from_numpy(preprocessed_img, binary_data=True) results = triton_client.infer(model_name="attr_retrieve_category", inputs=inputs) inference_output = torch.from_numpy(results.as_numpy(f'output__0')) scores = inference_output.detach().numpy() import pandas as pd attr_type = pd.read_csv(settings.CATEGORY_PATH) colattr = list(attr_type['labelName']) # attr_type['taskName'][0] maxsc = np.max(scores[0][:5]) indexs = np.argwhere(scores == maxsc)[:, 1] category = colattr[indexs[0]] if gender == "Male": if category == 'Trousers' or category == 'Skirt': category = 'Bottoms' elif category == 'Blouse' or category == 'Dress': category = 'Tops' else: category = 'Outwear' return category, scores, image def autoLevels(img, cutoff=0.1): channels = img.shape[2] # h,w,ch table = np.zeros((1, 256, 3), np.uint8) for ch in range(channels): # cutoff=0.1, 计算 0.1%, 99.9% 分位的灰度值 low = np.percentile(img[:, :, ch], q=cutoff) # ch 通道, cutoff=0.1, 0.1 分位的灰度值 high = np.percentile(img[:, :, ch], q=100 - cutoff) # 99.9 分位的灰度值, [0, high] 占比99.9% # 输入动态线性拉伸 Sin = min(max(low, 0), high - 2) # Sin, 黑场阈值, 0<=Sin 0: img_out = img * (1 - alpha) + alpha * 255.0 else: img_out = img * (1 + alpha) return np.array(img_out, dtype='uint8') # 14.14 Photoshop 自动色阶调整算法 def face_detect_pic(image, user_id, category, tasks_id): # 1、转灰度图 gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) # cv2.imshow("gray", gray) # 2、训练一组人脸 FACE_CLASSIFIER = "" face_detector = cv2.CascadeClassifier(FACE_CLASSIFIER) # 3、检测人脸（用灰度图检测，返回人脸矩形坐标(4个角)） faces_rect = face_detector.detectMultiScale(gray, 1.05, 3) if settings.DEBUG: 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) # 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) # 画出矩形框 upload_face_png_sd(dst, user_id=user_id, category=f"{category}", object_name=f"{tasks_id}.png") return len(faces_rect) if __name__ == '__main__': # Photoshop 自动色阶调整算法 img = cv2.imread("2.png", flags=1) # 读取彩色图像 gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 转换为灰度图像 print("cutoff={}, minG={}, maxG={}".format(0.0, gray.min(), gray.min())) # 色阶手动调整 # equManual = levelsAdjust(img, 63, 205, 0.8, 10, 245) # 手动调节 # 色阶自动调整 cutoff = 1.0 # 截断比例, 建议范围 [0.0,1.0] # cv2.imwrite("source.png", img) equAuto = autoLevels(img, cutoff) # cv2.imwrite("levels.png", equAuto) luminance = luminance_adjust(0.3, equAuto) # cv2.imwrite("luminance.png", luminance) # # # 将图像转换为灰度 # gray = cv2.cvtColor(luminance, cv2.COLOR_BGR2GRAY) # # # 使用Canny边缘检测算法检测图像的边缘 # edges = cv2.Canny(gray, 150, 200) # # # 对边缘进行膨胀操作，增强轮廓 # kernel = np.ones((1, 1), np.uint8) # dilated_edges = cv2.dilate(edges, kernel, iterations=1) # # # 创建一个与原始图像相同大小的空白图像 # # result = np.zeros_like(luminance) # # # 将增强后的轮廓叠加到原始图像上 # luminance[dilated_edges != 0] = (255, 255, 255) remove_bg_img = remove_background(luminance) # cv2.imwrite("remove_bg_img.png", remove_bg_img) # print(1) cv2.imshow("source", img) cv2.imshow("levels", equAuto) cv2.imshow("luminance", luminance) # cv2.imshow("dilated_edges", luminance) cv2.imshow("remove_bg_img", remove_bg_img) cv2.waitKey(0) image = cv2.imread("1.png") remove_background(image)