"""Gradio web app for Brushstroke Parameterized Style Transfer. Wraps the optimization pipeline from jupyter_notebooks/brushstrokes.ipynb in a Gradio Blocks UI. Streams intermediate canvases for live preview. """ import os import numpy as np import torch import torch.optim as optim import torchvision.transforms as transforms from PIL import Image import gradio as gr from renderer import BrushStrokeRenderer from losses import StyleTransferLosses, total_variation_loss, curvature_loss from utils import pick_device VGG_WEIGHTS = "vgg_weights/vgg19_weights_normalized.h5" CONTENT_LAYERS = ["conv4_2", "conv5_2"] STYLE_LAYERS = ["conv1_1", "conv2_1", "conv3_1", "conv4_1", "conv5_1"] DEVICE = pick_device() def _pil_to_tensor(img: Image.Image, size: int) -> torch.Tensor: transform = transforms.Compose([ transforms.Resize(size), transforms.ToTensor(), ]) return transform(img.convert("RGB")).unsqueeze(0).to(DEVICE, torch.float) def _canvas_to_pil(canvas: torch.Tensor) -> Image.Image: arr = canvas.detach().cpu().numpy().clip(0, 1) return Image.fromarray((arr * 255).astype(np.uint8)) def stylize( content_img: Image.Image, style_img: Image.Image, num_strokes: int, steps: int, img_size: int, canvas_color: str, progress=gr.Progress(), ): """Run brushstroke optimization. Yields intermediate canvases every 10 steps.""" if content_img is None or style_img is None: raise gr.Error("Please upload both a content image and a style image.") progress(0, desc="Loading images...") content_t = _pil_to_tensor(content_img, int(img_size)) style_t = _pil_to_tensor(style_img, 224) _, _, H, W = content_t.shape progress(0.02, desc="Loading VGG and initializing brushstrokes...") vgg_loss = StyleTransferLosses( VGG_WEIGHTS, content_t, style_t, CONTENT_LAYERS, STYLE_LAYERS, scale_by_y=True, ).to(DEVICE).eval() content_np = content_t[0].permute(1, 2, 0).cpu().numpy() renderer = BrushStrokeRenderer( H, W, num_strokes=int(num_strokes), samples_per_curve=10, strokes_per_pixel=20, canvas_color=canvas_color, length_scale=1.1, width_scale=0.1, content_img=content_np, ).to(DEVICE) optimizer_geom = optim.Adam( [renderer.location, renderer.curve_s, renderer.curve_e, renderer.curve_c, renderer.width], lr=1e-1, ) optimizer_color = optim.Adam([renderer.color], lr=1e-2) steps = int(steps) actual_strokes = renderer.location.shape[0] status = f"Optimizing {actual_strokes} brushstrokes on {H}x{W} canvas..." yield None, status for step in range(1, steps + 1): optimizer_geom.zero_grad() optimizer_color.zero_grad() canvas = renderer() canvas_img = canvas.unsqueeze(0).permute(0, 3, 1, 2).contiguous() content_loss, style_loss = vgg_loss(canvas_img) content_loss = content_loss * 1.0 style_loss = style_loss * 3.0 tv_loss = 0.008 * total_variation_loss( renderer.location, renderer.curve_s, renderer.curve_e, K=10 ) curv_loss = 4.0 * curvature_loss( renderer.curve_s, renderer.curve_e, renderer.curve_c ) loss = content_loss + style_loss + tv_loss + curv_loss loss.backward( inputs=[renderer.location, renderer.curve_s, renderer.curve_e, renderer.curve_c, renderer.width], retain_graph=True, ) optimizer_geom.step() style_loss.backward(inputs=[renderer.color]) optimizer_color.step() progress(step / steps, desc=f"Step {step}/{steps}") if step % 5 == 0 or step == steps: with torch.no_grad(): preview = _canvas_to_pil(renderer()) yield preview, ( f"Step {step}/{steps} | " f"content={content_loss.item():.3f} " f"style={style_loss.item():.3f} " f"tv={tv_loss.item():.3f} " f"curv={curv_loss.item():.3f}" ) with torch.no_grad(): final = _canvas_to_pil(renderer()) yield final, f"Done — {actual_strokes} strokes, {steps} steps on {DEVICE}." with gr.Blocks(title="Brushstroke Style Transfer") as demo: gr.Markdown( "# Brushstroke Parameterized Style Transfer\n" "Upload a **content image** and a **painting** (style). The app optimizes " "thousands of parametric brushstrokes to repaint your content in the style " "of the painting. Based on Kotovenko et al., " "[*Rethinking Style Transfer: From Pixels to Parameterized Brushstrokes*]" "(https://arxiv.org/abs/2103.17185).\n\n" f"_Running on: `{DEVICE}`. Optimization is iterative and takes ~30s–3min " "depending on settings and hardware._" ) with gr.Row(): with gr.Column(scale=1): content_in = gr.Image(label="Content image", type="pil", height=300) style_in = gr.Image(label="Style image (painting)", type="pil", height=300) with gr.Accordion("Settings", open=True): num_strokes = gr.Slider( minimum=500, maximum=10000, value=500, step=500, label="Number of brushstrokes", info="More strokes = finer detail, slower", ) steps = gr.Slider( minimum=20, maximum=300, value=20, step=10, label="Optimization steps", info="More steps = better convergence, slower", ) img_size = gr.Slider( minimum=256, maximum=768, value=256, step=64, label="Image size (shortest side)", ) canvas_color = gr.Radio( choices=["gray", "white", "black"], value="gray", label="Canvas background color", ) run_btn = gr.Button("Stylize", variant="primary") with gr.Column(scale=1): output_img = gr.Image(label="Stylized output", type="pil", height=500) status = gr.Textbox(label="Status", interactive=False) examples_dir = "images" content_examples = os.path.join(examples_dir, "content") style_examples = os.path.join(examples_dir, "style") if os.path.isdir(content_examples) and os.path.isdir(style_examples): content_files = sorted( os.path.join(content_examples, f) for f in os.listdir(content_examples) if f.lower().endswith((".jpg", ".jpeg", ".png")) )[:4] style_files = sorted( os.path.join(style_examples, f) for f in os.listdir(style_examples) if f.lower().endswith((".jpg", ".jpeg", ".png")) )[:4] if content_files and style_files: pairs = [ [content_files[i % len(content_files)], style_files[i % len(style_files)]] for i in range(min(4, max(len(content_files), len(style_files)))) ] gr.Examples( examples=pairs, inputs=[content_in, style_in], label="Example pairs", ) run_btn.click( fn=stylize, inputs=[content_in, style_in, num_strokes, steps, img_size, canvas_color], outputs=[output_img, status], show_progress="minimal", ) if __name__ == "__main__": demo.queue(max_size=4).launch()