Duplicate from MaxReimann/Whitebox-Style-Transfer-Editing

Co-authored-by: Max Reimann <MaxReimann@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.github/workflows/file_size_check.yml

@@ -0,0 +1,16 @@
+name: Check file size
+on:               # or directly `on: [push]` to run the action on every push on any branch
+  pull_request:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check large files
+        uses: ActionsDesk/lfs-warning@v2.0
+        with:
+          filesizelimit: 10485760 # this is 10MB so we can sync to HF Spaces

.github/workflows/push_to_huggingface.yml

@@ -0,0 +1,19 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push https://MaxReimann:$HF_TOKEN@huggingface.co/spaces/MaxReimann/Whitebox-Style-Transfer-Editing main

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "wise"]
+	path = wise
+	url = https://github.com/winfried-ripken/wise
+	branch = minimal_pipeline

README.md

@@ -0,0 +1,47 @@
+---
+title: White-box Style Transfer Editing (WISE)
+emoji: 🎨
+colorFrom: pink
+colorTo: red
+sdk: streamlit
+sdk_version: 1.10.0
+app_file: Whitebox_style_transfer.py
+tags:
+- Style Transfer
+- Image Synthesis
+- Editing
+- Painting
+pinned: false
+license: mit
+duplicated_from: MaxReimann/Whitebox-Style-Transfer-Editing
+---
+# White-box Style Transfer Editing (WISE) Demo
+
+This app demonstrates the editing capabilities of the [White-box Style Transfer Editing (WISE) framework](https://github.com/winfried-ripken/wise).
+It optimizes the parameters of classical image processing filters to match a given style image.
+After optimization, parameters can be tuned by hand to achieve a desired look.
+
+
+### How does it work?
+We provide a small stylization effect that contains several filters such as bump mapping or edge enhancement that can be optimized. The optimization yields so-called parameter masks, which contain per-pixel parameter settings for each filter.
+
+## 🚀 Try it out 🚀 
+**Our demo is now on huggingface: [huggingface/Whitebox-Style-Transfer-Editing](https://huggingface.co/spaces/MaxReimann/Whitebox-Style-Transfer-Editing)**
+
+![Streamlit Screenshot](images/screen_wise_demo.jpg?raw=true "WISE Editing Demo")
+
+To run **locally**, clone the repo recursively and install the dependencies in requirements.txt. Set HUGGINGFACE to false in demo_config.py. 
+Then run the streamlit app using `streamlit run Whitebox_style_transfer.py`
+
+
+
+## Links & Paper 
+[Project page](https://ivpg.hpi3d.de/wise/),
+[arxiv link](https://arxiv.org/abs/2207.14606),
+[framework code](https://github.com/winfried-ripken/wise)
+
+"WISE: Whitebox Image Stylization by Example-based Learning", by Winfried Lötzsch*, Max Reimann*, Martin Büßemeyer, Amir Semmo, Jürgen Döllner, Matthias Trapp, in ECCV 2022
+
+### Further notes
+Pull Requests and further improvements welcome.
+Please note that the shown effect is a minimal pipeline in terms of stylization capability, the much more feature-rich oilpaint and watercolor pipelines we show in our ECCV paper cannot be open-sourced due to IP reasons.

Whitebox_style_transfer.py

@@ -0,0 +1,312 @@
+import base64
+import datetime
+import os
+import sys
+from io import BytesIO
+from pathlib import Path
+import numpy as np
+import requests
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+PACKAGE_PARENT = 'wise'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
+
+import streamlit as st
+from streamlit.logger import get_logger
+from st_click_detector import click_detector
+import streamlit.components.v1 as components
+from streamlit.source_util import get_pages
+from streamlit_extras.switch_page_button import switch_page
+
+from demo_config import HUGGING_FACE
+from parameter_optimization.parametric_styletransfer import single_optimize
+from parameter_optimization.parametric_styletransfer import CONFIG as ST_CONFIG
+from parameter_optimization.strotss_org import strotss, pil_resize_long_edge_to
+import helpers.session_state as session_state
+from helpers import torch_to_np, np_to_torch
+from effects import get_default_settings, MinimalPipelineEffect 
+
+st.set_page_config(layout="wide")
+BASE_URL = "https://ivpg.hpi3d.de/wise/wise-demo/images/"
+LOGGER = get_logger(__name__)
+
+effect_type = "minimal_pipeline"
+
+if "click_counter" not in st.session_state:
+    st.session_state.click_counter = 1
+
+if "action" not in st.session_state:
+    st.session_state["action"] = ""
+
+content_urls = [
+    {
+        "name": "Portrait", "id": "portrait",
+        "src": BASE_URL + "/content/portrait.jpeg"
+    },
+    {
+        "name": "Tuebingen", "id": "tubingen",
+        "src": BASE_URL + "/content/tubingen.jpeg"
+    },
+    {
+        "name": "Colibri", "id": "colibri",
+        "src": BASE_URL + "/content/colibri.jpeg"
+    }
+]
+
+style_urls = [
+    {
+        "name": "Starry Night, Van Gogh", "id": "starry_night",
+        "src": BASE_URL + "/style/starry_night.jpg"
+    },
+    {
+        "name": "The Scream, Edward Munch", "id": "the_scream",
+        "src": BASE_URL + "/style/the_scream.jpg"
+    },
+    {
+        "name": "The Great Wave, Ukiyo-e", "id": "wave",
+        "src": BASE_URL + "/style/wave.jpg"
+    },
+    {
+        "name": "Woman with Hat, Henry Matisse", "id": "woman_with_hat",
+        "src": BASE_URL + "/style/woman_with_hat.jpg"
+    }
+]
+
+
+def last_image_clicked(type="content", action=None, ):
+    kw = "last_image_clicked" + "_" + type
+    if action:
+        session_state.get(**{kw: action})
+    elif kw not in session_state.get():
+        return None
+    else:
+        return session_state.get()[kw]
+
+
+@st.cache
+def _retrieve_from_id(clicked, urls):
+    src = [x["src"] for x in urls if x["id"] == clicked][0]
+    img = Image.open(requests.get(src, stream=True).raw)
+    return img, src
+
+
+def store_img_from_id(clicked, urls, imgtype):
+    img, src = _retrieve_from_id(clicked, urls)
+    session_state.get(**{f"{imgtype}_im": img, f"{imgtype}_render_src": src, f"{imgtype}_id": clicked})
+
+
+def img_choice_panel(imgtype, urls, default_choice, expanded):
+    with st.expander(f"Select {imgtype} image:", expanded=expanded):
+        html_code = '<div class="column" style="display: flex; flex-wrap: wrap; padding: 0 4px;">'
+        for url in urls:
+            html_code += f"<a href='#' id='{url['id']}' style='padding: 0px 5px'><img height='160px' style='margin-top: 8px;' src='{url['src']}'></a>"
+        html_code += "</div>"
+        clicked = click_detector(html_code)
+
+        if not clicked and st.session_state["action"] not in ("uploaded", "switch_page_from_local_edits", "switch_page_from_presets", "slider_change", "reset"):  # default val
+            store_img_from_id(default_choice, urls, imgtype)
+
+        st.write("OR:  ")
+
+        with st.form(imgtype + "-form", clear_on_submit=True):
+            uploaded_im = st.file_uploader(f"Load {imgtype} image:", type=["png", "jpg"], )
+            upload_pressed = st.form_submit_button("Upload")
+
+            if upload_pressed and uploaded_im is not None:
+                img = Image.open(uploaded_im)
+                buffered = BytesIO()
+                img.save(buffered, format="JPEG")
+                encoded = base64.b64encode(buffered.getvalue()).decode()
+                # session_state.get(uploaded_im=img, content_render_src=f"data:image/jpeg;base64,{encoded}")
+                session_state.get(**{f"{imgtype}_im": img, f"{imgtype}_render_src": f"data:image/jpeg;base64,{encoded}",
+                                     f"{imgtype}_id": "uploaded"})
+                st.session_state["action"] = "uploaded"
+                st.write("uploaded.")
+
+        last_clicked = last_image_clicked(type=imgtype)
+        print("last_clicked", last_clicked, "clicked", clicked, "action", st.session_state["action"] )
+        if not upload_pressed and clicked != "":  # trigger when no file uploaded
+            if last_clicked != clicked:  # only activate when content was actually clicked
+                store_img_from_id(clicked, urls, imgtype)
+                last_image_clicked(type=imgtype, action=clicked)
+                st.session_state["action"] = "clicked"
+                st.session_state.click_counter += 1  # hack to get page to reload at top
+
+        state = session_state.get()
+        st.sidebar.write(f'Selected {imgtype} image:')
+        st.sidebar.markdown(f'<img src="{state[f"{imgtype}_render_src"]}" width=240px></img>', unsafe_allow_html=True)
+
+
+def optimize(effect, preset, result_image_placeholder):
+    content = st.session_state["Content_im"]
+    style = st.session_state["Style_im"]
+    result_image_placeholder.text("<- Custom content/style needs to be style transferred")
+    optimize_button = st.sidebar.button("Optimize Style Transfer")
+    if optimize_button:
+        if HUGGING_FACE:
+            result_image_placeholder.warning("NST optimization is currently disabled in this HuggingFace Space because it takes ~5min to optimize. To try it out, please clone the repo and change the huggingface variable in demo_config.py")
+            st.stop()
+
+        result_image_placeholder.text("Executing NST to create reference image..")
+        base_dir = f"result/{datetime.datetime.now().strftime(r'%Y-%m-%d %H.%Mh %Ss')}"
+        os.makedirs(base_dir)
+        with st.spinner(text="Running NST"):
+            reference = strotss(pil_resize_long_edge_to(content, 1024),
+                                pil_resize_long_edge_to(style, 1024), content_weight=16.0,
+                                device=torch.device("cuda"), space="uniform")
+        progress_bar = result_image_placeholder.progress(0.0)
+        ref_save_path = os.path.join(base_dir, "reference.jpg")
+        content_save_path = os.path.join(base_dir, "content.jpg")
+        resize_to = 720
+        reference = pil_resize_long_edge_to(reference, resize_to)
+        reference.save(ref_save_path)
+        content.save(content_save_path)
+        ST_CONFIG["n_iterations"] = 300
+        with st.spinner(text="Optimizing parameters.."):
+            vp, content_img_cuda = single_optimize(effect, preset, "l1", content_save_path, str(ref_save_path),
+                                              write_video=False, base_dir=base_dir,
+                                              iter_callback=lambda i: progress_bar.progress(
+                                                  float(i) / ST_CONFIG["n_iterations"]))
+        return content_img_cuda.detach(), vp.cuda().detach()
+    else:
+        if not "result_vp" in st.session_state:
+            st.stop()
+        else:
+            return st.session_state["effect_input"], st.session_state["result_vp"]
+
+
+@st.cache(hash_funcs={MinimalPipelineEffect: id})
+def create_effect():
+    effect, preset, param_set = get_default_settings(effect_type)
+    effect.enable_checkpoints()
+    effect.cuda()
+    return effect, preset
+
+
+def load_visual_params(vp_path: str, img_org: Image, org_cuda: torch.Tensor, effect) -> torch.Tensor:
+    if Path(vp_path).exists():
+        vp = torch.load(vp_path).detach().clone()
+        vp = F.interpolate(vp, (img_org.height, img_org.width))
+        if len(effect.vpd.vp_ranges) == vp.shape[1]:
+            return vp
+    # use preset and save it
+    vp = effect.vpd.preset_tensor(preset, org_cuda, add_local_dims=True)
+    torch.save(vp, vp_path)
+    return vp
+
+
+# @st.cache(hash_funcs={torch.Tensor: id})
+@st.experimental_memo
+def load_params(content_id, style_id):#, effect):
+    preoptim_param_path = os.path.join("precomputed", effect_type, content_id, style_id)
+    img_org = Image.open(os.path.join(preoptim_param_path, "input.png"))
+    content_cuda = np_to_torch(img_org).cuda()
+    vp_path = os.path.join(preoptim_param_path, "vp.pt")
+    vp = load_visual_params(vp_path, img_org, content_cuda, effect)
+    return content_cuda, vp
+
+
+def render_effect(effect, content_cuda, vp):
+    with torch.no_grad():
+        result_cuda = effect(content_cuda, vp)
+    img_res = Image.fromarray((torch_to_np(result_cuda) * 255.0).astype(np.uint8))
+    return img_res
+
+
+result_container = st.container()
+coll1, coll2 = result_container.columns([3,2])
+coll1.header("Result")
+coll2.header("Global Edits")
+result_image_placeholder = coll1.empty()
+result_image_placeholder.markdown("## loading..")
+
+img_choice_panel("Content", content_urls, "portrait", expanded=True)
+img_choice_panel("Style", style_urls, "starry_night", expanded=True)
+
+state = session_state.get()
+content_id = state["Content_id"]
+style_id = state["Style_id"]
+
+effect, preset = create_effect()
+
+print("content id, style id", content_id, style_id  )
+if st.session_state["action"] == "uploaded":
+    content_img, _vp = optimize(effect, preset, result_image_placeholder)
+elif st.session_state["action"] in ("switch_page_from_local_edits", "switch_page_from_presets", "slider_change") or \
+      content_id == "uploaded" or style_id == "uploaded":
+    print("restore param")
+    _vp = st.session_state["result_vp"]
+    content_img = st.session_state["effect_input"]
+else:
+    print("load_params")
+    content_img, _vp = load_params(content_id, style_id)#, effect)
+
+vp = torch.clone(_vp)
+
+
+def reset_params(means, names):
+    for i, name in enumerate(names):
+        st.session_state["slider_" + name] = means[i]
+
+def on_slider():
+    st.session_state["action"] = "slider_change"
+
+
+with coll2:
+    show_params_names = [ 'bumpScale', "bumpOpacity", "contourOpacity"]
+    display_means = []
+    def create_slider(name):
+        mean = torch.mean(vp[:, effect.vpd.name2idx[name]]).item()
+        display_mean = mean + 0.5
+        display_means.append(display_mean)
+        if "slider_" + name not in st.session_state or st.session_state["action"] != "slider_change": 
+          st.session_state["slider_" + name] = display_mean
+        slider = st.slider(f"Mean {name}: ", 0.0, 1.0, step=0.05, key="slider_" + name, on_change=on_slider)
+        vp[:, effect.vpd.name2idx[name]] += slider - display_mean
+        vp.clamp_(-0.5, 0.5)
+    
+    for name in show_params_names:
+        create_slider(name)
+
+    others_idx = set(range(len(effect.vpd.vp_ranges))) - set([effect.vpd.name2idx[name] for name in show_params_names])
+    others_names = [effect.vpd.vp_ranges[i][0] for i in sorted(list(others_idx))]
+    other_param = st.selectbox("Other parameters: ", others_names)
+    create_slider(other_param)
+
+
+    reset_button = st.button("Reset Parameters", on_click=reset_params, args=(display_means, show_params_names))
+    if reset_button:
+        st.session_state["action"] = "reset"
+        st.experimental_rerun()
+
+    edit_locally_btn = st.button("Edit Local Parameter Maps")
+    if edit_locally_btn:
+        switch_page('️ local edits')
+
+    apply_presets = st.button("Paint Presets")
+    if apply_presets:
+        switch_page("Apply_preset")
+
+img_res = render_effect(effect, content_img, vp)
+
+st.session_state["result_vp"] = vp
+st.session_state["effect_input"] = content_img
+st.session_state["last_result"] = img_res
+
+with coll1:
+    # width = int(img_res.width * 500 / img_res.height)
+    result_image_placeholder.image(img_res)#, width=width)
+
+# a bit hacky way to return focus to top of page after clicking on images
+components.html(
+    f"""
+        <p>{st.session_state.click_counter}</p>
+        <script>
+            window.parent.document.querySelector('section.main').scrollTo(0, 0);
+        </script>
+    """,
+    height=0
+)

demo_config.py

@@ -0,0 +1 @@
+HUGGING_FACE=True # if run in hugging face. Disables some things like full NST optimization

pages/1_🎨_Apply_preset.py

@@ -0,0 +1,121 @@
+import os
+import sys
+import torch.nn.functional as F
+import torch
+
+PACKAGE_PARENT = '..'
+WISE_DIR = '../wise/'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, WISE_DIR)))
+
+
+import numpy as np
+from PIL import Image
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+
+from effects.minimal_pipeline import MinimalPipelineEffect
+from helpers.visual_parameter_def import minimal_pipeline_presets, minimal_pipeline_bump_mapping_preset, minimal_pipeline_xdog_preset
+from helpers import torch_to_np, np_to_torch
+from effects import get_default_settings
+from demo_config import HUGGING_FACE
+
+st.set_page_config(page_title="Preset Edit Demo", layout="wide")
+
+
+# @st.cache(hash_funcs={OilPaintEffect: id})
+@st.cache(hash_funcs={MinimalPipelineEffect: id})
+def local_edits_create_effect():
+    effect, preset, param_set = get_default_settings("minimal_pipeline")
+    effect.enable_checkpoints()
+    effect.cuda()
+    return effect, param_set
+
+
+effect, param_set = local_edits_create_effect()
+presets = {
+    "original": minimal_pipeline_presets,
+    "bump mapped": minimal_pipeline_bump_mapping_preset,
+    "contoured": minimal_pipeline_xdog_preset
+}
+
+st.session_state["action"] = "switch_page_from_presets" # on switchback, remember effect input
+
+active_preset = st.sidebar.selectbox("apply preset: ", ["original", "bump mapped", "contoured"])
+blend_strength = st.sidebar.slider("Parameter blending strength (non-hue) : ", 0.0, 1.0, 1.0, 0.05)
+hue_blend_strength = st.sidebar.slider("Hue-shift blending strength : ", 0.0, 1.0, 1.0, 0.05)
+
+st.sidebar.text("Drawing options:")
+stroke_width = st.sidebar.slider("Stroke width: ", 1, 80, 40)
+drawing_mode = st.sidebar.selectbox(
+    "Drawing tool:", ("freedraw", "line", "rect", "circle", "transform")
+)
+
+st.session_state["preset_canvas_key"] ="preset_canvas"
+
+vp = torch.clone(st.session_state["result_vp"])
+org_cuda = st.session_state["effect_input"]
+
+@st.experimental_memo
+def greyscale_original(_org_cuda, content_id): #content_id is used for hashing
+    if HUGGING_FACE:
+        wsize = 450
+        img_org_height, img_org_width = _org_cuda.shape[-2:]
+        wpercent = (wsize / float(img_org_width))
+        hsize = int((float(img_org_height) * float(wpercent)))
+    else:
+        longest_edge = 670
+        img_org_height, img_org_width = _org_cuda.shape[-2:]
+        max_width_height = max(img_org_width, img_org_height)
+        hsize = int((float(longest_edge) * float(float(img_org_height) / max_width_height)))
+        wsize = int((float(longest_edge) * float(float(img_org_width) / max_width_height)))
+
+    org_img = F.interpolate(_org_cuda, (hsize, wsize), mode="bilinear")
+    org_img = torch.mean(org_img, dim=1, keepdim=True) / 2.0
+    org_img = torch_to_np(org_img, multiply_by_255=True)[..., np.newaxis].repeat(3, axis=2)
+    org_img = Image.fromarray(org_img.astype(np.uint8))
+    return org_img, hsize, wsize
+
+greyscale_img, hsize, wsize = greyscale_original(org_cuda, st.session_state["Content_id"])
+
+coll1, coll2 = st.columns(2)
+coll1.header("Draw Mask")
+coll2.header("Live Result")
+
+with coll1:
+    # Create a canvas component
+    canvas_result = st_canvas(
+        fill_color="rgba(0, 0, 0, 1)",  # Fixed fill color with some opacity
+        stroke_width=stroke_width,
+        background_image=greyscale_img,
+        width=greyscale_img.width,
+        height=greyscale_img.height,
+        drawing_mode=drawing_mode,
+        key=st.session_state["preset_canvas_key"]
+    )
+    
+
+res_data = None
+if canvas_result.image_data is not None:
+    abc = np_to_torch(canvas_result.image_data.astype(np.float)).sum(dim=1, keepdim=True).cuda()
+
+    img_org_width = org_cuda.shape[-1]
+    img_org_height = org_cuda.shape[-2]
+    res_data = F.interpolate(abc, (img_org_height, img_org_width)).squeeze(1)
+
+    preset_tensor = effect.vpd.preset_tensor(presets[active_preset], org_cuda, add_local_dims=True)
+    hue = torch.clone(vp[:,effect.vpd.name2idx["hueShift"]])
+    vp[:] = preset_tensor * res_data * blend_strength + vp[:] * (1 - res_data * blend_strength)
+    vp[:, effect.vpd.name2idx["hueShift"]] = \
+        preset_tensor[:,effect.vpd.name2idx["hueShift"]] * res_data * hue_blend_strength + hue * (1 - res_data * hue_blend_strength)
+
+with torch.no_grad():
+    result_cuda = effect(org_cuda, vp)
+
+img_res = Image.fromarray((torch_to_np(result_cuda) * 255.0).astype(np.uint8))
+coll2.image(img_res)
+
+apply_btn = st.sidebar.button("Apply")
+if apply_btn:
+    st.session_state["result_vp"] = vp

pages/2_🖌️_Local_edits.py

@@ -0,0 +1,242 @@
+import os
+import sys
+
+import torch.nn.functional as F
+import torch
+import numpy as np
+import matplotlib
+from matplotlib import pyplot as plt
+import matplotlib.cm
+from PIL import Image
+
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+
+
+PACKAGE_PARENT = '..'
+WISE_DIR = '../wise/'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, WISE_DIR)))
+
+
+
+from effects.gauss2d_xy_separated import Gauss2DEffect
+from effects.minimal_pipeline import MinimalPipelineEffect
+from helpers import torch_to_np, np_to_torch
+from effects import get_default_settings
+from demo_config import HUGGING_FACE
+
+st.set_page_config(page_title="Editing Demo", layout="wide")
+
+# @st.cache(hash_funcs={OilPaintEffect: id})
+@st.cache(hash_funcs={MinimalPipelineEffect: id})
+def local_edits_create_effect():
+    effect, preset, param_set = get_default_settings("minimal_pipeline")
+    effect.enable_checkpoints()
+    effect.cuda()
+    return effect, param_set
+
+
+effect, param_set = local_edits_create_effect()
+
+@st.experimental_memo
+def gen_param_strength_fig():
+    cmap = matplotlib.cm.get_cmap('plasma')
+    # cmap show
+    gradient = np.linspace(0, 1, 256)
+    gradient = np.vstack((gradient, gradient))
+    fig, ax = plt.subplots(figsize=(3, 0.1))
+    fig.patch.set_alpha(0.0)
+    ax.set_title("parameter strength", fontsize=6.5, loc="left")
+    ax.imshow(gradient, aspect='auto', cmap=cmap)
+    ax.set_axis_off()
+    return fig, cmap
+
+cmap_fig, cmap = gen_param_strength_fig()
+
+st.session_state["canvas_key"] = "canvas"
+try:
+    vp = st.session_state["result_vp"]
+    org_cuda = st.session_state["effect_input"]
+except KeyError as e:
+    print("init run, certain keys not found. If this happens once its ok.")
+
+if st.session_state["action"] != "switch_page_from_local_edits":
+    st.session_state.local_edit_action = "init"
+
+st.session_state["action"] = "switch_page_from_local_edits" # on switchback, remember effect input
+
+if "mask_edit_counter" not in st.session_state:
+    st.session_state["mask_edit_counter"] = 1
+if "initial_drawing" not in st.session_state:
+    st.session_state["initial_drawing"] = {"random": st.session_state["mask_edit_counter"], "background": "#eee"}
+
+def on_slider_change():
+    if st.session_state.local_edit_action == "init":
+        st.stop()
+    st.session_state.local_edit_action = "slider"
+
+def on_param_change():
+    st.session_state.local_edit_action = "param_change"
+
+active_param = st.sidebar.selectbox("active parameter: ", param_set + ["smooth"], index=2, on_change=on_param_change)
+
+st.sidebar.text("Drawing options")
+if active_param != "smooth":
+    plus_or_minus = st.sidebar.slider("Increase or decrease param map: ", -1.0, 1.0, 0.8, 0.05,
+                                      on_change=on_slider_change)
+else:
+    sigma = st.sidebar.slider("Sigma: ", 0.1, 10.0, 0.5, 0.1, on_change=on_slider_change)
+
+stroke_width = st.sidebar.slider("Stroke width: ", 1, 50, 20, on_change=on_slider_change)
+drawing_mode = st.sidebar.selectbox(
+    "Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"), on_change=on_slider_change,
+)
+
+st.sidebar.text("Viewing options")
+if active_param != "smooth":
+    overlay = st.sidebar.slider("show parameter overlay: ", 0.0, 1.0, 0.8, 0.02, on_change=on_slider_change)
+    st.sidebar.pyplot(cmap_fig, bbox_inches='tight', pad_inches=0)
+
+st.sidebar.text("Update:")
+realtime_update = st.sidebar.checkbox("Update in realtime", True)
+clear_after_draw = st.sidebar.checkbox("Clear Canvas after each Stroke", False)
+invert_selection = st.sidebar.checkbox("Invert Selection", False)
+
+
+@st.experimental_memo
+def greyscale_org(_org_cuda, content_id): #content_id is used for hashing
+    if HUGGING_FACE:
+        wsize = 450
+        img_org_height, img_org_width = _org_cuda.shape[-2:]
+        wpercent = (wsize / float(img_org_width))
+        hsize = int((float(img_org_height) * float(wpercent)))
+    else:
+        longest_edge = 670
+        img_org_height, img_org_width = _org_cuda.shape[-2:]
+        max_width_height = max(img_org_width, img_org_height)
+        hsize = int((float(longest_edge) * float(float(img_org_height) / max_width_height)))
+        wsize = int((float(longest_edge) * float(float(img_org_width) / max_width_height)))
+
+    org_img = F.interpolate(_org_cuda, (hsize, wsize), mode="bilinear")
+    org_img = torch.mean(org_img, dim=1, keepdim=True) / 2.0
+    org_img = torch_to_np(org_img)[..., np.newaxis].repeat(3, axis=2)
+    return org_img, hsize, wsize
+
+def generate_param_mask(vp):
+    greyscale_img, hsize, wsize = greyscale_org(org_cuda, st.session_state["Content_id"])
+    if active_param != "smooth":
+        scaled_vp = F.interpolate(vp, (hsize, wsize))[:, effect.vpd.name2idx[active_param]]
+        param_cmapped = cmap((scaled_vp + 0.5).cpu().numpy())[...,:3][0]
+        greyscale_img = greyscale_img * (1 - overlay) + param_cmapped * overlay
+    return Image.fromarray((greyscale_img * 255).astype(np.uint8))
+
+def compute_results(_vp):
+    if "cached_canvas" in st.session_state and st.session_state["cached_canvas"].image_data is not None:
+        canvas_result = st.session_state["cached_canvas"]
+        abc = np_to_torch(canvas_result.image_data.astype(np.float32)).sum(dim=1, keepdim=True).cuda()
+
+        if invert_selection:
+            abc = abc * (- 1.0) + 1.0
+
+        img_org_width = org_cuda.shape[-1]
+        img_org_height = org_cuda.shape[-2]
+        res_data = F.interpolate(abc, (img_org_height, img_org_width)).squeeze(1)
+
+        if active_param != "smooth":
+            _vp[:, effect.vpd.name2idx[active_param]] += plus_or_minus * res_data
+            _vp.clamp_(-0.5, 0.5)
+        else:
+            gauss2dx = Gauss2DEffect(dxdy=[1.0, 0.0], dim_kernsize=5)
+            gauss2dy = Gauss2DEffect(dxdy=[0.0, 1.0], dim_kernsize=5)
+
+            vp_smoothed = gauss2dx(_vp, torch.tensor(sigma).cuda())
+            vp_smoothed = gauss2dy(vp_smoothed, torch.tensor(sigma).cuda())
+
+            print(res_data.shape)
+            print(_vp.shape)
+            print(vp_smoothed.shape)
+            _vp = torch.lerp(_vp, vp_smoothed, res_data.unsqueeze(1))
+
+    with torch.no_grad():
+        result_cuda = effect(org_cuda, _vp)
+
+    _, hsize, wsize = greyscale_org(org_cuda, st.session_state["Content_id"])
+    result_cuda = F.interpolate(result_cuda, (hsize, wsize), mode="bilinear")
+
+    return Image.fromarray((torch_to_np(result_cuda) * 255.0).astype(np.uint8)), _vp
+
+coll1, coll2 = st.columns(2)
+coll1.header("Draw Mask:")
+coll2.header("Live Result")
+
+# there is no way of removing the canvas history/state without rerunning the whole program.
+# therefore, giving the canvas a initial_drawing that differs from the canvas state will clear the background
+def mark_canvas_for_redraw():
+    print("mark for redraw")
+    st.session_state["mask_edit_counter"] += 1 # change state of initial drawing
+    initial_drawing = {"random": st.session_state["mask_edit_counter"], "background": "#eee"}
+    st.session_state["initial_drawing"] = initial_drawing
+
+
+with coll1:
+    print("edit action", st.session_state.local_edit_action)
+    if clear_after_draw and st.session_state.local_edit_action not in ("slider", "param_change", "init"):
+        if st.session_state.local_edit_action == "redraw":
+            st.session_state.local_edit_action = "draw"
+            mark_canvas_for_redraw()
+        else:
+            st.session_state.local_edit_action = "redraw"
+
+    mask = generate_param_mask(st.session_state["result_vp"])
+    st.session_state["last_mask"] = mask
+
+    # Create a canvas component
+    canvas_result = st_canvas(
+        fill_color="rgba(0, 0, 0, 1)",
+        stroke_width=stroke_width,
+        background_image=mask,
+        update_streamlit=realtime_update,
+        width=mask.width,
+        height=mask.height,
+        initial_drawing=st.session_state["initial_drawing"],
+        drawing_mode=drawing_mode,
+        key=st.session_state.canvas_key,
+    ) 
+
+    if canvas_result.json_data is  None:
+        print("stops")
+        st.stop()
+
+    st.session_state["cached_canvas"] = canvas_result
+
+    print("compute result")
+    img_res, vp = compute_results(vp)
+    st.session_state["last_result"] = img_res
+    st.session_state["result_vp"] = vp
+
+    st.markdown("### Mask: " + active_param)
+
+if st.session_state.local_edit_action in ("slider", "param_change", "init"):
+    print("set redraw")
+    st.session_state.local_edit_action = "redraw"
+
+
+print("plot masks")
+texts = []
+preview_masks = []
+img = st.session_state["last_mask"]
+for i, p in enumerate(param_set):
+    idx = effect.vpd.name2idx[p]
+    iii = F.interpolate(vp[:, idx:idx + 1] + 0.5, (int(img.height * 0.2), int(img.width * 0.2)))
+    texts.append(p[:15])
+    preview_masks.append(torch_to_np(iii))
+
+coll2.image(img_res)  # , use_column_width="auto")
+ppp = st.columns(len(param_set))
+for i, (txt, im) in enumerate(zip(texts, preview_masks)):
+    ppp[i].text(txt)
+    ppp[i].image(im, clamp=True)
+
+print("....")

