Update colors

app.py

@@ -10,6 +10,7 @@ import beartype
 import einops
 import einops.layers.torch
 import gradio as gr
+import matplotlib
 import numpy as np
 import saev.activations
 import saev.config
@@ -31,7 +32,7 @@ logger = logging.getLogger("app.py")
 ####################
 
 
-MAX_FREQ = 1e-2
+MAX_FREQ = 3e-2
 """Maximum frequency. Any feature that fires more than this is ignored."""
 
 RESIZE_SIZE = 512
@@ -46,12 +47,14 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 CWD = pathlib.Path(".")
 """Current working directory."""
 
-N_SAE_LATENTS = 2
+N_SAE_LATENTS = 4
 """Number of SAE latents to show."""
 
 N_LATENT_EXAMPLES = 4
 """Number of examples per SAE latent to show."""
 
+COLORMAP = matplotlib.colormaps.get_cmap("plasma")
+
 
 @beartype.beartype
 class Example(typing.TypedDict):
@@ -175,8 +178,11 @@ def add_highlights(
     overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
     draw = ImageDraw.Draw(overlay)
 
+    colors = np.zeros((len(patches), 3), dtype=np.uint8)
+    colors[:, 0] = ((patches / (upper + 1e-9)) * 255).astype(np.uint8)
+
     # Using semi-transparent red (255, 0, 0, alpha)
-    for p, val in enumerate(patches):
+    for p, (val, color) in enumerate(zip(patches, colors)):
         assert upper is not None
         val /= upper + 1e-9
         x_np, y_np = p % iw_np, p // ih_np
@@ -185,28 +191,13 @@ def add_highlights(
                 (x_np * pw_px, y_np * ph_px),
                 (x_np * pw_px + pw_px, y_np * ph_px + ph_px),
             ],
-            fill=(int(val * 256), 0, 0, int(opacity * val * 256)),
+            fill=(*color, int(opacity * val * 255)),
         )
 
     # Composite the original image and the overlay
     return Image.alpha_composite(img.convert("RGBA"), overlay)
 
 
-@jaxtyped(typechecker=beartype.beartype)
-@torch.inference_mode
-def upsample(
-    x_WH: Int[Tensor, "width_ps height_ps"],
-) -> UInt8[Tensor, "width_px height_px"]:
-    return (
-        torch.nn.functional.interpolate(
-            x_WH.view((1, 1, 16, 16)).float(),
-            scale_factor=28,
-        )
-        .view((448, 448))
-        .type(torch.uint8)
-    )
-
-
 #######################
 # Inference Functions #
 #######################
@@ -317,10 +308,10 @@ def get_orig_preds(img: Image.Image) -> Example:
     clf = load_clf()
     logits_WHC = clf(x_WHD)
 
-    pred_WH = logits_WHC.argmax(axis=-1)
+    pred_WH = logits_WHC[:, :, 1:].argmax(axis=-1) + 1
     return {
         "orig_url": data.img_to_base64(data.to_sized(img)),
-        "seg_url": data.img_to_base64(data.u8_to_img(upsample(pred_WH))),
+        "seg_url": data.img_to_base64(data.u8_to_overlay(pred_WH, img)),
         "classes": data.to_classes(pred_WH),
     }
 
@@ -384,11 +375,11 @@ def get_mod_preds(img: Image.Image, latents: dict[str, int | float]) -> Example:
     mod_WHD = einops.rearrange(mod_BPD, "() (w h) dim -> w h dim", w=16, h=16)
 
     logits_WHC = clf(mod_WHD)
-    pred_WH = logits_WHC.argmax(axis=-1)
+    pred_WH = logits_WHC[:, :, 1:].argmax(axis=-1) + 1
     # pred_WH = einops.rearrange(pred_P, "(w h) -> w h", w=16, h=16)
     return {
         "orig_url": data.img_to_base64(data.to_sized(img)),
-        "seg_url": data.img_to_base64(data.u8_to_img(upsample(pred_WH))),
+        "seg_url": data.img_to_base64(data.u8_to_overlay(pred_WH, img)),
         "classes": data.to_classes(pred_WH),
     }
 

