旷视low-level系列（一）：Bayer Pattern Unification and Bayer Preserving Augmentation

1. Motivation

对于RAW域去噪，通常会将单通道bayer格式的RAW图打包成4通道，然后送入神经网络。不同厂家生产的sensor出的RAW图可能具有不同的bayer模式，通常是RGGB，BGGR， GRBG和GBRG。
业内做AI-ISP的攻城狮们应该都会遇到这样一个问题，在适配不同sensor的过程中会积累大量具有不同Bayer模式的数据，然后在训练模型时都想用上，这时大家都会将这些异源的数据统一成相同的bayer模式，常用的操作有：① 在裁剪patch时根据目标bayer模式选择合适的起点；② 打包成4通道，然后交换通道顺序。论文作者发现第二种方式会产生伪影，而第一种方式不会。
另外，数据增强是训练神经网络时提升性能的一种常用手段，对于RAW数据，为了避免破坏bayer模式，通常会选择在打包成4通道后再做翻转和旋转等增强。然而作者发现这样也会产生伪影，并提出了相应的解决方案。

2. Contribution

• 提出了BayerUnify，将不同的bayer模式转换为一个统一的模式，充分利用异源数据，扩大训练集规模
• 提出了BayerAug，一种有效的RAW图像的数据增强方式

3. Methods

3.1 BayerUnify

训练阶段采用crop的方式将当前bayer模式转换为目标bayer模式

推理阶段采用先pad的方式转换bayer模式（crop会丢失信息），对神经网络的输出再做crop得到与原始图像格式一致的结果。

def bayer_unify(raw: np.ndarray, input_pattern: str, target_pattern: str, mode: str) -> Tuple:
    """
    Convert a bayer raw image from one bayer pattern to another.

    Parameters
    ----------
    raw : np.ndarray in shape (H, W)
        Bayer raw image to be unified.
    input_pattern : {"RGGB", "BGGR", "GRBG", "GBRG"}
        The bayer pattern of the input image.
    target_pattern : {"RGGB", "BGGR", "GRBG", "GBRG"}
        The expected output pattern.
    mode: {"crop", "pad"}
        The way to handle submosaic shift. "crop" abandons the outmost pixels,
        and "pad" introduces extra pixels. Use "crop" in training and "pad" in
        testing.
    """
    if input_pattern not in BAYER_PATTERNS:
        raise ValueError('Unknown input bayer pattern!')
    if target_pattern not in BAYER_PATTERNS:
        raise ValueError('Unknown target bayer pattern!')
    if mode not in NORMALIZATION_MODE:
        raise ValueError('Unknown normalization mode!')
    if not isinstance(raw, np.ndarray) or len(raw.shape) != 2:
        raise ValueError('raw should be a 2-dimensional numpy.ndarray!')

    if input_pattern == target_pattern:
        h_offset, w_offset = 0, 0
    elif input_pattern[0] == target_pattern[2] and input_pattern[1] == target_pattern[3]:
        h_offset, w_offset = 1, 0
    elif input_pattern[0] == target_pattern[1] and input_pattern[2] == target_pattern[3]:
        h_offset, w_offset = 0, 1
    elif input_pattern[0] == target_pattern[3] and input_pattern[1] == target_pattern[2]:
        h_offset, w_offset = 1, 1
    else:  # This is not happening in ["RGGB", "BGGR", "GRBG", "GBRG"]
        raise RuntimeError('Unexpected pair of input and target bayer pattern!')

    if mode == "pad":
        out = np.pad(raw, [[h_offset, h_offset], [w_offset, w_offset]], 'reflect')
    elif mode == "crop":
        h, w = raw.shape
        out = raw[h_offset:h - h_offset, w_offset:w - w_offset]
    else:
        raise ValueError('Unknown normalization mode!')

    return out, h_offset, w_offset

3.2 BayerAug

直接对RAW数据做翻转会改变bayer模式，BayerAug先翻转再执行BayerUnify，保证bayer模式不变。

def bayer_aug(raw: np.ndarray, flip_h: bool, flip_w: bool, transpose: bool, input_pattern: str) -> np.ndarray:
    """
    Apply augmentation to a bayer raw image.

    Parameters
    ----------
    raw : np.ndarray in shape (H, W)
        Bayer raw image to be augmented. H and W must be even numbers.
    flip_h : bool
        If True, do vertical flip.
    flip_w : bool
        If True, do horizontal flip.
    transpose : bool
        If True, do transpose.
    input_pattern : {"RGGB", "BGGR", "GRBG", "GBRG"}
        The bayer pattern of the input image.
    """

    if input_pattern not in BAYER_PATTERNS:
        raise ValueError('Unknown input bayer pattern!')
    if not isinstance(raw, np.ndarray) or len(raw.shape) != 2:
        raise ValueError('raw should be a 2-dimensional numpy.ndarray')
    if raw.shape[0] % 2 == 1 or raw.shape[1] % 2 == 1:
        raise ValueError('raw should have even number of height and width!')

    aug_pattern, target_pattern = input_pattern, input_pattern

    out = raw
    if flip_h:
        out = out[::-1, :]
        aug_pattern = aug_pattern[2] + aug_pattern[3] + aug_pattern[0] + aug_pattern[1]
    if flip_w:
        out = out[:, ::-1]
        aug_pattern = aug_pattern[1] + aug_pattern[0] + aug_pattern[3] + aug_pattern[2]
    if transpose:
        out = out.T
        aug_pattern = aug_pattern[0] + aug_pattern[2] + aug_pattern[1] + aug_pattern[3]

    out = bayer_unify(out, aug_pattern, target_pattern, "crop")
    return out

4. Comments

初看，就这？用起来，还挺香。没有很大的创新，胜在工程价值较高。

Reference

[1] Learning Raw Image Denoising with Bayer Pattern Unification and Bayer Preserving Augmentation
[2] 官方代码