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 argparse
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",
        title="Histogram"
    )
    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.]),
    contrast: float = (0., [-1., 1.]),
    global_params: dict = {}
) -> None:
    global_params["shadow_boost"] = shadow_boost
    global_params["highlight_boost"] = highlight_boost
    global_params["contrast"] = contrast
    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:
    shadow_boost = global_params.get("shadow_boost", 0.)
    highlight_boost = global_params.get("highlight_boost", 0.)
    contrast = global_params.get("contrast", 0.)
    exposure = global_params.get("exposure", 0.)
    if not apply_tone_curve_on_luma:
        return apply_s_curve_tone_mapping(
            img,
            shadow_boost,
            highlight_boost,
            contrast,
            exposure
        )
    hsv = rgb_to_hsv(img)
    luma_tone_mapped = apply_s_curve_tone_mapping(
        hsv[..., -1],
        shadow_boost,
        highlight_boost,
        contrast,
        exposure
    )
    hsv[..., -1] = luma_tone_mapped
    return hsv_to_rgb(hsv)


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


@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


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__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--backend", type=str,
                        choices=["gradio", "qt", "mpl"], default="gradio")
    args = parser.parse_args()
    img = Image.load_image("image_sample.jpg")
    markdown_description = r"# Tone mapping: S-curve" + "\n" \
        r"$$\text{shadow boost} \in [-1, 1]$$  $$\text{highlight boost} \in [-1, 1]$$  $$\text{contrast} \in [-1, 1]$$  $$\text{exposure} \in [-1, 1]$$" + "\n" \
        "If Tone curves are applied to RGB colors separately, you'll get serious color shifts. It is much better when tone curve is applied to the luminance channel." + "\n"
    interactive_pipeline(gui=args.backend, markdown_description=markdown_description)(
        image_editing_pipeline)(img)