Update app.py

app.py

@@ -2,19 +2,15 @@ import gradio as gr
 from datasets import load_dataset
 import numpy as np
 from PIL import Image
+from sentence_transformers import SentenceTransformer
 
+# reuse the same grayscale conversion
 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
+        axis = int(np.argmin(a.shape))
         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)
@@ -24,20 +20,28 @@ def flux_to_gray(flux_array):
     arr = (norm * 255).astype(np.uint8)
     return Image.fromarray(arr, mode="L")
 
-def show_first():
+# load a well-known CLIP model
+model = SentenceTransformer("clip-ViT-B-32")
+
+def test_single_embedding():
     ds = load_dataset("MultimodalUniverse/jwst", split="train", streaming=True)
     rec = next(iter(ds))
-    img = rec["image"]
-    pil = flux_to_gray(img["flux"])
-    caption = f"object_id: {rec.get('object_id')}, band: {img.get('band')}"
-    return pil, caption
+    pil = flux_to_gray(rec["image"]["flux"]).convert("RGB")  # CLIP expects RGB
+    emb = model.encode(pil, convert_to_numpy=True)
+    info = f"OK. Image embedding shape: {emb.shape}"
+    caption = f"object_id: {rec.get('object_id')}"
+    return pil, caption, info
 
 demo = gr.Interface(
-    fn=show_first,
+    fn=test_single_embedding,
     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 + minimal metadata."
+    outputs=[
+        gr.Image(type="pil", label="Preview"),
+        gr.Textbox(label="Info", lines=1),
+        gr.Textbox(label="Embedding", lines=1),
+    ],
+    title="JWST → CLIP embedding check",
+    description="Embeds one JWST image with CLIP to confirm the pipeline."
 )
 
 demo.launch()