tensor与list版本的iou比较

自己写的是list版本的，但是一般深度学习框架都是tensor版本的，所以做了两者对比，发现没有什么区别，但是个人感觉，tensor版本更灵活。

import random
import torchvision.transforms as transforms
import cv2
import torch
import os
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from skimage import io, transform
import xml.etree.ElementTree as ET
import torch.utils as utils


# https://www.pianshen.com/article/6419180043/
def list_iou(box1, box2):
    """Implement the intersection over union (IoU) between box1 and box2

    Arguments:
    box1 -- first box, list object with coordinates (x1, y1, x2, y2)
    box2 -- second box, list object with coordinates (x1, y1, x2, y2)
    """

    # Calculate the (y1, x1, y2, x2) coordinates of the intersection of box1 and box2. Calculate its Area
    xi1 = max(box1[0], box2[0])
    yi1 = max(box1[1], box2[1])
    xi2 = min(box1[2], box2[2])
    yi2 = min(box1[3], box2[3])
    inter_area = max(0, (yi2 - yi1)) * max(0, (xi2 - xi1))

    # Calculate the Union area by using Formula: Union(A,B) = A + B - Inter(A,B)
    box1_area = (box1[3] - box1[1]) * (box1[2] - box1[0])
    box2_area = (box2[3] - box2[1]) * (box2[2] - box2[0])
    union_area = box1_area + box2_area - inter_area

    # compute the IoU
    iou = inter_area / union_area

    return iou

def intersect(box_a, box_b):
    A = box_a.size(0)
    B = box_b.size(0)
    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
    inter = torch.clamp((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]

def tensorIou(box_a, box_b):

    inter = intersect(box_a, box_b)

    area_a = ((box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)
    area_b = ((box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1]).unsqueeze(0).expand_as(inter))

    area = inter / (area_a + area_b - inter)

    print(area)


if __name__ == '__main__':
    list_version_a = [[0.0000, 0.3057, 1.0000, 1.0000],
                      [0.0381, 0.2798, 0.8898, 1.0000]]
    list_version_b = [[-0.0367, -0.0367,  0.0633,  0.0633],
                      [-0.0574, -0.0574,  0.0840,  0.0840]]

    out_iou = []
    for box_a in list_version_a:
        for box_b in list_version_b:
            out_iou.append(list_iou(box_a, box_b))
    print(out_iou)

    tensor_version_a = torch.Tensor(list_version_a)
    # 也可以这样 tensor_version_a = torch.Tensor([[0.0000, 0.3057, 1.0000, 1.0000],
    #                       [0.0381, 0.2798, 0.8898, 1.0000])

    tensor_version_b = torch.Tensor(list_version_b)
    tensorIou(tensor_version_a, tensor_version_b)

从图上看，结果是一样的。