Upload 20 files

utils/io/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2021/5/17
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+
+from .image import read_color_array, read_gray_array
+from .params import load_weight, save_weight

utils/io/image.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2021/5/17
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+import cv2
+import numpy as np
+
+from utils.ops import minmax
+
+
+def read_gray_array(path, div_255=False, to_normalize=False, thr=-1, dtype=np.float32) -> np.ndarray:
+    """
+    1. read the binary image with the suffix `.jpg` or `.png`
+        into a grayscale ndarray
+    2. (to_normalize=True) rescale the ndarray to [0, 1]
+    3. (thr >= 0) binarize the ndarray with `thr`
+    4. return a gray ndarray (np.float32)
+    """
+    assert path.endswith(".jpg") or path.endswith(".png"), path
+    assert not div_255 or not to_normalize, path
+    gray_array = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+    assert gray_array is not None, f"Image Not Found: {path}"
+
+    if div_255:
+        gray_array = gray_array / 255
+
+    if to_normalize:
+        gray_array = minmax(gray_array, up_bound=255)
+
+    if thr >= 0:
+        gray_array = gray_array > thr
+
+    return gray_array.astype(dtype)
+
+
+def read_color_array(path: str):
+    assert path.endswith(".jpg") or path.endswith(".png")
+    bgr_array = cv2.imread(path, cv2.IMREAD_COLOR)
+    assert bgr_array is not None, f"Image Not Found: {path}"
+    rgb_array = cv2.cvtColor(bgr_array, cv2.COLOR_BGR2RGB)
+    return rgb_array

utils/io/params.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2020/12/19
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+
+import os
+
+import torch
+
+
+def save_weight(save_path, model):
+    print(f"Saving weight '{save_path}'")
+    if isinstance(model, dict):
+        model_state = model
+    else:
+        model_state = model.module.state_dict() if hasattr(model, "module") else model.state_dict()
+    torch.save(model_state, save_path)
+    print(f"Saved weight '{save_path}' " f"(only contain the net's weight)")
+
+
+def load_weight(load_path, model, *, strict=True, skip_unmatched_shape=False):
+    assert os.path.exists(load_path), load_path
+    model_params = model.state_dict()
+    for k, v in torch.load(load_path, map_location="cpu").items():
+        if k.endswith("module."):
+            k = k[7:]
+        if skip_unmatched_shape and k in model_params and v.shape != model_params[k].shape:
+            continue
+        model_params[k] = v
+    model.load_state_dict(model_params, strict=strict)

utils/ops/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2020/12/19
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+
+
+from .array_ops import *
+from .tensor_ops import *

utils/ops/array_ops.py

@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+import os
+
+import cv2
+import numpy as np
+
+
+def minmax(data_array: np.ndarray, up_bound: float = None) -> np.ndarray:
+    """
+    ::
+
+        data_array = (data_array / up_bound)
+        if min_value != max_value:
+            data_array = (data_array - min_value) / (max_value - min_value)
+
+    :param data_array:
+    :param up_bound: if is not None, data_array will devided by it before the minmax ops.
+    :return:
+    """
+    if up_bound is not None:
+        data_array = data_array / up_bound
+    max_value = data_array.max()
+    min_value = data_array.min()
+    if max_value != min_value:
+        data_array = (data_array - min_value) / (max_value - min_value)
+    return data_array
+
+
+def clip_to_normalize(data_array: np.ndarray, clip_range: tuple = None) -> np.ndarray:
+    clip_range = sorted(clip_range)
+    if len(clip_range) == 3:
+        clip_min, clip_mid, clip_max = clip_range
+        assert 0 <= clip_min < clip_mid < clip_max <= 1, clip_range
+        lower_array = data_array[data_array < clip_mid]
+        higher_array = data_array[data_array > clip_mid]
+        if lower_array.size > 0:
+            lower_array = np.clip(lower_array, a_min=clip_min, a_max=1)
+            max_lower = lower_array.max()
+            lower_array = minmax(lower_array) * max_lower
+            data_array[data_array < clip_mid] = lower_array
+        if higher_array.size > 0:
+            higher_array = np.clip(higher_array, a_min=0, a_max=clip_max)
+            min_lower = higher_array.min()
+            higher_array = minmax(higher_array) * (1 - min_lower) + min_lower
+            data_array[data_array > clip_mid] = higher_array
+    elif len(clip_range) == 2:
+        clip_min, clip_max = clip_range
+        assert 0 <= clip_min < clip_max <= 1, clip_range
+        if clip_min != 0 and clip_max != 1:
+            data_array = np.clip(data_array, a_min=clip_min, a_max=clip_max)
+        data_array = minmax(data_array)
+    elif clip_range is None:
+        data_array = minmax(data_array)
+    else:
+        raise NotImplementedError
+    return data_array
+
+
+def save_array_as_image(data_array: np.ndarray, save_name: str, save_dir: str, to_minmax: bool = False):
+    """
+    save the ndarray as a image
+
+    Args:
+        data_array: np.float32 the max value is less than or equal to 1
+        save_name: with special suffix
+        save_dir: the dirname of the image path
+        to_minmax: minmax the array
+    """
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    save_path = os.path.join(save_dir, save_name)
+    if data_array.dtype != np.uint8:
+        if data_array.max() > 1:
+            raise Exception("the range of data_array has smoe errors")
+        data_array = (data_array * 255).astype(np.uint8)
+    if to_minmax:
+        data_array = minmax(data_array, up_bound=255)
+        data_array = (data_array * 255).astype(np.uint8)
+    cv2.imwrite(save_path, data_array)
+
+
+def resize(image_array: np.ndarray, height, width, interpolation=cv2.INTER_LINEAR):
+    h, w = image_array.shape[:2]
+    if h == height and w == width:
+        return image_array
+
+    resized_image_array = cv2.resize(image_array, dsize=(width, height), interpolation=interpolation)
+    return resized_image_array
+
+
+def ms_resize(img, scales, base_h=None, base_w=None, interpolation=cv2.INTER_LINEAR):
+    assert isinstance(scales, (list, tuple))
+    if base_h is None:
+        base_h = img.shape[0]
+    if base_w is None:
+        base_w = img.shape[1]
+    return [resize(img, height=int(base_h * s), width=int(base_w * s), interpolation=interpolation) for s in scales]

utils/ops/tensor_ops.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2020
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+import torch
+import torch.nn.functional as F
+
+
+def rescale_2x(x: torch.Tensor, scale_factor=2):
+    return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False)
+
+
+def resize_to(x: torch.Tensor, tgt_hw: tuple):
+    return F.interpolate(x, size=tgt_hw, mode="bilinear", align_corners=False)
+
+
+def clip_grad(params, mode, clip_cfg: dict):
+    if mode == "norm":
+        if "max_norm" not in clip_cfg:
+            raise ValueError("`clip_cfg` must contain `max_norm`.")
+        torch.nn.utils.clip_grad_norm_(
+            params,
+            max_norm=clip_cfg.get("max_norm"),
+            norm_type=clip_cfg.get("norm_type", 2.0),
+        )
+    elif mode == "value":
+        if "clip_value" not in clip_cfg:
+            raise ValueError("`clip_cfg` must contain `clip_value`.")
+        torch.nn.utils.clip_grad_value_(params, clip_value=clip_cfg.get("clip_value"))
+    else:
+        raise NotImplementedError

