Update app.py

app.py

@@ -3,21 +3,32 @@ from datasets import load_dataset
 import numpy as np
 from PIL import Image
 
+def flux_to_gray(flux_array):
+    a = np.array(flux_array, dtype=np.float32)
+
+    # remove 1-length dimensions like (1,H,W) or (H,W,1)
+    a = np.squeeze(a)
+
+    # if still 3D (e.g., C,H,W or H,W,C), collapse the smallest axis
+    if a.ndim == 3:
+        axis = int(np.argmin(a.shape))  # pick the likely channel axis
+        a = np.nanmean(a, axis=axis)
+
+    # handle NaNs/infs and scale to 0..255
+    a = np.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0)
+    lo = np.nanpercentile(a, 1)
+    hi = np.nanpercentile(a, 99)
+    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
+        lo, hi = float(np.nanmin(a)), float(np.nanmax(a))
+    norm = np.clip((a - lo) / (hi - lo + 1e-9), 0, 1)
+    arr = (norm * 255).astype(np.uint8)
+    return Image.fromarray(arr, mode="L")
+
 def show_first():
     ds = load_dataset("MultimodalUniverse/jwst", split="train", streaming=True)
     rec = next(iter(ds))
     img = rec["image"]
-
-    # convert the flux array to a viewable grayscale image
-    flux = np.array(img["flux"], dtype=np.float32)
-    lo = np.percentile(flux, 1)
-    hi = np.percentile(flux, 99)
-    if hi <= lo:  # fallback if data is odd
-        lo, hi = float(flux.min()), float(flux.max())
-    norm = np.clip((flux - lo) / (hi - lo + 1e-9), 0, 1)
-    arr = (norm * 255).astype(np.uint8)
-    pil = Image.fromarray(arr, mode="L")
-
+    pil = flux_to_gray(img["flux"])
     caption = f"object_id: {rec.get('object_id')}, band: {img.get('band')}"
     return pil, caption
 
@@ -26,7 +37,7 @@ demo = gr.Interface(
     inputs=None,
     outputs=[gr.Image(type="pil", label="Preview"), gr.Textbox(label="Info", lines=2)],
     title="JWST sample preview",
-    description="First streamed record as image + a little metadata."
+    description="First streamed record as image + minimal metadata."
 )
 
 demo.launch()