Update settings of task_matting and fix the bug of pepper proproc.

fanchenhao · Sep 5, 2024 · 53f9f2d · 53f9f2d
1 parent a3e781d
commit 53f9f2d
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 18 deletions.
diff --git a/config.py b/config.py
@@ -136,7 +136,7 @@ def __init__(self) -> None:
 
         # PATH settings - inactive
         self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis')
-        self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights')
+        self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights/cv')
         self.weights = {
             'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'),
             'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]),
@@ -163,6 +163,7 @@ def __init__(self) -> None:
             with open(run_sh_file[0], 'r') as f:
                 lines = f.readlines()
                 self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0])
+                self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0])
 
     def print_task(self) -> None:
         # Return task for choosing settings in shell scripts.

diff --git a/eval_existingOnes.py b/eval_existingOnes.py
@@ -3,7 +3,7 @@
 from glob import glob
 import prettytable as pt
 
-from evaluation.evaluate import evaluator
+from evaluation.metrics import evaluator
 from config import Config
 
 
@@ -40,7 +40,7 @@ def do_eval(args):
             print('\t', 'Evaluating model: {}...'.format(_model_name))
             pred_paths = [p.replace(args.gt_root, os.path.join(args.pred_root, _model_name)).replace('/gt/', '/') for p in gt_paths]
             # print(pred_paths[:1], gt_paths[:1])
