参考：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))