utils/pipeline/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2021/5/31
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+
+from .optimizer import construct_optimizer
+from .scaler import Scaler
+from .scheduler import Scheduler

utils/pipeline/optimizer.py

@@ -0,0 +1,169 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2020/12/19
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+import types
+
+import torchvision.models
+from torch import nn
+from torch.optim import SGD, Adam, AdamW
+
+
+def get_optimizer(mode, params, initial_lr, optim_cfg):
+    if mode == "sgd":
+        optimizer = SGD(
+            params=params,
+            lr=initial_lr,
+            momentum=optim_cfg["momentum"],
+            weight_decay=optim_cfg["weight_decay"],
+            nesterov=optim_cfg.get("nesterov", False),
+        )
+    elif mode == "adamw":
+        optimizer = AdamW(
+            params=params,
+            lr=initial_lr,
+            betas=optim_cfg.get("betas", (0.9, 0.999)),
+            weight_decay=optim_cfg.get("weight_decay", 0),
+            amsgrad=optim_cfg.get("amsgrad", False),
+        )
+    elif mode == "adam":
+        optimizer = Adam(
+            params=params,
+            lr=initial_lr,
+            betas=optim_cfg.get("betas", (0.9, 0.999)),
+            weight_decay=optim_cfg.get("weight_decay", 0),
+            amsgrad=optim_cfg.get("amsgrad", False),
+        )
+    else:
+        raise NotImplementedError(mode)
+    return optimizer
+
+
+def group_params(model: nn.Module, group_mode: str, initial_lr: float, optim_cfg: dict):
+    if group_mode == "yolov5":
+        """
+        norm, weight, bias = [], [], []  # optimizer parameter groups
+        for k, v in model.named_modules():
+            if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
+                bias.append(v.bias)  # biases
+            if isinstance(v, nn.BatchNorm2d):
+                norm.append(v.weight)  # no decay
+            elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
+                weight.append(v.weight)  # apply decay
+
+        if opt.adam:
+            optimizer = optim.Adam(norm, lr=hyp["lr0"], betas=(hyp["momentum"], 0.999))  # adjust beta1 to momentum
+        else:
+            optimizer = optim.SGD(norm, lr=hyp["lr0"], momentum=hyp["momentum"], nesterov=True)
+
+        optimizer.add_param_group({"params": weight, "weight_decay": hyp["weight_decay"]})  # add weight with weight_decay
+        optimizer.add_param_group({"params": bias})  # add bias (biases)
+        """
+        norm, weight, bias = [], [], []  # optimizer parameter groups
+        for k, v in model.named_modules():
+            if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
+                bias.append(v.bias)  # conv bias and bn bias
+            if isinstance(v, nn.BatchNorm2d):
+                norm.append(v.weight)  # bn weight
+            elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
+                weight.append(v.weight)  # conv weight
+        params = [
+            {"params": filter(lambda p: p.requires_grad, bias), "weight_decay": 0.0},
+            {"params": filter(lambda p: p.requires_grad, norm), "weight_decay": 0.0},
+            {"params": filter(lambda p: p.requires_grad, weight)},
+        ]
+    elif group_mode == "r3":
+        params = [
+            # 不对bias参数执行weight decay操作，weight decay主要的作用就是通过对网络
+            # 层的参数（包括weight和bias）做约束（L2正则化会使得网络层的参数更加平滑）达
+            # 到减少模型过拟合的效果。
+            {
+                "params": [
+                    param for name, param in model.named_parameters() if name[-4:] == "bias" and param.requires_grad
+                ],
+                "lr": 2 * initial_lr,
+                "weight_decay": 0,
+            },
+            {
+                "params": [
+                    param for name, param in model.named_parameters() if name[-4:] != "bias" and param.requires_grad
+                ],
+                "lr": initial_lr,
+                "weight_decay": optim_cfg["weight_decay"],
+            },
+        ]
+    elif group_mode == "all":
+        params = model.parameters()
+    elif group_mode == "finetune":
+        if hasattr(model, "module"):
+            model = model.module
+        assert hasattr(model, "get_grouped_params"), "Cannot get the method get_grouped_params of the model."
+        params_groups = model.get_grouped_params()
+        params = [
+            {
+                "params": filter(lambda p: p.requires_grad, params_groups["pretrained"]),
+                "lr": optim_cfg.get("diff_factor", 0.1) * initial_lr,
+            },
+            {
+                "params": filter(lambda p: p.requires_grad, params_groups["retrained"]),
+                "lr": initial_lr,
+            },
+        ]
+    elif group_mode == "finetune2":
+        if hasattr(model, "module"):
+            model = model.module
+        assert hasattr(model, "get_grouped_params"), "Cannot get the method get_grouped_params of the model."
+        params_groups = model.get_grouped_params()
+        params = [
+            {
+                "params": filter(lambda p: p.requires_grad, params_groups["pretrained_backbone"]),
+                "lr": 0.1 * initial_lr,
+            },
+            {
+                "params": filter(lambda p: p.requires_grad, params_groups["pretrained_siamese"]),
+                "lr": 0.5 * initial_lr,
+            },
+            {
+                "params": filter(lambda p: p.requires_grad, params_groups["retrained"]),
+                "lr": initial_lr,
+            },
+        ]
+    else:
+        raise NotImplementedError
+    return params
+
+
+def construct_optimizer(model, initial_lr, mode, group_mode, cfg):
+    params = group_params(model, group_mode=group_mode, initial_lr=initial_lr, optim_cfg=cfg)
+    optimizer = get_optimizer(mode=mode, params=params, initial_lr=initial_lr, optim_cfg=cfg)
+    optimizer.lr_groups = types.MethodType(get_lr_groups, optimizer)
+    optimizer.lr_string = types.MethodType(get_lr_strings, optimizer)
+    return optimizer
+
+
+def get_lr_groups(self):
+    return [group["lr"] for group in self.param_groups]
+
+
+def get_lr_strings(self):
+    return ",".join([f"{group['lr']:.3e}" for group in self.param_groups])
+
+
+if __name__ == "__main__":
+    model = torchvision.models.vgg11_bn()
+    norm, weight, bias = [], [], []  # optimizer parameter groups
+    for k, v in model.named_modules():
+        if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
+            bias.append(v.bias)  # biases
+        if isinstance(v, nn.BatchNorm2d):
+            norm.append(v.weight)  # no decay
+        elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
+            weight.append(v.weight)  # apply decay
+
+    optimizer = Adam(norm, lr=0.001, betas=(0.98, 0.999))  # adjust beta1 to momentum
+    # optimizer = optim.SGD(norm, lr=hyp["lr0"], momentum=hyp["momentum"], nesterov=True)
+
+    optimizer.add_param_group({"params": weight, "weight_decay": 1e-4})  # add weight with weight_decay
+    optimizer.add_param_group({"params": bias})  # add bias (biases)
+
+    print(optimizer)

utils/pipeline/scaler.py

