update data and model

app.py

@@ -533,6 +533,9 @@ def make_dataset_images_section(open=False):
             labels = np.array(dataset['label'])
             unique_labels = np.unique(labels)
             valid_classes = [i for i in classes if i in unique_labels]
+            invalid_classes = [i for i in classes if i not in unique_labels]
+            if len(invalid_classes) > 0:
+                gr.Warning(f"Classes {invalid_classes} not found in the dataset.")
             if len(valid_classes) == 0:
                 gr.Error(f"Classes {classes} not found in the dataset.")
                 return None
@@ -580,11 +583,13 @@ def make_parameters_section():
         layer_dict = LAYER_DICT
         if model_name in layer_dict:
             value = layer_dict[model_name]
-            return gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True)
+            return (gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True),
+                    gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?"))
         else:
             value = 12
-            return gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True)
-    model_dropdown.change(fn=change_layer_slider, inputs=model_dropdown, outputs=layer_slider)
+            return (gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?"),
+                    gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True))
+    model_dropdown.change(fn=change_layer_slider, inputs=model_dropdown, outputs=[layer_slider, node_type_dropdown])
     
     with gr.Accordion("➡️ Click to expand: more parameters", open=False):
         affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")

backbone.py

@@ -1,17 +1,14 @@
+# Author: Huzheng Yang
+# %%
 from typing import Optional, Tuple
 from einops import rearrange
 import requests
 import torch
 import torch.nn.functional as F
 import timm
-from PIL import Image
 from torch import nn
 import numpy as np
 import os
-import time
-
-import gradio as gr
-
 from functools import partial
 
 MODEL_DICT = {}
@@ -613,7 +610,7 @@ class EVA02(nn.Module):
         super().__init__(**kwargs)
         
         model = timm.create_model(
-            'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k',
+            'eva02_base_patch14_448.mim_in22k_ft_in1k',
             pretrained=True,
             num_classes=0,  # remove classifier nn.Linear
         )
@@ -660,9 +657,9 @@ class EVA02(nn.Module):
             'block': block_outputs
         }
         
-MODEL_DICT["EVA-CLIP(eva02_large_patch14_448)"] = partial(EVA02)
-LAYER_DICT["EVA-CLIP(eva02_large_patch14_448)"] = 24
-RES_DICT["EVA-CLIP(eva02_large_patch14_448)"] = (448, 448)
+MODEL_DICT["CLIP(eva02_base_patch14_448.mim_in22k_ft_in1k)"] = partial(EVA02)
+LAYER_DICT["CLIP(eva02_base_patch14_448.mim_in22k_ft_in1k)"] = 12
+RES_DICT["CLIP(eva02_base_patch14_448.mim_in22k_ft_in1k)"] = (448, 448)
 
 class CLIPConvnext(nn.Module):
     def __init__(self):
@@ -862,6 +859,8 @@ def extract_features(images, model, model_name, node_type, layer, batch_size=8):
             inp = inp.cuda()
         out = model(inp)  # {'attn': [B, H, W, C], 'mlp': [B, H, W, C], 'block': [B, H, W, C]}
         out = out[node_type]
+        if out is None:
+            raise ValueError(f"Node type {node_type} not found in model {model_name}")
         out = out[layer]
         # normalize
         out = F.normalize(out, dim=-1)