Add Relight tab

app.py

@@ -1,4 +1,8 @@
-"""ShadeNet-3.2 Gradio Space (CPU-friendly: fp16 ONNX backend)."""
+"""ShadeNet-3.2 Gradio Space (CPU-friendly: fp16 ONNX backend).
+
+Tab 1: decompose a photo into albedo/depth/normal/shading.
+Tab 2: relight the albedo with a user-controlled Lambertian light.
+"""
 import os
 
 import gradio as gr
@@ -7,8 +11,7 @@ import onnxruntime as ort
 import spaces
 from PIL import Image
 
-from inference_utils import (build_grid, extract_maps, pil_to_np, resize_pad,
-                             run_ensemble_onnx)
+from inference_utils import build_grid, extract_maps, pil_to_np, resize_pad
 
 MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            "onnx", "model_fp16.onnx")
@@ -17,16 +20,61 @@ IMAGE_SIZE = 512
 _sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
 
 
+def _run_model(img_rgb):
+    x = pil_to_np(img_rgb).astype(np.float32)
+    return _sess.run(None, {"input_rgb": x})[0]  # [1, 8, H, W] in [-1, 1]
+
+
 @spaces.GPU(duration=60)
 def decompose(image: Image.Image):
     """Returns (grid, input, albedo, shading, depth, normal, recon)."""
     img_rgb = resize_pad(image.convert("RGB"), IMAGE_SIZE)
-    out = run_ensemble_onnx(_sess, img_rgb, IMAGE_SIZE)
+    out = _run_model(img_rgb)
     maps = extract_maps(out, img_rgb)
     return (build_grid(img_rgb, out), maps["input"], maps["albedo"],
             maps["shading"], maps["depth"], maps["normal"], maps["recon"])
 
 
+def _relight(albedo01, normal01, azimuth, elevation, intensity, ambient,
+             color):
+    """Lambertian relight. Arrays in [0, 1]; normal frame: x=right, y=up,
+    z=toward viewer (camera space). Returns relit RGB in [0, 1]."""
+    n = normal01.astype(np.float32) * 2.0 - 1.0
+    n = n / np.maximum(np.linalg.norm(n, axis=-1, keepdims=True), 1e-6)
+    # image y is down; flip to up-positive for the light math
+    n = n * np.array([1.0, -1.0, 1.0], dtype=np.float32)
+    az, el = np.deg2rad(float(azimuth)), np.deg2rad(float(elevation))
+    lv = np.array([np.cos(el) * np.cos(az), np.sin(el),
+                   np.cos(el) * np.sin(az)], dtype=np.float32)
+    lambert = np.clip((n * lv).sum(-1), 0.0, 1.0)
+    shade = float(ambient) + (1.0 - float(ambient)) * float(intensity) * lambert
+    col = np.array(color, dtype=np.float32) / 255.0
+    return np.clip(albedo01.astype(np.float32) * shade[..., None] *
+                   col[None, None, :], 0.0, 1.0)
+
+
+@spaces.GPU(duration=60)
+def relight_init(image):
+    """Run inference once; cache albedo+normal; show default relight."""
+    if image is None:
+        return None, None
+    img_rgb = resize_pad(image.convert("RGB"), IMAGE_SIZE)
+    out = _run_model(img_rgb)
+    alb = ((out[:, 0:3] + 1.0) / 2.0)[0].transpose(1, 2, 0).astype(np.float32)
+    nrm = ((out[:, 4:7] + 1.0) / 2.0)[0].transpose(1, 2, 0).astype(np.float32)
+    rel = _relight(alb, nrm, 135.0, 45.0, 1.0, 0.25, (255, 255, 255))
+    return (alb, nrm), Image.fromarray((rel * 255).astype("uint8"))
+
+
+def relight_render(state, azimuth, elevation, intensity, ambient, color):
+    """Pure-numpy re-render from cached maps (no inference)."""
+    if state is None:
+        return None
+    alb, nrm = state
+    rel = _relight(alb, nrm, azimuth, elevation, intensity, ambient, color)
+    return Image.fromarray((rel * 255).astype("uint8"))
+
+
 EXAMPLES = []
 _exdir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                       "assets", "examples")
@@ -34,7 +82,7 @@ if os.path.isdir(_exdir):
     EXAMPLES = sorted(os.path.join(_exdir, f) for f in os.listdir(_exdir)
                       if f.lower().endswith((".png", ".jpg", ".jpeg")))
 
-demo = gr.Interface(
+tab1 = gr.Interface(
     fn=decompose,
     inputs=gr.Image(type="pil", label="Input photo"),
     outputs=[
@@ -55,5 +103,30 @@ demo = gr.Interface(
     cache_examples=False,
 )
 
+with gr.Blocks() as tab2:
+    gr.Markdown("## Relight\nRun inference once, then move the light. "
+                "Single-bounce Lambertian — no shadows or interreflections.")
+    with gr.Row():
+        rin = gr.Image(type="pil", label="Input photo")
+        rout = gr.Image(type="pil", label="Relit")
+    state = gr.State(value=None)
+    with gr.Row():
+        az = gr.Slider(0, 360, value=135, step=1, label="Azimuth (deg)")
+        el = gr.Slider(0, 90, value=45, step=1, label="Elevation (deg)")
+    with gr.Row():
+        inten = gr.Slider(0.0, 2.0, value=1.0, step=0.01, label="Intensity")
+        amb = gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="Ambient")
+        col = gr.ColorPicker(value="#ffffff", label="Light color")
+    rin.change(relight_init, inputs=rin, outputs=[state, rout])
+    for comp in (az, el, inten, amb, col):
+        comp.change(relight_render,
+                    inputs=[state, az, el, inten, amb, col], outputs=rout)
+
+demo = gr.TabbedInterface(
+    [tab1, tab2],
+    ["Decompose", "Relight"],
+    title="ShadeNet-3.2 (5M): inverse rendering + relighting",
+)
+
 if __name__ == "__main__":
     demo.launch()