@@ -0,0 +1,59 @@
+from functools import partial
+from itertools import chain
+
+from torch.cuda.amp import GradScaler, autocast
+
+from .. import ops
+
+
+class Scaler:
+    def __init__(
+        self, optimizer, use_fp16=False, *, set_to_none=False, clip_grad=False, clip_mode=None, clip_cfg=None
+    ) -> None:
+        self.optimizer = optimizer
+        self.set_to_none = set_to_none
+        self.autocast = autocast(enabled=use_fp16)
+        self.scaler = GradScaler(enabled=use_fp16)
+
+        if clip_grad:
+            self.grad_clip_ops = partial(ops.clip_grad, mode=clip_mode, clip_cfg=clip_cfg)
+        else:
+            self.grad_clip_ops = None
+
+    def calculate_grad(self, loss):
+        self.scaler.scale(loss).backward()
+        if self.grad_clip_ops is not None:
+            self.scaler.unscale_(self.optimizer)
+            self.grad_clip_ops(chain(*[group["params"] for group in self.optimizer.param_groups]))
+
+    def update_grad(self):
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.optimizer.zero_grad(set_to_none=self.set_to_none)
+
+    def state_dict(self):
+        r"""
+        Returns the state of the scaler as a :class:`dict`.  It contains five entries:
+
+        * ``"scale"`` - a Python float containing the current scale
+        * ``"growth_factor"`` - a Python float containing the current growth factor
+        * ``"backoff_factor"`` - a Python float containing the current backoff factor
+        * ``"growth_interval"`` - a Python int containing the current growth interval
+        * ``"_growth_tracker"`` - a Python int containing the number of recent consecutive unskipped steps.
+
+        If this instance is not enabled, returns an empty dict.
+
+        .. note::
+           If you wish to checkpoint the scaler's state after a particular iteration, :meth:`state_dict`
+           should be called after :meth:`update`.
+        """
+        return self.scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        r"""
+        Loads the scaler state.  If this instance is disabled, :meth:`load_state_dict` is a no-op.
+
+        Args:
+           state_dict(dict): scaler state.  Should be an object returned from a call to :meth:`state_dict`.
+        """
+        self.scaler.load_state_dict(state_dict)

utils/pipeline/scheduler.py

