import cv2
import mmcv
import numpy as np
import tritonclient.http as httpclient
import torch


def preprocess(img):
    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, ori_shape


def get_attribute(model_save_name, sample):
    triton_client = httpclient.InferenceServerClient(url=f"10.1.1.240:10020")
    inputs = [
        httpclient.InferInput("input__0", sample.shape, datatype="FP32")
    ]
    inputs[0].set_data_from_numpy(sample, binary_data=True)
    results = triton_client.infer(model_name=model_save_name, inputs=inputs)
    inference_output = torch.from_numpy(results.as_numpy(f"output__0"))
    scores = inference_output.detach().numpy()
    print(f"{model} is ok")


image, shape = preprocess(cv2.imread("8365b04b-69d0-4f6d-a61b-a8cc136256263422b10ff45d4ae4e10be3fc25a9c36b.jpg"))
# get_attribute("attr_retrieve_D_sleeve_shape", image)

model_list = ['top_length',
              'top_type',
              'top_Sleeve_length',
              'top_Sleeve_shape',
              'top_Sleeve_shoulder',
              'top_Neckline',
              'top_print',
              'top_material',
              'top_Softness',
              'top_Design',
              'top_optype',
              'top_Silhouette',
              'top_Collar',
              'bottom_sub-Type',
              'bottom_length',
              'bottom_print',
              'bottom_material',
              'bottom_Softness_B',
              'bottom_Silhouette_B',
              'bottom_OPType_B',
              'bottom_design',
              'outwear_outer_length',
              #   'outwear_2_outer_type',
              'outwear_outer_sleeve_length',
              'outwear_outer_sleeve_shape',
              'outwear_outer_sleeve_shoulder',
              'outwear_outer_collar',
              'outwear_print',
              'outwear_material',
              'outwear_outer_softness',
              'outwear_outer_design',
              #   'outwear_11_outer_opening',
              'outwear_outer_optype',
              'outwear_outer_silhouette',
              'jumpsuit_length',
              'jumpsuit_sleeve_length',
              'jumpsuit_sleeve_shape',
              'jumpsuit_sleeve_shoulder',
              'jumpsuit_neckline',
              'jumpsuit_collar',
              'jumpsuit_print',
              'jumpsuit_material',
              'jumpsuit_softness',
              'jumpsuit_design',
              'jumpsuit_optype',
              'jumpsuit_subtype',
              'dress_length',
              'dress_sleeve_length',
              'dress_sleeve_shape',
              'dress_sleeve_shoulder',
              'dress_neckline',
              'dress_print',
              'dress_collar',
              'dress_material',
              'dress_design',
              'dress_softness',
              'dress_silohouette12',
              #   'dress_'
              'dress_type'
              ]
except_model = []
# 'attr_recong_B_optype', 'attr_recong_D_length', 'attr_recong_O_length', 'attr_recong_O_optype', 'attr_recong_material',
# 'attr_recong_print', 'attr_retrieve_D_collar', 'attr_retrieve_D_design', 'attr_retrieve_D_neckline', 'attr_retrieve_D_silohouette',
# 'attr_retrieve_D_sleeve_shape', 'attr_retrieve_D_sleeve_shoulder', 'attr_retrieve_D_type', 'attr_retrieve_T_length', 'attr_retrieve_T_optype'

# for model in model_list:
#     get_attribute(model, image)

for model in model_list:
    try:
        get_attribute(model, image)
    except Exception as e:
        print(e)
        except_model.append(model)
print(except_model)