"""Adapted from:
    @longcw faster_rcnn_pytorch: https://github.com/longcw/faster_rcnn_pytorch
    @rbgirshick py-faster-rcnn https://github.com/rbgirshick/py-faster-rcnn
    Licensed under The MIT License [see LICENSE for details]
"""
from __future__ import print_function

import argparse
import copy
import csv
import os
import pickle
import sys
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.utils.data as data
from pyheatmap.heatmap import HeatMap
from torch.autograd import Variable

from data import VOC_CLASSES as labelmap
from data import VOC_ROOT, BaseTransform, VOCAnnotationTransform, VOCDetection
from layers.box_utils import nms, py_cpu_nms, soft_nms
from ssd import build_ssd

if sys.version_info[0] == 2:
    import xml.etree.cElementTree as ET
else:
    import xml.etree.ElementTree as ET


def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")


parser = argparse.ArgumentParser(
    description='Single Shot MultiBox Detector Evaluation')
parser.add_argument('--trained_model',
                    default='weights/ssd300_mAP_77.43_v2.pth', type=str,
                    help='Trained state_dict file path to open')
parser.add_argument('--save_folder', default='eval/', type=str,
                    help='File path to save results')
parser.add_argument('--confidence_threshold', default=0.01, type=float,
                    help='Detection confidence threshold')
parser.add_argument('--top_k', default=200, type=int,
                    help='Further restrict the number of predictions to parse')
parser.add_argument('--cuda', default=True, type=str2bool,
                    help='Use cuda to train model')
parser.add_argument('--voc_root', default=VOC_ROOT,
                    help='Location of VOC root directory')
parser.add_argument('--cleanup', default=True, type=str2bool,
                    help='Cleanup and remove results files following eval')
parser.add_argument('--debug',
                    action='store_true',
                    help='small data for debug.')
parser.add_argument('--only_best_pred',
                    action='store_true',
                    help='only draw the best pred box.')
parser.add_argument('--pred_threshold',
                    default=0.0,
                    type=float,
                    help='set the threshold of pred boxes to draw')
parser.add_argument('--range',
                    required=True,
                    type=int,
                    help='set the range value when eval.')

args = parser.parse_args()

if not os.path.exists(args.save_folder):
    os.mkdir(args.save_folder)

if torch.cuda.is_available():
    if args.cuda:
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
    if not args.cuda:
        print("WARNING: It looks like you have a CUDA device, but aren't using \
              CUDA.  Run with --cuda for optimal eval speed.")
        torch.set_default_tensor_type('torch.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')

annopath = os.path.join(args.voc_root, 'Annotations', '%s.xml')
imgpath = os.path.join(args.voc_root, 'JPEGImages', '%s.jpg')
# imgsetpath = os.path.join(args.voc_root, 'VOC2007', 'ImageSets',
#                           'Main', '{:s}.txt')
imgsetpath = os.path.join(args.voc_root, 'ImageSets',
                          'Main', '{}.txt')
YEAR = '2007'
# devkit_path = args.voc_root + 'VOC' + YEAR
devkit_path = args.voc_root
dataset_mean = (104, 117, 123)
set_type = 'test'


class Timer(object):
    """A simple timer."""
    def __init__(self):
        self.total_time = 0.
        self.calls = 0
        self.start_time = 0.
        self.diff = 0.
        self.average_time = 0.

    def tic(self):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        self.start_time = time.time()

    def toc(self, average=True):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.average_time = self.total_time / self.calls
        if average:
            return self.average_time
        else:
            return self.diff


def parse_rec(filename):
    """ Parse a PASCAL VOC xml file """
    tree = ET.parse(filename)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        obj_struct['name'] = obj.find('name').text
        obj_struct['pose'] = obj.find('pose').text
        obj_struct['truncated'] = int(obj.find('truncated').text)
        obj_struct['difficult'] = int(obj.find('difficult').text)
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(bbox.find('xmin').text) - 1,
                              int(bbox.find('ymin').text) - 1,
                              int(bbox.find('xmax').text) - 1,
                              int(bbox.find('ymax').text) - 1]
        objects.append(obj_struct)

    return objects