-            em, sm, fm, mae, wfm, hce, mba, biou = evaluator(
+            em, sm, fm, mae, mse, wfm, hce, mba, biou = evaluator(
                 gt_paths=gt_paths,
                 pred_paths=pred_paths,
                 metrics=args.metrics.split('+'),
@@ -72,7 +72,7 @@ def do_eval(args):
                 ]
             elif config.task == 'Matting':
                 scores = [
-                    sm.round(3), fm['curve'].max().round(3), em['curve'].mean().round(3), mse.round(3),
+                    sm.round(3), fm['curve'].max().round(3), em['curve'].mean().round(3), mse.round(5),
                     em['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
                     mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
                 ]
@@ -110,7 +110,7 @@ def do_eval(args):
             'COD': '+'.join(['TE-COD10K', 'NC4K', 'TE-CAMO', 'CHAMELEON'][:]),
             'HRSOD': '+'.join(['DAVIS-S', 'TE-HRSOD', 'TE-UHRSD', 'TE-DUTS', 'DUT-OMRON'][:]),
             'General': '+'.join(['DIS-VD'][:]),
-            'Matting': '+'.join(['TE-P3M-500-P'][:]),
+            'Matting': '+'.join(['TE-AM2k'][:]),
         }[config.task])
     parser.add_argument(
         '--save_dir', type=str, help='candidate competitors',
@@ -120,9 +120,9 @@ def do_eval(args):
         default=False)
     parser.add_argument(
         '--metrics', type=str, help='candidate competitors',
-        default='+'.join(['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'HCE'][:100 if 'DIS5K' in config.task else -1]))
+        default='+'.join(['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'MSE', 'HCE'][:100 if 'DIS5K' in config.task else -1]))
     args = parser.parse_args()
-    args.metrics = '+'.join(['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'HCE'][:100 if sum(['DIS-' in _data for _data in args.data_lst.split('+')]) else -1])
+    args.metrics = '+'.join(['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'MSE', 'HCE'][:100 if sum(['DIS-' in _data for _data in args.data_lst.split('+')]) else -1])
 
     os.makedirs(args.save_dir, exist_ok=True)
     try:

diff --git a/evaluation/metrics.py b/evaluation/metrics.py
@@ -1,6 +1,7 @@
 import os
 from tqdm import tqdm
 import cv2
+from PIL import Image
 import numpy as np
 from scipy.ndimage import convolve, distance_transform_edt as bwdist
 from skimage.morphology import skeletonize
@@ -12,7 +13,7 @@
 _TYPE = np.float64
 
 
-def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'HCE'], verbose=False):
+def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'MSE', 'HCE'], verbose=False):
     # define measures
     if 'E' in metrics:
         EM = EMeasure()
@@ -22,6 +23,8 @@ def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA',
         FM = FMeasure()
     if 'MAE' in metrics:
         MAE = MAEMeasure()
+    if 'MSE' in metrics:
+        MSE = MSEMeasure()
     if 'WF' in metrics:
         WFM = WeightedFMeasure()
     if 'HCE' in metrics:
@@ -63,6 +66,8 @@ def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA',
             FM.step(pred=pred_ary, gt=gt_ary)
         if 'MAE' in metrics:
             MAE.step(pred=pred_ary, gt=gt_ary)
+        if 'MSE' in metrics:
+            MSE.step(pred=pred_ary, gt=gt_ary)
         if 'WF' in metrics:
             WFM.step(pred=pred_ary, gt=gt_ary)
         if 'HCE' in metrics:
@@ -99,6 +104,10 @@ def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA',
         mae = MAE.get_results()['mae']
     else:
         mae = np.float64(-1)
+    if 'MSE' in metrics:
+        mse = MSE.get_results()['mse']
+    else:
+        mse = np.float64(-1)
     if 'WF' in metrics:
         wfm = WFM.get_results()['wfm']
     else:
@@ -116,7 +125,7 @@ def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA',
     else:
         biou = {'curve': np.array([np.float64(-1)])}
 
-    return em, sm, fm, mae, wfm, hce, mba, biou
+    return em, sm, fm, mae, mse, wfm, hce, mba, biou
 
 
 def _prepare_data(pred: np.ndarray, gt: np.ndarray) -> tuple:
@@ -208,6 +217,25 @@ def get_results(self) -> dict:
         return dict(mae=mae)
 
 
+class MSEMeasure(object):
+    def __init__(self):
+        self.mses = []
+
+    def step(self, pred: np.ndarray, gt: np.ndarray):
+        pred, gt = _prepare_data(pred, gt)
+
+        mse = self.cal_mse(pred, gt)
+        self.mses.append(mse)
+
+    def cal_mse(self, pred: np.ndarray, gt: np.ndarray) -> float:
+        mse = np.mean((pred - gt) ** 2)
+        return mse
+
+    def get_results(self) -> dict:
+        mse = np.mean(np.array(self.mses, _TYPE))
+        return dict(mse=mse)
+
+
 class SMeasure(object):
     def __init__(self, alpha: float = 0.5):
         self.sms = []

diff --git a/image_proc.py b/image_proc.py
@@ -48,7 +48,7 @@ def preproc(image, label, preproc_methods=['flip']):
     if 'enhance' in preproc_methods:
         image = color_enhance(image)
     if 'pepper' in preproc_methods:
-        label = random_pepper(label)
+        image = random_pepper(image)
     return image, label
 
 
@@ -112,8 +112,5 @@ def random_pepper(img, N=0.0015):
     for i in range(noiseNum):
         randX = random.randint(0, img.shape[0] - 1)
         randY = random.randint(0, img.shape[1] - 1)
-        if random.randint(0, 1) == 0:
-            img[randX, randY] = 0
-        else:
-            img[randX, randY] = 255
+        img[randX, randY] = random.randint(0, 1) * 255
     return Image.fromarray(img)
diff --git a/inference.py b/inference.py
@@ -90,7 +90,7 @@ def main(args):
                             'COD': 'TE-COD10K+NC4K+TE-CAMO+CHAMELEON',
                             'HRSOD': 'DAVIS-S+TE-HRSOD+TE-UHRSD+TE-DUTS+DUT-OMRON',
                             'General': 'DIS-VD',
-                            'Matting': 'TE-P3M-500-P',
+                            'Matting': 'TE-AM2k',
                             'DIS5K-': 'DIS-VD',
                             'COD-': 'TE-COD10K',
                             'SOD-': 'DAVIS-S+TE-HRSOD+TE-UHRSD',

diff --git a/sub.sh b/sub.sh
@@ -1,8 +1,8 @@
 #!/bin/sh
 # Example: ./sub.sh tmp_proj 0,1,2,3 3 --> Use 0,1,2,3 for training, release GPUs, use GPU:3 for inference.
 
-# module load gcc/11.2.0 cuda/11.8 cudnn/8.6.0_cu11x && cpu_core_num=6
-module load compilers/cuda/11.8 compilers/gcc/12.2.0 cudnn/8.4.0.27_cuda11.x && cpu_core_num=32
+module load gcc/11.2.0 cuda/11.8 cudnn/8.6.0_cu11x && cpu_core_num=6
+# module load compilers/cuda/11.8 compilers/gcc/12.2.0 cudnn/8.4.0.27_cuda11.x && cpu_core_num=32
 
 export PYTHONUNBUFFERED=1
 

diff --git a/test.sh b/test.sh
@@ -16,7 +16,7 @@ case "${task}" in
     "COD") testsets='CHAMELEON,NC4K,TE-CAMO,TE-COD10K' ;;
     "HRSOD") testsets='DAVIS-S,TE-HRSOD,TE-UHRSD,DUT-OMRON,TE-DUTS' ;;
     "General") testsets='DIS-VD' ;;
-    "Matting") testsets='TE-P3M-500-P' ;;
+    "Matting") testsets='TE-AM2k' ;;
 esac
 testsets=(`echo ${testsets} | tr ',' ' '`) && testsets=${testsets[@]}