import torch_mlu
import torch_mlu.core.mlu_model as ct
import torch_mlu.core.mlu_quantize as mlu_quantize
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import numpy as np
import torchvision.transforms as transforms
import argparse
from collections import OrderedDict
from PIL import Image
import cv2
import os
from torchvision import datasets, transforms
import warnings
warnings.filterwarnings("ignore")
from darknet import Darknet
import dataset
from utils import *
from yolact_packing.data import COLORS, cfg
from yolact_packing.yolact import Yolact
from yolact_packing.utils.augmentations import FastBaseTransform
from yolact_packing.utils import timer
from yolact_packing.layers.output_utils import postprocess

if __name__ == '__main__':
    # Parse configuration files
    frame = torch.from_numpy(cv2.imread('./test_img/img/000086.png')).float()
    batch = FastBaseTransform()(frame.unsqueeze(0))
    yolact_weightfile = './yolact_packing/weights/yolact-mlu.pth'


    torch.set_grad_enabled(False)
    torch.set_default_tensor_type('torch.FloatTensor')
    yolact_net = Yolact()
    yolact_net.detect.use_fast_nms = True
    yolact_net.detect.use_cross_class_nms = False
    cfg.mask_proto_debug = False

    state_dict = torch.load(yolact_weightfile, map_location=torch.device('cpu'))
    # yolact_net.load_state_dict(state_dict, strict=False)
    yolact_net.load_state_dict(state_dict)
    yolact_net.eval()
    pred = yolact_net(batch)
    print(pred)

    # 配置量化参数
    qconfig={'iteration': 1, 'use_avg':False, 'data_scale':1.0, 'firstconv':False, 'per_channel': False}
    # 调用量化接口
    quantized_model = mlu_quantize.quantize_dynamic_mlu(yolact_net,qconfig_spec=None, dtype='int8', gen_quant=True)
    # 设置为推理模式
    quantized_model = quantized_model.eval()

    # 执行推理生成量化值
    print(batch.shape)
    print(batch)

    quantized_model(batch)
	# 保存量化模型
    torch.save(quantized_model.state_dict(), './yolact_packing/weights/yolact-mlu-int8.pth')