def get_output_dir(name, phase):
    """Return the directory where experimental artifacts are placed.
    If the directory does not exist, it is created.
    A canonical path is built using the name from an imdb and a network
    (if not None).
    """
    filedir = os.path.join(name, phase)
    if not os.path.exists(filedir):
        os.makedirs(filedir)
    return filedir


def get_voc_results_file_template(image_set, cls):
    # VOCdevkit/VOC2007/results/det_test_aeroplane.txt
    filename = 'det_' + image_set + '_%s.txt' % (cls)
    filedir = os.path.join(devkit_path, 'results')
    if not os.path.exists(filedir):
        os.makedirs(filedir)
    path = os.path.join(filedir, filename)
    return path


def write_voc_results_file(all_boxes, dataset):
    for cls_ind, cls in enumerate(labelmap):
        print('Writing {:s} VOC results file'.format(cls))
        filename = get_voc_results_file_template(set_type, cls)
        with open(filename, 'wt') as f:
            for im_ind, index in enumerate(dataset.ids):
                dets = all_boxes[cls_ind+1][im_ind]
                if dets == []:
                    continue
                # the VOCdevkit expects 1-based indices
                for k in range(dets.shape[0]):
                    f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.
                            format(index[1], dets[k, -1],
                                   dets[k, 0] + 1, dets[k, 1] + 1,
                                   dets[k, 2] + 1, dets[k, 3] + 1))

def write_csv(name_model, ap, recall, precision):
    f = open('performance.csv','a', newline='')
    csv_writer = csv.writer(f)
    list_performance = [name_model, ap, recall, precision]
    csv_writer.writerow(['name', 'ap', 'recall', 'precision'])
    csv_writer.writerow(list_performance)
    f.close()

def do_python_eval(output_dir='output', use_07=False):
    cachedir = os.path.join(devkit_path, 'annotations_cache')
    aps = []
    # The PASCAL VOC metric changed in 2010
    use_07_metric = use_07
    print('VOC07 metric? ' + ('Yes' if use_07_metric else 'No'))
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)
    for i, cls in enumerate(labelmap): # labelmap只有nodule
        filename = get_voc_results_file_template(set_type, cls)
        rec, prec, ap, recall, precision = voc_eval(
           filename, annopath, imgsetpath.format(set_type), cls, cachedir,
           ovthresh=0.5, use_07_metric=use_07_metric)

        aps += [ap]
        print('AP for {} = {:.4f}'.format(cls, ap))
        with open(os.path.join(output_dir, cls + '_pr.pkl'), 'wb') as f:
            pickle.dump({'rec': rec, 'prec': prec, 'ap': ap}, f)

    # print('Mean AP = {:.4f}'.format(np.mean(aps)))
    print('Recall = {:.4f}'.format(recall))
    print('Precision = {:.4f}'.format(precision))
    print('~~~~~~~~')
    print('Results:')
    for ap in aps:
        print('{:.3f}'.format(ap))
    # print('{:.3f}'.format(np.mean(aps)))
    print('~~~~~~~~')
    print('')
    print('--------------------------------------------------------------')
    print('Results computed with the **unofficial** Python eval code.')
    print('Results should be very close to the official MATLAB eval code.')
    print('--------------------------------------------------------------')
    write_csv(args.trained_model, ap, recall, precision)

