fix: constants.

app.py

@@ -9,12 +9,12 @@ from torchvision.models.feature_extraction import create_feature_extractor
 from torchvision.transforms import functional as F
 import glob
 
-cait_model = create_model("cait_xxs24_224.fb_dist_in1k", pretrained=True).eval()
-transform = timm.data.create_transform(
-    **timm.data.resolve_data_config(cait_model.pretrained_cfg)
+CAIT_MODEL = create_model("cait_xxs24_224.fb_dist_in1k", pretrained=True).eval()
+TRANSFORM = timm.data.create_transform(
+    **timm.data.resolve_data_config(CAIT_MODEL.pretrained_cfg)
 )
 
-patch_size = 16
+PATCH_SIZE = 16
 
 
 def create_attn_extractor(block_id=0):
@@ -23,7 +23,7 @@ def create_attn_extractor(block_id=0):
         https://github.com/huggingface/pytorch-image-models/discussions/926
     """
     feature_extractor = create_feature_extractor(
-        cait_model,
+        CAIT_MODEL,
         return_nodes=[f"blocks_token_only.{block_id}.attn.softmax"],
         tracer_kwargs={"leaf_modules": [PatchEmbed]},
     )
@@ -34,8 +34,8 @@ def get_cls_attention_map(
     image, attn_score_dict, block_key="blocks_token_only.0.attn.softmax"
 ):
     """Prepares attention maps so that they can be visualized."""
-    w_featmap = image.shape[3] // patch_size
-    h_featmap = image.shape[2] // patch_size
+    w_featmap = image.shape[3] // PATCH_SIZE
+    h_featmap = image.shape[2] // PATCH_SIZE
 
     attention_scores = attn_score_dict[block_key]
     nh = attention_scores.shape[1]  # Number of attention heads.
@@ -51,7 +51,7 @@ def get_cls_attention_map(
     # Resize the attention patches to 224x224 (224: 14x16)
     attentions = F.resize(
         attentions,
-        size=(h_featmap * patch_size, w_featmap * patch_size),
+        size=(h_featmap * PATCH_SIZE, w_featmap * PATCH_SIZE),
         interpolation=3,
     )
     print(attentions.shape)
@@ -85,8 +85,8 @@ def serialize_images(processed_map):
 def generate_class_attn_map(image, block_id=0):
     """Collates the above utilities together for generating
     a class attention map."""
-    image_tensor = transform(image).unsqueeze(0)
-    feature_extractor = create_attn_extractor(cait_model, block_id)
+    image_tensor = TRANSFORM(image).unsqueeze(0)
+    feature_extractor = create_attn_extractor(block_id)
 
     with torch.no_grad():
         out = feature_extractor(image_tensor)