Return depth as 16-bit PNG via gr.File (preserve full precision)

app.py

@@ -56,6 +56,8 @@ except Exception as _e:
     print(f"[warn] gradio_client patch failed: {_e}")
 # ─────────────────────────────────────────────────────────────────────────────
 
+import tempfile
+
 import numpy as np
 import spaces
 import torch
@@ -75,22 +77,40 @@ def _get_pipe(task: str):
     return _pipes[task]
 
 
-def _to_8bit(img: Image.Image) -> Image.Image:
-    """colorize_depth_map was modified to return 16-bit grayscale (uint16).
-    Gradio serialises via webp which drops the upper byte, yielding an almost-
-    white image.  Convert to proper 8-bit L before handing back to Gradio.
+def _save_16bit_png(img: Image.Image) -> str:
+    """Save the depth map as a true 16-bit PNG and return the path.
+
+    colorize_depth_map (modified previously) already returns a uint16 PIL
+    Image. The default gr.Image output would re-encode it as 8-bit webp,
+    which destroys the upper byte and yields visible stepping in the
+    fabricated relief (~20 µm per step over a 5 mm relief). Returning
+    a file path via gr.File bypasses Gradio's image processing entirely.
     """
     arr = np.array(img)
+    fd, tmp_path = tempfile.mkstemp(suffix='_depth16.png')
+    os.close(fd)
     if arr.dtype == np.uint16:
-        arr = (arr >> 8).astype(np.uint8)   # high byte → full 8-bit range
-        return Image.fromarray(arr, mode="L")
-    return img
+        # Native 16-bit path — preserve all 65,536 levels
+        Image.fromarray(arr, mode="I;16").save(tmp_path, format="PNG")
+    elif arr.dtype == np.uint8:
+        # Up-promote to 16-bit (no extra precision, but consistent output)
+        arr16 = arr.astype(np.uint16) * 257  # 0->0, 255->65535
+        Image.fromarray(arr16, mode="I;16").save(tmp_path, format="PNG")
+    else:
+        # Float / other — normalise to 16-bit
+        a = arr.astype(np.float32)
+        rng = max(a.max() - a.min(), 1e-9)
+        a = (a - a.min()) / rng
+        Image.fromarray((a * 65535).astype(np.uint16),
+                         mode="I;16").save(tmp_path, format="PNG")
+    return tmp_path
 
 
 @spaces.GPU
 def infer_depth(image_path):
     pipe = _get_pipe("depth")
-    return _to_8bit(infer_pipe(pipe, image_path, "depth", SEED, device))
+    img = infer_pipe(pipe, image_path, "depth", SEED, device)
+    return _save_16bit_png(img)
 
 
 @spaces.GPU
@@ -101,11 +121,12 @@ def infer_normal(image_path):
 
 with gr.Blocks(title="Lotus 1 - Depth + Normal") as demo:
     gr.Markdown("# Lotus 1 - Depth & Normal")
-    gr.Markdown("API: `/depth` returns depth map, `/normal` returns normal map.")
+    gr.Markdown("API: `/depth` returns a **16-bit grayscale PNG** depth map, "
+                "`/normal` returns an RGB normal map.")
 
     with gr.Tab("Depth"):
         d_in = gr.Image(label="Input", type="filepath")
-        d_out = gr.Image(label="Depth (g)")
+        d_out = gr.File(label="Depth (16-bit PNG)")
         d_btn = gr.Button("Run depth")
         d_btn.click(infer_depth, inputs=d_in, outputs=d_out, api_name="depth")