@@ -0,0 +1,410 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2020/12/19
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+
+import copy
+import math
+import os.path
+import warnings
+from bisect import bisect_right
+
+import matplotlib
+import numpy as np
+import torch.optim
+from adjustText import adjust_text
+
+matplotlib.use("Agg")
+from matplotlib import pyplot as plt
+
+# helper function ----------------------------------------------------------------------
+
+
+def linear_increase(low_bound, up_bound, percentage):
+    """low_bound + [0, 1] * (up_bound - low_bound)"""
+    assert 0 <= percentage <= 1, f"percentage({percentage}) must be in [0, 1]"
+    return low_bound + (up_bound - low_bound) * percentage
+
+
+def cos_anneal(low_bound, up_bound, percentage):
+    assert 0 <= percentage <= 1, f"percentage({percentage}) must be in [0, 1]"
+    cos_percentage = (1 + math.cos(math.pi * percentage)) / 2.0
+    return linear_increase(low_bound, up_bound, percentage=cos_percentage)
+
+
+def poly_anneal(low_bound, up_bound, percentage, lr_decay):
+    assert 0 <= percentage <= 1, f"percentage({percentage}) must be in [0, 1]"
+    poly_percentage = pow((1 - percentage), lr_decay)
+    return linear_increase(low_bound, up_bound, percentage=poly_percentage)
+
+
+def linear_anneal(low_bound, up_bound, percentage):
+    assert 0 <= percentage <= 1, f"percentage({percentage}) must be in [0, 1]"
+    return linear_increase(low_bound, up_bound, percentage=1 - percentage)
+
+
+# coefficient function ----------------------------------------------------------------------
+
+
+def get_f3_coef_func(num_iters):
+    """
+    F3Net
+
+    :param num_iters: The number of iterations for the total process.
+    :return:
+    """
+
+    def get_f3_coef(curr_idx):
+        assert 0 <= curr_idx <= num_iters
+        return 1 - abs((curr_idx + 1) / (num_iters + 1) * 2 - 1)
+
+    return get_f3_coef
+
+
+def get_step_coef_func(gamma, milestones):
+    """
+    lr = baselr * gamma ** 0    if curr_idx < milestones[0]
+    lr = baselr * gamma ** 1   if milestones[0] <= epoch < milestones[1]
+    ...
+
+    :param gamma:
+    :param milestones:
+    :return: The function for generating the coefficient.
+    """
+    if isinstance(milestones, (tuple, list)):
+        milestones = list(sorted(milestones))
+        return lambda curr_idx: gamma ** bisect_right(milestones, curr_idx)
+    elif isinstance(milestones, int):
+        return lambda curr_idx: gamma ** ((curr_idx + 1) // milestones)
+    else:
+        raise ValueError(f"milestones only can be list/tuple/int, but now it is {type(milestones)}")
+
+
+def get_cos_coef_func(half_cycle, min_coef, max_coef=1):
+    """
+    :param half_cycle: The number of iterations in a half cycle.
+    :param min_coef: The minimum coefficient of the learning rate.
+    :param max_coef: The maximum coefficient of the learning rate.
+    :return: The function for generating the coefficient.
+    """
+
+    def get_cos_coef(curr_idx):
+        recomputed_idx = curr_idx % (half_cycle + 1)
+        # recomputed \in [0, half_cycle]
+        return cos_anneal(low_bound=min_coef, up_bound=max_coef, percentage=recomputed_idx / half_cycle)
+
+    return get_cos_coef
+
+
+def get_fatcos_coef_func(start_iter, half_cycle, min_coef, max_coef=1):
+    """
+    :param half_cycle: The number of iterations in a half cycle.
+    :param min_coef: The minimum coefficient of the learning rate.
+    :param max_coef: The maximum coefficient of the learning rate.
+    :return: The function for generating the coefficient.
+    """
+
+    def get_cos_coef(curr_idx):
+        curr_idx = max(0, curr_idx - start_iter)
+        recomputed_idx = curr_idx % (half_cycle + 1)
+        # recomputed \in [0, half_cycle]
+        return cos_anneal(low_bound=min_coef, up_bound=max_coef, percentage=recomputed_idx / half_cycle)
+
+    return get_cos_coef
+
+
+def get_poly_coef_func(num_iters, lr_decay, min_coef, max_coef=1):
+    """
+    :param num_iters: The number of iterations for the polynomial descent process.
+    :param lr_decay: The decay item of the polynomial descent process.
+    :param min_coef: The minimum coefficient of the learning rate.
+    :param max_coef: The maximum coefficient of the learning rate.
+    :return: The function for generating the coefficient.
+    """
+
+    def get_poly_coef(curr_idx):
+        assert 0 <= curr_idx <= num_iters, (curr_idx, num_iters)
+        return poly_anneal(low_bound=min_coef, up_bound=max_coef, percentage=curr_idx / num_iters, lr_decay=lr_decay)
+
+    return get_poly_coef
+
+
+# coefficient entry function ----------------------------------------------------------------------
+
+
+def get_scheduler_coef_func(mode, num_iters, cfg):
+    """
+    the region is a closed interval: [0, num_iters]
+    """
+    assert num_iters > 0
+    min_coef = cfg.get("min_coef", 1e-6)
+    if min_coef is None or min_coef == 0:
+        warnings.warn(f"The min_coef ({min_coef}) of the scheduler will be replaced with 1e-6")
+        min_coef = 1e-6
+
+    if mode == "step":
+        coef_func = get_step_coef_func(gamma=cfg["gamma"], milestones=cfg["milestones"])
+    elif mode == "cos":
+        if half_cycle := cfg.get("half_cycle"):
+            half_cycle -= 1
+        else:
+            half_cycle = num_iters
+        if (num_iters - half_cycle) % (half_cycle + 1) != 0:
+            # idx starts from 0
+            percentage = ((num_iters - half_cycle) % (half_cycle + 1)) / (half_cycle + 1) * 100
+            warnings.warn(
+                f"The final annealing process ({percentage:.3f}%) is not complete. "
+                f"Please pay attention to the generated 'lr_coef_curve.png'."
+            )
+        coef_func = get_cos_coef_func(half_cycle=half_cycle, min_coef=min_coef)
+    elif mode == "fatcos":
+        assert 0 <= cfg.start_percent < 1, cfg.start_percent
+        start_iter = int(cfg.start_percent * num_iters)
+
+        num_iters -= start_iter
+        if half_cycle := cfg.get("half_cycle"):
+            half_cycle -= 1
+        else:
+            half_cycle = num_iters
+        if (num_iters - half_cycle) % (half_cycle + 1) != 0:
+            # idx starts from 0
+            percentage = ((num_iters - half_cycle) % (half_cycle + 1)) / (half_cycle + 1) * 100
+            warnings.warn(
+                f"The final annealing process ({percentage:.3f}%) is not complete. "
+                f"Please pay attention to the generated 'lr_coef_curve.png'."
+            )
+        coef_func = get_fatcos_coef_func(start_iter=start_iter, half_cycle=half_cycle, min_coef=min_coef)
+    elif mode == "poly":
+        coef_func = get_poly_coef_func(num_iters=num_iters, lr_decay=cfg["lr_decay"], min_coef=min_coef)
+    elif mode == "constant":
+        coef_func = lambda x: cfg.get("coef", 1)
+    elif mode == "f3":
+        coef_func = get_f3_coef_func(num_iters=num_iters)
+    else:
+        raise NotImplementedError(f"{mode} must be in {Scheduler.supported_scheduler_modes}")
+    return coef_func
+
+
+def get_warmup_coef_func(num_iters, min_coef, max_coef=1, mode="linear"):
+    """
+    the region is a closed interval: [0, num_iters]
+    """
+    assert num_iters > 0
+    if mode == "cos":
+        anneal_func = cos_anneal
+    elif mode == "linear":
+        anneal_func = linear_anneal
+    else:
+        raise NotImplementedError(f"{mode} must be in {Scheduler.supported_warmup_modes}")
+
+    def get_warmup_coef(curr_idx):
+        return anneal_func(low_bound=min_coef, up_bound=max_coef, percentage=1 - curr_idx / num_iters)
+
+    return get_warmup_coef
+
+
+# main class ----------------------------------------------------------------------
+
+
+class Scheduler:
+    supported_scheduler_modes = ("step", "cos", "fatcos", "poly", "constant", "f3")
+    supported_warmup_modes = ("cos", "linear")
+
+    def __init__(self, optimizer, num_iters, epoch_length, scheduler_cfg, step_by_batch=True):
+        """A customized wrapper of the scheduler.
+
+        Args:
+            optimizer (): Optimizer.
+            num_iters (int): The total number of the iterations.
+            epoch_length (int): The number of the iterations of one epoch.
+            scheduler_cfg (dict): The config of the scheduler.
+            step_by_batch (bool, optional): The mode of updating the scheduler. Defaults to True.
+
+        Raises:
+            NotImplementedError:
+        """
+        self.optimizer = optimizer
+        self.num_iters = num_iters
+        self.epoch_length = epoch_length
+        self.step_by_batch = step_by_batch
+
+        self.scheduler_cfg = copy.deepcopy(scheduler_cfg)
+        self.mode = scheduler_cfg["mode"]
+        if self.mode not in self.supported_scheduler_modes:
+            raise NotImplementedError(
+                f"{self.mode} is not implemented. Has been supported: {self.supported_scheduler_modes}"
+            )
+        warmup_cfg = scheduler_cfg.get("warmup", None)
+
+        num_warmup_iters = 0
+        if warmup_cfg is not None and isinstance(warmup_cfg, dict):
+            num_warmup_iters = warmup_cfg["num_iters"]
+            if num_warmup_iters > 0:
+                print("Will using warmup")
+                self.warmup_coef_func = get_warmup_coef_func(
+                    num_warmup_iters,
+                    min_coef=warmup_cfg.get("initial_coef", 0.01),
+                    mode=warmup_cfg.get("mode", "linear"),
+                )
+        self.num_warmup_iters = num_warmup_iters
+
+        if step_by_batch:
+            num_scheduler_iters = num_iters - num_warmup_iters
+        else:
+            num_scheduler_iters = (num_iters - num_warmup_iters) // epoch_length
+        # the region is a closed interval
+        self.lr_coef_func = get_scheduler_coef_func(
+            mode=self.mode, num_iters=num_scheduler_iters - 1, cfg=scheduler_cfg["cfg"]
+        )
+        self.num_scheduler_iters = num_scheduler_iters
+
+        self.last_lr_coef = 0
+        self.initial_lrs = None
+
+    def __repr__(self):
+        formatted_string = [
+            f"{self.__class__.__name__}: (\n",
+            f"num_iters: {self.num_iters}\n",
+            f"epoch_length: {self.epoch_length}\n",
+            f"warmup_iter: [0, {self.num_warmup_iters})\n",
+            f"scheduler_iter: [{self.num_warmup_iters}, {self.num_iters - 1}]\n",
+            f"mode: {self.mode}\n",
+            f"scheduler_cfg: {self.scheduler_cfg}\n",
+            f"initial_lrs: {self.initial_lrs}\n",
+            f"step_by_batch: {self.step_by_batch}\n)",
+        ]
+        return "    ".join(formatted_string)
+
+    def record_lrs(self, param_groups):
+        self.initial_lrs = [g["lr"] for g in param_groups]
+
+    def update(self, coef: float):
+        assert self.initial_lrs is not None, "Please run .record_lrs(optimizer) first."
+        for curr_group, initial_lr in zip(self.optimizer.param_groups, self.initial_lrs):
+            curr_group["lr"] = coef * initial_lr
+
+    def step(self, curr_idx):
+        if curr_idx < self.num_warmup_iters:
+            # get maximum value (1.0) when curr_idx == self.num_warmup_iters
+            self.update(coef=self.get_lr_coef(curr_idx))
+        else:
+            # Start from a value lower than 1 (curr_idx == self.num_warmup_iters)
+            if self.step_by_batch:
+                self.update(coef=self.get_lr_coef(curr_idx))
+            else:
+                if curr_idx % self.epoch_length == 0:
+                    self.update(coef=self.get_lr_coef(curr_idx))
+
+    def get_lr_coef(self, curr_idx):
+        coef = None
+        if curr_idx < self.num_warmup_iters:
+            coef = self.warmup_coef_func(curr_idx)
+        else:
+            # when curr_idx == self.num_warmup_iters, coef == 1.0
+            # down from the largest coef (1.0)
+            if self.step_by_batch:
+                coef = self.lr_coef_func(curr_idx - self.num_warmup_iters)
+            else:
+                if curr_idx % self.epoch_length == 0 or curr_idx == self.num_warmup_iters:
+                    # warmup结束后尚未开始按照epoch进行调整的学习率调整，此时需要将系数调整为最大。
+                    coef = self.lr_coef_func((curr_idx - self.num_warmup_iters) // self.epoch_length)
+        if coef is not None:
+            self.last_lr_coef = coef
+        return self.last_lr_coef
+
+    def plot_lr_coef_curve(self, save_path=""):
+        plt.rc("xtick", labelsize="small")
+        plt.rc("ytick", labelsize="small")
+
+        fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(8, 4), dpi=600)
+        # give plot a title
+        ax.set_title("Learning Rate Coefficient Curve")
+        # make axis labels
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("Coefficient")
+
+        x_data = np.arange(self.num_iters)
+        y_data = np.array([self.get_lr_coef(x) for x in x_data])
+
+        # set lim
+        x_min, x_max = 0, self.num_iters - 1
+        dx = self.num_iters * 0.1
+        ax.set_xlim(x_min - dx, x_max + 2 * dx)
+
+        y_min, y_max = y_data.min(), y_data.max()
+        dy = (y_data.max() - y_data.min()) * 0.1
+        ax.set_ylim((y_min - dy, y_max + dy))
+
+        if self.step_by_batch:
+            marker_on = [0, -1]
+            key_point_xs = [0, self.num_iters - 1]
+            for idx in range(1, len(y_data) - 1):
+                prev_y = y_data[idx - 1]
+                curr_y = y_data[idx]
+                next_y = y_data[idx + 1]
+                if (
+                    (curr_y > prev_y and curr_y >= next_y)
+                    or (curr_y >= prev_y and curr_y > next_y)
+                    or (curr_y <= prev_y and curr_y < next_y)
+                    or (curr_y < prev_y and curr_y <= next_y)
+                ):
+                    marker_on.append(idx)
+                    key_point_xs.append(idx)
+
+            marker_on = sorted(set(marker_on))
+            key_point_xs = sorted(set(key_point_xs))
+            key_point_ys = []
+
+            texts = []
+            for x in key_point_xs:
+                y = y_data[x]
+                key_point_ys.append(y)
+
+                texts.append(ax.text(x=x, y=y, s=f"({x:d},{y:.3e})"))
+            adjust_text(texts, arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=-0.3"))
+
+            # set ticks
+            ax.set_xticks(key_point_xs)
+            # ax.set_yticks(key_point_ys)
+
+            ax.plot(x_data, y_data, marker="o", markevery=marker_on)
+        else:
+            ax.plot(x_data, y_data)
+
+        ax.spines["right"].set_visible(False)
+        ax.spines["top"].set_visible(False)
+        ax.spines["left"].set_visible(True)
+        ax.spines["bottom"].set_visible(True)
+
+        plt.tight_layout()
+        if save_path:
+            fig.savefig(os.path.join(save_path, "lr_coef.png"))
+        plt.close()
+
+
+if __name__ == "__main__":
+    model = torch.nn.Conv2d(10, 10, 3, 1, 1)
+    sche = Scheduler(
+        optimizer=torch.optim.SGD(model.parameters(), lr=0.1),
+        num_iters=30300,
+        epoch_length=505,
+        scheduler_cfg=dict(
+            warmup=dict(
+                num_iters=6060,
+                initial_coef=0.01,
+                mode="cos",
+            ),
+            mode="cos",
+            cfg=dict(
+                half_cycle=6060,
+                lr_decay=0.9,
+                min_coef=0.001,
+            ),
+        ),
+        step_by_batch=True,
+    )
+    print(sche)
+    sche.plot_lr_coef_curve(
+        # save_path="/home/lart/Coding/SOD.torch",
+        show=True,
+    )

utils/pt_utils.py

@@ -0,0 +1,66 @@
+import logging
+import os
+import random
+
+import numpy as np
+import torch
+from torch import nn
+from torch.backends import cudnn
+
+LOGGER = logging.getLogger("main")
+
+
+def customized_worker_init_fn(worker_id):
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+
+
+def set_seed_for_lib(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def initialize_seed_cudnn(seed, deterministic):
+    assert isinstance(deterministic, bool) and isinstance(seed, int)
+    if seed >= 0:
+        LOGGER.info(f"We will use a fixed seed {seed}")
+    else:
+        seed = np.random.randint(2**32)
+        LOGGER.info(f"We will use a random seed {seed}")
+    set_seed_for_lib(seed)
+    if not deterministic:
+        LOGGER.info("We will use `torch.backends.cudnn.benchmark`")
+    else:
+        LOGGER.info("We will not use `torch.backends.cudnn.benchmark`")
+    cudnn.enabled = True
+    cudnn.benchmark = not deterministic
+    cudnn.deterministic = deterministic
+
+
+def to_device(data, device="cuda"):
+    if isinstance(data, (tuple, list)):
+        return [to_device(item, device) for item in data]
+    elif isinstance(data, dict):
+        return {name: to_device(item, device) for name, item in data.items()}
+    elif isinstance(data, torch.Tensor):
+        return data.to(device=device, non_blocking=True)
+    else:
+        raise TypeError(f"Unsupported type {type(data)}. Only support Tensor or tuple/list/dict containing Tensors.")
+
+
+def frozen_bn_stats(model, freeze_affine=False):
+    """
+    Set all the bn layers to eval mode.
+    Args:
+        model (model): model to set bn layers to eval mode.
+    """
+    for m in model.modules():
+        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
+            m.eval()
+            if freeze_affine:
+                m.requires_grad_(False)

utils/py_utils.py

@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+import copy
+import logging
+import os
+import shutil
+from collections import OrderedDict, abc
+from datetime import datetime
+
+LOGGER = logging.getLogger("main")
+
+
+def construct_path(output_dir: str, exp_name: str) -> dict:
+    proj_root = os.path.join(output_dir, exp_name)
+    exp_idx = 0
+    exp_output_dir = os.path.join(proj_root, f"exp_{exp_idx}")
+    while os.path.exists(exp_output_dir):
+        exp_idx += 1
+        exp_output_dir = os.path.join(proj_root, f"exp_{exp_idx}")
+
+    tb_path = os.path.join(exp_output_dir, "tb")
+    save_path = os.path.join(exp_output_dir, "pre")
+    pth_path = os.path.join(exp_output_dir, "pth")
+
+    final_full_model_path = os.path.join(pth_path, "checkpoint_final.pth")
+    final_state_path = os.path.join(pth_path, "state_final.pth")
+
+    log_path = os.path.join(exp_output_dir, f"log_{str(datetime.now())[:10]}.txt")
+    cfg_copy_path = os.path.join(exp_output_dir, f"config.py")
+    trainer_copy_path = os.path.join(exp_output_dir, f"trainer.txt")
+    excel_path = os.path.join(exp_output_dir, f"results.xlsx")
+
+    path_config = {
+        "output_dir": output_dir,
+        "pth_log": exp_output_dir,
+        "tb": tb_path,
+        "save": save_path,
+        "pth": pth_path,
+        "final_full_net": final_full_model_path,
+        "final_state_net": final_state_path,
+        "log": log_path,
+        "cfg_copy": cfg_copy_path,
+        "excel": excel_path,
+        "trainer_copy": trainer_copy_path,
+    }
+
+    return path_config
+
+
+def construct_exp_name(model_name: str, cfg: dict):
+    # bs_16_lr_0.05_e30_noamp_2gpu_noms_352
+    focus_item = OrderedDict(
+        {
+            "train/batch_size": "bs",
+            "train/lr": "lr",
+            "train/num_epochs": "e",
+            "train/num_iters": "i",
+            "train/data/shape/h": "h",
+            "train/data/shape/w": "w",
+            "train/optimizer/mode": "opm",
+            "train/optimizer/group_mode": "opgm",
+            "train/scheduler/mode": "sc",
+            "train/scheduler/warmup/num_iters": "wu",
+            "train/use_amp": "amp",
+        }
+    )
+    config = copy.deepcopy(cfg)
+
+    def _format_item(_i):
+        if isinstance(_i, bool):
+            _i = "" if _i else "false"
+        elif isinstance(_i, (int, float)):
+            if _i == 0:
+                _i = "false"
+        elif isinstance(_i, (list, tuple)):
+            _i = "" if _i else "false"  # 只是判断是否非空
+        elif isinstance(_i, str):
+            if "_" in _i:
+                _i = _i.replace("_", "").lower()
+        elif _i is None:
+            _i = "none"
+        # else: other types and values will be returned directly
+        return _i
+
+    if (epoch_based := config.train.get("epoch_based", None)) is not None and (not epoch_based):
+        focus_item.pop("train/num_epochs")
+    else:
+        # 默认基于epoch
+        focus_item.pop("train/num_iters")
+
+    exp_names = [model_name]
+    for key, alias in focus_item.items():
+        item = get_value_recurse(keys=key.split("/"), info=config)
+        formatted_item = _format_item(item)
+        if formatted_item == "false":
+            continue
+        exp_names.append(f"{alias.upper()}{formatted_item}")
+
+    info = config.get("info", None)
+    if info:
+        exp_names.append(f"INFO{info.lower()}")
+
+    return "_".join(exp_names)
+
+
+def pre_mkdir(path_config):
+    # 提前创建好记录文件，避免自动创建的时候触发文件创建事件
+    check_mkdir(path_config["pth_log"])
+    make_log(path_config["log"], f"=== log {datetime.now()} ===")
+
+    # 提前创建好存储预测结果和存放模型的文件夹
+    check_mkdir(path_config["save"])
+    check_mkdir(path_config["pth"])
+
+
+def check_mkdir(dir_name, delete_if_exists=False):
+    if not os.path.exists(dir_name):
+        os.makedirs(dir_name)
+    else:
+        if delete_if_exists:
+            print(f"{dir_name} will be re-created!!!")
+            shutil.rmtree(dir_name)
+            os.makedirs(dir_name)
+
+
+def make_log(path, context):
+    with open(path, "a") as log:
+        log.write(f"{context}\n")
+
+
+def iterate_nested_sequence(nested_sequence):
+    """
+    当前支持list/tuple/int/float/range()的多层嵌套，注意不要嵌套的太深，小心超出python默认的最大递归深度
+
+    例子
+    ::
+
+        for x in iterate_nested_sequence([[1, (2, 3)], range(3, 10), 0]):
+            print(x)
+
+        1
+        2
+        3
+        3
+        4
+        5
+        6
+        7
+        8
+        9
+        0
+
+    :param nested_sequence: 多层嵌套的序列
+    :return: generator
+    """
+    for item in nested_sequence:
+        if isinstance(item, (int, float)):
+            yield item
+        elif isinstance(item, (list, tuple, range)):
+            yield from iterate_nested_sequence(item)
+        else:
+            raise NotImplementedError
+
+
+def get_value_recurse(keys: list, info: dict):
+    curr_key, sub_keys = keys[0], keys[1:]
+
+    if (sub_info := info.get(curr_key, "NoKey")) == "NoKey":
+        raise KeyError(f"{curr_key} must be contained in {info}")
+
+    if sub_keys:
+        return get_value_recurse(keys=sub_keys, info=sub_info)
+    else:
+        return sub_info
+
+
+def mapping_to_str(mapping: abc.Mapping, *, prefix: str = "    ", lvl: int = 0, max_lvl: int = 1) -> str:
+    """
+    Print the structural information of the dict.
+    """
+    sub_lvl = lvl + 1
+    cur_prefix = prefix * lvl
+    sub_prefix = prefix * sub_lvl
+
+    if lvl == max_lvl:
+        sub_items = str(mapping)
+    else:
+        sub_items = ["{"]
+        for k, v in mapping.items():
+            sub_item = sub_prefix + k + ": "
+            if isinstance(v, abc.Mapping):
+                sub_item += mapping_to_str(v, prefix=prefix, lvl=sub_lvl, max_lvl=max_lvl)
+            else:
+                sub_item += str(v)
+            sub_items.append(sub_item)
+        sub_items.append(cur_prefix + "}")
+        sub_items = "\n".join(sub_items)
+    return sub_items

utils/recorder/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+
+from .counter import TrainingCounter
+from .group_metric_caller import GroupedMetricRecorder
+from .logger import TBLogger
+from .meter_recorder import AvgMeter, HistoryBuffer
+from .visualize_results import plot_results

utils/recorder/counter.py

@@ -0,0 +1,75 @@
+import math
+
+
+class TrainingCounter:
+    def __init__(self, epoch_length, epoch_based=True, *, num_epochs=None, num_total_iters=None) -> None:
+        self.num_inner_iters = epoch_length
+        self._iter_counter = 0
+        self._epoch_counter = 0
+
+        if epoch_based:
+            assert num_epochs is not None
+            self.num_epochs = num_epochs
+            self.num_total_iters = num_epochs * epoch_length
+        else:
+            assert num_total_iters is not None
+            self.num_total_iters = num_total_iters
+            self.num_epochs = math.ceil(num_total_iters / epoch_length)
+
+    def set_start_epoch(self, start_epoch):
+        self._epoch_counter = start_epoch
+        self._iter_counter = start_epoch * self.num_inner_iters
+
+    def set_start_iterations(self, start_iteration):
+        self._iter_counter = start_iteration
+        self._epoch_counter = start_iteration // self.num_inner_iters
+
+    def every_n_epochs(self, n: int) -> bool:
+        return (self._epoch_counter + 1) % n == 0 if n > 0 else False
+
+    def every_n_iters(self, n: int) -> bool:
+        return (self._iter_counter + 1) % n == 0 if n > 0 else False
+
+    def is_first_epoch(self) -> bool:
+        return self._epoch_counter == 0
+
+    def is_last_epoch(self) -> bool:
+        return self._epoch_counter == self.num_epochs - 1
+
+    def is_first_inner_iter(self) -> bool:
+        return self._iter_counter % self.num_inner_iters == 0
+
+    def is_last_inner_iter(self) -> bool:
+        return (self._iter_counter + 1) % self.num_inner_iters == 0
+
+    def is_first_total_iter(self) -> bool:
+        return self._iter_counter == 0
+
+    def is_last_total_iter(self) -> bool:
+        return self._iter_counter == self.num_total_iters - 1
+
+    def update_iter_counter(self):
+        self._iter_counter += 1
+
+    def update_epoch_counter(self):
+        self._epoch_counter += 1
+
+    def reset_iter_all_counter(self):
+        self._iter_counter = 0
+        self._epoch_counter = 0
+
+    @property
+    def curr_iter(self):
+        return self._iter_counter
+
+    @property
+    def next_iter(self):
+        return self._iter_counter + 1
+
+    @property
+    def curr_epoch(self):
+        return self._epoch_counter
+
+    @property
+    def curr_percent(self):
+        return self._iter_counter / self.num_total_iters

utils/recorder/group_metric_caller.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2021/1/4
+# @Author  : Lart Pang
+# @GitHub  : https://github.com/lartpang
+
+from collections import OrderedDict
+
+import numpy as np
+import py_sod_metrics
+
+
+def ndarray_to_basetype(data):
+
+
+    def _to_list_or_scalar(item):
+        listed_item = item.tolist()
+        if isinstance(listed_item, list) and len(listed_item) == 1:
+            listed_item = listed_item[0]
+        return listed_item
+
+    if isinstance(data, (tuple, list)):
+        results = [_to_list_or_scalar(item) for item in data]
+    elif isinstance(data, dict):
+        results = {k: _to_list_or_scalar(item) for k, item in data.items()}
+    else:
+        assert isinstance(data, np.ndarray)
+        results = _to_list_or_scalar(data)
+    return results
+
+
+def round_w_zero_padding(x, bit_width):
+    x = str(round(x, bit_width))
+    x += "0" * (bit_width - len(x.split(".")[-1]))
+    return x
+
+
+INDIVADUAL_METRIC_MAPPING = {
+    "sm": py_sod_metrics.Smeasure,
+    "wfm": py_sod_metrics.WeightedFmeasure,
+    "mae": py_sod_metrics.MAE,
+    "em": py_sod_metrics.Emeasure,
+}
+BINARY_CLASSIFICATION_METRIC_MAPPING = {
+    "fmeasure": {
+        "handler": py_sod_metrics.FmeasureHandler,
+        "kwargs": dict(with_dynamic=True, with_adaptive=True, with_binary=False, beta=0.3),
+    },
+    "iou": {
+        "handler": py_sod_metrics.IOUHandler,
+        "kwargs": dict(with_dynamic=True, with_adaptive=True, with_binary=False),
+    },
+    "dice": {
+        "handler": py_sod_metrics.DICEHandler,
+        "kwargs": dict(with_dynamic=True, with_adaptive=False, with_binary=False),
+    },
+}
+
+
+class ImageMetricRecorder:
+    supported_metrics = sorted(INDIVADUAL_METRIC_MAPPING.keys()) + sorted(BINARY_CLASSIFICATION_METRIC_MAPPING.keys())
+
+    def __init__(self, metric_names=("sm", "wfm", "mae", "fmeasure", "em")):
+        
+        if not metric_names:
+            metric_names = self.supported_metrics
+        assert all([m in self.supported_metrics for m in metric_names]), f"Only support: {self.supported_metrics}"
+
+        self.metric_objs = {}
+        has_existed = False
+        for metric_name in metric_names:
+            if metric_name in INDIVADUAL_METRIC_MAPPING:
+                self.metric_objs[metric_name] = INDIVADUAL_METRIC_MAPPING[metric_name]()
+            else:  # metric_name in BINARY_CLASSIFICATION_METRIC_MAPPING
+                if not has_existed:  # only init once
+                    self.metric_objs["fmeasurev2"] = py_sod_metrics.FmeasureV2()
+                    has_existed = True
+                metric_handler = BINARY_CLASSIFICATION_METRIC_MAPPING[metric_name]
+                self.metric_objs["fmeasurev2"].add_handler(
+                    handler_name=metric_name, metric_handler=metric_handler["handler"](**metric_handler["kwargs"])
+                )
+
+    def step(self, pre: np.ndarray, gt: np.ndarray, gt_path: str):
+        assert pre.shape == gt.shape, (pre.shape, gt.shape, gt_path)
+        assert pre.dtype == gt.dtype == np.uint8, (pre.dtype, gt.dtype, gt_path)
+
+        for m_obj in self.metric_objs.values():
+            m_obj.step(pre, gt)
+
+    def get_all_results(self, num_bits: int = 3, return_ndarray: bool = False) -> dict:
+        sequential_results = {}
+        numerical_results = {}
+        for m_name, m_obj in self.metric_objs.items():
+            info = m_obj.get_results()
+            if m_name == "fmeasurev2":
+                for _name, results in info.items():
+                    dynamic_results = results.get("dynamic")
+                    adaptive_results = results.get("adaptive")
+                    if dynamic_results is not None:
+                        sequential_results[_name] = np.flip(dynamic_results)
+                        numerical_results[f"max{_name}"] = dynamic_results.max()
+                        numerical_results[f"avg{_name}"] = dynamic_results.mean()
+                    if adaptive_results is not None:
+                        numerical_results[f"adp{_name}"] = adaptive_results
+            else:
+                results = info[m_name]
+                if m_name in ("wfm", "sm", "mae"):
+                    numerical_results[m_name] = results
+                elif m_name == "em":
+                    sequential_results[m_name] = np.flip(results["curve"])
+                    numerical_results.update(
+                        {
+                            "maxem": results["curve"].max(),
+                            "avgem": results["curve"].mean(),
+                            "adpem": results["adp"],
+                        }
+                    )
+                else:
+                    raise NotImplementedError(m_name)
+
+        if num_bits is not None and isinstance(num_bits, int):
+            numerical_results = {k: v.round(num_bits) for k, v in numerical_results.items()}
+        if not return_ndarray:
+            sequential_results = ndarray_to_basetype(sequential_results)
+            numerical_results = ndarray_to_basetype(numerical_results)
+        return {"sequential": sequential_results, "numerical": numerical_results}
+
+    def show(self, num_bits: int = 3, return_ndarray: bool = False) -> dict:
+        return self.get_all_results(num_bits=num_bits, return_ndarray=return_ndarray)["numerical"]
+
+
+class GroupedMetricRecorder(object):
+    supported_metrics = ["mae", "em", "sm", "wfm"] + sorted(BINARY_CLASSIFICATION_METRIC_MAPPING.keys())
+
+    def __init__(self, group_names=None, metric_names=("sm", "wfm", "mae", "fmeasure", "em")):
+        self.group_names = group_names
+        self.metric_names = metric_names
+        self.zero()
+
+    def zero(self):
+        self.metric_recorders = {}
+        if self.group_names is not None:
+            self.metric_recorders.update(
+                {n: ImageMetricRecorder(metric_names=self.metric_names) for n in self.group_names}
+            )
+
+    def step(self, group_name: str, pre: np.ndarray, gt: np.ndarray, gt_path: str):
+        if group_name not in self.metric_recorders:
+            self.metric_recorders[group_name] = ImageMetricRecorder(metric_names=self.metric_names)
+        self.metric_recorders[group_name].step(pre, gt, gt_path)
+
+    def show(self, num_bits: int = 3, return_group: bool = False):
+        groups_metrics = {
+            n: r.get_all_results(num_bits=None, return_ndarray=True) for n, r in self.metric_recorders.items()
+        }
+
+        results = {}
+        for group_metrics in groups_metrics.values():
+            for metric_type, metric_group in group_metrics.items():  # sequential and numerical
+                results.setdefault(metric_type, {})
+                for metric_name, metric_array in metric_group.items():
+                    results[metric_type].setdefault(metric_name, []).append(metric_array)
+
+        numerical_results = {}
+        for metric_type, metric_group in results.items():
+            for metric_name, metric_array in metric_group.items():
+                metric_array = np.mean(np.vstack(metric_array), axis=0)  # average over all groups
+
+                if metric_name in BINARY_CLASSIFICATION_METRIC_MAPPING or metric_name == "em":
+                    if metric_type == "sequential":
+                        numerical_results[f"max{metric_name}"] = metric_array.max()
+                        numerical_results[f"avg{metric_name}"] = metric_array.mean()
+                else:
+                    if metric_type == "numerical":
+                        if metric_name.startswith(("max", "avg")):
+                            # these metrics (maxfm, avgfm, maxem, avgem) will be recomputed within the group
+                            continue
+                        numerical_results[metric_name] = metric_array
+
+        numerical_results = ndarray_to_basetype(numerical_results)
+        numerical_results = {k: round(v, num_bits) for k, v in numerical_results.items()}
+        numerical_results = self.sort_results(numerical_results)
+        if not return_group:
+            return numerical_results
+
+        group_numerical_results = {}
+        for group_name, group_metric in groups_metrics.items():
+            group_metric = {k: v.round(num_bits) for k, v in group_metric["numerical"].items()}
+            group_metric = ndarray_to_basetype(group_metric)
+            group_numerical_results[group_name] = self.sort_results(group_metric)
+        return numerical_results, group_numerical_results
+
+    def sort_results(self, results: dict) -> OrderedDict:
+        """for a single group of metrics"""
+        sorted_results = OrderedDict()
+        all_keys = sorted(results.keys(), key=lambda item: item[::-1])
+        for name in self.metric_names:
+            for key in all_keys:
+                if key.endswith(name):
+                    sorted_results[key] = results[key]
+        return sorted_results

utils/recorder/logger.py

@@ -0,0 +1,23 @@
+from torch.utils.tensorboard import SummaryWriter
+
+
+class TBLogger:
+    def __init__(self, tb_root):
+        self.tb_root = tb_root
+        self.tb = None
+
+    def write_to_tb(self, name, data, curr_iter):
+        assert self.tb_root is not None
+
+        if self.tb is None:
+            self.tb = SummaryWriter(self.tb_root)
+
+        if not isinstance(data, (tuple, list)):
+            self.tb.add_scalar(f"data/{name}", data, curr_iter)
+        else:
+            for idx, data_item in enumerate(data):
+                self.tb.add_scalar(f"data/{name}_{idx}", data_item, curr_iter)
+
+    def close_tb(self):
+        if self.tb is not None:
+            self.tb.close()

utils/recorder/meter_recorder.py

@@ -0,0 +1,91 @@
+# -*- coding: utf-8 -*-
+from collections import deque
+
+
+class AvgMeter(object):
+    __slots__ = ["value", "sum", "count"]
+
+    def __init__(self):
+        self.value = 0
+        self.sum = 0
+        self.count = 0
+
+    def reset(self):
+        self.value = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, value, num=1):
+        self.value = value
+        self.sum += value * num
+        self.count += num
+
+    @property
+    def avg(self):
+        return self.sum / self.count
+
+    def __repr__(self) -> str:
+        return f"{self.avg:.5f}"
+
+
+class HistoryBuffer:
+    """The class tracks a series of values and provides access to the smoothed
+    value over a window or the global average / sum of the sequence.
+
+    Args:
+        window_size (int): The maximal number of values that can
+            be stored in the buffer. Defaults to 20.
+
+    Example::
+
+        >>> his_buf = HistoryBuffer()
+        >>> his_buf.update(0.1)
+        >>> his_buf.update(0.2)
+        >>> his_buf.avg
+        0.15
+    """
+
+    def __init__(self, window_size: int = 20) -> None:
+        self._history = deque(maxlen=window_size)
+        self._count: int = 0
+        self._sum: float = 0
+        self.reset()
+
+    def reset(self):
+        self._history.clear()
+        self._count = 0
+        self._sum = 0
+
+    def update(self, value: float, num: int = 1) -> None:
+        """Add a new scalar value. If the length of queue exceeds ``window_size``,
+        the oldest element will be removed from the queue.
+        """
+        self._history.append(value)
+        self._count += num
+        self._sum += value * num
+
+    @property
+    def latest(self) -> float:
+        """The latest value of the queue."""
+        return self._history[-1]
+
+    @property
+    def avg(self) -> float:
+        """The average over the window."""
+        if len(self._history) == 0:
+            return 0
+        else:
+            return sum(self._history) / len(self._history)
+
+    @property
+    def global_avg(self) -> float:
+        """The global average of the queue."""
+        if self._count == 0:
+            return 0
+        else:
+            return self._sum / self._count
+
+    @property
+    def global_sum(self) -> float:
+        """The global sum of the queue."""
+        return self._sum

utils/recorder/visualize_results.py

@@ -0,0 +1,43 @@
+import os
+
+import cv2
+import matplotlib
+
+matplotlib.use("Agg")
+import numpy as np
+import torchvision.transforms.functional as tv_tf
+from torchvision.utils import make_grid
+
+
+def plot_results(data_container, save_path, base_size=256, is_rgb=True):
+    """Plot the results conresponding to the batched images based on the `make_grid` method from `torchvision`.
+
+    Args:
+        data_container (dict): Dict containing data you want to plot.
+        save_path (str): Path of the exported image.
+    """
+    font_cfg = dict(fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, thickness=2)
+
+    grids = []
+    for subplot_id, (name, data) in enumerate(data_container.items()):
+        if data.ndim == 3:
+            data = data.unsqueeze(1)
+
+        grid = make_grid(data, nrow=data.shape[0], padding=2, normalize=False)
+        grid = np.array(tv_tf.to_pil_image(grid.float()))
+        h, w = grid.shape[:2]
+        ratio = base_size / h
+        grid = cv2.resize(grid, dsize=None, fx=ratio, fy=ratio, interpolation=cv2.INTER_LINEAR)
+
+        (text_w, text_h), baseline = cv2.getTextSize(text=name, **font_cfg)
+        text_xy = 20, 20 + text_h // 2 + baseline
+        cv2.putText(grid, text=name, org=text_xy, color=(255, 255, 255), **font_cfg)
+
+        grids.append(grid)
+    grids = np.concatenate(grids, axis=0)  # H,W,C
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+    if is_rgb:
+        grids = cv2.cvtColor(grids, cv2.COLOR_RGB2BGR)
+    cv2.imwrite(save_path, grids)