更新图像和视频的评估代码，以及一些小的修复。

Author: lartpang

.gitignore

@@ -280,6 +280,8 @@ gen
 /configs/
 # /*.py
 /*.sh
+/*.ps1
+/*.bat
 /results/rgb_sod.md
 /results/htmls/*.html
 !/.github/assets/*.jpg

eval.py

@@ -4,7 +4,7 @@
 import textwrap
 import warnings
 
-from metrics import cal_sod_matrics
+from metrics import image_metrics, video_metrics
 from utils.generate_info import get_datasets_info, get_methods_info
 from utils.recorders import SUPPORTED_METRICS
 
@@ -39,7 +39,7 @@ def get_args():
 
     EXAMPLES:
 
-    python eval_image.py \
+    python eval.py \
         --dataset-json configs/datasets/rgbd_sod.json \
         --method-json \
             configs/methods/json/rgbd_other_methods.json \
@@ -113,12 +113,41 @@ def get_args():
         "--metric-names",
         type=str,
         nargs="+",
-        default=["mae", "fmeasure", "precision", "recall", "em", "sm", "wfm"],
+        default=["sm", "wfm", "mae", "fmeasure", "em", "precision", "recall", "msiou"],
         choices=SUPPORTED_METRICS,
         help="Names of metrics",
     )
+    parser.add_argument(
+        "--data-type",
+        type=str,
+        default="image",
+        choices=["image", "video"],
+        help="Type of data.",
+    )
+
+    known_args = parser.parse_known_args()[0]
+    if known_args.data_type == "video":
+        parser.add_argument(
+            "--valid-frame-start",
+            type=int,
+            default=0,
+            help="Valid start index of the frame in each gt video. Defaults to 1, it will skip the first frame. If it is set to None, the code will not skip frames.",
+        )
+        parser.add_argument(
+            "--valid-frame-end",
+            type=int,
+            default=0,
+            help="Valid end index of the frame in each gt video. Defaults to -1, it will skip the last frame. If it is set to 0, the code will not skip frames.",
+        )
+
     args = parser.parse_args()
 
+    if args.data_type == "video":
+        args.valid_frame_start = max(args.valid_frame_start, 0)
+        args.valid_frame_end = min(args.valid_frame_end, 0)
+        if args.valid_frame_end == 0:
+            args.valid_frame_end = None
+
     if args.metric_npy:
         os.makedirs(os.path.dirname(args.metric_npy), exist_ok=True)
     if args.curves_npy:
@@ -149,22 +178,39 @@ def main():
         exclude_methods=args.exclude_methods,
     )
 
-    # 确保多进程在windows上也可以正常使用
-    cal_sod_matrics.cal_image_matrics(
-        sheet_name="Results",
-        to_append=not args.to_overwrite,
-        txt_path=args.record_txt,
-        xlsx_path=args.record_xlsx,
-        methods_info=methods_info,
-        datasets_info=datasets_info,
-        curves_npy_path=args.curves_npy,
-        metrics_npy_path=args.metric_npy,
-        num_bits=args.num_bits,
-        num_workers=args.num_workers,
-        metric_names=args.metric_names,
-        ncols_tqdm=119,
-    )
-
-
+    if args.data_type == "image":
+        image_metrics.cal_metrics(
+            sheet_name="Results",
+            to_append=not args.to_overwrite,
+            txt_path=args.record_txt,
+            xlsx_path=args.record_xlsx,
+            methods_info=methods_info,
+            datasets_info=datasets_info,
+            curves_npy_path=args.curves_npy,
+            metrics_npy_path=args.metric_npy,
+            num_bits=args.num_bits,
+            num_workers=args.num_workers,
+            metric_names=args.metric_names,
+        )
+    else:
+        video_metrics.cal_metrics(
+            sheet_name="Results",
+            to_append=not args.to_overwrite,
+            txt_path=args.record_txt,
+            xlsx_path=args.record_xlsx,
+            methods_info=methods_info,
+            datasets_info=datasets_info,
+            curves_npy_path=args.curves_npy,
+            metrics_npy_path=args.metric_npy,
+            num_bits=args.num_bits,
+            num_workers=args.num_workers,
+            metric_names=args.metric_names,
+            return_group=False,
+            start_idx=args.valid_frame_start,
+            end_idx=args.valid_frame_end,
+        )
+
+
+# 确保多进程在windows上也可以正常使用
 if __name__ == "__main__":
     main()

metrics/cal_sod_matrics.py renamed to metrics/image_metrics.py

@@ -101,7 +101,7 @@ def export(self):
                     )
 
 
-def cal_image_matrics(
+def cal_metrics(
     sheet_name: str = "results",
     txt_path: str = "",
     to_append: bool = True,
@@ -112,7 +112,6 @@ def cal_image_matrics(
     metrics_npy_path: str = "./metrics.npy",
     num_bits: int = 3,
     num_workers: int = 2,
-    ncols_tqdm: int = 79,
     metric_names: tuple = ("sm", "wfm", "mae", "fmeasure", "em"),
 ):
     """Save the results of all models on different datasets in a `npy` file in the form of a
@@ -129,7 +128,6 @@ def cal_image_matrics(
         metrics_npy_path (str, optional): The npy file path for saving metric values. Defaults to "./metrics.npy".
         num_bits (int, optional): The number of bits used to format results. Defaults to 3.
         num_workers (int, optional): The number of workers of multiprocessing or multithreading. Defaults to 2.
-        ncols_tqdm (int, optional): Number of columns for tqdm. Defaults to 79.
         metric_names (tuple, optional): Names of metrics. Defaults to ("sm", "wfm", "mae", "fmeasure", "em").
 
     Returns:
@@ -167,16 +165,12 @@ def cal_image_matrics(
         sheet_name=sheet_name,
     )
 
-    # multi-process mode
-    # tqdm.set_lock(RLock())
-    # pool_cls = pool.Pool
-    # multi-threading mode
     tqdm.set_lock(TRLock())
-    pool_cls = pool.ThreadPool
-    procs = pool_cls(processes=num_workers, initializer=tqdm.set_lock, initargs=(tqdm.get_lock(),))
+    procs = pool.ThreadPool(
+        processes=num_workers, initializer=tqdm.set_lock, initargs=(tqdm.get_lock(),)
+    )
     print(f"Create a {procs}).")
 
-    procs_idx = 0
     for dataset_name, dataset_path in datasets_info.items():
         # 获取真值图片信息
         gt_info = dataset_path["mask"]
@@ -187,16 +181,13 @@ def cal_image_matrics(
         gt_index_file = dataset_path.get("index_file")
         if gt_index_file:
             gt_name_list = get_name_list(
-                data_path=gt_index_file,
-                name_prefix=gt_prefix,
-                name_suffix=gt_suffix,
+                data_path=gt_index_file, name_prefix=gt_prefix, name_suffix=gt_suffix
             )
         else:
             gt_name_list = get_name_list(
-                data_path=gt_root,
-                name_prefix=gt_prefix,
-                name_suffix=gt_suffix,
+                data_path=gt_root, name_prefix=gt_prefix, name_suffix=gt_suffix
             )
+        gt_info_pair = (gt_root, gt_prefix, gt_suffix)
         assert len(gt_name_list) > 0, "there is not ground truth."
 
         # ==>> test the intersection between pre and gt for each method <<==
@@ -214,33 +205,28 @@ def cal_image_matrics(
             pre_name_list = get_name_list(
                 data_path=pre_root, name_prefix=pre_prefix, name_suffix=pre_suffix
             )
+            pre_info_pair = (pre_root, pre_prefix, pre_suffix)
 
             # get the intersection
-            eval_name_list = sorted(list(set(gt_name_list).intersection(pre_name_list)))
+            eval_name_list = sorted(set(gt_name_list).intersection(pre_name_list))
             if len(eval_name_list) == 0:
                 tqdm.write(f"{method_name} does not have results on {dataset_name}")
                 continue
 
+            desc = f"[{dataset_name}({len(gt_name_list)}):{method_name}({len(pre_name_list)})]"
             kwargs = dict(
                 names=eval_name_list,
                 num_bits=num_bits,
-                pre_root=pre_root,
-                pre_prefix=pre_prefix,
-                pre_suffix=pre_suffix,
-                gt_root=gt_root,
-                gt_prefix=gt_prefix,
-                gt_suffix=gt_suffix,
-                desc=f"[{dataset_name}({len(gt_name_list)}):{method_name}({len(pre_name_list)})]",
-                proc_idx=procs_idx,
+                pre_info_pair=pre_info_pair,
+                gt_info_pair=gt_info_pair,
                 metric_names=metric_names,
-                ncols_tqdm=ncols_tqdm,
                 metric_class=metric_class,
+                desc=desc,
             )
             callback = partial(recorder.record, dataset_name=dataset_name, method_name=method_name)
             procs.apply_async(func=evaluate, kwds=kwargs, callback=callback)
-            # for debugging
+            # print(" -------------------- [DEBUG] -------------------- ")
             # callback(evaluate(**kwargs), dataset_name=dataset_name, method_name=method_name)
-            procs_idx += 1
     procs.close()
     procs.join()
 
@@ -257,38 +243,22 @@ def cal_image_matrics(
     tqdm.write(f"All methods have been evaluated:\n{formatted_string}")
 
 
-def evaluate(
-    names,
-    num_bits,
-    gt_root,
-    gt_prefix,
-    gt_suffix,
-    pre_root,
-    pre_prefix,
-    pre_suffix,
-    metric_class,
-    desc="",
-    proc_idx=None,
-    metric_names=None,
-    ncols_tqdm=79,
-):
+def evaluate(names, num_bits, pre_info_pair, gt_info_pair, metric_class, metric_names, desc=""):
     metric_recoder = metric_class(metric_names=metric_names)
     # https://github.com/tqdm/tqdm#parameters
     # https://github.com/tqdm/tqdm/blob/master/examples/parallel_bars.py
-    for name in tqdm(
-        names, total=len(names), desc=desc, position=proc_idx, ncols=ncols_tqdm, lock_args=(False,)
-    ):
+    for name in tqdm(names, total=len(names), desc=desc, ncols=79, lock_args=(False,)):
         gt, pre = get_gt_pre_with_name(
             img_name=name,
-            pre_root=pre_root,
-            pre_prefix=pre_prefix,
-            pre_suffix=pre_suffix,
-            gt_root=gt_root,
-            gt_prefix=gt_prefix,
-            gt_suffix=gt_suffix,
+            pre_root=pre_info_pair[0],
+            pre_prefix=pre_info_pair[1],
+            pre_suffix=pre_info_pair[2],
+            gt_root=gt_info_pair[0],
+            gt_prefix=gt_info_pair[1],
+            gt_suffix=gt_info_pair[2],
             to_normalize=False,
         )
-        metric_recoder.step(pre=pre, gt=gt, gt_path=os.path.join(gt_root, name))
+        metric_recoder.step(pre=pre, gt=gt, gt_path=os.path.join(gt_info_pair[0], name))
 
     method_results = metric_recoder.show(num_bits=num_bits, return_ndarray=False)
     return method_results