Include original predictions

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

app.py

@@ -125,7 +125,6 @@ def add_highlights(
     upper: int | None = None,
     opacity: float = 0.9,
 ) -> Image.Image:
-    breakpoint()
     if not len(patches):
         return img
 
@@ -198,18 +197,22 @@ class SaeActivation(typing.TypedDict):
 
 
 @beartype.beartype
-def get_image(i: int) -> tuple[str, str, int]:
-    img_sized = data.to_sized(data.get_image(i))
+def get_img(i: int) -> dict[str, object]:
+    img_sized = data.to_sized(data.get_img(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), i
+    return {
+        "index": i,
+        "orig_url": data.img_to_base64(img_sized),
+        "seg_url": data.img_to_base64(seg_img_sized),
+    }
 
 
 @beartype.beartype
 @torch.inference_mode
-def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]:
+def get_sae_latents(img_i: int, patches: list[int]) -> list[SaeActivation]:
     """
     Given a particular cell, returns some highlighted images showing what feature fires most on this cell.
     """
@@ -219,7 +222,7 @@ def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]
     split_vit, vit_transform = modeling.load_vit(DEVICE)
     sae = load_sae(DEVICE)
 
-    img = data.get_image(image_i)
+    img = data.get_img(img_i)
 
     x_BCWH = vit_transform(img)[None, ...].to(DEVICE)
 
@@ -261,7 +264,7 @@ def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]
         examples = []
         for i_im, values_p in pairs:
             seg_sized = data.to_sized(data.get_seg(i_im))
-            img_sized = data.to_sized(data.get_image(i_im))
+            img_sized = data.to_sized(data.get_img(i_im))
 
             seg_u8_sized = data.to_u8(seg_sized)
             seg_img_sized = data.u8_to_img(seg_u8_sized)
@@ -286,26 +289,27 @@ def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]
 
 
 @torch.inference_mode
-def get_true_labels(image_i: int) -> Image.Image:
-    seg = human_dataset[image_i]["segmentation"]
-    image = seg_to_img(seg)
-    return image
+def get_preds(i: int) -> dict[str, object]:
+    img = data.get_img(i)
+    split_vit, vit_transform = modeling.load_vit(DEVICE)
 
+    x_BCWH = vit_transform(img)[None, ...].to(DEVICE)
 
-@torch.inference_mode
-def get_pred_labels(i: int) -> list[Image.Image | list[int]]:
-    sample = vit_dataset[i]
-    x = sample["image"][None, ...].to(device)
-    x_BPD = rest_of_vit.forward_start(x)
-    x_BPD = rest_of_vit.forward_end(x_BPD)
+    x_BPD = split_vit.forward_start(x_BCWH)
+    x_BPD = split_vit.forward_end(x_BPD)
 
     x_WHD = einops.rearrange(x_BPD, "() (w h) dim -> w h dim", w=16, h=16)
 
-    logits_WHC = head(x_WHD)
+    clf = load_clf()
+    logits_WHC = clf(x_WHD)
 
     pred_WH = logits_WHC.argmax(axis=-1)
-    preds = einops.rearrange(pred_WH, "w h -> (w h)").tolist()
-    return [seg_to_img(upsample(pred_WH)), preds]
+    # preds = einops.rearrange(pred_WH, "w h -> (w h)").tolist()
+    return {
+        "index": i,
+        "orig_url": data.img_to_base64(data.to_sized(img)),
+        "seg_url": data.img_to_base64(data.u8_to_img(upsample(pred_WH))),
+    }
 
 
 @beartype.beartype
