134 lines
5.0 KiB
Python
134 lines
5.0 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import tritonclient.http as httpclient
|
|
import requests
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from app.service.outfit_matcher.foco import extract_main_colors
|
|
|
|
|
|
class OutfitMatcherHon:
|
|
def __init__(self, outfits):
|
|
self.outfits = outfits
|
|
self.tritonclient = httpclient.InferenceServerClient(url="localhost:8000")
|
|
|
|
@staticmethod
|
|
def imnormalize(img, mean, std, to_rgb=True):
|
|
"""Normalize an image with mean and std.
|
|
|
|
Args:
|
|
img (ndarray): Image to be normalized.
|
|
mean (ndarray): The mean to be used for normalize.
|
|
std (ndarray): The std to be used for normalize.
|
|
to_rgb (bool): Whether to convert to rgb.
|
|
|
|
Returns:
|
|
ndarray: The normalized image.
|
|
"""
|
|
img = img.copy().astype(np.float32)
|
|
assert img.dtype != np.uint8
|
|
mean = np.float64(mean.reshape(1, -1))
|
|
stdinv = 1 / np.float64(std.reshape(1, -1))
|
|
if to_rgb:
|
|
cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) # inplace
|
|
cv2.subtract(img, mean, img) # inplace
|
|
cv2.multiply(img, stdinv, img) # inplace
|
|
return img
|
|
|
|
@staticmethod
|
|
def load_image(img_path):
|
|
if 'http' in img_path:
|
|
file = requests.get(img_path)
|
|
image = cv2.imdecode(np.fromstring(file.content, np.uint8), 1)
|
|
image = Image.fromarray(image.astype('uint8'), 'RGB')
|
|
else:
|
|
image = Image.open(img_path).convert('RGB')
|
|
return np.array(image)
|
|
|
|
@staticmethod
|
|
def resize_image(img):
|
|
"""
|
|
Args:
|
|
img: ndarray (height, width, channel)
|
|
"""
|
|
resized_img = cv2.resize(img, (224, 224), dst=None, interpolation=1)
|
|
return resized_img
|
|
|
|
@staticmethod
|
|
def pad_array(input_value):
|
|
"""pad List of Array into same batch size
|
|
|
|
Args:
|
|
input_value: List of numpy arrary need to be padded
|
|
|
|
Returns:
|
|
Tensor: [batch_dim, max_dim, original_tensor_size]
|
|
"""
|
|
max_dim = max([len(x) for x in input_value])
|
|
mask = np.zeros((len(input_value), max_dim), dtype=np.float32)
|
|
|
|
# Pad each array
|
|
padded_arrays = []
|
|
for i, array in enumerate(input_value):
|
|
# Compute padding amount along the pad dimension
|
|
pad_dim = max_dim - array.shape[0]
|
|
consistent_shape = array.shape[1:]
|
|
pad_widths = [(0, pad_dim)] + [(0, 0)] * len(consistent_shape)
|
|
padded_array = np.pad(array, pad_widths, mode='constant', constant_values=0)
|
|
padded_arrays.append(padded_array)
|
|
|
|
mask[i, array.shape[0]:] = float("-inf")
|
|
|
|
# Stack the padded arrays and change the dimension
|
|
batched_arrays = np.stack(padded_arrays, axis=0)
|
|
return batched_arrays, mask
|
|
|
|
def preprocess(self):
|
|
outfit_images = []
|
|
outfit_colors = []
|
|
for outfit in self.outfits:
|
|
images = []
|
|
colors = []
|
|
for item in outfit["items"]:
|
|
image = self.load_image(item["image_path"])
|
|
image = self.resize_image(image)
|
|
normalized_image = self.imnormalize(image,
|
|
mean=np.array([208.32996145, 201.28227452, 198.47047691], dtype=np.float32),
|
|
std=np.array([75.48939648, 80.47423057, 82.21144189], dtype=np.float32))
|
|
images.append(normalized_image.transpose(2, 0, 1))
|
|
color = extract_main_colors(image)
|
|
colors.append(color)
|
|
images = np.stack(images, axis=0)
|
|
outfit_images.append(images) # List[(items, 3, 224, 224)]
|
|
colors = np.stack(colors, axis=0)
|
|
outfit_colors.append(colors)
|
|
outfit_images, mask = self.pad_array(outfit_images)
|
|
outfit_colors, _ = self.pad_array(outfit_colors)
|
|
return outfit_images, outfit_colors, mask
|
|
|
|
def get_result(self, outfits):
|
|
# start = time.time()
|
|
image, color, mask = self.preprocess()
|
|
# print(start - time.time())
|
|
# transformed_img = image.astype(np.float32)
|
|
# 输入集
|
|
inputs = [
|
|
httpclient.InferInput("input__0", image.shape, datatype="FP32"),
|
|
httpclient.InferInput("input__1", color.shape, datatype="FP32"),
|
|
httpclient.InferInput("input__2", mask.shape, datatype="FP32"),
|
|
]
|
|
inputs[0].set_data_from_numpy(image.astype(np.float32), binary_data=True)
|
|
inputs[1].set_data_from_numpy(color.astype(np.float32), binary_data=True)
|
|
inputs[2].set_data_from_numpy(mask.astype(np.float32), binary_data=True)
|
|
# 输出集
|
|
outputs = [
|
|
httpclient.InferRequestedOutput("output__0", binary_data=True),
|
|
]
|
|
results = self.tritonclient.infer(model_name="outfit_matcher", inputs=inputs, outputs=outputs)
|
|
# 推理
|
|
# 取结果
|
|
inference_output1 = torch.from_numpy(results.as_numpy("output__0"))
|
|
return inference_output1 # Shape (N, 1)
|