【计算点云语义分割中的iou】IoU Calculator

参考：https://github.com/mit-han-lab/spvnas/blob/master/core/callbacks.py

'''
Description: 
Author: suyunzheng
Date: 2022-04-06 17:59:29
LastEditTime: 2022-04-06 19:15:34
LastEditors: maple
'''
import numpy as np
import torch
from torch import Tensor
from plyfile import PlyData
from sklearn.metrics import confusion_matrix

def getPointCloud(filepath):
    plydata = PlyData.read(filepath)
    data = plydata.elements[0].data
    coords = np.array([data['x'], data['y'], data['z']], dtype=np.float32).T
    feats = np.array([data['red'], data['green'], data['blue']], dtype=np.float32).T/255
    labels = np.array(data['label'], dtype=np.int32)

    normalized_coords = coords-coords.mean(0)
    feats = np.concatenate((feats-0.5, normalized_coords), 1)   # dim = 6

    return coords, feats, labels, None


def generate_pc():
    plyfilename = "/media/xiaokeai1/Program_dataset/suyunzzz/data/S3DIS/sparse_conv_dataset/Stanford3D/Area_1/conferenceRoom_1.ply"
    return getPointCloud(plyfilename)


class IoUCalculator:
    def __init__(self, class_num, ignore_labe) -> None:
        self.class_num = class_num
        self.ignore_label = ignore_labe
        # self.gt_classes = [0 for _ in range(self.class_num)]
        # self.positive_classes = [0 for _ in range(self.class_num)]
        # self.true_positive_classes = [0 for _ in range(self.class_num)]
    
    
        self.total_seen = np.zeros(self.class_num)
        self.total_correct = np.zeros(self.class_num)
        self.total_positive = np.zeros(self.class_num)
        
    def add_data(self, outputs, targets):
        if isinstance(outputs, list):
            outputs = np.array(outputs)
        if isinstance(targets, list):
            targets = np.array(targets)
        outputs = outputs[targets != self.ignore_label]
        targets = targets[targets != self.ignore_label]
    
        for i in range(self.class_num):
            self.total_seen[i] += np.sum(targets == i).item()
            self.total_correct[i] += np.sum((targets == i)
                                                & (outputs == targets)).item()
            self.total_positive[i] += np.sum(outputs == i).item()
            
            
    def compute_iou(self):
        ious = []

        for i in range(self.class_num):
            if self.total_seen[i] == 0:
                ious.append(1)
            else:
                cur_iou = self.total_correct[i] / (self.total_seen[i]
                                                   + self.total_positive[i]
                                                   - self.total_correct[i])
                ious.append(cur_iou)

        miou = np.mean(ious)
        
        return miou, ious

def iou(y_pred, y_true):
    # y_pred = y_pred.flatten()
    # y_true = y_true.flatten()

    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2])
    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2,3,4,5,6,7,8,9,10,11,12])
    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2,3,4,5,6,7,8,9,10,11,12])
    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2])
    current = confusion_matrix(y_true, y_pred)
    print(current)
    intersection =  np.diag(current)
    gt = current.sum(axis = 1)
    pred = current.sum(axis = 0)
    union = gt + pred - intersection
    ioU = intersection / union.astype(np.float32)+1e-8
    return np.mean(ioU), ioU


if __name__ == '__main__':
    a = [1,0,2,0,1,2]
    pred = [0,0,0,0,0,0]
    
    iou_cla = IoUCalculator(class_num=3, ignore_labe=3)
    iou_cla.add_data(pred, a)
    re = iou_cla.compute_iou()
    print("re:{}".format(re))
    
    # a  = np.array(a)
    # pred = np.array(pred)
    iou_value, iou_list = iou(pred, a)
    print("miou:{}, iou_list:{}".format(iou_value, iou_list))