Merge pull request #18 from MHGL/Test_Add_IoU_Eval

Test add io u eval

Author: gemfield

test.py

@@ -1,3 +1,4 @@
+# -*- coding:utf-8 -*-
 import sys
 import os
 import time
@@ -7,12 +8,15 @@
 import torch.nn.functional as F
 import deepvac
 from deepvac import LOG, Deepvac
+from utils.utils_IOU_eval import IOUEval
 from data.dataloader import OsWalkDataset2
 
 class ESPNetTest(Deepvac):
     def __init__(self, deepvac_config):
         super(ESPNetTest, self).__init__(deepvac_config)
         os.makedirs(self.config.show_output_dir, exist_ok=True)
+        if self.config.test_label_path is not None:
+            self.config.iou_eval = IOUEval(self.config.cls_num)
 
     def preIter(self):
         assert len(self.config.target) == 1, 'config.core.test_batch_size must be set to 1 in current test mode.'
@@ -35,6 +39,11 @@ def postIter(self):
         savepath = os.path.join(self.config.show_output_dir, filename)
         mask_savepath = os.path.join(self.config.show_output_dir, mask_filename)
 
+        if self.config.test_label_path:
+            label_file = os.path.join(self.config.test_label_path, filename.replace(".jpg", ".png"))
+            self.config.label = cv2.imread(label_file, 0)
+            self.config.iou_eval.addBatch(self.config.mask, self.config.label)
+
         classMap_numpy_color = np.zeros((h, w, 3), dtype=np.uint8)
         for idx in np.unique(self.config.mask):
             [r, g, b] = self.config.pallete[idx]
@@ -44,6 +53,17 @@ def postIter(self):
         cv2.imwrite(mask_savepath, classMap_numpy_color)
         LOG.logI('{}: [out cv image save to {}] [{}/{}]\n'.format(self.config.phase, savepath, self.config.test_step + 1, len(self.config.test_loader)))
 
+    def testFly(self):
+        if self.config.test_loader:
+            self.test()
+            if self.config.test_label_path is None:
+                return
+            *_, self.config.mIOU = self.config.iou_eval.getMetric()
+            LOG.logI(">>> {}: [dataset: {}, mIOU: {:.3f}]".format(self.config.phase, self.config.filepath.split('/')[-2], self.config.mIOU))
+            return
+
+        LOG.logE("You have to reimplement testFly() in subclass {} if you didn't set any valid input, e.g. config.core.test_loader.".format(self.name()), exit=True)
+
 
 if __name__ == "__main__":
     from config import config
@@ -55,13 +75,16 @@ def check_args(idx, argv):
         config.core.model_path = sys.argv[1]
     if check_args(2, sys.argv):
         config.test_sample_path = sys.argv[2]
+    if check_args(3, sys.argv):
+        config.core.test_label_path = sys.argv[3]
 
     if (config.core.model_path is None) or (config.test_sample_path is None):
         helper = '''model_path or test_sample_path not found, please check:
                 config.core.model_path or sys.argv[1] to init model path
                 config.test_sample_path or sys.argv[2] to init test sample path
+                config.test_label_path or sys.argv[3] to init test sample path (not required)
                 for example:
-                python3 test.py <trained-model-path> <test sample path>'''
+                python3 test.py <trained-model-path> <test sample path> [test label path(not required)]'''
         print(helper)
         sys.exit(1)
 

train.py

@@ -21,7 +21,7 @@ def train(self):
         for i, loader in enumerate(self.config.train_loader_list):
             self.config.train_loader = loader
             super(ESPNetTrain, self).train()
-            
+
     #only save model for last loader
     def doSave(self):
         if not self.config.train_loader.is_last_loader:
@@ -34,9 +34,9 @@ def postIter(self):
 
         self.config.epoch_loss.append(self.config.loss.item())
         if self.config.phase == 'TRAIN':
-            self.iou_eval_train.addBatch(self.config.output[0].max(1)[1].data, self.config.target.data)
+            self.iou_eval_train.addBatch(self.config.output[0].max(1)[1].data.cpu().numpy(), self.config.target.data.cpu().numpy())
         else:
-            self.iou_eval_val.addBatch(self.config.output[0].max(1)[1].data, self.config.target.data)
+            self.iou_eval_val.addBatch(self.config.output[0].max(1)[1].data.cpu().numpy(), self.config.target.data.cpu().numpy())
 
     def preEpoch(self):
         self.config.epoch_loss = []

utils/utils_IOU_eval.py

@@ -21,8 +21,8 @@ def compute_hist(self, predict, gth):
         return hist
 
     def addBatch(self, predict, gth):
-        predict = predict.cpu().numpy().flatten()
-        gth = gth.cpu().numpy().flatten()
+        predict = predict.flatten()
+        gth = gth.flatten()
 
         epsilon = 0.00000001
         hist = self.compute_hist(predict, gth)