data.py

@@ -8,8 +8,9 @@ import beartype
 import einops.layers.torch
 import numpy as np
 import requests
-from jaxtyping import Integer, UInt8, jaxtyped
-from PIL import Image
+import torch
+from jaxtyping import Int, Integer, UInt8, jaxtyped
+from PIL import Image, ImageDraw
 from torch import Tensor
 from torchvision.transforms import v2
 
@@ -50,6 +51,7 @@ def make_colors() -> UInt8[np.ndarray, "n 3"]:
 
     # Fixed colors. Must be synced with Segmentation.elm.
     colors[2] = np.array([201, 249, 255], dtype=np.uint8)
+    colors[2] = np.array([201, 249, 255], dtype=np.uint8)
     colors[4] = np.array([151, 204, 4], dtype=np.uint8)
     colors[13] = np.array([104, 139, 88], dtype=np.uint8)
     colors[16] = np.array([54, 48, 32], dtype=np.uint8)
@@ -89,6 +91,50 @@ def to_u8(seg_raw: Image.Image) -> UInt8[Tensor, "width height"]:
     return u8_transform(seg_raw)
 
 
+@jaxtyped(typechecker=beartype.beartype)
+def upsample(
+    x_WH: Int[Tensor, "width_ps height_ps"],
+) -> UInt8[Tensor, "width_px height_px"]:
+    return (
+        torch.nn.functional.interpolate(
+            x_WH.view((1, 1, 16, 16)).float(),
+            scale_factor=28,
+        )
+        .view((448, 448))
+        .type(torch.uint8)
+    )
+
+
+@jaxtyped(typechecker=beartype.beartype)
+def u8_to_overlay(
+    map: Integer[Tensor, "width_ps height_ps"],
+    img: Image.Image,
+    *,
+    opacity: float = 0.5,
+) -> Image.Image:
+    iw_np, ih_np = map.shape
+    iw_px, ih_px = img.size
+    pw_px, ph_px = iw_px // iw_np, ih_px // ih_np
+
+    # Create a transparent overlay
+    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
+    draw = ImageDraw.Draw(overlay)
+
+    # Using semi-transparent red (255, 0, 0, alpha)
+    for p, i in enumerate(map.view(-1).tolist()):
+        x_np, y_np = p % iw_np, p // ih_np
+        draw.rectangle(
+            [
+                (x_np * pw_px, y_np * ph_px),
+                (x_np * pw_px + pw_px, y_np * ph_px + ph_px),
+            ],
+            fill=(*colors[i - 1], int(opacity * 256)),
+        )
+
+    # Composite the original image and the overlay
+    return Image.alpha_composite(img.convert("RGBA"), overlay)
+
+
 @jaxtyped(typechecker=beartype.beartype)
 def u8_to_img(map: UInt8[Tensor, "width height"]) -> Image.Image:
     map = map.cpu().numpy()
@@ -109,7 +155,7 @@ def to_classes(map: Integer[Tensor, "width height"]) -> list[int]:
 @beartype.beartype
 def img_to_base64(img: Image.Image) -> str:
     buf = io.BytesIO()
-    img.save(buf, format="webp")
+    img.save(buf, format="webp", lossless=True)
     b64 = base64.b64encode(buf.getvalue())
     s64 = b64.decode("utf8")
     return "data:image/webp;base64," + s64

requirements.txt

@@ -2,9 +2,9 @@
 #    uv pip compile pyproject.toml
 aiofiles==23.2.1
     # via gradio
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.12
     # via
     #   datasets
     #   fsspec
@@ -18,46 +18,22 @@ anyio==4.8.0
     # via
     #   gradio
     #   httpx
