from __future__ import print_function

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from PIL import Image, ImageFont, ImageDraw
from torch import nn
from torchvision.models import resnet18

from config import parser
from datalist import data_transform_test as data_transform

# from pytorch_network.ShuffleNet import ShuffleNetG2


'''
细度分类，用BCNN
'''
font = ImageFont.truetype(font='simhei.ttf', size=50)


class Inference(object):
    def __init__(self, image_path):
        self.args = parser.parse_args()

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        torch.manual_seed(self.args.seed)
        self.device = torch.device("cuda" if use_cuda else "cpu")

        self.model = resnet18(pretrained=True)
        self.model.fc = nn.Linear(self.model.fc.in_features, 10)
        self.model = self.model.to(self.device)

        self.model = torch.nn.DataParallel(self.model).cuda()
        if self.args.resume == False:
            try:
                print("load the weight from pretrained-weight file")
                model_dict = self.model.state_dict()
                checkpoint = torch.load("weights/best.pth",
                                        map_location=self.device)
                pretrained_dict = checkpoint['model_state_dict']
                pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
                model_dict.update(pretrained_dict)
                self.model.load_state_dict(model_dict, strict=True)

                print("Finished to load the weight")
            except Exception as e:
                print("can not load weight")
                raise e

        if use_cuda:
            cudnn.benchmark = True

        '''
        数据准备
        '''
        self.image = Image.open(image_path).convert("RGB")
        self.image2 = self.image
        self.image = self.image.resize((224, 224), Image.BICUBIC)
        self.image = data_transform(self.image)
        self.image = self.image.unsqueeze(0)

        self.predict()

    def predict(self):

        with torch.no_grad():
            # self.image = torch.from_numpy(self.image)
            '''
            加一个维度之后，才可以送入模型
            '''
            # self.image = Variable(torch.unsqueeze(self.image, dim=0).float(), requires_grad=False)
            if self.args.use_cuda:
                self.image = self.image.cuda()

            output = self.model(self.image)
            pred = torch.argmax(output, 1)
            print(int(pred.item()))

        '''
        预测类别
        '''
        self.classes = ["beagle", "dachshund", "dalmatian", "jindo", "maltese", "pomeranian", "retriever",
                        "ShihTzu", "toypoodle", "Yorkshireterrier"]
        print(self.classes[pred.item()])

        draw = ImageDraw.Draw(self.image2)
        draw.text((0, 0), str(self.classes[pred.item()]), fill=(255, 0, 0), font=font)
        self.image2.show()
        self.image2.save("classification.jpg")


if __name__ == "__main__":
    path = "G:/datasets/分类/dog_class_num=10/val\Yorkshireterrier\9.116879290-origpic-672ff5.jpg"
    Inference(path)