import torch
from torchvision.models import resnet101
from fastapi import FastAPI
# 引入gradio
import gradio as gr
from PIL import Image
import numpy as np
import pickle
from base import BaseUI
# 测试集
from torchvision.transforms import transforms
class Test(BaseUI):
    def __init__(self):
        super(Test, self).__init__()

    def clear_gr(self):
        image = None
        text = ""
        return image, text

    def load_model(self):
        # 读取模型
        model = resnet101()
        state_dict = torch.load(self.model_path)
        model.load_state_dict(state_dict)
        model.eval()
        return model
    def prediciton_model(self,model,img):
        # 预测一个batchsize的数据
        pre_y = model(img)
        print(pre_y.shape)
        _, pre_y = torch.max(pre_y, axis=1)
        print("预测的的label下标是: ", pre_y)
        # 打开并读取文件
        with open(self.contrast_path, 'r', encoding='utf-8') as f:
            meta = f.readlines()

        for i in meta:
            list = i.split()
            if pre_y == int(list[0]):
                return list[2].rstrip(',')
            # 输出标签对应的分类名称
            # return str(class_names[pre_y.item()]).split("'")[1]

    def model_VGG16(self,value):
        # 读取模型
        model = self.load_model()
        if len(self.transforms_Resize) == 2:
            # 修改数据格式
            # 将三维数组转换为 PIL.Image 对象
            image_pil = Image.fromarray((value * 255).astype(np.uint8))
            transform = transforms.Compose([
                transforms.Resize((self.transforms_Resize[0],self.transforms_Resize[1])),
                transforms.ToTensor(),
                transforms.Normalize(mean=self.transforms_Normalize_mean,
                                     std=self.transforms_Normalize_std)
            ])
            img = transform(image_pil)
            img = torch.unsqueeze(img, 0)
            return self.prediciton_model(model, img)

        elif len(self.transforms_Resize) == 1:
            # 修改数据格式
            # 将三维数组转换为 PIL.Image 对象
            image_pil = Image.fromarray((value * 255).astype(np.uint8))
            transform = transforms.Compose([
                transforms.Resize(self.transforms_Resize[0]),
                transforms.CenterCrop(self.transforms_centercrop_value),
                transforms.ToTensor(),
                transforms.Normalize(mean=self.transforms_Normalize_mean,
                                     std=self.transforms_Normalize_std)
            ])
            img = transform(image_pil)
            img = torch.unsqueeze(img, 0)
            return self.prediciton_model(model, img)


    def webui_framework(self):
        # 构建gradio组件
        with gr.Blocks() as demo:
            with gr.Column():
                input = gr.Image()
                with gr.Row():
                    clear = gr.Button(value="Clear")
                    submit = gr.Button(value="Submit")
                text = gr.TextArea()
                submit.click(self.model_VGG16, input, text)
                clear.click(self.clear_gr, None, [input, text])


if __name__ == '__main__':
    with gr.Blocks() as demo:
        Test().webui_framework()
    demo.launch(share=True)