Spaces:
Running on Zero
Running on Zero
| """FlashReFocus: interactive image refocusing (Hugging Face Space, ZeroGPU). | |
| 1. Restore: Stage I (one-step diffusion with detail weight omega and the decoder refiner) restores the all-in-focus | |
| image once, and Depth Anything V2 estimates its depth once. | |
| 2. Click the restored image (or the depth map) to focus; only the 1.3M-parameter renderer runs per edit. | |
| Modes: one focal plane; two focal planes (first click A, second click B); tilt-shift (a click sets the sharp row). | |
| 3. The aperture (blur strength K), the depth curve and an optional user depth map re-render the current edit. | |
| """ | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| from PIL import Image, ImageDraw, ImageFont | |
| import engine | |
| SINGLE, DUAL, TILT = "π One focal plane", "π₯ Two focal planes", "ποΈ Tilt-shift" | |
| EXAMPLES = ["examples/portrait_city.jpg", "examples/portrait_aquarium.jpg", "examples/flower.jpg", | |
| "examples/dock.jpg", "examples/sculpture.jpg"] | |
| FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" | |
| try: | |
| ImageFont.truetype(FONT, 12) | |
| except OSError: | |
| FONT = None | |
| def fit(img, long_side): | |
| """Long side at most `long_side`, both sides multiples of 16.""" | |
| w, h = img.size | |
| s = min(1.0, long_side / max(w, h)) | |
| w2, h2 = max(16, int(w * s) // 16 * 16), max(16, int(h * s) // 16 * 16) | |
| return img.convert("RGB").resize((w2, h2), Image.LANCZOS) | |
| def gpu_restore(rgb, omega): | |
| t0 = time.perf_counter() | |
| aif, depth = engine.restore(rgb, omega) | |
| return aif, depth, time.perf_counter() - t0 | |
| def gpu_render(aif, depth, focal_depth, strength): | |
| t0 = time.perf_counter() | |
| out = engine.render(aif, depth, focal_depth, strength) | |
| return out, time.perf_counter() - t0 | |
| def depth_image(depth): | |
| return Image.fromarray((np.clip(1 - depth, 0, 1) * 255).astype(np.uint8)) # near is bright | |
| def scene_depth(state, gamma): | |
| """The depth the edits start from: the estimate (or the user's map), with the depth curve D^gamma applied.""" | |
| d = state["user_depth"] if state.get("user_depth") is not None else state["depth"] | |
| return np.power(np.clip(d, 0, 1), float(gamma)).astype(np.float32) | |
| def depth_at(depth, point): | |
| h, w = depth.shape | |
| return float(depth[min(max(point[1], 0), h - 1), min(max(point[0], 0), w - 1)]) | |
| def edit(state, mode, gamma, band): | |
| """(depth to render, focal depth D_f) for the current mode, or None when a focal point is still missing.""" | |
| d = scene_depth(state, gamma) | |
| h, w = d.shape | |
| if mode == TILT: | |
| return engine.tilt_shift(h, w, state.get("tilt_row", 0.5), float(band)), 0.5 | |
| pts = state["pts"] | |
| if not pts: | |
| return None | |
| if mode == DUAL and len(pts) == 2: | |
| da, db = depth_at(d, pts[0]), depth_at(d, pts[1]) | |
| return engine.two_focal_planes(d, da, db), da | |
| return d, depth_at(d, pts[0]) | |
| def restore(img, long_side, omega, gamma): | |
| if img is None: | |
| return None, None, None, None, "Upload an image or pick an example first." | |
| rgb = np.asarray(fit(img, int(long_side))) | |
| aif, depth, t = gpu_restore(rgb, float(omega)) | |
| state = {"aif": aif, "depth": depth, "user_depth": None, "pts": [], "tilt_row": 0.5} | |
| msg = (f"Restored {aif.shape[1]}Γ{aif.shape[0]} with Ο = {float(omega):.2f} and estimated its depth in {t:.2f} s. " | |
| "**Click the restored image or the depth map to focus.**") | |
| return Image.fromarray(aif), depth_image(scene_depth(state, gamma)), None, state, msg | |
| def render(state, strength, mode, gamma, band, show_marker): | |
| """(refocused image, depth view, message).""" | |
| edited = edit(state, mode, gamma, band) | |
| if edited is None: | |
| return None, depth_image(scene_depth(state, gamma)), "Click the restored image or the depth map to focus." | |
| depth, focal_depth = edited | |
| out, dt = gpu_render(state["aif"], depth, focal_depth, float(strength)) | |
| out = Image.fromarray(out) | |
| if show_marker: | |
| out = draw_markers(out, state, mode) | |
| k = f"aperture K = {float(strength):.2f}" + (" (extrapolated)" if float(strength) > 1 else "") | |
| if mode == TILT: | |
| msg = f"Tilt-shift band at {state.get('tilt_row', 0.5):.0%} of the height, {k}: rendered in {dt:.2f} s." | |
| elif mode == DUAL and len(state["pts"]) == 2: | |
| msg = f"Two focal planes A {state['pts'][0]} and B {state['pts'][1]}, {k}: rendered in {dt:.2f} s." | |
| elif mode == DUAL: | |
| msg = f"Plane A {state['pts'][0]} set (rendered in {dt:.2f} s). Click a second point for plane B." | |
| else: | |
| msg = f"Focus {state['pts'][0]} (D_f = {focal_depth:.2f}), {k}: rendered in {dt:.2f} s." | |
| return out, depth_image(depth), msg | |
| def draw_markers(im, state, mode): | |
| """Focal points (A / B with two planes), or the sharp row of tilt-shift, on a 3x supersampled overlay.""" | |
| im = im.convert("RGBA"); S = 3 | |
| ov = Image.new("RGBA", (im.width * S, im.height * S), (0, 0, 0, 0)); d = ImageDraw.Draw(ov) | |
| r = max(9, round(min(im.size) * 0.022)) * S; lw = max(2, round(r / 6)) | |
| if mode == TILT: | |
| y = state.get("tilt_row", 0.5) * ov.height | |
| for x in range(0, ov.width, 12 * lw): | |
| d.line([x, y, x + 6 * lw, y], fill=(255, 255, 255, 220), width=lw) | |
| return Image.alpha_composite(im, ov.resize(im.size, Image.LANCZOS)).convert("RGB") | |
| font = ImageFont.truetype(FONT, round(r * 1.25)) if FONT else ImageFont.load_default() | |
| for i, (x, y) in enumerate(state["pts"]): | |
| cx, cy = x * S, y * S | |
| d.ellipse([cx - r - lw, cy - r - lw, cx + r + lw, cy + r + lw], outline=(0, 0, 0, 170), width=lw * 3) | |
| d.ellipse([cx - r, cy - r, cx + r, cy + r], outline=(255, 255, 255, 255), width=lw) | |
| dot = lw * 1.3 | |
| d.ellipse([cx - dot, cy - dot, cx + dot, cy + dot], fill=(255, 255, 255, 255)) | |
| if mode == DUAL: | |
| tx = cx + r + lw * 2 if cx + r * 3 < ov.width else cx - r * 2.2 | |
| d.text((tx, max(0, cy - r * 1.9)), "AB"[i], font=font, fill=(255, 255, 255, 255), | |
| stroke_width=max(2, lw), stroke_fill=(0, 0, 0, 200)) | |
| return Image.alpha_composite(im, ov.resize(im.size, Image.LANCZOS)).convert("RGB") | |
| def aif_view(state, mode): | |
| return draw_markers(Image.fromarray(state["aif"]), state, mode) | |
| def on_click(state, strength, mode, gamma, band, show_marker, evt: gr.SelectData): | |
| if not state: | |
| return None, None, None, state, "Restore an image first." | |
| p = (int(evt.index[0]), int(evt.index[1])) | |
| if mode == TILT: | |
| state = dict(state, tilt_row=p[1] / state["aif"].shape[0]) | |
| else: | |
| state = dict(state, pts=state["pts"] + [p] if mode == DUAL and len(state["pts"]) == 1 else [p]) | |
| out, depth_view, msg = render(state, strength, mode, gamma, band, show_marker) | |
| return aif_view(state, mode), out, depth_view, state, msg | |
| def on_change(state, strength, mode, gamma, band, show_marker): | |
| if not state: | |
| return gr.update(), None, gr.update(), state, "" | |
| if mode != DUAL and len(state["pts"]) > 1: | |
| state = dict(state, pts=state["pts"][:1]) | |
| out, depth_view, msg = render(state, strength, mode, gamma, band, show_marker) | |
| return aif_view(state, mode), out, depth_view, state, msg | |
| def on_user_depth(state, user_depth, strength, mode, gamma, band, show_marker): | |
| if not state: | |
| return gr.update(), None, gr.update(), state, "Restore an image first, then upload its depth map." | |
| h, w = state["aif"].shape[:2] | |
| state = dict(state, user_depth=None if user_depth is None else engine.depth_from_image(user_depth, h, w)) | |
| out, depth_view, msg = render(state, strength, mode, gamma, band, show_marker) | |
| note = "Using your depth map. " if user_depth is not None else "Using the estimated depth again. " | |
| return aif_view(state, mode), out, depth_view, state, note + msg | |
| def on_clear(state, gamma): | |
| if not state: | |
| return gr.update(), None, gr.update(), state, "" | |
| state = dict(state, pts=[], tilt_row=0.5) | |
| return (Image.fromarray(state["aif"]), None, depth_image(scene_depth(state, gamma)), state, | |
| "Focal points cleared. Click the restored image or the depth map to focus.") | |
| RECOMMENDED = {"long_side": 1024, "omega": 1.25, "gamma": 1.0, "band": 0.06, "strength": 0.5, "mode": SINGLE, | |
| "show_marker": False} | |
| def on_reset(state): | |
| """Recommended settings (the paper's defaults), the estimated depth again; focal points are kept.""" | |
| r = RECOMMENDED | |
| controls = [r["long_side"], r["omega"], r["gamma"], r["band"], r["strength"], r["mode"], r["show_marker"], None] | |
| if not state: | |
| return controls + [gr.update(), None, gr.update(), state, "Settings reset to the recommended values."] | |
| state = dict(state, user_depth=None, pts=state["pts"][:1]) | |
| out, depth_view, msg = render(state, r["strength"], r["mode"], r["gamma"], r["band"], r["show_marker"]) | |
| note = "Settings reset to the recommended values (Ο applies on the next Restore). " | |
| return controls + [aif_view(state, r["mode"]), out, depth_view, state, note + msg] | |
| LOGO = """<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 64 64" width="64" height="64"><circle cx="32" cy="32" r="30" fill="#FFFFFF" stroke="#2E7D32" stroke-width="3"/><g fill="#48A43F" stroke="#FFFFFF" stroke-width="1.6" stroke-linejoin="round"><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(0 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(60 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(120 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(180 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(240 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(300 32 32)"/></g><path d="M35.5 19 L25.5 34 H31.5 L28.5 45 L39 29 H33 Z" fill="#FFD21E" stroke="#363636" stroke-width="1.4" stroke-linejoin="round"/></svg> | |
| """ | |
| HEADER = f""" | |
| <div id="header"> | |
| <div class="title">{LOGO}<div><h1>FlashReFocus</h1><p>Interactive image refocusing in seconds: restore a photo once, | |
| then click to refocus, open the aperture, or keep two subjects sharp.</p></div></div> | |
| <div class="authors">Ching-Heng Cheng<sup>1*</sup> Β· <a href="https://ming053l.github.io/" target="_blank">Chia-Ming Lee</a><sup>2,3*</sup> Β· | |
| <a href="https://www.albany.edu/faculty/mchang2/" target="_blank">Ming-Ching Chang</a><sup>2</sup> Β· | |
| <a href="https://scholar.google.com/citations?user=gMBvzGoAAAAJ&hl=en" target="_blank">Xin Li</a><sup>2</sup> Β· | |
| <a href="https://yulunalexliu.github.io/" target="_blank">Yu-Lun Liu</a><sup>3</sup> Β· | |
| <a href="https://cchsu.info/" target="_blank">Chih-Chung Hsu</a><sup>1,3†</sup><br> | |
| <small><sup>1</sup>National Cheng Kung University Β· <sup>2</sup>State University of New York at Albany Β· | |
| <sup>3</sup>National Yang Ming Chiao Tung University Β· * equal contribution Β· † corresponding author</small></div> | |
| <div class="links"> | |
| <a href="https://ming053l.github.io/FlashReFocus/" target="_blank">π Project page</a> | |
| <a href="https://ming053l.github.io/FlashReFocus/assets/FlashReFocus.pdf" target="_blank">π Paper</a> | |
| <a href="https://github.com/ming053l/FlashReFocus" target="_blank">π» Code</a> | |
| <a href="https://youtu.be/g032M7moWQ4" target="_blank">π¬ Video</a> | |
| <a href="https://huggingface.co/datasets/ming0531/3CReal" target="_blank">ποΈ 3CReal dataset</a> | |
| </div> | |
| <div class="tip">π <b>Try this first:</b> pick the aquarium example β <b>Restore</b> β choose <b>π₯ Two focal planes</b> | |
| β click the face, then the shark. Every click after Restore runs only the 1.3M-parameter renderer.</div> | |
| </div> | |
| """ | |
| CSS = """ | |
| #header {margin-bottom: 4px} | |
| #header .title {display: flex; align-items: center; justify-content: center; gap: 16px} | |
| #header h1 {margin: 0; font-size: 2.1em; color: #2E7D32} | |
| #header .title p {margin: 2px 0 0; color: var(--body-text-color-subdued)} | |
| #header .authors {text-align: center; margin-top: 10px; line-height: 1.6} | |
| #header .authors a {color: #2E7D32; text-decoration: none; font-weight: 600} | |
| #header .authors small {color: var(--body-text-color-subdued)} | |
| #header .links {display: flex; flex-wrap: wrap; justify-content: center; gap: 8px; margin: 12px 0} | |
| #header .links a {padding: 5px 14px; border-radius: 999px; border: 1px solid #48A43F; color: #2E7D32; | |
| text-decoration: none; font-weight: 600} | |
| #header .links a:hover {background: #48A43F; color: white} | |
| #header .tip {background: #f1f8ef; border-left: 4px solid #48A43F; padding: 10px 14px; border-radius: 8px; | |
| color: #1f3d1f} | |
| .step h3 {margin: 6px 0 2px; color: #2E7D32} | |
| #status {background: #f1f8ef; border-left: 4px solid #48A43F; padding: 8px 14px; border-radius: 8px; min-height: 2.4em} | |
| #status p {margin: 0; color: #1f3d1f} | |
| """ | |
| with gr.Blocks(title="FlashReFocus", theme=gr.themes.Soft(primary_hue="green"), css=CSS) as demo: | |
| gr.HTML(HEADER) | |
| state = gr.State(None) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### β Upload a photo and restore it once", elem_classes="step") | |
| inp = gr.Image(label="Input photograph", type="pil", sources=["upload", "clipboard"], height=320) | |
| gr.Examples(EXAMPLES, inputs=inp, label="Examples (Unsplash portraits and our own photographs)") | |
| with gr.Accordion("Restoration settings", open=False): | |
| long_side = gr.Slider(512, 1536, value=1024, step=64, label="Long side (pixels)") | |
| omega = gr.Slider(0.0, 1.5, value=1.25, step=0.25, label="Detail weight Ο", | |
| info="0: pixel LoRA only (most faithful); higher restores more detail.") | |
| btn = gr.Button("β¨ Restore", variant="primary", size="lg") | |
| gr.Markdown("### β’ Depth map", elem_classes="step") | |
| depth = gr.Image(label="Depth used for rendering (near is bright): click to focus", type="pil", | |
| interactive=False, height=280) | |
| with gr.Accordion("Depth controls", open=False): | |
| gamma = gr.Slider(0.4, 2.5, value=1.0, step=0.05, label="Depth curve Ξ³ (D β D^Ξ³)", | |
| info="< 1 spreads out the near range, > 1 the far range; blur follows the new spacing.") | |
| band = gr.Slider(0.02, 0.3, value=0.06, step=0.01, label="Tilt-shift: half-height of the sharp band") | |
| user_depth = gr.Image(label="Your own depth map (optional; near is bright, any size)", type="numpy", | |
| sources=["upload", "clipboard"], height=150) | |
| with gr.Column(): | |
| gr.Markdown("### β‘ Click the restored image to focus", elem_classes="step") | |
| aif = gr.Image(label="Restored all-in-focus image (click to focus)", type="pil", interactive=False, | |
| height=320) | |
| mode = gr.Radio([SINGLE, DUAL, TILT], value=SINGLE, label="Focus mode", | |
| info="Two planes: the first click sets A, the second B. Tilt-shift: click the row to keep sharp.") | |
| strength = gr.Slider(0.0, 1.5, value=0.5, step=0.05, label="π Aperture (blur strength K)", | |
| info="0: everything sharp; 1: the widest aperture seen in training; above 1 extrapolates.") | |
| with gr.Row(): | |
| show_marker = gr.Checkbox(value=False, label="Show focal markers") | |
| clear = gr.Button("Clear focal points", size="sm") | |
| reset = gr.Button("βΊ Recommended settings", size="sm") | |
| status = gr.Markdown("Upload a photo or pick an example, then press β¨ Restore.", elem_id="status") | |
| out = gr.Image(label="Refocused result", type="pil", interactive=False, height=430, format="png") | |
| controls = [state, strength, mode, gamma, band, show_marker] | |
| outputs = [aif, out, depth, state, status] | |
| btn.click(restore, [inp, long_side, omega, gamma], [aif, depth, out, state, status]) | |
| aif.select(on_click, controls, outputs) | |
| depth.select(on_click, controls, outputs) | |
| for slider in (strength, gamma, band): | |
| slider.release(on_change, controls, outputs) | |
| mode.change(on_change, controls, outputs) | |
| show_marker.change(on_change, controls, outputs) | |
| user_depth.change(on_user_depth, [state, user_depth, strength, mode, gamma, band, show_marker], outputs) | |
| clear.click(on_clear, [state, gamma], outputs) | |
| reset.click(on_reset, [state], [long_side, omega, gamma, band, strength, mode, show_marker, user_depth] + outputs) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=20).launch(show_api=False) | |