Use cloudflare for ade20k images

app.py

@@ -143,9 +143,9 @@ def load_tensor(path: str | pathlib.Path) -> Tensor:
     return torch.load(path, weights_only=True, map_location="cpu")
 
 
-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")
+# 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)
@@ -231,14 +231,13 @@ class SaeActivation(typing.TypedDict):
 
 
 @beartype.beartype
-def get_image(image_i: int) -> tuple[str, str, int]:
-    sample = data.get_sample(image_i)
-    img_sized = data.to_sized(sample["image"])
-    seg_sized = data.to_sized(sample["segmentation"])
+def get_image(i: int) -> tuple[str, str, int]:
+    img_sized = data.to_sized(data.get_image(i))
+    seg_sized = data.to_sized(data.get_seg(i))
     seg_u8_sized = data.to_u8(seg_sized)
     seg_img_sized = data.u8_to_img(seg_u8_sized)
 
-    return data.img_to_base64(img_sized), data.img_to_base64(seg_img_sized), image_i
+    return data.img_to_base64(img_sized), data.img_to_base64(seg_img_sized), i
 
 
 @beartype.beartype
@@ -253,9 +252,9 @@ def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]
     vit, vit_transform = load_vit()
     sae = load_sae()
 
-    sample = data.get_sample(image_i)
+    img = data.get_image(image_i)
 
-    x = vit_transform(sample["image"])[None, ...].to(DEVICE)
+    x = vit_transform(img)[None, ...].to(DEVICE)
 
     _, vit_acts_BLPD = vit(x)
     vit_acts_PD = (
@@ -268,8 +267,6 @@ def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]
     acts_SP = einops.rearrange(f_x_PS, "patches n_latents -> n_latents patches")
     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)
 

constants.py

@@ -1,6 +1,5 @@
 import torch
 
-
 DINOV2_IMAGENET1K_SCALAR = 2.0181241035461426
 
 

data.py

@@ -1,15 +1,13 @@
 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
+import requests
 from jaxtyping import UInt8, jaxtyped
 from PIL import Image
 from torch import Tensor
@@ -17,104 +15,25 @@ from torchvision.transforms import v2
 
 logger = logging.getLogger("data.py")
 
+R2_URL = "https://pub-129e98faed1048af94c4d4119ea47be7.r2.dev"
+
 
 @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",
-    )
+@functools.lru_cache(maxsize=512)
+def get_image(i: int) -> Image.Image:
+    fpath = f"/images/ADE_val_{i + 1:08}.jpg"
+    url = R2_URL + fpath
+    logger.info("Getting image from '%s'.", url)
+    return Image.open(requests.get(url, stream=True).raw)
 
 
 @beartype.beartype
-def get_sample(i: int) -> dict[str, object]:
-    dataset = get_dataset()
-    return dataset[i]
+@functools.lru_cache(maxsize=512)
+def get_seg(i: int) -> Image.Image:
+    fpath = f"/annotations/ADE_val_{i + 1:08}.png"
+    url = R2_URL + fpath
+    logger.info("Getting annotations from '%s'.", url)
+    return Image.open(requests.get(url, stream=True).raw)
 
 
 @jaxtyped(typechecker=beartype.beartype)