Update app.py

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
 from datasets import load_dataset
+from itertools import islice
 import numpy as np
 from PIL import Image
 import torch
 from transformers import CLIPModel, CLIPProcessor
+import torch.nn.functional as F
 
+# ---------- utils ----------
 def flux_to_gray(flux_array):
     a = np.array(flux_array, dtype=np.float32)
     a = np.squeeze(a)
@@ -12,32 +15,60 @@ def flux_to_gray(flux_array):
         axis = int(np.argmin(a.shape))
         a = np.nanmean(a, axis=axis)
     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)
+    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")
 
+# ---------- model ----------
 model_id = "openai/clip-vit-base-patch32"
 model = CLIPModel.from_pretrained(model_id)
 processor = CLIPProcessor.from_pretrained(model_id)
+model.eval()
 
-def test_clip():
+# ---------- in-memory index ----------
+INDEX = {
+    "feats": None,   # torch.Tensor [N, 512]
+    "ids": [],       # list[str]
+    "thumbs": [],    # list[PIL.Image]
+    "bands": []      # list[str]
+}
+
+def build_index(n=200):
     ds = load_dataset("MultimodalUniverse/jwst", split="train", streaming=True)
-    rec = next(iter(ds))
-    pil = flux_to_gray(rec["image"]["flux"]).convert("RGB")
-    with torch.no_grad():
-        image_inputs = processor(images=pil, return_tensors="pt")
-        image_feats = model.get_image_features(**image_inputs)  # [1, 512]
-    return pil, f"image_features shape: {tuple(image_feats.shape)}", f"object_id: {rec.get('object_id')}"
-
-demo = gr.Interface(
-    fn=test_clip,
-    inputs=None,
-    outputs=[gr.Image(type="pil", label="Preview"),
-             gr.Textbox(label="Shape", lines=1),
-             gr.Textbox(label="Info", lines=1)],
-    title="JWST → CLIP embedding check (transformers)"
-)
+    feats, ids, thumbs, bands = [], [], [], []
+    for rec in islice(ds, int(n)):
+        pil = flux_to_gray(rec["image"]["flux"]).convert("RGB")
+        t = pil.copy(); t.thumbnail((128, 128))
+        with torch.no_grad():
+            inp = processor(images=pil, return_tensors="pt")
+            f = model.get_image_features(**inp)           # [1, 512]
+            f = F.normalize(f, p=2, dim=-1)[0]            # [512]
+        feats.append(f)
+        ids.append(str(rec.get("object_id")))
+        bands.append(str(rec["image"].get("band")))
+        thumbs.append(t)
+
+    if not feats:
+        return "No records indexed."
+
+    INDEX["feats"] = torch.stack(feats)                   # [N, 512]
+    INDEX["ids"] = ids
+    INDEX["thumbs"] = thumbs
+    INDEX["bands"] = bands
+    return f"Index built: {len(ids)} images."
+
+# ---------- UI ----------
+with gr.Blocks() as demo:
+    gr.Markdown("JWST multimodal search — build the index")
+
+    n = gr.Slider(50, 1000, value=200, step=10, label="How many images to index")
+    build_btn = gr.Button("Build index")
+    status = gr.Textbox(label="Status", lines=2)
+
+    build_btn.click(build_index, inputs=n, outputs=status)
+
 demo.launch()