-    #   jupyter-server
     #   pycrdt
     #   starlette
-argon2-cffi==23.1.0
-    # via jupyter-server
-argon2-cffi-bindings==21.2.0
-    # via argon2-cffi
-arrow==1.3.0
-    # via isoduration
-asttokens==3.0.0
-    # via stack-data
-async-lru==2.0.4
-    # via jupyterlab
 attrs==25.1.0
     # via
     #   aiohttp
     #   jsonschema
     #   referencing
-babel==2.17.0
-    # via jupyterlab-server
 beartype==0.19.0
     # via
     #   saev-semantic-segmentation (pyproject.toml)
     #   saev
-beautifulsoup4==4.13.1
-    # via nbconvert
-bleach==6.2.0
-    # via nbconvert
-braceexpand==0.1.7
-    # via webdataset
 certifi==2025.1.31
     # via
     #   httpcore
     #   httpx
     #   requests
-    #   sentry-sdk
-cffi==1.17.1
-    # via
-    #   argon2-cffi-bindings
-    #   pyvips
 charset-normalizer==3.4.1
     # via requests
 click==8.1.8
@@ -65,43 +41,28 @@ click==8.1.8
     #   marimo
     #   typer
     #   uvicorn
-    #   wandb
 cloudpickle==3.1.1
     # via submitit
-comm==0.2.2
-    # via ipykernel
 contourpy==1.3.1
     # via matplotlib
 cycler==0.12.1
     # via matplotlib
-datasets==3.2.0
+datasets==3.3.2
     # via saev
-debugpy==1.8.12
-    # via ipykernel
-decorator==5.1.1
-    # via ipython
-defusedxml==0.7.1
-    # via nbconvert
 dill==0.3.8
     # via
     #   datasets
     #   multiprocess
-docker-pycreds==0.4.0
-    # via wandb
 docstring-parser==0.16
     # via tyro
 docutils==0.21.2
     # via marimo
-einops==0.8.0
+einops==0.8.1
     # via
     #   saev-semantic-segmentation (pyproject.toml)
     #   saev
-executing==2.2.0
-    # via stack-data
 fastapi==0.115.8
     # via gradio
-fastjsonschema==2.21.1
-    # via nbformat
 ffmpy==0.5.0
     # via gradio
 filelock==3.17.0
@@ -109,15 +70,13 @@ filelock==3.17.0
     #   datasets
     #   huggingface-hub
     #   torch
-fonttools==4.55.8
+fonttools==4.56.0
     # via matplotlib
-fqdn==1.5.1
-    # via jsonschema
 frozenlist==1.5.0
     # via
     #   aiohttp
     #   aiosignal
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
     #   datasets
     #   gradio-client
@@ -125,13 +84,9 @@ fsspec==2024.9.0
     #   torch
 ftfy==6.3.1
     # via open-clip-torch
-gitdb==4.0.12
-    # via gitpython
-gitpython==3.1.44
-    # via wandb
-gradio==5.14.0
+gradio==5.16.2
     # via saev-semantic-segmentation (pyproject.toml)
-gradio-client==1.7.0
+gradio-client==1.7.1
     # via gradio
 h11==0.14.0
     # via
@@ -143,9 +98,8 @@ httpx==0.28.1
     # via
     #   gradio
     #   gradio-client
-    #   jupyterlab
     #   safehttpx
-huggingface-hub==0.28.1
+huggingface-hub==0.29.1
     # via
     #   datasets
     #   gradio
@@ -156,88 +110,30 @@ idna==3.10
     # via
     #   anyio
     #   httpx
-    #   jsonschema
     #   requests
     #   yarl
-ipykernel==6.29.5
-    # via jupyterlab
-ipython==8.32.0
-    # via ipykernel
-isoduration==20.11.0
-    # via jsonschema
 itsdangerous==2.2.0
     # via marimo
-jaxtyping==0.2.37
+jaxtyping==0.2.38
     # via saev
 jedi==0.19.2
