From 3fa3903438fc29b3b0ccf13f47039d8425428dd4 Mon Sep 17 00:00:00 2001
From: chenguowei01 <chenguowei01@baidu.com>
Date: Thu, 24 Sep 2020 15:13:37 +0800
Subject: [PATCH] update confusion matrix calculation

---
 dygraph/paddleseg/core/val.py      |  3 +--
 dygraph/paddleseg/utils/metrics.py | 28 +++++++++++++++++++++-------
 dygraph/val.py                     |  4 ++--
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/dygraph/paddleseg/core/val.py b/dygraph/paddleseg/core/val.py
index b0f408c3..25627033 100644
--- a/dygraph/paddleseg/core/val.py
+++ b/dygraph/paddleseg/core/val.py
@@ -49,7 +49,6 @@ def evaluate(model,
     for iter, (im, im_info, label) in tqdm.tqdm(
             enumerate(eval_dataset), total=total_iters):
         im = to_variable(im)
-        # pred, _ = model(im)
         logits = model(im)
         pred = paddle.argmax(logits[0], axis=1)
         pred = pred.numpy().astype('float32')
@@ -68,7 +67,7 @@ def evaluate(model,
         pred = pred.astype('int64')
         mask = label != ignore_index
         # To-DO Test Execution Time
-        conf_mat.calculate(pred=pred, label=label, ignore=mask)
+        conf_mat.calculate(pred=pred, label=label, mask=mask)
         _, iou = conf_mat.mean_iou()
 
         time_iter = timer.elapsed_time()
diff --git a/dygraph/paddleseg/utils/metrics.py b/dygraph/paddleseg/utils/metrics.py
index 44eaf4ed..e365f7f7 100644
--- a/dygraph/paddleseg/utils/metrics.py
+++ b/dygraph/paddleseg/utils/metrics.py
@@ -29,18 +29,32 @@ class ConfusionMatrix(object):
         self.num_classes = num_classes
         self.streaming = streaming
 
-    def calculate(self, pred, label, ignore=None):
+    def calculate(self, pred, label, mask):
+        """
+        Calculate confusion matrix
+
+        Args:
+            pred (np.ndarray): The prediction of input image by model.
+            label (np.ndarray): The ground truth of input image.
+            mask (np.ndarray): The mask which pixel is valid. The dtype should be bool.
+        """
         # If not in streaming mode, clear matrix everytime when call `calculate`
         if not self.streaming:
             self.zero_matrix()
 
-        label = np.transpose(label, (0, 2, 3, 1))
-        ignore = np.transpose(ignore, (0, 2, 3, 1))
-        mask = np.array(ignore) == 1
+        pred = np.squeeze(pred)
+        label = np.squeeze(label)
+        mask = np.squeeze(mask)
 
-        label = np.asarray(label)[mask]
-        pred = np.asarray(pred)[mask]
-        one = np.ones_like(pred)
+        if not pred.shape == label.shape == mask.shape:
+            raise ValueError(
+                'Shape of `pred`, `label` and `mask` should be equal, '
+                'but there are {}, {} and {}.'.format(pred.shape, label.shape,
+                                                      mask.shape))
+
+        label = label[mask]
+        pred = pred[mask]
+        one = np.ones_like(pred).astype('int64')
         # Accumuate ([row=label, col=pred], 1) into sparse
         spm = csr_matrix((one, (label, pred)),
                          shape=(self.num_classes, self.num_classes))
diff --git a/dygraph/val.py b/dygraph/val.py
index 4d168dbf..e36fe2ff 100644
--- a/dygraph/val.py
+++ b/dygraph/val.py
@@ -18,8 +18,8 @@ import paddle
 from paddle.distributed import ParallelEnv
 
 import paddleseg
-from paddleseg.cvlibs import manager
-from paddleseg.utils import get_environ_info, Config, logger
+from paddleseg.cvlibs import manager, Config
+from paddleseg.utils import get_environ_info, logger
 from paddleseg.core import evaluate
 
 
-- 
GitLab