from __future__ import division
# 导入精确除法
from model import *
from utils.utils import *
from utils.datasets import *

import os
import time
import datetime
import argparse
import random
import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import NullLocator

parser = argparse.ArgumentParser()
parser.add_argument('--image_folder', type=str, default='data/samples', help='data sample')
parser.add_argument('--batch_size', type=int, default=1, help='size of each image batch')
parser.add_argument('--config_path', type=str, default='cfg/WPOD_NET.cfg', help='path to model config file')
parser.add_argument('--weights_path', type=str, default='weights/WPDF_NET.weights', help='path to weights file')
parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads')
parser.add_argument('--img_size', type=int, default=416, help='size of each images')
parser.add_argument('--use_cuda', type=bool, default=True, help='whether to use cuda')

opt = parser.parse_args()
print(opt)

os.makedirs('output', exist_ok=True)

cuda = opt.use_cuda and torch.cuda.is_available()


model = WPOD_NET(opt.config_path)
model.load_weights(opt.weights_path)
if cuda:
    model.cuda()
model.eval()
dataloader = DataLoader(ImageFolder(opt.image_folder, img_size=opt.img_size),
                        batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_cpu)
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
imgs = []
img_detections = []
print("\nPerforming object detection: ")
prev_time = time.time()
for batch_i, (img_paths, input_imgs) in enumerate(dataloader):
    input_imgs = Variable(input_imgs.type(Tensor))
    with torch.no_grad():
        detections = model(input_imgs)
        detections = non_max_suppression(detections)
        # Log progress
        current_time = time.time()
        inference_time = datetime.timedelta(seconds=current_time - prev_time)
        prev_time = current_time
        print('\t+ Batch %d, Inference Time: %s' % (batch_i, inference_time))  # 处理一张图片所需的时间
        imgs.extend(img_paths)
        img_detections.extend(detections)

# Bounding-box colors
cmap = plt.get_cmap('tab20b')
colors = [cmap(i) for i in np.linspace(0, 1, 20)]
print('\nSavinng images: ')
for img_i, (path, detections) in enumerate(zip(imgs, img_detections)):

    print("(%d) Image: '%s'" % (img_i, path))

    img = np.array(Image.open(path))
    plt.figure()
    fig, ax = plt.subplots(1)
    ax.imshow(img)

    # The amount of padding that was added
    pad_x = max(img.shape[0] - img.shape[1], 0) * (opt.img_size / max(img.shape))
    pad_y = max(img.shape[1] - img.shape[0], 0) * (opt.img_size / max(img.shape))
    # Image height and width after padding is removed
    unpad_h = opt.img_size - pad_y
    unpad_w = opt.img_size - pad_x
    if detections is not None:
        
        bbox_colors = random.sample(colors, 1)
        v1 = detections[0].item()
        v2 = detections[1].item()
        v3 = detections[2].item()
        v4 = detections[3].item()
        v5 = detections[4].item()
        v6 = detections[5].item()
        v7 = detections[6].item()
        v8 = detections[7].item()
        print('\t conf: %.5f' % v1)
        q = np.array([[-0.5, -0.5, 0.5, 0.5], [-0.5, 0.5, 0.5, -0.5]])
        pol = (np.dot(np.array([[[max(0, v3), v4], [v5, max(v6,0)]]]), q) + np.array([[v7],[v8]])).reshape(4,2)
        bbox = patches.Polygon(pol)
        ax.add_patch(bbox)
    plt.axis('off')
    plt.gca().xaxis.set_major_locator(NullLocator())  # gca获取当前绘图区域 xaxis.set_major_locator(NullLocator())删除x坐标的刻度显示
    plt.gca().yaxis.set_major_locator(NullLocator())
    plt.savefig('output/%d.png' % (img_i), bbox_inches='tight', pad_inches=0.0)
    plt.close()