from __future__ import division

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import cv2

# 获取任意给定图像中存在的类别
def unique(tensor):
    tensor_np =tensor.cpu().numpy()
    unique_np = np.unique(tensor_np)
    unique_tensor = torch.from_numpy(unique_np)

    tensor_res = tensor.new(unique_tensor.shape)
    tensor_res.copy_(unique_tensor)

    return tensor_res

# 计算两个边界框的IoU
def bbox_iou(box1, box2):
    # 获取边框的坐标
    b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
    b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]

    # 获取交叉矩形的坐标
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)

    # 交叉面积
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
   
    # 合并面积
    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)
    union_area = b1_area + b2_area - inter_area

    # IoU
    iou = inter_area / union_area

    return iou

# 把检测特征图转换成二维张量，张量的每一行对应边界框的属性，5个参数：输出，输入图像的维度……
def predict_transform(prediction, inp_dim, anchors, num_classes, CUDA=True):
    
    batch_size = prediction.size(0)
    stride = inp_dim // prediction.size(2)
    grid_size = inp_dim // stride
    bbox_attrs = 5 + num_classes
    num_anchors = len(anchors)

    prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size)
    prediction = prediction.transpose(1,2).contiguous()
    prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs)
    
    # 锚点的维度与net块的h和w属性一致，输入图像的维度和检测图的维度之商就是步长，用检测特征图的步长分割锚点
    anchors = [(a[0]/stride, a[1]/stride) for a in anchors]

    # 对（x，y）坐标和objectness分数执行Sigmoid函数操作
    prediction[:,:,0] = torch.sigmoid(prediction[:,:,0])
    prediction[:,:,1] = torch.sigmoid(prediction[:,:,1])
    prediction[:,:,4] = torch.sigmoid(prediction[:,:,4])

    # 将网格偏移添加到中心坐标预测中
    grid = np.arange(grid_size)
    a,b = np.meshgrid(grid, grid)

    x_offset = torch.FloatTensor(a).view(-1,1)
    y_offset = torch.FloatTensor(b).view(-1,1)

    if CUDA:
        x_offset = x_offset.cuda()
        y_offset = y_offset.cuda()
    
    x_y_offset = torch.cat((x_offset, y_offset), 1).repeat(1,num_anchors).view(-1,2).unsqueeze(0)

    prediction[:,:,:2] += x_y_offset

    # 将锚点应用到边界框维度中
    anchors = torch.FloatTensor(anchors)

    if CUDA:
        anchors = anchors.cuda()
    
    anchors = anchors.repeat(grid_size*grid_size, 1).unsqueeze(0)
    prediction[:,:,2:4] = torch.exp(prediction[:,:,2:4])*anchors

    # 将sigmoid激活函数应用到类别分数中
    prediction[:,:,5:5 + num_classes] = torch.sigmoid((prediction[:,:,5:5 + num_classes]))

    # 将检测图的大小调整到与输入图像大小一致，乘以stride变量（边界框属性根据特征图大小而定）
    prediction[:,:,:4] *= stride

    return prediction

# 加载类别，返回字典——将每个类别的索引映射到其名称的字符串
def load_classes(namesfile):
    fp = open(namesfile, "r")
    names = fp.read().split("\n")[:-1]
    return names

