An example implementation of stlimg.filter_method

hsv_filter.py

"""
Copyright 2023 AstroArts Inc. All rights reserved.
"""

import stlimg
import numpy


@stlimg.filter_method
def demo_filter(img: stlimg.Image) -> stlimg.Image:
    arr = numpy.asarray(img)
    hsv = rgb2hsv(arr)
    v = hsv[2]
    v_min = numpy.min(v)
    v_max = numpy.max(v)
    v_range = v_max - v_min
    nv = v / v_range

    y = curve(nv)
    hsv[1] *= y
    return stlimg.Image.fromarray(hsv2rgb(hsv))


def curve(x: float) -> float:
    return numpy.sin((x * numpy.pi) / 2)


def rgb2hsv(arr: numpy.array) -> numpy.array:
    assert len(arr.shape) == 3
    assert arr.shape[0] == 3
    hsv = numpy.zeros_like(arr)
    r = arr[0]
    g = arr[1]
    b = arr[2]
    max0 = numpy.where(r > g, r, g)
    min0 = numpy.where(r > g, g, r)
    rgb_max = numpy.where(max0 > b, max0, b)
    rgb_min = numpy.where(min0 < b, min0, b)
    rgb_range = rgb_max - rgb_min
    domain = 6.0
    step1 = domain / 3.0
    step2 = step1 * 2.0
    # R
    h_r = numpy.where(rgb_range > 0.0, (g - b) / rgb_range, 0.0)
    h_r = numpy.where(h_r < 0.0, h_r + domain, h_r)
    # G
    h_g = numpy.where(rgb_range > 0.0, step1 + (b - r) / rgb_range, 0.0)
    h_b = numpy.where(rgb_range > 0.0, step2 + (r - g) / rgb_range, 0.0)
    h = numpy.where(
        rgb_range > 0.0,
        numpy.where(
            rgb_max == r,
            h_r,
            numpy.where(
                rgb_max == g,
                h_g,
                h_b
            )
        ),
        0.0
    )
    hsv[0] = h / domain
    hsv[1] = numpy.where(rgb_max != 0.0, rgb_range / rgb_max, rgb_range)
    hsv[2] = rgb_max
    return hsv


def hsv2rgb(arr: numpy.array) -> numpy.array:
    assert len(arr.shape) == 3
    assert arr.shape[0] == 3
    rgb = numpy.zeros_like(arr)
    h = arr[0]
    s = arr[1]
    v = arr[2]

    f = h * 6.0
    i = numpy.array(f, dtype=numpy.int32)
    f = f - i
    w = v * (1.0 - s)
    x = v * (1.0 - s * f)
    y = v * (1.0 - s * (1.0 - f))

    # vxwwyv
    rgb[0] = numpy.where(
        s > 0.0,
        numpy.where(i == 0, v,
            numpy.where(i == 1, x,
                numpy.where(i == 2, w,
                    numpy.where(i == 3, w,
                        numpy.where(i == 4, y,
                            numpy.where(i == 5, v, v)))))), v)
    # yvvxww
    rgb[1] = numpy.where(
        s > 0.0,
        numpy.where(i == 0, y,
            numpy.where(i == 1, v,
                numpy.where(i == 2, v,
                    numpy.where(i == 3, x,
                        numpy.where(i == 4, w,
                            numpy.where(i == 5, w, v)))))), v)
    # wwyvvx
    rgb[2] = numpy.where(
        s > 0.0,
        numpy.where(i == 0, w,
            numpy.where(i == 1, w,
                numpy.where(i == 2, y,
                    numpy.where(i == 3, v,
                        numpy.where(i == 4, v,
                            numpy.where(i == 5, x, v)))))), v)

    return rgb