-    # via
-    #   ipython
-    #   marimo
+    # via marimo
 jinja2==3.1.5
     # via
     #   altair
     #   gradio
-    #   jupyter-server
-    #   jupyterlab
-    #   jupyterlab-server
-    #   nbconvert
     #   torch
-joblib==1.4.2
-    # via scikit-learn
-json5==0.10.0
-    # via jupyterlab-server
-jsonpointer==3.0.0
-    # via jsonschema
 jsonschema==4.23.0
-    # via
-    #   altair
-    #   jupyter-events
-    #   jupyterlab-server
-    #   nbformat
+    # via altair
 jsonschema-specifications==2024.10.1
     # via jsonschema
-jupyter-client==8.6.3
-    # via
-    #   ipykernel
-    #   jupyter-server
-    #   nbclient
-jupyter-core==5.7.2
-    # via
-    #   ipykernel
-    #   jupyter-client
-    #   jupyter-server
-    #   jupyterlab
-    #   nbclient
-    #   nbconvert
-    #   nbformat
-jupyter-events==0.12.0
-    # via jupyter-server
-jupyter-lsp==2.2.5
-    # via jupyterlab
-jupyter-server==2.15.0
-    # via
-    #   jupyter-lsp
-    #   jupyterlab
-    #   jupyterlab-server
-    #   notebook-shim
-jupyter-server-terminals==0.5.3
-    # via jupyter-server
-jupyterlab==4.3.5
-    # via saev
-jupyterlab-pygments==0.3.0
-    # via nbconvert
-jupyterlab-server==2.27.3
-    # via jupyterlab
 kiwisolver==1.4.8
     # via matplotlib
-mako==1.3.9
-    # via pdoc3
-marimo==0.10.19
+marimo==0.11.7
     # via saev
 markdown==3.7
     # via
     #   marimo
-    #   pdoc3
     #   pymdown-extensions
 markdown-it-py==3.0.0
     # via rich
@@ -245,18 +141,10 @@ markupsafe==2.1.5
     # via
     #   gradio
     #   jinja2
-    #   mako
-    #   nbconvert
 matplotlib==3.10.0
     # via saev
-matplotlib-inline==0.1.7
-    # via
-    #   ipykernel
-    #   ipython
 mdurl==0.1.2
     # via markdown-it-py
-mistune==3.1.1
-    # via nbconvert
 mpmath==1.3.0
     # via sympy
 multidict==6.1.0
@@ -265,26 +153,13 @@ multidict==6.1.0
     #   yarl
 multiprocess==0.70.16
     # via datasets
-narwhals==1.25.0
+narwhals==1.27.1
     # via
     #   altair
     #   marimo
-nbclient==0.10.2
-    # via nbconvert
-nbconvert==7.16.6
-    # via jupyter-server
-nbformat==5.10.4
-    # via
-    #   jupyter-server
-    #   nbclient
-    #   nbconvert
-nest-asyncio==1.6.0
-    # via ipykernel
 networkx==3.4.2
     # via torch
-notebook-shim==0.2.4
-    # via jupyterlab
-numpy==2.2.2
+numpy==2.2.3
     # via
     #   saev-semantic-segmentation (pyproject.toml)
     #   contourpy
@@ -292,10 +167,7 @@ numpy==2.2.2
     #   gradio
     #   matplotlib
     #   pandas
-    #   scikit-learn
-    #   scipy
     #   torchvision
-    #   webdataset
 nvidia-cublas-cu12==12.4.5.8
     # via
     #   nvidia-cudnn-cu12
@@ -334,8 +206,6 @@ open-clip-torch==2.30.0
     # via saev
 orjson==3.10.15
     # via gradio
-overrides==7.7.0
-    # via jupyter-server
 packaging==24.2
     # via
     #   altair
@@ -343,81 +213,43 @@ packaging==24.2
     #   gradio
     #   gradio-client
     #   huggingface-hub
