#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# Copyright (c) Megvii Inc. All rights reserved.

import torch
import torch.nn as nn

# IOU loss表示的是两个框的重叠部分与两个框的并集的比值，IOU越大，表示两个框越接近，越小表示两个框越远离。
class IOUloss(nn.Module):
    def __init__(self, reduction="none", loss_type="iou"):
        super(IOUloss, self).__init__()
        self.reduction = reduction
        # reduction表示的是计算loss的方式，有mean和sum两种方式，mean表示的是对loss求平均，sum表示的是对loss求和。
        self.loss_type = loss_type
        # loss_type表示的是计算loss的方式，有iou和giou两种方式，iou表示的是计算两个框的IOU，giou表示的是计算两个框的GIoU。
        # 其中GIoU是IOU的改进版，GIoU的计算方式是先计算两个框的并集，然后计算两个框的IOU，最后用并集的面积减去IOU的面积，再除以并集的面积，得到的结果就是GIoU的值。

    def forward(self, pred, target):
        assert pred.shape[0] == target.shape[0]
        # pred 表示的是预测的框，target表示的是真实的框，两者的shape都是[batch_size, 4]，4表示的是x,y,w,h。

        pred = pred.view(-1, 4)
        target = target.view(-1, 4)
        # view()函数表示的是把目标Tensor变成指定的形状，-1表示的是自动推断，这里的-1表示的是自动推断出来的结果是batch_size。
        # 目标列数是4，表示的是x,y,w,h。
        tl = torch.max(
            (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
        )
        # tl的结果是个二维数组，表示的是两个框的左上角的坐标。
        br = torch.min(
            (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
        )

        area_p = torch.prod(pred[:, 2:], 1)
        area_g = torch.prod(target[:, 2:], 1)
        # prod()函数表示的是对数组中的元素进行乘积运算，这里的乘积运算是指对数组中的每个元素进行乘积运算，最后得到的结果是一个数。
        # area_p表示的是预测框的面积，area_g表示的是真实框的面积。
        # area_g表示的是真实框的面积，area_g表示的是真实框的面积。
        en = (tl < br).type(tl.type()).prod(dim=1)
        # tl < br 表示的是两个框的左上角的坐标是否小于右下角的坐标，如果小于，那么就表示两个框有重叠部分，如果大于，那么就表示两个框没有重叠部分。
        # type(tl.type()) 表示的是把tl的数据类型转换成tl.type()的数据类型。
        # prod(dim=1) 表示的是对数组中的元素进行乘积运算，这里的乘积运算是指对数组中的每一行进行乘积运算，最后得到的结果是一个数组。
        # en表示的是两个框是否有重叠部分，如果有重叠部分，那么en的值为1，如果没有重叠部分，那么en的值为0。
        area_i = torch.prod(br - tl, 1) * en
        # area_i表示的是两个框的重叠部分的面积。
        area_u = area_p + area_g - area_i
        # area_u表示的是两个框的并集的面积。
        iou = (area_i) / (area_u + 1e-16)
        # iou表示的是两个框的重叠部分与两个框的并集的比值。
        if self.loss_type == "iou":
            loss = 1 - iou ** 2
            # loss表示的是两个框的重叠部分与两个框的并集的比值的平方。
        elif self.loss_type == "giou":
            c_tl = torch.min(
                (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
            )
            # c_tl表示的是两个框的左上角的坐标的最小值。
            c_br = torch.max(
                (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
            )
            # c_br表示的是两个框的右下角的坐标的最大值。
            area_c = torch.prod(c_br - c_tl, 1)
            # area_c表示的是两个框的并集的面积。
            giou = iou - (area_c - area_u) / area_c.clamp(1e-16)
            # giou表示的是两个框的重叠部分与两个框的并集的比值的平方。
            # clamp(min=-1.0, max=1.0)表示的是把giou的值限制在[-1.0, 1.0]之间。
            loss = 1 - giou.clamp(min=-1.0, max=1.0)
            # loss表示的是两个框的重叠部分与两个框的并集的比值的平方。

        if self.reduction == "mean":
            loss = loss.mean()
            # loss类似于一个数组，loss.mean()表示的是对loss求平均。
        elif self.reduction == "sum":
            loss = loss.sum()
            # loss类似于一个数组，loss.sum()表示的是对loss求和。

        return loss