目标检测-斜框IOU,nms计算

1 旋转IOU

"""
2019.7.4 计算旋转的iou
"""
#coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy as np
import cv2
import tensorflow as tf
import time
import torch

def rbbx_overlaps(boxes1, boxes2, gpu_id):
    pass

"""
计算旋转面积
boxes1,boxes2格式为x,y,w,h,theta
"""
def iou_rotate_calculate(boxes1, boxes2, use_gpu=False, gpu_id=1):
    print("####boxes2:", boxes1.shape)
    print("####boxes2:", boxes2.shape)
    os._exit()
    # start = time.time()
    if use_gpu:
        print("暂时用不了")
        ious = rbbx_overlaps(boxes1, boxes2, gpu_id)
    else:
        area1 = boxes1[0] * boxes1[1]
        area2 = boxes2[2] * boxes2[3]
        r1 = ((boxes2[0], boxes2[1]), (boxes1[0], boxes1[1]), boxes1[2])
        r2 = ((boxes2[0], boxes2[1]), (boxes2[2], boxes2[3]), boxes2[4])
        int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
        if int_pts is not None:
            order_pts = cv2.convexHull(int_pts, returnPoints=True)
            int_area = cv2.contourArea(order_pts)
            # 计算出iou
            ious = int_area * 1.0 / (area1 + area2 - int_area)
        else:
            ious=0
    return ious


"""
计算旋转面积
boxes1,boxes2格式为x,y,w,h,theta
输出为张量格式
"""
def iou_rotate_calculate1(boxes1, boxes2, use_gpu=False, gpu_id=1):
    #将gpu的tensor转成cpu的
    boxes1=boxes1.cpu()
    boxes2 = boxes2.cpu()
    # 将tensor转成numpy的
    boxes1 = boxes1.numpy()
    boxes2 = boxes2.numpy()
    ious_total=[]
    # start = time.time()
    if use_gpu:
        print("@@@@暂时用不了")
        ious = rbbx_overlaps(boxes1, boxes2, gpu_id)
    else:
        for num in range(0,len(boxes2)):
            area1 = boxes1[0] * boxes1[1]
            area2 = boxes2[num,2] * boxes2[num,3]
            r1 = ((boxes2[num,0], boxes2[num,1]), (boxes1[0], boxes1[1]), boxes1[2])
            r2 = ((boxes2[num,0], boxes2[num,1]), (boxes2[num,2], boxes2[num,3]), boxes2[num,4])
            int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
            if int_pts is not None:
                order_pts = cv2.convexHull(int_pts, returnPoints=True)
                int_area = cv2.contourArea(order_pts)
                # 计算出iou
                ious = int_area * 1.0 / (area1 + area2 - int_area)
                ious_total.append(ious)
            else:
                ious = 0
                ious_total.append(ious)
        # numpy转为CPU tensor
        total_ious = torch.from_numpy(np.array(ious_total))
        total_ious = total_ious.type(torch.FloatTensor)  # 转Float
        # CPU tensor转GPU tensor
        ious = total_ious.cuda()
    return ious

"""
计算旋转面积
boxes1,boxes2格式为x,y,w,h,theta
输出为张量格式
"""
def bbox_iou_rotate_calculate1(boxes1, boxes2, use_gpu=False, gpu_id=1):
    #将gpu的tensor转成cpu的
    boxes1=boxes1.cpu()
    boxes2 = boxes2.cpu()
    # 将tensor转成numpy的
    boxes1 = boxes1.numpy()
    boxes2 = boxes2.numpy()
    ious_total=[]
    # start = time.time()
    if use_gpu:
        print("暂时用不了")
        ious = rbbx_overlaps(boxes1, boxes2, gpu_id)
    else:
        for num in range(0,len(boxes2)):
            area1 = boxes1[num,2] * boxes1[num,3]
            area2 = boxes2[num,2] * boxes2[num,3]
            r1 = ((boxes1[num,0], boxes1[num,1]), (boxes1[num,2], boxes1[num,3]), boxes1[num,4])
            r2 = ((boxes2[num,0], boxes2[num,1]), (boxes2[num,2], boxes2[num,3]), boxes2[num,4])
            int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
            if int_pts is not None:
                order_pts = cv2.convexHull(int_pts, returnPoints=True)
                int_area = cv2.contourArea(order_pts)
                # 计算出iou
                ious = int_area * 1.0 / (area1 + area2 - int_area)
                ious_total.append(ious)
            else:
                ious = 0
                ious_total.append(ious)
        # numpy转为CPU tensor
        total_ious = torch.from_numpy(np.array(ious_total))
        total_ious = total_ious.type(torch.FloatTensor)  # 转Float
        # CPU tensor转GPU tensor
        ious = total_ious.cuda()
    return ious

if __name__ == '__main__':
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = '1'
    boxes1 = np.array([50, 50, 100, 300, 0],np.float32)

    boxes2 = np.array([50, 50, 100, 300, -45.], np.float32)

    ious_area = iou_rotate_calculate(boxes1, boxes2, use_gpu=False)
    print('###ious_area:',ious_area)

