目标检测之3维合成

现在有一系列的图片，图片之间可以按照z轴方向进行排列。图片经过了目标检测，输出了一系列的检测框，现在的需求是将检测框按类别进行合成，以在3维上生成检测结果。

思路：将图片按照z轴方向排列，以z轴索引作为检测框的z值。等同于输入为（x, y, w, h, z, class_id），可以计算得到每个检测框的中心点来标定这个框（x_center, y_center, z, class_id）。

然后可以通过聚类算法在4维空间上进行聚类，最后取出聚类出的每一类的点的xyz的最大值与最小值和class_id来生成聚类结果[x_min, y_min, z_min, x_max, y_max, z_max, class_id]。

代码展示：

from sklearn.cluster import DBSCAN
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

class BBoxClusterv3:
    def __init__(self, bbox_list):
        self.bbox_list = bbox_list
        self.clustering = None
        # self.color_map = plt.cm.get_cmap('hsv', len(set([bbox[5] for bbox in bbox_list])))

    def cluster(self, eps=100, min_samples=2):
        X = [[bbox[0]+bbox[2]/2, bbox[1]+bbox[3]/2, bbox[4], bbox[5]] for bbox in self.bbox_list]  # 中心点的x，y，z坐标和类别ID
        # breakpoint()
        self.clustering = DBSCAN(eps=eps, min_samples=min_samples).fit(X)

    def get_new_bbox_list(self):
        # self.cluster()
        labels = self.clustering.labels_
        print("聚类出的类别：",labels)
        new_bbox_list = []
        for label in set(labels):
            if label != -1:  # Ignore noise
                idxs = np.where(labels == label)[0]
                print("每一类的bboxes索引: ",idxs)
                bboxes = np.array(self.bbox_list)[idxs]
                print("每一类的bboxes集合: ",bboxes)
                x_min = np.min(bboxes[:, 0])
                y_min = np.min(bboxes[:, 1])
                x_max = np.max(bboxes[:, 0] + bboxes[:, 2])
                y_max = np.max(bboxes[:, 1] + bboxes[:, 3])
                z_min = np.min(bboxes[:, 4])
                z_max = np.max(bboxes[:, 4])
                class_id = bboxes[0, 5]
                new_bbox_list.append([x_min, y_min, z_min, x_max, y_max, z_max, class_id])
        return new_bbox_list

    def draw_bbox_2d(self, bbox, ax):
        x_min, y_min, w, h, z, class_id = bbox
        color = self.color_map(class_id)
        # print(color)
        for xi in [x_min, x_min+w]:
            for yi in [y_min, y_min+h]:
                ax.plot([xi, xi], [yi, yi], [z, z], color=color, linestyle='dashed')
        for xi in [x_min, x_min+w]:
            ax.plot([xi, xi], [y_min, y_min+h], [z, z], color=color, linestyle='dashed')
        for yi in [y_min, y_min+h]:
            ax.plot([x_min, x_min+w], [yi, yi], [z, z], color=color, linestyle='dashed')

    def draw_bbox_3d(self, bbox, ax):
        x_min, y_min, z_min, x_max, y_max, z_max, class_id = bbox
        color = self.color_map(class_id)
        for xi in [x_min, x_max]:
            for yi in [y_min, y_max]:
                ax.plot([xi, xi], [yi, yi], [z_min, z_max], color=color)
        for xi in [x_min, x_max]:
            for zi in [z_min, z_max]:
                ax.plot([xi, xi], [y_min, y_max], [zi, zi], color=color)
        for yi in [y_min, y_max]:
            for zi in [z_min, z_max]:
                ax.plot([x_min, x_max], [yi, yi], [zi, zi], color=color)

    def visualize(self, bbox_list=None, new_bbox_list=None):
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        for bbox in bbox_list:
            self.draw_bbox_2d(bbox, ax)

        for bbox in new_bbox_list:
            self.draw_bbox_3d(bbox, ax)

        plt.show()

    def draw(self):
        new_bbox_list = self.get_new_bbox_list()
        print(bbox_list,new_bbox_list)
        self.visualize(bbox_list, new_bbox_list)

    def color_map(self, class_id):
        # 假设这里使用映射字典将类别 ID 映射到不同的颜色
        color_mapping = {0: 'red', 1: 'blue', 2: 'green'}
        return color_mapping.get(class_id, 'black')  # 默认为黑色

if __name__ == "__main__":
    bbox_list = [
        #-------------------------#
        #[x, y, w, h, z, class_id]#
        #-------------------------#
        [100, 200, 50, 50, 0, 0],
        [110, 210, 50, 50, 1, 0],
        [120, 220, 50, 50, 2, 0],
        [130, 230, 50, 50, 3, 0],
        [140, 240, 50, 50, 4, 0],
        [200, 300, 60, 60, 0, 1],
        [210, 310, 60, 60, 1, 1],
        [220, 320, 60, 60, 2, 1],
        [300, 400, 70, 70, 6, 0],
        [310, 410, 70, 70, 7, 0],
        [320, 420, 70, 70, 8, 0],
        [400, 500, 80, 80, 9, 1],
        [410, 510, 80, 80, 10, 1],
        [420, 520, 80, 80, 11, 2]
    ]
    bbox_cluster = BBoxClusterv3(bbox_list)
    bbox_cluster.cluster()
    bbox_cluster.draw()

假如有以下几类框