app.py

from interactive_pipe import interactive_pipeline, interactive
from interactive_pipe.data_objects.curves import SingleCurve, Curve
from interactive_pipe.data_objects.image import Image
from global_tone_mapping import apply_s_curve_tone_mapping
from color_conversions import rgb_to_hsv, hsv_to_rgb
from synthetic_charts import generate_color_wheel
import numpy as np


def histogram(img: np.ndarray) -> Curve:
    hist_curves = []
    for ch in range(img.shape[-1]):
        hist, bins = np.histogram(
            img[..., ch].flatten(),
            bins=128,
            range=(0, 1)
        )
        hist_curves.append(
            SingleCurve(bins[:-1], hist, style='rgb'[ch]+"-")
        )
    hist_curve = Curve(
        hist_curves,
        xlabel="Intensity",
        ylabel="Frequency"
    )
    return hist_curve


@interactive()
def set_tone_mapping_params(
    shadow_boost: float = (0., [-1., 1.]),
    highlight_boost: float = (0., [-1., 1.]),
    exposure: float = (0., [-1., 1.]),
    global_params: dict = {}
) -> None:
    global_params["shadow_boost"] = shadow_boost
    global_params["highlight_boost"] = highlight_boost
    global_params["exposure"] = exposure


@interactive(apply_tone_curve_on_luma=(True,))
def s_curve_tone_mapping(
    img: np.ndarray,
    apply_tone_curve_on_luma: bool = True,
    global_params={}
) -> np.ndarray:
    exposure = global_params.get("exposure", 0.)
    shadow_boost = global_params.get("shadow_boost", 0.)
    highlight_boost = global_params.get("highlight_boost", 0.)
    if not apply_tone_curve_on_luma:
        return apply_s_curve_tone_mapping(img, shadow_boost, highlight_boost, exposure)
    hsv = rgb_to_hsv(img)
    luma_tone_mapped = apply_s_curve_tone_mapping(
        hsv[..., -1],
        shadow_boost,
        highlight_boost,
        exposure
    )
    hsv[..., -1] = luma_tone_mapped
    return hsv_to_rgb(hsv)


def s_curve_visualization(global_params={}) -> Curve:
    exposure = global_params.get("exposure", 0.)
    shadow_boost = global_params.get("shadow_boost", 0.)
    highlight_boost = global_params.get("highlight_boost", 0.)
    x = np.linspace(0, 1, 256)
    y = apply_s_curve_tone_mapping(x, shadow_boost,
                                   highlight_boost,
                                   exposure)
    return Curve(
        [
            SingleCurve(x, y, style='m-'),
            SingleCurve(x, x, style='k--'),
        ],
        xlabel="Input", ylabel="Output"
    )


@interactive(
    input_image=("sample_image", ["sample_image", "color_wheel"]),
)
def pick_image(
    default_image: np.ndarray,
    input_image: str = "sample_image",
    global_params={}
) -> np.ndarray:
    if input_image == "sample_image":
        return default_image
    elif input_image == "color_wheel":
        color_wheel = global_params.get("color_wheel", None)
        if color_wheel is None:
            color_wheel = generate_color_wheel(resolution=496)
            global_params["color_wheel"] = color_wheel
        return color_wheel
    else:
        return default_image


@interactive_pipeline(gui="gradio")
def image_editing_pipeline(sample_image):
    input_image = pick_image(sample_image)
    set_tone_mapping_params()
    tc_image = s_curve_tone_mapping(input_image)
    tone_curve = s_curve_visualization()
    histogram_curve = histogram(tc_image)
    return tc_image, histogram_curve, tone_curve


if __name__ == "__main__":
    img = Image.load_image("image_sample.jpg")
    image_editing_pipeline(img)