cleaning code up

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+.venv/
+.hypothesis/
+.aider*
+.env
+.DS_Store
+.coverage
+saev.egg-info/

README.md

@@ -0,0 +1,26 @@
+---
+title: SAEs for Semantic Segmentation
+emoji: 🐨
+colorFrom: gray
+colorTo: blue
+sdk: gradio
+sdk_version: 5.9.1
+python_version: 3.12.8
+app_file: app.py
+pinned: false
+license: mit
+short_description: Interpret semantic segmentation models using SAEs.
+---
+
+
+I used [s5cmd](https://github.com/peak/s5cmd) to upload ADE20K to Cloudflare R2.
+
+```sh
+# in images/
+s5cmd --credentials-file ~/.local/etc/cloudflare/r2-credentials --endpoint-url https://6391ae4399fb354a41cab96372935a6e.r2.cloudflarestorage.com \
+  cp validation/ s3://saev-ade20k/images/
+
+# in annotations/
+s5cmd --credentials-file ~/.local/etc/cloudflare/r2-credentials --endpoint-url https://6391ae4399fb354a41cab96372935a6e.r2.cloudflarestorage.com \
+  cp validation/ s3://saev-ade20k/annotations/```
+```

app.py

@@ -1,23 +1,28 @@
+import functools
+import io
+import json
+import logging
 import os.path
+import pathlib
 import typing
-import functools
 
 import beartype
 import einops
 import einops.layers.torch
 import gradio as gr
+import saev.activations
+import saev.config
+import saev.nn
+import saev.visuals
 import torch
 from jaxtyping import Float, Int, UInt8, jaxtyped
 from PIL import Image
 from torch import Tensor
 
-import saev.activations
-import saev.config
-import saev.nn
-import saev.visuals
+import constants
+import data
 
-from .. import training
-from . import data
+logger = logging.getLogger("app.py")
 
 ####################
 # Global Constants #
@@ -30,9 +35,6 @@ DEBUG = False
 max_frequency = 1e-2
 """Maximum frequency. Any feature that fires more than this is ignored."""
 
-ckpt = "oebd6e6i"
-"""Which SAE checkpoint to use."""
-
 n_sae_latents = 3
 """Number of SAE latents to show."""
 
@@ -51,14 +53,8 @@ CROP_SIZE = (448, 448)
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 """Hardware accelerator, if any."""
 
-####################
-# Helper Functions #
-####################
-
-
-@beartype.beartype
-def load_tensor(path: str) -> Tensor:
-    return torch.load(path, weights_only=True, map_location="cpu")
+CWD = pathlib.Path(".")
+"""Current working directory."""
 
 
 ##########
@@ -67,43 +63,49 @@ def load_tensor(path: str) -> Tensor:
 
 
 @functools.cache
-def load_vit(
-    model_cfg: modeling.Config,
-) -> tuple[
-    activations.WrappedVisionTransformer,
-    typing.Callable,
-    float,
-    Float[Tensor, " d_vit"],
-]:
+def load_vit() -> tuple[saev.activations.WrappedVisionTransformer, typing.Callable]:
     vit = (
-        saev.activations.WrappedVisionTransformer(model_cfg.wrapped_cfg)
+        saev.activations.WrappedVisionTransformer(
+            saev.config.Activations(
+                model_family="dinov2",
+                model_ckpt="dinov2_vitb14_reg",
+                layers=[-2],
+                n_patches_per_img=256,
+            )
+        )
         .to(DEVICE)
         .eval()
     )
-    vit_transform = saev.activations.make_img_transform(
-        model_cfg.vit_family, model_cfg.vit_ckpt
-    )
-    logger.info("Loaded ViT: %s.", model_cfg.key)
+    vit_transform = saev.activations.make_img_transform("dinov2", "dinov2_vitb14_reg")
+    logger.info("Loaded ViT.")
 
-    try:
-        # Normalizing constants
-        acts_dataset = saev.activations.Dataset(model_cfg.acts_cfg)
-        logger.info("Loaded dataset norms: %s.", model_cfg.key)
-    except RuntimeError as err:
-        logger.warning("Error loading ViT: %s", err)
-        return None, None, None, None
+    return vit, vit_transform
 
-    return vit, vit_transform, acts_dataset.scalar.item(), acts_dataset.act_mean
 
+@functools.cache
+def load_sae() -> saev.nn.SparseAutoencoder:
+    """
+    Loads a sparse autoencoder from disk.
+    """
+    sae_ckpt_fpath = CWD / "assets" / "sae.pt"
+    sae = saev.nn.load(str(sae_ckpt_fpath))
+    sae.to(device).eval()
+    return sae
 
-sae_ckpt_fpath = f"/home/stevens.994/projects/saev/checkpoints/{ckpt}/sae.pt"
-sae = saev.nn.load(sae_ckpt_fpath)
-sae.to(device).eval()
 
+@functools.cache
+def load_clf() -> torch.nn.Module:
+    # /home/stevens.994/projects/saev/checkpoints/contrib/semseg/lr_0_001__wd_0_001/model_step8000.pt
+    head_ckpt_fpath = CWD / "assets" / "clf.pt"
+    with open(head_ckpt_fpath, "rb") as fd:
+        kwargs = json.loads(fd.readline().decode())
+        buffer = io.BytesIO(fd.read())
 
-head_ckpt_fpath = "/home/stevens.994/projects/saev/checkpoints/contrib/semseg/lr_0_001__wd_0_001/model_step8000.pt"
-head = training.load(head_ckpt_fpath)
-head = head.to(device).eval()
+    model = torch.nn.Linear(**kwargs)
+    state_dict = torch.load(buffer, weights_only=True, map_location=device)
+    model.load_state_dict(state_dict)
+    model = model.to(device).eval()
+    return model
 
 
 class RestOfDinoV2(torch.nn.Module):
@@ -136,17 +138,18 @@ rest_of_vit = rest_of_vit.to(device)
 ####################
 
 
-ckpt_data_root = (
-    f"/research/nfs_su_809/workspace/stevens.994/saev/features/{ckpt}/sort_by_patch"
-)
+@beartype.beartype
+def load_tensor(path: str | pathlib.Path) -> Tensor:
+    return torch.load(path, weights_only=True, map_location="cpu")
 
-top_img_i = load_tensor(os.path.join(ckpt_data_root, "top_img_i.pt"))
-top_values = load_tensor(os.path.join(ckpt_data_root, "top_values.pt"))
-sparsity = load_tensor(os.path.join(ckpt_data_root, "sparsity.pt"))
 
+top_img_i = load_tensor(CWD / "assets" / "top_img_i.pt")
+top_values = load_tensor(CWD / "assets" / "top_values_uint8.pt")
+sparsity = load_tensor(CWD / "assets" / "sparsity.pt")
 
-mask = torch.ones((sae.cfg.d_sae), dtype=bool)
-mask = mask & (sparsity < max_frequency)
+
+# mask = torch.ones((sae.cfg.d_sae), dtype=bool)
+# mask = mask & (sparsity < max_frequency)
 
 
 ############
@@ -229,9 +232,13 @@ class SaeActivation(typing.TypedDict):
 
 @beartype.beartype
 def get_image(image_i: int) -> tuple[str, str, int]:
-    img_sized, labels_sized = data.get_sample(image_i)
+    sample = data.get_sample(image_i)
+    img_sized = data.to_sized(sample["image"])
+    seg_sized = data.to_sized(sample["segmentation"])
+    seg_u8_sized = data.to_u8(seg_sized)
+    seg_img_sized = data.u8_to_img(seg_u8_sized)
 
-    return data.pil_to_base64(img_sized), data.pil_to_base64(labels_sized), image_i
+    return data.img_to_base64(img_sized), data.img_to_base64(seg_img_sized), image_i
 
 
 @beartype.beartype
@@ -243,27 +250,29 @@ def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]
     if not patches:
         return []
 
-    vit, vit_transform, scalar, mean = load_vit(model_cfg)
-    if vit is None:
-        logger.warning("Skipping ViT '%s'", model_name)
-        return []
-    sae = load_sae(model_cfg)
+    vit, vit_transform = load_vit()
+    sae = load_sae()
 
-    mean = mean.to(DEVICE)
-    x = vit_transform(img_p)[None, ...].to(DEVICE)
+    sample = data.get_sample(image_i)
+
+    x = vit_transform(sample["image"])[None, ...].to(DEVICE)
 
     _, vit_acts_BLPD = vit(x)
-    vit_acts_PD = (vit_acts_BLPD[0, 0, 1:].to(DEVICE).clamp(-1e-5, 1e5) - mean) / scalar
+    vit_acts_PD = (
+        vit_acts_BLPD[0, 0, 1:].to(DEVICE).clamp(-1e-5, 1e5)
+        - (constants.DINOV2_IMAGENET1K_MEAN).to(DEVICE)
+    ) / constants.DINOV2_IMAGENET1K_SCALAR
 
     _, f_x_PS, _ = sae(vit_acts_PD)
     # Ignore [CLS] token and get just the requested latents.
     acts_SP = einops.rearrange(f_x_PS, "patches n_latents -> n_latents patches")
-    logger.info("Got SAE activations for '%s'.", model_name)
-    top_img_i, top_values = load_tensors(model_cfg)
-    logger.info("Loaded top SAE activations for '%s'.", model_name)
+    logger.info("Got SAE activations.")
 
     breakpoint()
 
+    top_img_i, top_values = load_tensors(model_cfg)
+    logger.info("Loaded top SAE activations for '%s'.", model_name)
+
     vit_acts_MD = torch.stack([
         acts_dataset[image_i * acts_dataset.metadata.n_patches_per_img + i]["act"]
         for i in patches

constants.py

@@ -0,0 +1,776 @@
+import torch
+
+
+DINOV2_IMAGENET1K_SCALAR = 2.0181241035461426
+
+
+DINOV2_IMAGENET1K_MEAN = torch.tensor([
+    0.1450997292995453,
+    -1.0630134344100952,
+    -0.3518574833869934,
+    -0.38624095916748047,
+    -0.4866980314254761,
+    -0.28983384370803833,
+    0.9997676014900208,
+    -1.231179118156433,
+    -0.7889889478683472,
+    -0.4450306296348572,
+    -0.09231726080179214,
+    0.13243812322616577,
+    0.09571082890033722,
+    -0.29342857003211975,
+    0.05933428555727005,
+    -0.21923032402992249,
+    0.08959043025970459,
+    -0.6981018781661987,
+    0.4853704869747162,
+    -0.29948222637176514,
+    0.3107207119464874,
+    -0.3812718093395233,
+    -0.5013473033905029,
+    2.88395094871521,
+    -0.5611682534217834,
+    -0.3514024615287781,
+    0.025546086952090263,
+    -0.24438244104385376,
+    -0.23365195095539093,
+    -0.2533780336380005,
+    0.4445696473121643,
+    1.1176759004592896,
+    -0.4188934564590454,
+    0.09051182866096497,
+    -0.04133417829871178,
+    -0.008052834309637547,
+    -0.5118610858917236,
+    0.22084011137485504,
+    -0.7333402633666992,
+    0.8644523620605469,
+    -0.43727627396583557,
+    -0.22333095967769623,
+    -1.5415295362472534,
+    -0.24187016487121582,
+    -0.33239033818244934,
+    -1.2828021049499512,
+    -0.21485395729541779,
+    0.6667488813400269,
+    -0.25890952348709106,
+    -0.8630414009094238,
+    1.5059994459152222,
+    -0.00952776987105608,
+    0.18695995211601257,
+    0.0200128685683012,
+    -0.221832275390625,
+    1.2800148725509644,
+    -0.1416555792093277,
+    0.61446613073349,
+    0.053658585995435715,
+    -0.08877403289079666,
+    1.0190010070800781,
+    -0.308927446603775,
+    -0.3903353214263916,
+    -0.35504740476608276,
+    -0.7907304763793945,
+    -0.18439480662345886,
+    -0.1797204464673996,
+    0.8199827075004578,
+    -0.1736353039741516,
+    -0.16373644769191742,
+    0.7541728019714355,
+    -0.3236996829509735,
+    0.8245170712471008,
+    0.3411649167537689,
+    -0.21873517334461212,
+    -0.7620954513549805,
+    -0.10635858029127121,
+    -0.592278003692627,
+    0.8314691781997681,
+    -0.2021609991788864,
+    -0.24301563203334808,
+    -0.03504444658756256,
+    -0.061244938522577286,
+    -0.36000630259513855,
+    -0.38578882813453674,
+    -1.2314008474349976,
+    -0.3416382968425751,
+    0.5925644636154175,
+    0.32259607315063477,
+    0.13169726729393005,
+    -0.131134033203125,
+    0.05763484537601471,
+    -0.7130515575408936,
+    -0.5685354471206665,
+    0.04428980499505997,
+    0.9245452880859375,
+    0.37724241614341736,
+    -0.4426809549331665,
+    0.5091503262519836,
+    -0.08006338775157928,
+    -0.18945513665676117,
+    -0.770736575126648,
+    -0.3588047921657562,
+    0.04727765917778015,
+    -0.16137081384658813,
+    -0.021555813029408455,
+    0.6381930708885193,
+    0.30161890387535095,
+    -0.0710706040263176,
+    -0.13884945213794708,
+    -0.22726555168628693,
+    -0.6134527921676636,
+    0.2969088852405548,
+    -0.2334780991077423,
+    -0.46334928274154663,
+    -0.3058214485645294,
+    0.5196799039840698,
+    0.6341780424118042,
+    0.12271945178508759,
+    -1.0072089433670044,
+    -0.1198473796248436,
+    -0.24667270481586456,
+    -0.19228138029575348,
+    -0.3955901861190796,
+    -0.19902971386909485,
+    0.7407659292221069,
+    2.3908257484436035,
+    0.02820657566189766,
+    0.07064329087734222,
+    -0.2637694776058197,
+    0.2560977339744568,
+    0.3973558247089386,
+    -0.17345857620239258,
+    -0.9541534185409546,
+    -0.21434728801250458,
+    0.41178393363952637,
+    -0.008175228722393513,
+    0.5115303993225098,
+    -0.9667210578918457,
+    1.6499103307724,
+    -1.8320564031600952,
+    1.1143667697906494,
+    0.24006624519824982,
+    -0.02112947776913643,
+    -0.4952388405799866,
+    1.1000680923461914,
+    -0.4901401102542877,
+    0.22758258879184723,
+    -0.6699370741844177,
+    0.6926363706588745,
+    -0.5719613432884216,
+    0.008403707295656204,
+    2.0220773220062256,
+    -0.1789812445640564,
+    -0.8777256011962891,
+    0.3709064722061157,
+    -0.2629733681678772,
+    0.08407248556613922,
+    -0.27063870429992676,
+    0.09993340820074081,
+    -0.3755860924720764,
+    0.07000888139009476,
+    0.3775370419025421,
+    0.5653945207595825,
+    -0.11404427886009216,
+    -0.06088113784790039,
+    -0.0898045226931572,
+    0.19868576526641846,
+    0.14287644624710083,
+    -0.669394314289093,
+    -0.07882463932037354,
+    -0.12379930168390274,
+    -0.010277876630425453,
+    -0.5625343918800354,
+    -0.6508009433746338,
+    0.06929764896631241,
+    -2.0470166206359863,
+    1.0193544626235962,
+    -0.9747569561004639,
+    -0.25624850392341614,
+    -0.04412469267845154,
+    -0.01941649615764618,
+    0.04781557247042656,
+    -0.2561051845550537,
+    -0.09596704691648483,
+    -1.0529744625091553,
+    -0.32774603366851807,
+    -0.1931363344192505,
+    -0.36885082721710205,
+    -0.9351740479469299,
+    -0.47905397415161133,
+    -0.678762674331665,
+    2.336048126220703,
+    0.26323413848876953,
+    -0.36512619256973267,
+    -0.3650853633880615,
+    -0.8287989497184753,
+    0.5866581201553345,
+    -0.420742005109787,
+    0.008546118624508381,
+    -0.7811568975448608,
+    -0.34993329644203186,
+    -0.373068243265152,
+    0.028424998745322227,
+    -0.537581205368042,
+    -0.15937983989715576,
+    -0.5638740062713623,
+    -0.4413940906524658,
+    -0.05887509509921074,
+    -0.12291032075881958,
+    -0.26565149426460266,
+    -0.23059803247451782,
+    -0.2925986349582672,
+    0.04849022254347801,
+    -0.4770037531852722,
+    0.040383752435445786,
+    -0.8186637759208679,
+    -0.062463242560625076,
+    -0.3251510262489319,
+    -0.4319412112236023,
+    -0.34569647908210754,
+    0.9713658690452576,
+    -0.25668394565582275,
+    -0.37531179189682007,
+    0.5259386301040649,
+    -0.06112021207809448,
+    0.06980857998132706,
+    -0.38363778591156006,
+    -0.1948518007993698,
+    -0.7897586822509766,
+    -0.600932776927948,
+    -0.4269576072692871,
+    -0.32002967596054077,
+    0.08897170424461365,
+    -0.3079395294189453,
+    -0.05779555067420006,
+    -0.782086968421936,
+    1.9608103036880493,
+    0.1145739033818245,
+    0.06164107844233513,
+    -0.3024725317955017,
+    -0.6308553218841553,
+    -0.7640243172645569,
+    -4.433685302734375,
+    -0.31690648198127747,
+    -0.019084235653281212,
+    -0.09761863201856613,
+    -0.029514605179429054,
+    -0.5096182823181152,
+    1.112805962562561,
+    -0.3302820324897766,
+    -0.23730400204658508,
+    0.044646695256233215,
+    -0.805400013923645,
+    -7.766678333282471,
+    -0.2016162872314453,
+    -0.5018128752708435,
+    0.6819560527801514,
+    -0.2735823392868042,
+    -2.2288968563079834,
+    -0.36170846223831177,
+    -0.7745882868766785,
+    0.4644778370857239,
+    0.2525951564311981,
+    -0.22642317414283752,
+    -0.5394997596740723,
+    -0.5064775347709656,
+    -0.5716705918312073,
+    0.19713695347309113,
+    -0.5411649942398071,
+    -0.17092496156692505,
+    0.45778003334999084,
+    0.6894896030426025,
+    -0.21671152114868164,
+    -0.9160588383674622,
+    -0.10307890176773071,
+    0.11703722178936005,
+    -0.7433905601501465,
+    -1.5170584917068481,
+    2.163774013519287,
+    -1.542649507522583,
+    -0.1601075381040573,
+    -0.5249155163764954,
+    0.44509291648864746,
+    -0.5261067152023315,
+    -0.02273540571331978,
+    -0.28311043977737427,
+    0.9144242405891418,
+    0.43954336643218994,
+    -0.2469814419746399,
+    0.18752114474773407,
+    -0.6066163778305054,
+    -0.14480441808700562,
+    -0.3546217679977417,
+    -0.11870954185724258,
+    -0.09891107678413391,
+    -0.377458781003952,
+    0.33304381370544434,
+    -0.156569704413414,
+    -0.9730328321456909,
+    -0.5034677386283875,
+    0.042613230645656586,
+    0.08271210640668869,
+    -0.2368200123310089,
+    -0.07397157698869705,
+    0.011974042281508446,
+    -0.2115129977464676,
+    -0.3752884566783905,
+    -0.24985794723033905,
+    -0.25223013758659363,
+    1.8311675786972046,
+    -0.1650543361902237,
+    -0.031050190329551697,
+    0.10702164471149445,
+    0.8963613510131836,
+    -0.9483885169029236,
+    -0.8156309723854065,
+    -1.7132004499435425,
+    0.08163392543792725,
+    0.4886241555213928,
+    -0.016470594331622124,
+    -0.37671732902526855,
+    -0.025105634704232216,
+    -0.2695018947124481,
+    -0.8450148701667786,
+    -0.9802296757698059,
+    -0.21868866682052612,
+    -0.5872927308082581,
+    1.019242763519287,
+    0.01872517168521881,
+    0.5087792873382568,
+    0.06771136820316315,
+    1.4142885208129883,
+    0.13146139681339264,
+    -0.36489933729171753,
+    0.37572142481803894,
+    -0.3490581810474396,
+    -0.13830198347568512,
+    -1.8019393682479858,
+    1.5129766464233398,
+    0.07059808075428009,
+    1.7206473350524902,
+    0.02890164405107498,
+    0.3628808557987213,
+    0.3914141058921814,
+    0.4993101954460144,
+    0.3969678580760956,
+    -0.058554816991090775,
+    -0.3434300422668457,
+    -0.4157616198062897,
+    -0.7624511122703552,
+    -0.3997197449207306,
+    1.4573990106582642,
+    -0.3363801836967468,
+    -0.46490129828453064,
+    -0.7445303797721863,
+    -0.3460237979888916,
+    -0.6315308809280396,
+    0.8536337018013,
+    -0.08939796686172485,
+    -0.21093742549419403,
+    -0.08742645382881165,
+    -0.020040960982441902,
+    0.09354449808597565,
+    -0.809800386428833,
+    -0.0018062496092170477,
+    -1.0083088874816895,
+    0.3428219258785248,
+    0.012708818539977074,
+    -0.3535612225532532,
+    1.9481208324432373,
+    0.013826621696352959,
+    -0.026771225035190582,
+    0.18734635412693024,
+    0.9365230798721313,
+    1.247671025339514e-05,
+    -0.4420109987258911,
+    0.10769690573215485,
+    -0.6858118176460266,
+    -0.24754805862903595,
+    1.0027467012405396,
+    -0.26436665654182434,
+    -0.33883318305015564,
+    0.38209766149520874,
+    0.479579895734787,
+    -0.5910238027572632,
+    0.1890297830104828,
+    -0.29854580760002136,
+    -0.5636696219444275,
+    -0.504091739654541,
+    -0.32814571261405945,
+    -0.748496949672699,
+    -0.3217906653881073,
+    -0.12439341843128204,
+    -0.3949342668056488,
+    0.09739203751087189,
+    -0.4254276752471924,
+    0.8690429329872131,
+    -0.26380032300949097,
+    -1.2738139629364014,
+    -0.12694764137268066,
+    -0.7331164479255676,
+    0.11337947845458984,
+    -0.7573927640914917,
+    -0.41507089138031006,
+    -0.18960340321063995,
+    1.2390563488006592,
+    -0.10859012603759766,
+    -0.021934548392891884,
+    -0.05041227489709854,
+    -0.055214136838912964,
+    0.20024456083774567,
+    -0.2689618766307831,
+    -0.3135489821434021,
+    -0.07520166784524918,
+    -0.5906742811203003,
+    0.2828388512134552,
+    0.05117213353514671,
+    1.4600849151611328,
+    -0.1967628449201584,
+    0.011182722635567188,
+    0.028878701850771904,
+    -0.12146933376789093,
+    0.6056286096572876,
+    0.22920559346675873,
+    -0.008979334495961666,
+    -0.2874019742012024,
+    -0.4887332320213318,
+    0.8754663467407227,
+    -0.05393843352794647,
+    -0.2956174910068512,
+    -0.18953847885131836,
+    -0.19063766300678253,
+    -0.8141281008720398,
+    0.11052622646093369,
+    -0.020359158515930176,
+    -0.1262499988079071,
+    -1.7762614488601685,
+    -0.4864279627799988,
+    -0.8644945621490479,
+    0.1278448849916458,
+    1.1127605438232422,
+    -0.595068097114563,
+    -0.06630692631006241,
+    1.5608118772506714,
+    -0.9473971724510193,
+    -0.1827543079853058,
+    -0.25564679503440857,
+    -0.4378860294818878,
+    -0.8285927176475525,
+    -1.1397618055343628,
+    -0.06226593255996704,
+    -0.09025824069976807,
+    -0.518083393573761,
+    -0.893482506275177,
+    0.5022943615913391,
+    -0.5922176837921143,
+    0.2571451961994171,
+    0.25571396946907043,
+    0.832092821598053,
+    -0.061823680996894836,
+    -0.08963754773139954,
+    -0.42173218727111816,
+    -0.4375287890434265,
+    -0.43921560049057007,
+    0.5626742243766785,
+    -0.011294233612716198,
+    0.626301646232605,
+    -0.28029197454452515,
+    0.15464802086353302,
+    -0.7071759700775146,
+    -0.0337684191763401,
+    -0.20901329815387726,
+    -0.29788798093795776,
+    0.6644192934036255,
+    -0.049459852278232574,
+    0.039552830159664154,
+    -0.2790898084640503,
+    0.3250356614589691,
+    -0.12668772041797638,
+    -0.46142634749412537,
+    -0.35542988777160645,
+    -1.1817448139190674,
+    0.007615066133439541,
+    -0.43865758180618286,
+    -0.16142761707305908,
+    -0.37852972745895386,
+    -0.582589328289032,
+    0.4371003210544586,
+    -0.2603273391723633,
+    -0.03284638375043869,
+    0.8895729184150696,
+    -0.025997856631875038,
+    0.5761443376541138,
+    -0.28437164425849915,
+    -0.11191761493682861,
+    -0.07794637233018875,
+    0.02127309888601303,
+    -0.10069284588098526,
+    -0.2177346795797348,
+    -1.029278039932251,
+    -0.5014596581459045,
+    -0.5774326920509338,
+    -0.2856050431728363,
+    -0.24715296924114227,
+    0.1243511438369751,
+    0.042631667107343674,
+    -0.846584677696228,
+    -0.7308683395385742,
+    -0.09307371079921722,
+    -0.35250845551490784,
+    0.12801845371723175,
+    -0.5423708558082581,
+    -0.22422067821025848,
+    1.574460744857788,
+    -0.27640238404273987,
+    -0.37266722321510315,
+    -0.12533603608608246,
+    0.3177711069583893,
+    -0.4530303478240967,
+    0.24940718710422516,
+    -0.1272897720336914,
+    0.6882254481315613,
+    -0.2153051793575287,
+    -0.6189695000648499,
+    -0.38704702258110046,
+    -0.14360225200653076,
+    -0.08159925043582916,
+    0.4714410603046417,
+    -0.16035029292106628,
+    0.005880486220121384,
+    -0.5742312669754028,
+    -0.33733850717544556,
+    -0.39702731370925903,
+    -0.14614750444889069,
+    -0.06936132907867432,
+    0.2528288662433624,
+    -0.25900882482528687,
+    0.45907658338546753,
+    -0.20694994926452637,
+    0.4083366394042969,
+    -0.9925484657287598,
+    -0.17098328471183777,
+    0.3215583860874176,
+    -0.33823585510253906,
+    -0.07112737745046616,
+    -0.05322866141796112,
+    0.19237284362316132,
+    -0.6257429122924805,
+    0.23328493535518646,
+    -0.17247024178504944,
+    -0.3362499177455902,
+    -0.17041970789432526,
+    -0.014526017010211945,
+    -0.12138030678033829,
+    0.0698552280664444,
+    -0.609315037727356,
+    0.8142863512039185,
+    -2.295081615447998,
+    -0.07903101295232773,
+    -0.48268306255340576,
+    -0.2097805291414261,
+    -0.4481655955314636,
+    -1.059373378753662,
+    0.17675237357616425,
+    -0.5335419774055481,
+    0.7713444232940674,
+    0.6341530084609985,
+    1.1411781311035156,
+    -0.18365903198719025,
+    -0.4029919505119324,
+    -0.34328755736351013,
+    -1.1935101747512817,
+    -0.4249494671821594,
+    0.10720300674438477,
+    -0.13509584963321686,
+    -0.610278844833374,
+    -0.1007867231965065,
+    -0.13094481825828552,
+    0.3319343030452728,
+    -0.22466504573822021,
+    -0.33384865522384644,
+    -0.3001727759838104,
+    -0.48621413111686707,
+    0.10271137952804565,
+    -0.3953743577003479,
+    -0.3412061631679535,
+    -1.3808176517486572,
+    -0.3035687804222107,
+    0.27737119793891907,
+    -0.10266303271055222,
+    -0.472690224647522,
+    0.03376518189907074,
+    -0.2053908109664917,
+    -0.46477705240249634,
+    -0.0046875146217644215,
+    0.8462978005409241,
+    -0.7554765343666077,
+    -0.9736349582672119,
+    -0.14118513464927673,
+    -0.2665828466415405,
+    -0.9371470212936401,
+    -0.007497116923332214,
+    0.6816821098327637,
+    0.20980679988861084,
+    -0.5602611303329468,
+    -0.7874919176101685,
+    -0.01479698158800602,
+    -0.45345690846443176,
+    -0.12117742747068405,
+    -0.5790822505950928,
+    -0.27737149596214294,
+    0.08818025887012482,
+    -0.25239622592926025,
+    1.1271374225616455,
+    0.0044799973256886005,
+    0.2183203548192978,
+    -2.0634095668792725,
+    -0.007129574194550514,
+    0.32677894830703735,
+    0.019878007471561432,
+    0.060301825404167175,
+    -0.6844122409820557,
+    0.35185739398002625,
+    -0.0028550554998219013,
+    -0.5629953145980835,
+    0.06621643155813217,
+    -0.043473124504089355,
+    -0.3398932218551636,
+    -0.1782192587852478,
+    -0.24575252830982208,
+    -0.20299431681632996,
+    -0.3652290999889374,
+    -0.9888001680374146,
+    -0.30628740787506104,
+    0.6184420585632324,
+    -0.33409008383750916,
+    0.20486755669116974,
+    -0.8251897692680359,
+    -0.08471876382827759,
+    -0.5613390803337097,
+    0.057765014469623566,
+    0.5359746813774109,
+    -0.7063419818878174,
+    0.28122395277023315,
+    -0.004502696450799704,
+    -0.6543170213699341,
+    0.04663177207112312,
+    -0.05775964632630348,
+    -6.37779594399035e-05,
+    0.46121329069137573,
+    -0.004464420489966869,
+    1.4332563877105713,
+    0.20597098767757416,
+    -0.17879879474639893,
+    0.4316228926181793,
+    -1.2352955341339111,
+    -0.19363455474376678,
+    -0.32174810767173767,
+    -0.23037514090538025,
+    0.17044368386268616,
+    0.13070613145828247,
+    1.2171069383621216,
+    -1.171966314315796,
+    0.04596511274576187,
+    -0.1690378040075302,
+    -0.030221890658140182,
+    0.3216114342212677,
+    -0.08577033132314682,
+    -0.26656001806259155,
+    -0.4321160316467285,
+    -0.22010475397109985,
+    -0.6187731623649597,
+    -0.4711909890174866,
+    -0.3499036431312561,
+    0.13558903336524963,
+    -0.2124641239643097,
+    -0.28327351808547974,
+    0.12788993120193481,
+    -1.3083688020706177,
+    -0.0332779586315155,
+    -0.4718656837940216,
+    1.031941533088684,
+    -0.07811620831489563,
+    -0.5331435799598694,
+    -0.2602376341819763,
+    -0.8461449146270752,
+    0.18593788146972656,
+    0.5763140320777893,
+    -0.45714831352233887,
+    -0.1056162416934967,
+    0.2665534019470215,
+    -0.4580163061618805,
+    -0.25224190950393677,
+    -0.2334505170583725,
+    -0.6723064184188843,
+    0.12331533432006836,
+    0.054681699723005295,
+    -0.14116793870925903,
+    -0.10254379361867905,
+    2.0082550048828125,
+    -1.4980225563049316,
+    0.00379346776753664,
+    -0.8470208644866943,
+    0.06866040825843811,
+    -0.3133383095264435,
+    -0.20381635427474976,
+    -0.03295162320137024,
+    1.1624072790145874,
+    -1.2590479850769043,
+    -0.5051106810569763,
+    -0.5310556292533875,
+    0.11350126564502716,
+    -0.5141156315803528,
+    1.0333826541900635,
+    -0.5528491735458374,
+    -0.6508246064186096,
+    -1.0594176054000854,
+    -0.03546600416302681,
+    -0.0008655009442009032,
+    0.06422116607427597,
+    -0.5845358371734619,
+    -0.049052149057388306,
+    -0.578079104423523,
+    -0.46709108352661133,
+    -0.6544204354286194,
+    -0.13105393946170807,
+    -0.12359122931957245,
+    0.19125737249851227,
+    -0.9108084440231323,
+    -0.24640944600105286,
+    -0.5813102126121521,
+    -0.2342103123664856,
+    0.645296573638916,
+    0.4200597405433655,
+    1.030412197113037,
+    0.026015933603048325,
+    0.03929654508829117,
+    -0.18394766747951508,
+    -0.2946997582912445,
+    0.029773380607366562,
+    -1.1292797327041626,
+    -0.3272054195404053,
+    -0.19441728293895721,
+    -0.8372487425804138,
+    0.5765964984893799,
+    -0.28797629475593567,
+    -0.6211466789245605,
+    0.09933445602655411,
+    -0.5617806911468506,
+    1.163861870765686,
+    0.1421220600605011,
+    -0.790323793888092,
+    -0.4003753960132599,
+    -0.6941299438476562,
+    -0.5033494830131531,
+    -0.2234964221715927,
+    -0.12398113310337067,
+    -0.26237404346466064,
+    -0.4991702139377594,
+    -0.7963886260986328,
+    -0.012063371017575264,
+    -1.1415417194366455,
+    0.40668150782585144,
+    0.33048388361930847,
+    1.3195141553878784,
+    -0.0008099540136754513,
+    -0.06793856620788574,
+])

data.py

@@ -0,0 +1,187 @@
+import base64
+import dataclasses
+import functools
+import io
+import logging
+import os.path
+import random
+
+import beartype
+import einops.layers.torch
+import numpy as np
+import torchvision.datasets.folder
+from jaxtyping import UInt8, jaxtyped
+from PIL import Image
+from torch import Tensor
+from torchvision.transforms import v2
+
+logger = logging.getLogger("data.py")
+
+
+@beartype.beartype
+class Ade20k:
+    @beartype.beartype
+    @dataclasses.dataclass(frozen=True)
+    class Sample:
+        img_path: str
+        seg_path: str
+        label: str
+        target: int
+
+    samples: list[Sample]
+
+    def __init__(self, root: str, split: str):
+        self.logger = logging.getLogger("ade20k")
+        self.root = root
+        self.split = split
+        self.img_dir = os.path.join(root, "images")
+        self.seg_dir = os.path.join(root, "annotations")
+
+        # Check that we have the right path.
+        for subdir in ("images", "annotations"):
+            if not os.path.isdir(os.path.join(root, subdir)):
+                # Something is missing.
+                if os.path.realpath(root).endswith(subdir):
+                    self.logger.warning(
+                        "The ADE20K root should contain 'images/' and 'annotations/' directories."
+                    )
+                raise ValueError(f"Can't find path '{os.path.join(root, subdir)}'.")
+
+        _, split_mapping = torchvision.datasets.folder.find_classes(self.img_dir)
+        split_lookup: dict[int, str] = {
+            value: key for key, value in split_mapping.items()
+        }
+        self.loader = torchvision.datasets.folder.default_loader
+
+        err_msg = f"Split '{split}' not in '{set(split_lookup.values())}'."
+        assert split in set(split_lookup.values()), err_msg
+
+        # Load all the image paths.
+        imgs: list[str] = [
+            path
+            for path, s in torchvision.datasets.folder.make_dataset(
+                self.img_dir,
+                split_mapping,
+                extensions=torchvision.datasets.folder.IMG_EXTENSIONS,
+            )
+            if split_lookup[s] == split
+        ]
+
+        segs: list[str] = [
+            path
+            for path, s in torchvision.datasets.folder.make_dataset(
+                self.seg_dir,
+                split_mapping,
+                extensions=torchvision.datasets.folder.IMG_EXTENSIONS,
+            )
+            if split_lookup[s] == split
+        ]
+
+        # Load all the targets, classes and mappings
+        with open(os.path.join(root, "sceneCategories.txt")) as fd:
+            img_labels: list[str] = [line.split()[1] for line in fd.readlines()]
+
+        label_set = sorted(set(img_labels))
+        label_to_idx = {label: i for i, label in enumerate(label_set)}
+
+        self.samples = [
+            self.Sample(img_path, seg_path, label, label_to_idx[label])
+            for img_path, seg_path, label in zip(imgs, segs, img_labels)
+        ]
+
+    def __getitem__(self, index: int) -> dict[str, object]:
+        # Convert to dict.
+        sample = dataclasses.asdict(self.samples[index])
+
+        sample["image"] = self.loader(sample.pop("img_path"))
+        sample["segmentation"] = Image.open(sample.pop("seg_path")).convert("L")
+        sample["index"] = index
+
+        return sample
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+
+@functools.cache
+def get_dataset() -> Ade20k:
+    return Ade20k(
+        root="/research/nfs_su_809/workspace/stevens.994/datasets/ade20k/",
+        split="validation",
+    )
+
+
+@beartype.beartype
+def get_sample(i: int) -> dict[str, object]:
+    dataset = get_dataset()
+    return dataset[i]
+
+
+@jaxtyped(typechecker=beartype.beartype)
+def make_colors() -> UInt8[np.ndarray, "n 3"]:
+    values = (0, 51, 102, 153, 204, 255)
+    colors = []
+    for r in values:
+        for g in values:
+            for b in values:
+                colors.append((r, g, b))
+    # Fixed seed
+    random.Random(42).shuffle(colors)
+    colors = np.array(colors, dtype=np.uint8)
+
+    # Fixed colors for example 3122
+    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)
+    colors[26] = np.array([45, 125, 210], dtype=np.uint8)
+    colors[46] = np.array([238, 185, 2], dtype=np.uint8)
+    colors[52] = np.array([88, 91, 86], dtype=np.uint8)
+    colors[72] = np.array([76, 46, 5], dtype=np.uint8)
+    colors[94] = np.array([12, 15, 10], dtype=np.uint8)
+
+    return colors
+
+
+colors = make_colors()
+
+resize_transform = v2.Compose([
+    v2.Resize((512, 512), interpolation=v2.InterpolationMode.NEAREST),
+    v2.CenterCrop((448, 448)),
+])
+
+
+@beartype.beartype
+def to_sized(img_raw: Image.Image) -> Image.Image:
+    return resize_transform(img_raw)
+
+
+u8_transform = v2.Compose([
+    v2.ToImage(),
+    einops.layers.torch.Rearrange("() width height -> width height"),
+])
+
+
+@beartype.beartype
+def to_u8(seg_raw: Image.Image) -> UInt8[Tensor, "width height"]:
+    return u8_transform(seg_raw)
+
+
+@jaxtyped(typechecker=beartype.beartype)
+def u8_to_img(map: UInt8[Tensor, "width height"]) -> Image.Image:
+    map = map.cpu().numpy()
+    width, height = map.shape
+    colored = np.zeros((width, height, 3), dtype=np.uint8)
+    for i, color in enumerate(colors):
+        colored[map == i + 1, :] = color
+
+    return Image.fromarray(colored)
+
+
+@beartype.beartype
+def img_to_base64(img: Image.Image) -> str:
+    buf = io.BytesIO()
+    img.save(buf, format="webp")
+    b64 = base64.b64encode(buf.getvalue())
+    s64 = b64.decode("utf8")
+    return "data:image/webp;base64," + s64

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "einops>=0.8.0",
     "gradio>=5.3.0",
     "numpy>=2.2.2",
+    "saev",
     "torch>=2.6.0",
     "torchvision>=0.21.0",
 ]