def voc_ap(rec, prec, use_07_metric=True):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:True).
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], rec, [1.]))
        mpre = np.concatenate(([0.], prec, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap


def voc_eval(detpath,
             annopath,
             imagesetfile,
             classname,
             cachedir,
             ovthresh=0.5,
             use_07_metric=True):
    """rec, prec, ap = voc_eval(detpath,
                           annopath,
                           imagesetfile,
                           classname,
                           [ovthresh],
                           [use_07_metric])
Top level function that does the PASCAL VOC evaluation.
detpath: Path to detections
   detpath.format(classname) should produce the detection results file.
annopath: Path to annotations
   annopath.format(imagename) should be the xml annotations file.
imagesetfile: Text file containing the list of images, one image per line.
classname: Category name (duh)
cachedir: Directory for caching the annotations
[ovthresh]: Overlap threshold (default = 0.5)
[use_07_metric]: Whether to use VOC07's 11 point AP computation
   (default True)
"""
# assumes detections are in detpath.format(classname)
# assumes annotations are in annopath.format(imagename)
# assumes imagesetfile is a text file with each line an image name
# cachedir caches the annotations in a pickle file
# first load gt
    if not os.path.isdir(cachedir):
        os.mkdir(cachedir)
    cachefile = os.path.join(cachedir, 'annots.pkl')
    # read list of images
    with open(imagesetfile, 'r') as f:
        lines = f.readlines()
    imagenames = [x.strip() for x in lines]
    if not os.path.isfile(cachefile):
        # load annots
        recs = {}
        for i, imagename in enumerate(imagenames):
            recs[imagename] = parse_rec(annopath % (imagename))
            if i % 100 == 0:
                print('Reading annotation for {:d}/{:d}'.format(
                   i + 1, len(imagenames)))
        # save
        print('Saving cached annotations to {:s}'.format(cachefile))
        with open(cachefile, 'wb') as f:
            pickle.dump(recs, f)
    else:
        # load
        with open(cachefile, 'rb') as f:
            recs = pickle.load(f)

    # extract gt objects for this class
    class_recs = {}
    npos = 0
    for imagename in imagenames:
        R = [obj for obj in recs[imagename] if obj['name'] == classname]
        bbox = np.array([x['bbox'] for x in R])
        difficult = np.array([x['difficult'] for x in R]).astype(np.bool)
        det = [False] * len(R)
        npos = npos + sum(~difficult)
        class_recs[imagename] = {'bbox': bbox,
                                 'difficult': difficult,
                                 'det': det}

    num_positive = 0
    for each_class_recs in class_recs:
        num_positive = num_positive + class_recs[each_class_recs]['bbox'].shape[0]

    # read dets
    detfile = detpath.format(classname)
    with open(detfile, 'r') as f:
        lines = f.readlines()
    if any(lines) == 1:

        splitlines = [x.strip().split(' ') for x in lines]
        image_ids = [x[0] for x in splitlines]
        confidence = np.array([float(x[1]) for x in splitlines])
        BB = np.array([[float(z) for z in x[2:]] for x in splitlines])

        # sort by confidence
        sorted_ind = np.argsort(-confidence)
        sorted_scores = np.sort(-confidence)
        BB = BB[sorted_ind, :]
        image_ids = [image_ids[x] for x in sorted_ind]

        # go down dets and mark TPs and FPs
        nd = len(image_ids)
        tp = np.zeros(nd)
        fp = np.zeros(nd)
        for d in range(nd):
            R = class_recs[image_ids[d]]
            bb = BB[d, :].astype(float)
            ovmax = -np.inf
            BBGT = R['bbox'].astype(float)
            if BBGT.size > 0:
                # compute overlaps
                # intersection
                ixmin = np.maximum(BBGT[:, 0], bb[0])
                iymin = np.maximum(BBGT[:, 1], bb[1])
                ixmax = np.minimum(BBGT[:, 2], bb[2])
                iymax = np.minimum(BBGT[:, 3], bb[3])
                iw = np.maximum(ixmax - ixmin, 0.)
                ih = np.maximum(iymax - iymin, 0.)
                inters = iw * ih
                uni = ((bb[2] - bb[0]) * (bb[3] - bb[1]) +
                       (BBGT[:, 2] - BBGT[:, 0]) *
                       (BBGT[:, 3] - BBGT[:, 1]) - inters)
                overlaps = inters / uni
                ovmax = np.max(overlaps)
                jmax = np.argmax(overlaps)

            if ovmax > ovthresh:
                if not R['difficult'][jmax]:
                    if not R['det'][jmax]:
                        tp[d] = 1.
                        R['det'][jmax] = 1
                    else:
                        fp[d] = 1.
            else:
                fp[d] = 1.

        # compute precision recall
        fp = np.cumsum(fp)
        tp = np.cumsum(tp)
        rec = tp / float(npos)
        # avoid divide by zero in case the first detection matches a difficult
        # ground truth
        prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
        ap = voc_ap(rec, prec, use_07_metric)

        # save the rec and prec
        with open('performance_rec.csv', 'w') as f:
            f_csv = csv.writer(f)
            f_csv.writerow(list(rec))

        with open('performance_prec.csv', 'w') as f:
            f_csv = csv.writer(f)
            f_csv.writerow(list(prec))

    else:
        rec = -1.
        prec = -1.
        ap = -1.

    count_tp = max(tp)
    count_fp = max(fp)
    recall = count_tp / (num_positive)
    precision = count_tp / (count_tp + count_fp)

    return rec, prec, ap, recall, precision


def test_net(save_folder, net, cuda, dataset, transform, top_k,
             im_size=512, thresh=0.05):
    num_images = len(dataset)
    # all detections are collected into:
    #    all_boxes[cls][image] = N x 5 array of detections in
    #    (x1, y1, x2, y2, score)
    all_boxes = [[[] for _ in range(num_images)]
                 for _ in range(len(labelmap)+1)]

    # timers
    _t = {'im_detect': Timer(), 'misc': Timer()}
    output_dir = get_output_dir('ssd300_120000', set_type)
    det_file = os.path.join(output_dir, 'detections.pkl')

    context_list = ['', 'u', 'd']
    img_range = args.range

    if img_range == 0:
        img_range = 1
        context_list = ['']

    for i in range(0, num_images):
        index, img_name, img_original, im, gt, h, w = dataset.pull_item(i)
        list_heat_map = []
        list_heat_map_weighted = []

        num_gt = gt.shape[0]
        list_gt_x_1 = []
        list_gt_y_1 = []
        list_gt_x_2 = []
        list_gt_y_2 = []
        for gt_index in range(num_gt):
            list_gt_x_1.append(int(gt[gt_index][0] * 512))
            list_gt_y_1.append(int(gt[gt_index][1] * 512))
            list_gt_x_2.append(int(gt[gt_index][2] * 512))
            list_gt_y_2.append(int(gt[gt_index][3] * 512))

        cls_dets_all = None
        flag_cls_dets_all = True
        for each_context in context_list:
            for each_range in range(1, img_range+1):

                if not each_context == '':
                    img_name_range = '{}_{}_{}.png'.format(img_name, each_context, each_range)
                    img_original, im, h, w = dataset.pull_item_eval(index, img_name, img_name_range)
                    
                if each_context == '':
                    img_original_center = copy.deepcopy(img_original)

                x = Variable(im.unsqueeze(0))
                if args.cuda and torch.cuda.is_available():
                    x = x.cuda()
                _t['im_detect'].tic()
                detections = net(x).data
                detect_time = _t['im_detect'].toc(average=False)

                back_colour = int(img_original[0,0,0] + 0.5)
                # skip j = 0, because it's the background class
                for j in range(1, detections.size(1)):
                    dets = detections[0, j, :]
                    mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
                    dets = torch.masked_select(dets, mask).view(-1, 5)
                    if dets.size(0) == 0:
                        continue
                    boxes = dets[:, 1:]
                    boxes[:, 0] *= w
                    boxes[:, 2] *= w
                    boxes[:, 1] *= h
                    boxes[:, 3] *= h
                    scores = dets[:, 0].cpu().numpy()
                    cls_dets = np.hstack((boxes.cpu().numpy(),
                                        scores[:, np.newaxis])).astype(np.float32,
                                                                        copy=False)
                    # clean up the dets
                    ## clean up background boxes
                    num_box = cls_dets.shape[0]
                    list_clean_up = []
                    for each_box_index in range(num_box):
                        each_box = cls_dets[each_box_index]
                        x_1 = int(each_box[0] + 0.5)
                        y_1 = int(each_box[1] + 0.5)
                        x_2 = int(each_box[2] + 0.5)
                        y_2 = int(each_box[3] + 0.5)
                        image_box = img_original[y_1: y_2, x_1: x_2]

                        range_index_0 = image_box.shape[0]
                        range_index_1 = image_box.shape[1]
                        
                        count_back_colour = 0
                        for index_0 in range(range_index_0):
                            for index_1 in range(range_index_1):
                                if image_box[index_0][index_1][0] == back_colour:
                                    count_back_colour += 1

                        condition_1 = x_1 < 0 or x_2 < 0 or y_1 < 0 or y_2 < 0
                        condition_2 = range_index_0 == 0 or range_index_1 == 0
                        if condition_2:
                            list_clean_up.append(each_box_index)
                            continue
                        if count_back_colour / (range_index_0 * range_index_1) > 0.5 or condition_1:
                            list_clean_up.append(each_box_index)

                    cls_dets = np.delete(cls_dets, list_clean_up, axis=0)

                    ## hard nms
                    # num_box = boxes.shape[0]
                    # keep = py_cpu_nms(boxes.cpu().numpy(), scores)
                    # list_clean_up = []
                    # for index_clean in range(num_box):
                    #     if not index_clean in keep:
                    #         list_clean_up.append(index_clean)
                    # cls_dets = np.delete(cls_dets, list_clean_up, axis=0)


                    ## soft nms
                    # cls_dets = soft_nms(torch.Tensor(cls_dets), score_threshold=0)
                    # cls_dets = cls_dets.cpu().numpy()

                    ## top-k
                    list_clean_up = list(range(cls_dets.shape[0]))
                    del(list_clean_up[0: top_k])
                    cls_dets = np.delete(cls_dets, list_clean_up, axis=0)


                    ## set thresh
                    num_box = cls_dets.shape[0]
                    list_clean_up = []
                    for each_box_index in range(num_box):
                        each_box = cls_dets[each_box_index]
                        if each_box[4] < thresh:
                            list_clean_up.append(each_box_index)
                    cls_dets = np.delete(cls_dets, list_clean_up, axis=0)

                    # if each_context == '':
                    #     cls_dets_center = cls_dets

                    # stock cls_dets
                    if flag_cls_dets_all:
                        cls_dets_all = cls_dets
                        flag_cls_dets_all = False
                        
                    else:
                        # cls_dets[:, 4] = cls_dets[:, 4] / 3 * (3 - each_range)
                        try:
                            cls_dets_all = np.concatenate((cls_dets_all, cls_dets), axis=0)
                        except ValueError:
                            pass

                    # show cleaned boxes
                    if args.debug:
                        num_box = cls_dets.shape[0]
                        for each_box_index in range(num_box):
                            each_box = cls_dets[each_box_index]

                            x_1 = int(each_box[0] + 0.5)
                            y_1 = int(each_box[1] + 0.5)
                            x_2 = int(each_box[2] + 0.5)
                            y_2 = int(each_box[3] + 0.5)
                            point_left_up = (x_1, y_1)
                            point_right_down = (x_2, y_2)

                            print(each_box[4])
                            lung_location = img_original[y_1: y_2, x_1: x_2, :]
                            cv2.imwrite('test/test_lung_location_{}.png'.format(each_box_index), lung_location)

                            cv2.rectangle(img_original, point_left_up, point_right_down, (0, 0, 255), 1)
                            font = cv2.FONT_HERSHEY_SIMPLEX
                            img_original = cv2.putText(img_original, ' {:.4f}'.format(each_box[4]), point_right_down, font, 0.4, (0, 0, 255))

                        for each_gt_x_1, each_gt_y_1, each_gt_x_2, each_gt_y_2, in zip(list_gt_x_1, list_gt_y_1, list_gt_x_2, list_gt_y_2):
                            point_left_up_gt = (each_gt_x_1, each_gt_y_1)
                            point_right_down_gt = (each_gt_x_2, each_gt_y_2)
                            cv2.rectangle(img_original, point_left_up_gt, point_right_down_gt, (0, 255, 0), 1)
                            cv2.imwrite('test/test_boxes_{}_{}.png'.format(each_context, each_range),
                                img_original, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])

                    # heat map
                    # heat_data = np.zeros((512, 512))
                    # num_box = cls_dets.shape[0]
                    # for each_box_index in range(num_box):
                    #     each_box = cls_dets[each_box_index]
                    #     x_1 = int(each_box[0] + 0.5)
                    #     y_1 = int(each_box[1] + 0.5)
                    #     x_2 = int(each_box[2] + 0.5)
                    #     y_2 = int(each_box[3] + 0.5)

                    #     heat_data[y_1: y_2, x_1: x_2] = heat_data[y_1: y_2, x_1: x_2] + each_box[4]

                    # if each_context == '' or each_context == 'd':
                    #     list_heat_map.append(heat_data)
                    # elif each_context == 'u':
                    #     list_heat_map.insert(0, heat_data)

                    # if each_context in ['d', 'u']:
                    #     weighted = (3 - each_range) / 3
                    # elif each_context == '':
                    #     weighted = 1
                    # list_heat_map_weighted.append(heat_data * weighted)


                    # if args.debug:
                    #     # sns.heatmap(heat_data, vmin=0, vmax=1, cmap='PuBuGn', cbar=False)
                    #     print('max heat: {}'.format(np.max(heat_data)))
                    #     sns.heatmap(heat_data, vmin=0, vmax=1)
                    #     plt.savefig('test/heatmap_{}_{}.png'.format(each_context, each_range))
                    #     plt.close()
                    #     print(img_name)

                if each_context == '':
                    break

        # sort cls_dets_all
        sort_index = np.argsort(cls_dets_all[:,4])[::-1]
        cls_dets_all = cls_dets_all[sort_index]

        # get average heatmap
        heat_data_average = np.zeros((512, 512))
        for each_heat_data in list_heat_map:
            heat_data_average = heat_data_average + each_heat_data
        heat_data_average = heat_data_average / len(list_heat_map)
        if args.debug:
            print('max heart average: {}'.format(np.max(heat_data_average)))
            sns.heatmap(heat_data_average, vmin=0, vmax=1)
            plt.savefig('test/heatmap_average.png')
            plt.close()

        # get weighted heatmap
        heat_data_weighted = np.zeros((512, 512))
        for each_heat_data_weighted in list_heat_map_weighted:
            heat_data_weighted = heat_data_weighted + each_heat_data_weighted
        heat_data_weighted = heat_data_weighted / 3
        if args.debug:
            print('max heart weighted: {}'.format(np.max(heat_data_weighted)))
            sns.heatmap(heat_data_weighted, vmin=0, vmax=1)
            plt.savefig('test/heatmap_weighted.png')
            plt.close()

        ## nms
        # num_box = cls_dets_all.shape[0]
        # keep = py_cpu_nms(cls_dets_all)
        # list_clean_up = []
        # for index_clean in range(num_box):
        #     if not index_clean in keep:
        #         list_clean_up.append(index_clean)
        # cls_dets_all = np.delete(cls_dets_all, list_clean_up, axis=0)

        # soft nms
        # cls_dets_all = soft_nms(torch.Tensor(cls_dets_all), score_threshold=0)
        # cls_dets_all = cls_dets_all.cpu().numpy()

        # show all range cls_dets
        num_box = cls_dets_all.shape[0]
        for each_box_index in range(num_box):
            each_box = cls_dets_all[each_box_index]

            x_1 = int(each_box[0] + 0.5)
            y_1 = int(each_box[1] + 0.5)
            x_2 = int(each_box[2] + 0.5)
            y_2 = int(each_box[3] + 0.5)
            point_left_up = (x_1, y_1)
            point_right_down = (x_2, y_2)
            cv2.rectangle(img_original_center, point_left_up, point_right_down, (0, 0, 255), 1)

        for each_gt_x_1, each_gt_y_1, each_gt_x_2, each_gt_y_2, in zip(list_gt_x_1, list_gt_y_1, list_gt_x_2, list_gt_y_2):
            point_left_up_gt = (each_gt_x_1, each_gt_y_1)
            point_right_down_gt = (each_gt_x_2, each_gt_y_2)
            cv2.rectangle(img_original_center, point_left_up_gt, point_right_down_gt, (0, 255, 0), 1)
            cv2.imwrite('test/test_boxes_all.png', img_original_center, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])

        num_det = cls_dets_all.shape[0]
        heatmap_cls_dets_all = np.zeros((512, 512))
        for index_det in range(num_det):
            x_1 = int(cls_dets_all[index_det][0] + 0.5)
            y_1 = int(cls_dets_all[index_det][1] + 0.5)
            x_2 = int(cls_dets_all[index_det][2] + 0.5)
            y_2 = int(cls_dets_all[index_det][3] + 0.5)

            heatmap_cls_dets_all[y_1: y_2, x_1: x_2] = heatmap_cls_dets_all[y_1: y_2, x_1: x_2] + cls_dets_all[index_det][4]
        sns.heatmap(heatmap_cls_dets_all, vmin=0, vmax=1)
        plt.savefig('test/heatmap_cls_dets_all.png')
        plt.close()

        ## set all_boxes[j][i] = cls_dets
        # cls_dets = cls_dets_center
        # num_det = cls_dets.shape[0]
        # for index_det in range(num_det):
        #     x_1 = int(cls_dets[index_det][0] + 0.5)
        #     y_1 = int(cls_dets[index_det][1] + 0.5)
        #     x_2 = int(cls_dets[index_det][2] + 0.5)
        #     y_2 = int(cls_dets[index_det][3] + 0.5)

        #     image_box = heat_data_weighted[y_1: y_2, x_1: x_2]
        #     # sns.heatmap(image_box, vmin=0, vmax=1)
        #     # plt.savefig('test/test_heat_box.png')
        #     # plt.close()

        #     image_box_average = np.mean(image_box)
        #     cls_dets[index_det][4] = image_box_average

        # ##
        # heat_all_boxes = np.zeros((512, 512))
        # for index_det in range(num_det):
        #     x_1 = int(cls_dets[index_det][0] + 0.5)
        #     y_1 = int(cls_dets[index_det][1] + 0.5)
        #     x_2 = int(cls_dets[index_det][2] + 0.5)
        #     y_2 = int(cls_dets[index_det][3] + 0.5)

        #     heat_all_boxes[y_1: y_2, x_1: x_2] = heat_all_boxes[y_1: y_2, x_1: x_2] + cls_dets[index_det][4]
        # sns.heatmap(heat_all_boxes, vmin=0, vmax=1)
        # plt.savefig('test/heatmap_all_boxes.png')
        # plt.close()

        #
        all_boxes[j][i] = cls_dets_all

        print('\rim_detect: {:d}/{:d} {:.3f}s'.format(i + 1,
                                                    num_images, detect_time), end='', flush=True)
    print('')

    with open(det_file, 'wb') as f:
        pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

    print('Evaluating detections')
    evaluate_detections(all_boxes, output_dir, dataset)


def evaluate_detections(box_list, output_dir, dataset):
    write_voc_results_file(box_list, dataset)
    do_python_eval(output_dir)


if __name__ == '__main__':
    # load net
    num_classes = len(labelmap) + 1                      # +1 for background
    net = build_ssd('test', 512, num_classes)            # initialize SSD

    if torch.cuda.is_available():
        net.load_state_dict(torch.load(args.trained_model))
    else:
        net.load_state_dict(torch.load(args.trained_model, map_location=torch.device('cpu')))

    net.eval()
    print('Finished loading model!')
    # load data
    dataset = VOCDetection(args.voc_root, set_type,
                           BaseTransform(512, dataset_mean),
                           VOCAnnotationTransform())
    if args.cuda and torch.cuda.is_available():
        net = net.cuda()
        cudnn.benchmark = True
    # evaluation
    test_net(args.save_folder, net, args.cuda, dataset,
             BaseTransform(net.size, dataset_mean), args.top_k, 512,
             thresh=args.confidence_threshold)