@@ -393,64 +397,55 @@ def upsample(
 
 
 with gr.Blocks() as demo:
-    image_number = gr.Number(label="Validation Example")
+    ###########
+    # get-img #
+    ###########
+
+    # Inputs
+    img_number = gr.Number(label="Example Index")
 
-    input_image_base64 = gr.Text(label="Image in Base64")
-    true_labels_base64 = gr.Text(label="Labels in Base64")
+    # Outputs
+    get_img_out = gr.JSON(label="get_img_out", value={})
 
-    get_input_image_btn = gr.Button(value="Get Input Image")
-    get_input_image_btn.click(
-        get_image,
-        inputs=[image_number],
-        outputs=[input_image_base64, true_labels_base64, image_number],
-        api_name="get-image",
+    get_input_img_btn = gr.Button(value="Get Input Image")
+    get_input_img_btn.click(
+        get_img, inputs=[img_number], outputs=[get_img_out], api_name="get-img"
     )
 
-    # input_image = gr.Image(
-    #     label="Input Image",
-    #     sources=["upload", "clipboard"],
-    #     type="pil",
-    #     interactive=True,
-    # )
-    # patch_numbers = gr.CheckboxGroup(label="Image Patch", choices=list(range(256)))
-    # top_latent_numbers = gr.CheckboxGroup(label="Top Latents")
-    # top_latent_numbers = [
-    #     gr.Number(label="Top Latents #{j+1}") for j in range(n_sae_latents)
-    # ]
-    # sae_example_images = [
-    #     gr.Image(label=f"Latent #{j}, Example #{i + 1}", format="png")
-    #     for i in range(n_sae_examples)
-    #     for j in range(n_sae_latents)
-    # ]
+    ###################
+    # get-sae-latents #
+    ###################
 
+    # Inputs
     patches_json = gr.JSON(label="Patches", value=[])
-    activations_json = gr.JSON(label="Activations", value=[])
-
-    get_sae_activations_btn = gr.Button(value="Get SAE Activations")
-    get_sae_activations_btn.click(
-        get_sae_activations,
-        inputs=[image_number, patches_json],
-        outputs=[activations_json],
-        api_name="get-sae-examples",
+    # Outputs
+    get_sae_latents_out = gr.JSON(label="get_sae_latents_out", value=[])
+
+    get_sae_latents_btn = gr.Button(value="Get SAE Latents")
+    get_sae_latents_btn.click(
+        get_sae_latents,
+        inputs=[img_number, patches_json],
+        outputs=[get_sae_latents_out],
+        api_name="get-sae-latents",
     )
-    # semseg_image = gr.Image(label="Semantic Segmentaions", format="png")
-    # semseg_colors = gr.CheckboxGroup(
-    #     label="Sem Seg Colors", choices=list(range(1, 151))
-    # )
 
-    # get_pred_labels_btn = gr.Button(value="Get Pred. Labels")
-    # get_pred_labels_btn.click(
-    #     get_pred_labels,
-    #     inputs=[image_number],
-    #     outputs=[semseg_image, semseg_colors],
-    #     api_name="get-pred-labels",
-    # )
+    #############
+    # get-preds #
+    #############
+
+    # Outputs
+    get_preds_out = gr.JSON(label="get_preds_out", value=[])
+
+    get_pred_labels_btn = gr.Button(value="Get Predictions")
+    get_pred_labels_btn.click(
+        get_preds, inputs=[img_number], outputs=[get_preds_out], api_name="get-preds"
+    )
 
     # get_true_labels_btn = gr.Button(value="Get True Label")
     # get_true_labels_btn.click(
     #     get_true_labels,
-    #     inputs=[image_number],
-    #     outputs=semseg_image,
+    #     inputs=[img_number],
+    #     outputs=semseg_img,
     #     api_name="get-true-labels",
     # )
 
@@ -462,8 +457,8 @@ with gr.Blocks() as demo:
     # get_modified_labels_btn = gr.Button(value="Get Modified Label")
     # get_modified_labels_btn.click(
     #     get_modified_labels,
-    #     inputs=[image_number] + latent_numbers + value_sliders,
-    #     outputs=[semseg_image, semseg_colors],
+    #     inputs=[img_number] + latent_numbers + value_sliders,
+    #     outputs=[semseg_img, semseg_colors],
     #     api_name="get-modified-labels",
     # )
 

data.py

@@ -20,7 +20,7 @@ R2_URL = "https://pub-129e98faed1048af94c4d4119ea47be7.r2.dev"
 
 @beartype.beartype
 @functools.lru_cache(maxsize=512)
-def get_image(i: int) -> Image.Image:
+def get_img(i: int) -> Image.Image:
     fpath = f"/images/ADE_val_{i + 1:08}.jpg"
     url = R2_URL + fpath
     logger.info("Getting image from '%s'.", url)

modeling.py

@@ -21,7 +21,7 @@ class SplitDinov2(torch.nn.Module):
 
     def forward_start(
         self, x: Float[Tensor, "batch channels width height"]
-    ) -> Float[Tensor, "batch patches dim"]:
+    ) -> Float[Tensor, "batch total_patches dim"]:
         x_BPD = self.vit.prepare_tokens_with_masks(x)
         for blk in self.vit.blocks[: self.split_at]:
             x_BPD = blk(x_BPD)
@@ -29,7 +29,7 @@ class SplitDinov2(torch.nn.Module):
         return x_BPD
 
     def forward_end(
-        self, x_BPD: Float[Tensor, "batch n_patches dim"]
+        self, x_BPD: Float[Tensor, "batch total_patches dim"]
     ) -> Float[Tensor, "batch patches dim"]:
         for blk in self.vit.blocks[-self.split_at :]:
             x_BPD = blk(x_BPD)