kuwahara.py

import logging

import cv2
import numpy as np

logger = logging.getLogger(__name__)


class Kuwahara:
    @staticmethod
    def transform(im: np.ndarray, r: int = 15):
        h, w, _ = im.shape
        img = np.pad(im, ((r, r), (r, r), (0, 0)), "edge")
        logger.info("Start processing (rate: {})".format(r))
        ave, var = cv2.integral2(img)
        ave = (
            ave[: -r - 1, : -r - 1]
            + ave[r + 1 :, r + 1 :]
            - ave[r + 1 :, : -r - 1]
            - ave[: -r - 1, r + 1 :]
        )
        ave = ave / (r + 1) ** 2

        var = (
            var[: -r - 1, : -r - 1]
            + var[r + 1 :, r + 1 :]
            - var[r + 1 :, : -r - 1]
            - var[: -r - 1, r + 1 :]
        )
        var = var / (r + 1) ** 2 - ave ** 2
        var = var.sum(axis=2)

        def filt(i: np.ndarray, j: np.ndarray) -> np.ndarray:
            return (
                np.array([ave[i, j], ave[i + r, j], ave[i, j + r], ave[i + r, j + r]])[
                    (
                        np.array(
                            [var[i, j], var[i + r, j], var[i, j + r], var[i + r, j + r]]
                        )
                        .argmin(axis=0)
                        .flatten(),
                        j.flatten(),
                        i.flatten(),
                    )
                ]
                .reshape(w, h, _)
                .transpose(1, 0, 2)
            )

        return filt(*np.meshgrid(np.arange(h), np.arange(w))).astype(img.dtype)


if __name__ == "__main__":
    import argparse
    import os

    logging.basicConfig(level=logging.DEBUG)

    def get_args():
        def is_ok(filename: str):
            assert os.path.isfile(filename)
            assert filename.endswith((".jpg", ".png"))
            return filename

        parser = argparse.ArgumentParser()
        parser.add_argument("filename", type=is_ok)
        parser.add_argument("-r", "--rate", default=10, type=int)
        return parser.parse_args()

    def until_esc(kc: int = 27, max_call: int = 10):
        # TODO: cv2のwindowがないときはすぐ抜けるようにしたい
        if max_call < 1:
            return False
        key = cv2.waitKey(0)
        if key != kc:
            logger.debug("key code {} ('{}') was pressed".format(key, chr(key)))
            return until_esc(kc=kc, max_call=max_call - 1)
        return True

    def process1(_im, r: int = 10):
        im = Kuwahara.transform(_im, r)
        return np.hstack((_im, im))

    def process2(_im, w: int = 10, h: int = 5):
        return np.vstack(
            np.hstack(Kuwahara.transform(_im, j * w + i) for i in range(w))
            for j in range(h)
        )

    args = get_args()
    im = np.array(cv2.imread(filename=args.filename))

    cv2.imshow("process1", process1(im, r=args.rate))
    cv2.imshow("process2", process2(im, w=6, h=3))
    until_esc()