pages/3_📖_Readme.py

@@ -0,0 +1,32 @@
+import streamlit as st
+
+st.title("White-box Style Transfer Editing")
+
+st.markdown("""
+    This app demonstrates the editing capabilities of the White-box Style Transfer Editing (WISE) framework.
+    It optimizes the parameters of classical image processing filters to match a given style image.
+    After optimization, parameters can be tuned by hand to achieve a desired look.
+
+    ### How does it work?
+    We provide a small stylization effect that contains several filters such as bump mapping or edge enhancement that can be optimized. The optimization yields so-called parameter masks, which contain per pixel parameter settings of each filter.
+
+    ### How to use the app ?
+    - On the first page select existing content/style combinations or upload images to optimize. 
+    - After the effect has been applied, use the parameter sliders to adjust a parameter value globally
+    - On the "apply preset" page, we defined several parameter presets that can be drawn on the image. Press "Apply" to make the changes permanent
+    - On the " local editing" page, individual parameter masks can be edited regionally. Choose the parameter on the left sidebar, and use the parameter strength slider to either increase or decrease the strength of the drawn strokes
+    - Strokes on the drawing canvas (left column) are updated in real-time on the result in the right column. 
+    - Strokes stay on the canvas unless manually deleted by clicking the trash button. To remove them from the canvas after each stroke, tick the corresponding checkbox in the sidebar.
+
+    ### Links & Paper 
+    [Project page](https://ivpg.hpi3d.de/wise/),
+    [arxiv link](https://arxiv.org/abs/2207.14606)
+    [demo code](https://github.com/MaxReimann/WISE-Editing)
+
+    "WISE: Whitebox Image Stylization by Example-based Learning", by Winfried Lötzsch*, Max Reimann*, Martin Büßemeyer, Amir Semmo, Jürgen Döllner, Matthias Trapp, in ECCV 2022
+
+    ### Further notes
+    Pull Requests and further improvements are very welcome.
+    On huggingface, optimization is currently disabled because it takes ~5min to optimize. If you want to try out your own images, please check out the app on [github](https://github.com/MaxReimann/WISE-Editing).
+    Please note that the shown effect is a minimal pipeline in terms of stylization capability, the much more feature-rich oilpaint and watercolor pipelines we show in our ECCV paper cannot be open-sourced due to IP reasons.
+""")