2 旋转框NMS

"""
2019.7.4 : 对斜框计算来计算nms
"""
#coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import cv2
import tensorflow as tf

def nms_rotate(decode_boxes, scores, iou_threshold, max_output_size,
               use_angle_condition=False, angle_threshold=0, use_gpu=False, gpu_id=1):
    """
    :param boxes: format [x_c, y_c, w, h, theta]
    :param scores: scores of boxes
    :param threshold: iou threshold (0.7 or 0.5)
    :param max_output_size: max number of output
    :return: the remaining index of boxes
    """
    if use_gpu:
        #采用gpu方式
        keep = nms_rotate_gpu(boxes_list=decode_boxes,
                              scores=scores,
                              iou_threshold=iou_threshold,
                              angle_gap_threshold=angle_threshold,
                              use_angle_condition=use_angle_condition,
                              device_id=gpu_id)

        keep = tf.cond(
            tf.greater(tf.shape(keep)[0], max_output_size),
            true_fn=lambda: tf.slice(keep, [0], [max_output_size]),
            false_fn=lambda: keep)
    else: #采用cpu方式
        keep = tf.py_func(nms_rotate_cpu,
                          inp=[decode_boxes, scores, iou_threshold, max_output_size],
                          Tout=tf.int64)
    return keep

def nms_rotate_cpu(boxes, scores, iou_threshold, max_output_size):
    print("###运行cpu")
    keep = [] #保留框的结果集合
    order = scores.argsort()[::-1] #对检测结果得分进行降序排序
    num = boxes.shape[0] #获取检测框的个数

    suppressed = np.zeros((num), dtype=np.int)
    for _i in range(num):
        # 若当前保留框集合中的个数大于max_output_size时,直接返回
        if len(keep) >= max_output_size:
            break
        i = order[_i]
        # 对于抑制的检测框直接跳过
        if suppressed[i] == 1:
            continue
        keep.append(i) #保留当前框的索引
        #根据box信息组合成opencv中的旋转bbox
        r1 = ((boxes[i, 0], boxes[i, 1]), (boxes[i, 2], boxes[i, 3]), boxes[i, 4])
        #计算当前检测框的面积
        area_r1 = boxes[i, 2] * boxes[i, 3]
        #对剩余的而进行遍历
        for _j in range(_i + 1, num):
            j = order[_j]
            if suppressed[i] == 1:
                continue
            r2 = ((boxes[j, 0], boxes[j, 1]), (boxes[j, 2], boxes[j, 3]), boxes[j, 4])
            area_r2 = boxes[j, 2] * boxes[j, 3]
            inter = 0.0
            #求两个旋转矩形的交集,并返回相交的点集合
            int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
            if int_pts is not None:
                #求点集的凸边形
                order_pts = cv2.convexHull(int_pts, returnPoints=True)
                # 计算当前点集合组成的凸边形的面积
                int_area = cv2.contourArea(order_pts)
                #计算出iou
                inter = int_area * 1.0 / (area_r1 + area_r2 - int_area + 0.0000001)
            # 对大于设定阈值的检测框进行滤除
            if inter >= iou_threshold:
                suppressed[j] = 1

    return np.array(keep, np.int64)

# gpu的实现方式
def nms_rotate_gpu(boxes_list, scores, iou_threshold, use_angle_condition=False, angle_gap_threshold=0, device_id=0):
    if use_angle_condition:
        y_c, x_c, h, w, theta = tf.unstack(boxes_list, axis=1)
        boxes_list = tf.transpose(tf.stack([x_c, y_c, w, h, theta]))
        det_tensor = tf.concat([boxes_list, tf.expand_dims(scores, axis=1)], axis=1)
        keep = tf.py_func(rotate_gpu_nms,
                          inp=[det_tensor, iou_threshold, device_id],
                          Tout=tf.int64)
        return keep
    else:
        y_c, x_c, h, w, theta = tf.unstack(boxes_list, axis=1)
        boxes_list = tf.transpose(tf.stack([x_c, y_c, w, h, theta]))
        det_tensor = tf.concat([boxes_list, tf.expand_dims(scores, axis=1)], axis=1)
        keep = tf.py_func(rotate_gpu_nms,
                          inp=[det_tensor, iou_threshold, device_id],
                          Tout=tf.int64)
        keep = tf.reshape(keep, [-1])
        return keep

if __name__ == '__main__':
    boxes = np.array([[50, 40, 100, 100, 0],
                      [60, 50, 100, 100, 0],
                      [50, 30, 100, 100, -45.],
                      [200, 190, 100, 100, 0.]])

    scores = np.array([ 0.77,0.79, 0.88, 0.66,])
    keep = nms_rotate(tf.convert_to_tensor(boxes, dtype=tf.float32), tf.convert_to_tensor(scores, dtype=tf.float32),
                      0.7, 5)

    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = '1'
    with tf.Session() as sess:
        print(sess.run(keep))

 

上一篇:盒模型学习(一)


下一篇:IoU、NMS、bounding box regression、