提交 c849198a 编写于 作者: W wuzewu

Print the miou in the segmentation training demo.

上级 030daa23
import paddle
import numpy as np
import paddlehub as hub
from paddlehub.finetune.trainer import Trainer
from paddlehub.datasets import OpticDiscSeg
from paddlehub.vision.segmentation_transforms import Compose, Resize, Normalize
from paddlehub.vision.utils import ConfusionMatrix
if __name__ == "__main__":
transform = Compose([Resize(target_size=(512, 512)), Normalize()])
train_reader = OpticDiscSeg(transform)
train_transforms = Compose([Resize(target_size=(512, 512)), Normalize()])
eval_transforms = Compose([Normalize()])
train_reader = OpticDiscSeg(train_transforms)
eval_reader = OpticDiscSeg(eval_transforms, mode='val')
model = hub.Module(name='ocrnet_hrnetw18_voc', num_classes=2)
scheduler = paddle.optimizer.lr.PolynomialDecay(learning_rate=0.01, decay_steps=1000, power=0.9, end_lr=0.0001)
optimizer = paddle.optimizer.Adam(learning_rate=scheduler, parameters=model.parameters())
trainer = Trainer(model, optimizer, checkpoint_dir='test_ckpt_img_ocr', use_gpu=True)
trainer.train(train_reader, epochs=20, batch_size=4, eval_dataset=train_reader, log_interval=10, save_interval=4)
\ No newline at end of file
optimizer = paddle.optimizer.Momentum(learning_rate=scheduler, parameters=model.parameters())
trainer = Trainer(model, optimizer, checkpoint_dir='test_ckpt_img_seg', use_gpu=True)
trainer.train(train_reader, epochs=10, batch_size=4, log_interval=10, save_interval=4)
cfm = ConfusionMatrix(eval_reader.num_classes, streaming=True)
model.eval()
for imgs, labels in eval_reader:
imgs = imgs[np.newaxis, :, :, :]
preds = model(paddle.to_tensor(imgs))[0]
preds = paddle.argmax(preds, axis=1, keepdim=True).numpy()
labels = labels[np.newaxis, :, :, :]
ignores = labels != eval_reader.ignore_index
cfm.calculate(preds, labels, ignores)
_, miou = cfm.mean_iou()
print('miou: {:.4f}'.format(miou))
......@@ -21,6 +21,7 @@ import PIL
import numpy as np
import matplotlib as plt
import paddle.nn.functional as F
from scipy.sparse import csr_matrix
def is_image_file(filename: str) -> bool:
......@@ -290,3 +291,129 @@ def reverse_transform(pred: paddle.Tensor, ori_shape: List[int], transforms: Lis
else:
raise Exception("Unexpected info '{}' in im_info".format(item[0]))
return pred
class ConfusionMatrix(object):
"""
Confusion Matrix for segmentation evaluation.
Args:
num_classes (int): Number of categories of the confusion matrix.
streaming (bool): Whether to use streaming mode. If the value is set to True, the data will be
accumulated every time the `calculate` interface is called. Default to False.
"""
def __init__(self, num_classes: int, streaming: bool = False):
self.confusion_matrix = np.zeros([num_classes, num_classes], dtype='int64')
self.num_classes = num_classes
self.streaming = streaming
def calculate(self, pred, label, ignore=None):
# If not in streaming mode, clear matrix everytime when call `calculate`
if not self.streaming:
self.zero_matrix()
mask = np.array(ignore) == 1
label = np.asarray(label)[mask]
pred = np.asarray(pred)[mask]
one = np.ones_like(pred)
# Accumuate ([row=label, col=pred], 1) into sparse matrix
spm = csr_matrix((one, (label, pred)), shape=(self.num_classes, self.num_classes))
spm = spm.todense()
self.confusion_matrix += spm
def zero_matrix(self):
""" Clear confusion matrix """
self.confusion_matrix = np.zeros([self.num_classes, self.num_classes], dtype='int64')
def mean_iou(self) -> float:
iou_list = []
avg_iou = 0
# TODO: use numpy sum axis api to simpliy
vji = np.zeros(self.num_classes, dtype=int)
vij = np.zeros(self.num_classes, dtype=int)
for j in range(self.num_classes):
v_j = 0
for i in range(self.num_classes):
v_j += self.confusion_matrix[j][i]
vji[j] = v_j
for i in range(self.num_classes):
v_i = 0
for j in range(self.num_classes):
v_i += self.confusion_matrix[j][i]
vij[i] = v_i
for c in range(self.num_classes):
total = vji[c] + vij[c] - self.confusion_matrix[c][c]
if total == 0:
iou = 0
else:
iou = float(self.confusion_matrix[c][c]) / total
avg_iou += iou
iou_list.append(iou)
avg_iou = float(avg_iou) / float(self.num_classes)
return np.array(iou_list), avg_iou
def accuracy(self) -> float:
total = self.confusion_matrix.sum()
total_right = 0
for c in range(self.num_classes):
total_right += self.confusion_matrix[c][c]
if total == 0:
avg_acc = 0
else:
avg_acc = float(total_right) / total
vij = np.zeros(self.num_classes, dtype=int)
for i in range(self.num_classes):
v_i = 0
for j in range(self.num_classes):
v_i += self.confusion_matrix[j][i]
vij[i] = v_i
acc_list = []
for c in range(self.num_classes):
if vij[c] == 0:
acc = 0
else:
acc = self.confusion_matrix[c][c] / float(vij[c])
acc_list.append(acc)
return np.array(acc_list), avg_acc
def kappa(self) -> float:
vji = np.zeros(self.num_classes)
vij = np.zeros(self.num_classes)
for j in range(self.num_classes):
v_j = 0
for i in range(self.num_classes):
v_j += self.confusion_matrix[j][i]
vji[j] = v_j
for i in range(self.num_classes):
v_i = 0
for j in range(self.num_classes):
v_i += self.confusion_matrix[j][i]
vij[i] = v_i
total = self.confusion_matrix.sum()
# avoid spillovers
# TODO: is it reasonable to hard code 10000.0?
total = float(total) / 10000.0
vji = vji / 10000.0
vij = vij / 10000.0
tp = 0
tc = 0
for c in range(self.num_classes):
tp += vji[c] * vij[c]
tc += self.confusion_matrix[c][c]
tc = tc / 10000.0
pe = tp / (total * total)
po = tc / total
kappa = (po - pe) / (1 - pe)
return kappa
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册