Chainer で Bicubic 補間による画像二倍拡大層の実装

chainer-bicubic2x.py

import numpy as np
from chainer import Link
from chainer.functions import pad, convolution_2d, depth2space

def cubic(b, c):
    def k(x):
        if abs(x) < 1:
            return 1 / 6 * ((12 - 9 * b - 6 * c) * abs(x) ** 3 + (-18 + 12 * b + 6 * c) * abs(x) ** 2 + (6 - 2 * b))
        elif 1 <= abs(x) < 2:
            return 1 / 6 * ((-b - 6 * c) * abs(x) ** 3 + (6 * b + 30 * c) * abs(x) ** 2 + (-12 * b - 48 * c) * abs(x) + (8 * b + 24 * c))
        else:
            return 0
    return k

class Bicubic2xUpsampler(Link):

    def __init__(self, b=1/3, c=1/3):
        super().__init__()
        window = cubic(b, c)
        s = np.array([window(i + 0.25) for i in range(-2, 2)])
        e = np.array([window(i + 0.75) for i in range(-2, 2)])
        k1 = np.pad(s.reshape(1, 4) * s.reshape(4, 1), ((0, 1), (0, 1)))
        k2 = np.pad(e.reshape(1, 4) * s.reshape(4, 1), ((0, 1), (1, 0)))
        k3 = np.pad(s.reshape(1, 4) * e.reshape(4, 1), ((1, 0), (0, 1)))
        k4 = np.pad(e.reshape(1, 4) * e.reshape(4, 1), ((1, 0), (1, 0)))
        self.w = np.array([[k1], [k2], [k3], [k4]], dtype=np.float32)

    def __call__(self, x):
        b, c, h, w = x.shape
        h1 = x.reshape(b * c, 1, h, w)
        h2 = pad(h1, ((0, 0), (0, 0), (2, 2), (2, 2)), mode="symmetric")
        h3 = convolution_2d(h2, self.xp.array(self.w))
        h4 = depth2space(h3, 2)
        return h4.reshape(b, c, h * 2, w * 2)