from pprint import pprint import cv2 import mmcv import numpy as np import pandas as pd import torch import tritonclient.http as httpclient from minio import Minio from skimage import transform from app.core.config import MINIO_IP, MINIO_PORT, MINIO_SECURE, MINIO_ACCESS, MINIO_SECRET, ATT_TRITON_IP, ATT_TRITON_PORT from app.service.attribute_recognition.config.category_mapping import category_mapping def Merge(dict1, dict2): res = {**dict1, **dict2} return res class Rescale(object): """Rescale the image in a sample to a given size. Args: output_size (tuple or int): Desired output size. If tuple, output is matched to output_size. If int, smaller of image edges is matched to output_size keeping aspect ratio the same. """ def __init__(self, output_size): assert isinstance(output_size, (int, tuple)) self.output_size = output_size def __call__(self, image, landmarks): h, w = image.shape[:2] if isinstance(self.output_size, int): if h > w: new_h, new_w = self.output_size * h / w, self.output_size else: new_h, new_w = self.output_size, self.output_size * w / h else: new_h, new_w = self.output_size new_h, new_w = int(new_h), int(new_w) img = transform.resize(image, (new_h, new_w), mode='constant') # h and w are swapped for landmarks because for images, # x and y axes are axis 1 and 0 respectively landmarks = landmarks * [new_w / w, new_h / h] return img, landmarks class AttributeRecognition: def __init__(self): self.httpclient = httpclient.InferenceServerClient(url=f"{ATT_TRITON_IP}:{ATT_TRITON_PORT}") self.minio_client = Minio( f"{MINIO_IP}:{MINIO_PORT}", access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE) def get_image(self, url): # Get data of an object. # Read data from response. response = self.minio_client.get_object(url.split("/", 1)[0], url.split("/", 1)[1]) img = np.frombuffer(response.data, np.uint8) # 转成8位无符号整型 img = cv2.imdecode(img, cv2.IMREAD_COLOR) # 解码 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return img def preprocess(self, img_url): img = self.get_image(img_url) img = mmcv.imread(img) ori_shape = img.shape[:2] img_scale = (224, 224) scale_factor = [] img, x, y = mmcv.imresize(img, img_scale, return_scale=True) scale_factor.append(x) scale_factor.append(y) 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 def attribute(self, args, request_data): final_dict = {} for l in range(len(request_data.upload_img_path)): img_path = request_data.upload_img_path[l] sample = self.preprocess(img_path) category_model_path = args.category_model category_description = args.category_discription # 如果category为all-body 则通过category获取类别 ， 主要用去区分为连体裤和连衣裙 mapping_category = category_mapping[request_data.upload_img_category[l]] if mapping_category not in ["tops", "bottoms", "skirt", "dress", "outerwear"]: category_list = self.get_attribute(category_model_path, category_description, sample)['category'] else: category_list = [mapping_category] attr_dict = {} if len(category_list) >= 1: category = category_list[0] print(category) if category == 'tops' or category == "top": attr_dict = {'Item': "top"} for i in range(len(args.top_discription_list)): attr_description = args.top_discription_list[i] attr_model_path = args.top_model_list[i] present_dict = self.get_attribute(attr_model_path, attr_description, sample) attr_dict = Merge(attr_dict, present_dict) elif category == 'bottoms' or category == 'bottom': attr_dict = {'Item': "bottom", 'Type': ['Pants']} for i in range(len(args.bottom_discription_list)): attr_description = args.bottom_discription_list[i] attr_model_path = args.bottom_model_list[i] present_dict = self.get_attribute(attr_model_path, attr_description, sample) attr_dict = Merge(attr_dict, present_dict) elif category == 'skirt': attr_dict = {'Type': ['Skirt']} category = 'bottom' attr_dict['Item'] = category for i in range(len(args.bottom_discription_list)): attr_description = args.bottom_discription_list[i] attr_model_path = args.bottom_model_list[i] present_dict = self.get_attribute(attr_model_path, attr_description, sample) attr_dict = Merge(attr_dict, present_dict) elif category == 'dress': attr_dict = {'Item': category} for i in range(len(args.dress_discription_list)): attr_description = args.dress_discription_list[i] attr_model_path = args.dress_model_list[i] present_dict = self.get_attribute(attr_model_path, attr_description, sample) attr_dict = Merge(attr_dict, present_dict) elif category == 'outwear' or category == 'outerwear': attr_dict = {'Item': 'outer'} for i in range(len(args.outwear_discription_list)): attr_description = args.outwear_discription_list[i] attr_model_path = args.outwear_model_list[i] present_dict = self.get_attribute(attr_model_path, attr_description, sample) attr_dict = Merge(attr_dict, present_dict) elif category == 'jumpsuit': attr_dict = {'Item': 'jumpsuit'} for i in range(len(args.jumpsuit_discription_list)): attr_description = args.jumpsuit_discription_list[i] attr_model_path = args.jumpsuit_model_list[i] present_dict = self.get_attribute(attr_model_path, attr_description, sample) attr_dict = Merge(attr_dict, present_dict) final_dict[request_data.upload_img_id[l]] = attr_dict return final_dict def get_attribute(self, model_name, description, image): attr_type = pd.read_csv(description) inputs = [ httpclient.InferInput("input__0", image.shape, datatype="FP32") ] inputs[0].set_data_from_numpy(image, binary_data=True) results = self.httpclient.infer(model_name=model_name, inputs=inputs) inference_output = torch.from_numpy(results.as_numpy(f"output__0")) scores = inference_output.detach().numpy() colattr = list(attr_type['labelName']) task = attr_type['taskName'][0] if task == 'category': maxsc = np.max(scores) if maxsc > 0: indexs = np.argwhere(scores == maxsc)[:, 1] else: attr = {'category': []} return attr elif task[:-2] == 'Print' or task[:-2] == 'Material': # maxsc = np.max(scores) scores[scores > 0] = 1 scores[scores <= 0] = 0 indexs = np.argwhere(scores == 1)[:, 1] # if maxsc> 0: # indexs = np.argwhere(scores == maxsc)[:,1] else: maxsc = np.max(scores) if maxsc > 0: indexs = np.argwhere(scores == maxsc)[:, 1] else: indexs = [] # scores[scores > 0] = 1 # scores[scores <= 0] = 0 # indexs = np.argwhere(scores == 1)[:,1] attr = {} attr[task] = [] for i in range(len(indexs)): # print('i and indexs: ', i, indexs, indexs[i], type(indexs[i]), colattr[indexs[i]]) atr = colattr[indexs[i]] # print('corresponeding atr: ', atr) attr[task].append(atr) return attr if __name__ == '__main__': from app.service.attribute_recognition import const_debug request_data = {'upload_img_path': ['./test_top1.jpg'], 'upload_img_id': ["2"], "update_img_category": ["TOP/ONE PIECE"]} service = AttributeRecognition() pprint(service.attribute(const_debug, request_data))