precomputed/minimal_pipeline/colibri/starry_night/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f40c444b91434d3372267e267b7668a93966d79b29592f3cc8de1dd2bc756c4
+size 12753387

precomputed/minimal_pipeline/colibri/the_scream/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ec0b6a86555481aca68c4a20f12980d827fe69754bc8ffba6a2b6cf09c95211
+size 12753387

precomputed/minimal_pipeline/colibri/wave/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:013a9a3654664bdb201ab7a48e98711014c889fee5b05e1237fafa08a75dec67
+size 12753387

precomputed/minimal_pipeline/colibri/woman_with_hat/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3995afcba2a055d190b08e18661e6ca809cb9755d426696130c793793ec2c548
+size 12753387

precomputed/minimal_pipeline/portrait/starry_night/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:102c79861638e382fb79fbe203e7b643c500b93d400b63c06bebdad22150a490
+size 14904747

precomputed/minimal_pipeline/portrait/the_scream/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24e4dde0fbaae2eded6c760b1ad15b56c1babe4ba70fb2dd2bc49beca95eee15
+size 14904747

precomputed/minimal_pipeline/portrait/wave/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51a9fe9033f2b77fc21b5f1c546e3a9af1079bc60c0f2d4b61a0ad7b35b7034f
+size 14904747