-    #   ipykernel
-    #   jupyter-events
-    #   jupyter-server
-    #   jupyterlab
-    #   jupyterlab-server
     #   marimo
     #   matplotlib
-    #   nbconvert
 pandas==2.2.3
     # via
     #   datasets
     #   gradio
-pandocfilters==1.5.1
-    # via nbconvert
 parso==0.8.4
     # via jedi
-pdoc3==0.11.5
-    # via saev
-pexpect==4.9.0
-    # via ipython
 pillow==11.1.0
     # via
     #   gradio
     #   matplotlib
     #   saev
     #   torchvision
-pkgconfig==1.5.5
-    # via pyvips
-platformdirs==4.3.6
-    # via
-    #   jupyter-core
-    #   wandb
-polars==1.21.0
+polars==1.22.0
     # via saev
-prometheus-client==0.21.1
-    # via jupyter-server
-prompt-toolkit==3.0.50
-    # via ipython
-propcache==0.2.1
+propcache==0.3.0
     # via
     #   aiohttp
     #   yarl
-protobuf==5.29.3
-    # via wandb
-psutil==6.1.1
-    # via
-    #   ipykernel
-    #   marimo
-    #   wandb
-ptyprocess==0.7.0
-    # via
-    #   pexpect
-    #   terminado
-pure-eval==0.2.3
-    # via stack-data
-pyarrow==19.0.0
+psutil==7.0.0
+    # via marimo
+pyarrow==19.0.1
     # via datasets
-pycparser==2.22
-    # via cffi
 pycrdt==0.11.1
     # via marimo
 pydantic==2.10.6
     # via
     #   fastapi
     #   gradio
-    #   wandb
 pydantic-core==2.27.2
     # via pydantic
 pydub==0.25.1
     # via gradio
 pygments==2.19.1
     # via
-    #   ipython
     #   marimo
-    #   nbconvert
     #   rich
 pymdown-extensions==10.14.3
     # via marimo
@@ -425,68 +257,43 @@ pyparsing==3.2.1
     # via matplotlib
 python-dateutil==2.9.0.post0
     # via
-    #   arrow
-    #   jupyter-client
     #   matplotlib
     #   pandas
-python-json-logger==3.2.1
-    # via jupyter-events
 python-multipart==0.0.20
     # via gradio
 pytz==2025.1
     # via pandas
-pyvips==2.2.3
-    # via saev
 pyyaml==6.0.2
     # via
     #   datasets
     #   gradio
     #   huggingface-hub
-    #   jupyter-events
     #   marimo
     #   pymdown-extensions
     #   timm
-    #   wandb
-    #   webdataset
-pyzmq==26.2.1
-    # via
-    #   ipykernel
-    #   jupyter-client
-    #   jupyter-server
 referencing==0.36.2
     # via
     #   jsonschema
     #   jsonschema-specifications
-    #   jupyter-events
 regex==2024.11.6
     # via open-clip-torch
 requests==2.32.3
     # via
     #   datasets
     #   huggingface-hub
-    #   jupyterlab-server
-    #   wandb
-rfc3339-validator==0.1.4
-    # via
-    #   jsonschema
-    #   jupyter-events
-rfc3986-validator==0.1.1
-    # via
-    #   jsonschema
-    #   jupyter-events
 rich==13.9.4
     # via
     #   typer
     #   tyro
-rpds-py==0.22.3
+rpds-py==0.23.0
     # via
     #   jsonschema
     #   referencing
-ruff==0.9.4
+ruff==0.9.7
     # via
     #   gradio
     #   marimo
-saev @ git+https://github.com/samuelstevens/saev@44f1aa334828ed994d4d670629ff639c655db38b
+saev @ git+https://github.com/samuelstevens/saev@298cabdb6b771c76b402d0fdddab6907d1941d7a
     # via saev-semantic-segmentation (pyproject.toml)
 safehttpx==0.1.6
     # via gradio
