Upload 2 files

Added functionality to cropping tab

app.py

@@ -10,6 +10,7 @@ import threading
 import gradio as gr
 import spaces
 import torch
+import numpy as np
 from PIL import Image
 
 # ── Local modules — single source of truth for each concern ─────────────────
@@ -48,6 +49,8 @@ from image_utils import (
     build_pnginfo,
     push_pil_to_base,
     push_pil_to_reference,
+    extend_editor_canvas,
+    render_extend_schematic,
 )
 from control_tools import (
     generate_depthmap,
@@ -639,6 +642,33 @@ with gr.Blocks() as demo:
                     send_to_base_btn = gr.Button("→ Send to Base Image", variant="primary")
                     send_to_ref_btn  = gr.Button("→ Add to Reference Images")
 
+                # ── Extend canvas section ────────────────────────────────────────────────
+                # Uses the editor's current composite as the source so cropping + painting
+                # happen first, then we grow the canvas around the result. Output is
+                # loaded back into the same editor — Send → Base / Reference from there.
+                with gr.Accordion("📐 Extend canvas (add padding around image)", open=True):
+                    gr.Markdown(
+                        "Grow the editor image's canvas by a percentage in any combination "
+                        "of directions. Percentages are relative to the *current* image "
+                        "size — `Down = 100` doubles the height with the image on top. "
+                        "The result replaces the editor contents so you can crop again or "
+                        "send it to Base / Reference with the buttons above."
+                    )
+                    with gr.Row():
+                        ext_up    = gr.Number(label="Up %",    value=0,   minimum=0, precision=2)
+                        ext_down  = gr.Number(label="Down %",  value=0,   minimum=0, precision=2)
+                        ext_left  = gr.Number(label="Left %",  value=0,   minimum=0, precision=2)
+                        ext_right = gr.Number(label="Right %", value=0,   minimum=0, precision=2)
+                    with gr.Row():
+                        ext_fill = gr.ColorPicker(label="Fill colour", value="#000000")
+                        extend_btn = gr.Button("📐 Extend canvas", variant="primary")
+                    with gr.Row():
+                        ext_schematic = gr.Image(
+                            label="Layout preview (red outline = current image)",
+                            type="pil", interactive=False, height=220,
+                        )
+                        ext_info = gr.Markdown("*Upload something into the editor first.*")
+
             # ── Bulk processing tab ──────────────────────────────────────────
             with gr.Tab("📦 Bulk Process", id="tab_bulk"):
                 gr.Markdown(
@@ -815,6 +845,42 @@ with gr.Blocks() as demo:
         outputs=[reference_images],
     ).then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info])
 
