DOTA格式的8个坐标值转换为角度倾斜矩形框 (cx,cy,w,h,ang)

素颜马尾好姑娘i 2022-10-15 07:55 40阅读 0赞

def polygonToRotRectangle_batch(bbox, with_module=True):
        """
        :param bbox: The polygon stored in format [x1, y1, x2, y2, x3, y3, x4, y4]
                shape [num_boxes, 8]
        :return: Rotated Rectangle in format [cx, cy, w, h, theta]
                shape [num_rot_recs, 5]
        """
        # print('bbox: ', bbox)
        bbox = np.array(bbox,dtype=np.float32)
        bbox = np.reshape(bbox,newshape=(-1, 2, 4),order='F')
        # angle = math.atan2(-(bbox[0,1]-bbox[0,0]),bbox[1,1]-bbox[1,0])
        # print('bbox: ', bbox)
        angle = np.arctan2(-(bbox[:, 0,1]-bbox[:, 0,0]),bbox[:, 1,1]-bbox[:, 1,0])
        # angle = np.arctan2(-(bbox[:, 0,1]-bbox[:, 0,0]),bbox[:, 1,1]-bbox[:, 1,0])
        # center = [[0],[0]] ## shape [2, 1]
        # print('angle: ', angle)
        center = np.zeros((bbox.shape[0], 2, 1))
        for i in range(4):
            center[:, 0, 0] += bbox[:, 0,i]
            center[:, 1, 0] += bbox[:, 1,i]
     
        center = np.array(center,dtype=np.float32)/4.0
     
        # R = np.array([[math.cos(angle), -math.sin(angle)], [math.sin(angle), math.cos(angle)]], dtype=np.float32)
        R = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]], dtype=np.float32)
     
        normalized = np.matmul(R.transpose((2, 1, 0)),bbox-center)
     
     
        xmin = np.min(normalized[:, 0, :], axis=1)
        # print('diff: ', (xmin - normalized[:, 0, 3]))
        # assert sum((abs(xmin - normalized[:, 0, 3])) > eps) == 0
        xmax = np.max(normalized[:, 0, :], axis=1)
        # assert sum(abs(xmax - normalized[:, 0, 1]) > eps) == 0
        # print('diff2: ', xmax - normalized[:, 0, 1])
        ymin = np.min(normalized[:, 1, :], axis=1)
        # assert sum(abs(ymin - normalized[:, 1, 3]) > eps) == 0
        # print('diff3: ', ymin - normalized[:, 1, 3])
        ymax = np.max(normalized[:, 1, :], axis=1)
        # assert sum(abs(ymax - normalized[:, 1, 1]) > eps) == 0
        # print('diff4: ', ymax - normalized[:, 1, 1])
     
        w = xmax - xmin + 1
        h = ymax - ymin + 1
     
        w = w[:, np.newaxis]
        h = h[:, np.newaxis]
        # TODO: check it
        if with_module:
            angle = angle[:, np.newaxis] % ( 2 * np.pi)
        else:
            angle = angle[:, np.newaxis]
        dboxes = np.concatenate((center[:, 0].astype(np.float), center[:, 1].astype(np.float), w, h, angle), axis=1)
        return dboxes