@@ -494,40 +301,18 @@ safetensors==0.5.2
     # via
     #   open-clip-torch
     #   timm
-scikit-learn==1.6.1
-    # via saev
-scipy==1.15.1
-    # via scikit-learn
 semantic-version==2.10.0
     # via gradio
-send2trash==1.8.3
-    # via jupyter-server
-sentry-sdk==2.20.0
-    # via wandb
-setproctitle==1.3.4
-    # via wandb
 setuptools==75.8.0
-    # via
-    #   jupyterlab
-    #   torch
-    #   wandb
+    # via torch
 shellingham==1.5.4
     # via typer
 shtab==1.7.1
     # via tyro
 six==1.17.0
-    # via
-    #   docker-pycreds
-    #   python-dateutil
-    #   rfc3339-validator
-smmap==5.0.2
-    # via gitdb
+    # via python-dateutil
 sniffio==1.3.1
     # via anyio
-soupsieve==2.6
-    # via beautifulsoup4
-stack-data==0.6.3
-    # via ipython
 starlette==0.45.3
     # via
     #   fastapi
@@ -537,16 +322,8 @@ submitit==1.5.2
     # via saev
 sympy==1.13.1
     # via torch
-terminado==0.18.1
-    # via
-    #   jupyter-server
-    #   jupyter-server-terminals
-threadpoolctl==3.5.0
-    # via scikit-learn
 timm==1.0.14
     # via open-clip-torch
-tinycss2==1.4.0
-    # via bleach
 tomlkit==0.13.2
     # via
     #   gradio
@@ -563,46 +340,22 @@ torchvision==0.21.0
     #   saev-semantic-segmentation (pyproject.toml)
     #   open-clip-torch
     #   timm
-tornado==6.4.2
-    # via
-    #   ipykernel
-    #   jupyter-client
-    #   jupyter-server
-    #   jupyterlab
-    #   terminado
 tqdm==4.67.1
     # via
     #   datasets
     #   huggingface-hub
     #   open-clip-torch
     #   saev
-traitlets==5.14.3
-    # via
-    #   comm
-    #   ipykernel
-    #   ipython
-    #   jupyter-client
-    #   jupyter-core
-    #   jupyter-events
-    #   jupyter-server
-    #   jupyterlab
-    #   matplotlib-inline
-    #   nbclient
-    #   nbconvert
-    #   nbformat
 triton==3.2.0
     # via torch
-typeguard==4.4.1
+typeguard==4.4.2
     # via tyro
 typer==0.15.1
     # via gradio
-types-python-dateutil==2.9.0.20241206
-    # via arrow
 typing-extensions==4.12.2
     # via
     #   altair
     #   anyio
-    #   beautifulsoup4
     #   fastapi
     #   gradio
     #   gradio-client
@@ -615,40 +368,20 @@ typing-extensions==4.12.2
     #   typeguard
     #   typer
     #   tyro
-tyro==0.9.13
+tyro==0.9.16
     # via saev
 tzdata==2025.1
     # via pandas
-uri-template==1.3.0
-    # via jsonschema
 urllib3==2.3.0
-    # via
-    #   requests
-    #   sentry-sdk
+    # via requests
 uvicorn==0.34.0
     # via
     #   gradio
     #   marimo
-vl-convert-python==1.7.0
-    # via saev
 wadler-lindig==0.1.3
     # via jaxtyping
-wandb==0.19.5
-    # via saev
 wcwidth==0.2.13
-    # via
-    #   ftfy
-    #   prompt-toolkit
-webcolors==24.11.1
-    # via jsonschema
-webdataset==0.2.100
-    # via saev
-webencodings==0.5.1
-    # via
-    #   bleach
-    #   tinycss2
-websocket-client==1.8.0
-    # via jupyter-server
+    # via ftfy
 websockets==14.2
     # via
     #   gradio-client