+    # Extend-canvas wiring
+    # The schematic previews the *current editor composite* so users see live
+    # feedback as they nudge the percentages / fill colour.
+    def _editor_source_for_preview(editor_value):
+        if not editor_value or editor_value.get("composite") is None:
+            return None
+        comp = editor_value["composite"]
+        if isinstance(comp, np.ndarray):
+            from PIL import Image as _Image
+            comp = _Image.fromarray(comp)
+        return comp
+
+    def _update_extend_preview(editor_value, up, down, left, right, fill):
+        return render_extend_schematic(
+            _editor_source_for_preview(editor_value), up, down, left, right, fill,
+        )
+
+    _extend_preview_inputs = [editor, ext_up, ext_down, ext_left, ext_right, ext_fill]
+    _extend_preview_outputs = [ext_schematic, ext_info]
+    for _c in (ext_up, ext_down, ext_left, ext_right, ext_fill):
+        _c.change(fn=_update_extend_preview,
+                  inputs=_extend_preview_inputs, outputs=_extend_preview_outputs)
+    # Also refresh whenever the editor content itself changes (upload, HEIC load,
+    # a fresh crop) — that way the schematic scale updates too.
+    editor.change(fn=_update_extend_preview,
+                  inputs=_extend_preview_inputs, outputs=_extend_preview_outputs)
+
+    # Run: extend, then hand the new PIL back to the editor. `render_extend_
+    # schematic` re-runs via editor.change once the new image lands, so no
+    # extra .then() is needed for the preview.
+    extend_btn.click(
+        fn=extend_editor_canvas,
+        inputs=[editor, ext_up, ext_down, ext_left, ext_right, ext_fill],
+        outputs=[editor],
+    )
+
     # ── Bulk tab wiring ──────────────────────────────────────────────────────
     bulk_event = bulk_run_btn.click(
         fn=bulk_infer,

image_utils.py

@@ -10,7 +10,7 @@ import tempfile
 from typing import Any
 
 import numpy as np
-from PIL import Image, ImageOps, ImageFilter
+from PIL import Image, ImageOps, ImageFilter, ImageDraw, ImageColor
 from PIL.PngImagePlugin import PngInfo
 import gradio as gr
 
@@ -316,3 +316,91 @@ def push_pil_to_reference(img, current_gallery):
     current.append(img)
     gr.Info("Added to Reference Images.")
     return current
+
+# ── Canvas extension (used by the Crop / Fix tab) ───────────────────────────
+# Percentages are relative to the *original* dimensions. Padding is filled
+# with a flat colour. Returns a new PIL image; does NOT save to disk — the
+# result is loaded straight back into the ImageEditor so the existing
+# Send → Base / Send → Reference buttons handle export.
+
+from PIL import ImageColor  # add near the other PIL imports at top of file
+
+
+def _parse_fill(color_str: str, mode: str):
+    rgba = ImageColor.getcolor(color_str or "#000000", "RGBA")
+    if mode == "RGB":
+        return rgba[:3]
+    if mode == "L":
+        r, g, b, _ = rgba
+        return int(0.299 * r + 0.587 * g + 0.114 * b)
+    return rgba  # RGBA
+
+
+def compute_extend_padding(w: int, h: int, up, down, left, right) -> tuple[int, int, int, int]:
+    """Return (pad_left, pad_top, pad_right, pad_bottom) in pixels for the
+    given percentage inputs. Percentages are relative to original W/H so
+    Down=100 on 960×960 → 960×1920 with the source at the top."""
+    pl = int(round(w * (float(left  or 0) / 100.0)))
+    pr = int(round(w * (float(right or 0) / 100.0)))
+    pt = int(round(h * (float(up    or 0) / 100.0)))
+    pb = int(round(h * (float(down  or 0) / 100.0)))
+    return pl, pt, pr, pb
+
+
+def extend_canvas(img: Image.Image, up, down, left, right, fill: str) -> Image.Image:
+    """Extend `img`'s canvas by the given per-side percentages, filling the
+    new area with `fill`. Original image mode is preserved so RGBA stays
+    RGBA (no transparency loss)."""
+    if img is None:
+        raise gr.Error("Nothing to extend — upload an image into the editor first.")
+    if img.mode not in ("RGB", "RGBA", "L"):
+        img = img.convert("RGBA")
+    for name, v in (("Up", up), ("Down", down), ("Left", left), ("Right", right)):
+        if v is None or float(v) < 0:
+            raise gr.Error(f"'{name} %' must be ≥ 0.")
+    w, h = img.size
+    pl, pt, pr, pb = compute_extend_padding(w, h, up, down, left, right)
+    if pl == pt == pr == pb == 0:
+        gr.Info("All percentages are 0 — image unchanged.")
+        return img
+    new_size = (w + pl + pr, h + pt + pb)
+    canvas = Image.new(img.mode, new_size, _parse_fill(fill, img.mode))
+    canvas.paste(img, (pl, pt))
+    return canvas
+
+
+def render_extend_schematic(
+    img: Image.Image | None, up, down, left, right, fill: str,
+    max_dim: int = 320,
+) -> tuple[Image.Image | None, str]:
+    """Live, to-scale preview of what extend_canvas() will produce, without
+    doing the full render. Returns (schematic PIL, info markdown)."""
+    if img is None:
+        return None, "*Upload something into the editor first.*"
+    w, h = img.size
+    pl, pt, pr, pb = compute_extend_padding(w, h, up, down, left, right)
+    nw, nh = w + pl + pr, h + pt + pb
+    scale = min(max_dim / nw, max_dim / nh, 1.0)
+    sw, sh = max(1, int(nw * scale)), max(1, int(nh * scale))
+    spl, spt = int(pl * scale), int(pt * scale)
+    sow, soh = max(1, int(w * scale)), max(1, int(h * scale))
+
+    schem = Image.new("RGB", (sw, sh), _parse_fill(fill, "RGB"))
+    d = ImageDraw.Draw(schem)  # ImageDraw already imported at top of module
+    d.rectangle([spl, spt, spl + sow - 1, spt + soh - 1],
+                fill=(200, 200, 200), outline=(255, 0, 0), width=2)
+    info = (f"**Original:** {w} × {h}  \n"
+            f"**Padding (L, T, R, B):** {pl}, {pt}, {pr}, {pb}  \n"
+            f"**Final canvas:** {nw} × {nh}")
+    return schem, info
+
+
+def extend_editor_canvas(editor_value, up, down, left, right, fill) -> Image.Image:
+    """Take the current editor composite, extend it, and return the result so
+    it can be loaded straight back into the same ImageEditor. Any crop marks
+    the user had placed are baked in before extending (that's what
+    `_editor_composite` already does)."""
+    composite = fix_orientation(_editor_composite(editor_value))
+    out = extend_canvas(composite, up, down, left, right, fill)
+    gr.Info(f"Canvas extended to {out.size[0]} × {out.size[1]}.")
+    return out