precomputed/minimal_pipeline/portrait/woman_with_hat/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4384dc543accfec73dec64f5c02cc8f11ed7ac281edcbaea6db2bddc35718acb
+size 14904747

precomputed/minimal_pipeline/tubingen/starry_night/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75a63bfc9e71641e1f7a7b99127aaa8f16c1ed40636a210b523fb5c82f73abd6
+size 12416427

precomputed/minimal_pipeline/tubingen/the_scream/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f0118f3aae6559ae743255026cc6d81019bd8b907e2fa59a93f3c2a54176eb4
+size 12416427

precomputed/minimal_pipeline/tubingen/wave/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6412df16c03f5d48d9c2fe5e91449c99cc2b6fc240fd7c385edc5ab941526d4
+size 12416427

precomputed/minimal_pipeline/tubingen/woman_with_hat/vp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52a9e9c1f716bb3dbc6f5adbf7274356f3f06545f404244a930ecf7155389f99
+size 12416427

requirements.txt

@@ -0,0 +1,13 @@
+imageio
+imageio-ffmpeg
+matplotlib
+Pillow
+numpy
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
+torchvision
+streamlit==1.10.0
+streamlit_drawable_canvas==0.8.0
+streamlit_extras==0.1.5
+st_click_detector
+scipy

wise

@@ -0,0 +1 @@
+Subproject commit e4d085e22a081738667b32167f7901fa5febe25a