# 输出满足objectness分数阈值和非极大值抑制(NMS)，得到真实检测结果
def write_results(prediction, confidence, num_classes, nms_conf=0.4):
    # 输入为预测结果，置信度，类别数，NMS阈值

    # 低于objectness分数的每个边界框，其每个属性值都置0，即一整行。
    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
    prediction = prediction*conf_mask

    # 每个框的两个对焦坐标更容易计算两个框的IoU，故将(中心x，中心y，高度，宽度)属性转化成(左上角x，左上角y，右下角x，右下角y)
    box_a = prediction.new(prediction.shape)
    box_a[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
    box_a[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
    box_a[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2)
    box_a[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
    prediction[:,:,:4] = box_a[:,:,:4]

    batch_size = prediction.size(0)

    #output = prediction.new(1, prediction.size(2) + 1)
    write = False # 标识尚未初始化输出

    # 在第一个维度即bacth上循环，一次完成一个图像的置信度阈值和NMS
    for ind in range(batch_size):
        # 获取图像,10647x85
        image_pred = prediction[ind]

        # 每个边界框行有85个属性，其中80个类别分数，只取最大值的类别分数
        # 获取具有最高分数的类及其索引
        max_conf, max_conf_score = torch.max(image_pred[:,5:5+num_classes], 1)
        max_conf = max_conf.float().unsqueeze(1)
        max_conf_score = max_conf_score.float().unsqueeze(1)
        # 删除80个分类分数，增加最高分数类别的索引及最高分数
        seq = (image_pred[:,:5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)

        # 删除objectness置信度小于阈值的置0条目,try-except处理无检测结果的情况，continue跳过对本图像的循环
        non_zero_ind = torch.nonzero(image_pred[:,4])
        try:
            image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7) # 7列
        except:
            continue
        
        # PyTorch 0.4兼容
        if image_pred_.shape[0] == 0:
            continue

        # 获得一个图像的所有种类
        img_classes = unique(image_pred_[:,-1])
        
        # 按类别执行NMS
        for cls in img_classes:
            # 得到一个类别的所有检测
            cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1) 
            class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()

            image_pred_class = image_pred_[class_mask_ind].view(-1,7)

            # 对所有检测排序,按照objectness置信度
            conf_sort_index = torch.sort(image_pred_class[:,4], descending=True)[1]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)

            # 对于每一个检测,执行NMS
            for i in range(idx):
                # 获取正在查看的box之后所有boxes的IoUs
                try:
                    ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
                except ValueError: # image_pred_class[i+1,:]返回空张量
                    break
                except IndexError: # image_pred_class移除部分后，idx索引越界
                    break
                    
                # 清除IoU>阈值的检测
                iou_mask = (ious < nms_conf).float().unsqueeze(1)
                image_pred_class[i+1:] *= iou_mask

                # 移除0条目
                non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                image_pred_class = image_pred_class[non_zero_ind].view(-1,7)

            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
            seq = batch_ind, image_pred_class

            if not write:
                output = torch.cat(seq, 1)
                write = True
            else:
                out = torch.cat(seq, 1)
                output = torch.cat((output, out))
    # 输出一个形状为Dx8的张量；其中D是所有图像中的「真实」检测结果，每个都用一行表示。
    # 每一个检测结果都有8个属性，即该检测结果所属的batch中图像的索引、4个对角的坐标、objectness分数、有最大置信度的类别的分数、该类别的索引。
    try:
        return output
    except:
        return 0

# 使用填充调整具有不变长宽性的图像
def letterbox_image(img, inp_dim):
    
    img_w, img_h = img.shape[1], img.shape[0]
    w, h = inp_dim
    new_w = int(img_w * min(w/img_w, h/img_h))
    new_h = int(img_h * min(w/img_w, h/img_h))
    resized_image = cv2.resize(img, (new_w,new_h), interpolation = cv2.INTER_CUBIC)
    
    canvas = np.full((inp_dim[1], inp_dim[0], 3), 128)

    canvas[(h-new_h)//2:(h-new_h)//2 + new_h,(w-new_w)//2:(w-new_w)//2 + new_w,  :] = resized_image
    
    return canvas

# 将numpy数组转换成PyTorch的的输入格式
# OpenCV将图像载入成numpy数组，颜色通道为BGR。
# PyTorch的图像输入格式是(batch x 通道 x 高度 x 宽度)，通道顺序RGB。
def prep_image(img, inp_dim):

    img = letterbox_image(img, (inp_dim, inp_dim)) # 转换格式大小
    img = img[:,:,::-1].transpose((2,0,1)).copy() # BGR -> RGB(起止位置省略，步长为-1，负:从右往左)) | H x W x C -> C x H x W
    img = torch.from_numpy(img).float().div(255.0).unsqueeze(0)

    return img