Added logic in timm table generation for adding ViT encoders for DPT

Author: ved

misc/generate_table_timm.py

@@ -15,32 +15,62 @@ def has_dilation_support(name):
         return True
     except Exception:
         return False
+    
+def valid_vit_encoder_for_dpt(name):
+    if "vit" not in name:
+        return False
+    encoder = timm.create_model(name)
+    feature_info = encoder.feature_info
+    feature_info_obj = timm.models.FeatureInfo(
+                                                feature_info=feature_info, out_indices=[0,1,2,3]
+                                            )
+    reduction_scales = list(feature_info_obj.reduction())
+
+    if len(set(reduction_scales)) > 1:
+        return False
+    
+    output_stride = reduction_scales[0]
+    if bin(output_stride).count("1") != 1:
+        return False
+    
+    return True
 
 
 def make_table(data):
     names = data.keys()
     max_len1 = max([len(x) for x in names]) + 2
     max_len2 = len("support dilation") + 2
+    max_len3 = len("Supported for DPT") + 2
 
-    l1 = "+" + "-" * max_len1 + "+" + "-" * max_len2 + "+\n"
-    l2 = "+" + "=" * max_len1 + "+" + "=" * max_len2 + "+\n"
+    l1 = "+" + "-" * max_len1 + "+" + "-" * max_len2 + "+" + "-" * max_len3 + "+\n"
+    l2 = "+" + "=" * max_len1 + "+" + "=" * max_len2 + "+" + "-" * max_len3 + "+\n"
     top = (
         "| "
         + "Encoder name".ljust(max_len1 - 2)
         + " | "
         + "Support dilation".center(max_len2 - 2)
+        + " | "
+        + "Supported for DPT".center(max_len3 - 2)
         + " |\n"
     )
 
     table = l1 + top + l2
 
     for k in sorted(data.keys()):
-        support = (
-            "✅".center(max_len2 - 3)
-            if data[k]["has_dilation"]
-            else " ".center(max_len2 - 2)
-        )
-        table += "| " + k.ljust(max_len1 - 2) + " | " + support + " |\n"
+
+        if "has_dilation" in data[k] and data[k]["has_dilation"]:
+            support = ("✅".center(max_len2 - 3))
+
+        else:
+            support = (" ".center(max_len2 - 2))
+
+        if "supported_only_for_dpt" in data[k]:
+            supported_for_dpt = ("✅".center(max_len3 - 3))
+
+        else:
+            supported_for_dpt = (" ".center(max_len3 - 2))
+
+        table += "| " + k.ljust(max_len1 - 2) + " | " + support + " | " + supported_for_dpt + " |\n"
         table += l1
 
     return table
@@ -55,8 +85,15 @@ def make_table(data):
                 check_features_and_reduction(name)
                 has_dilation = has_dilation_support(name)
                 supported_models[name] = dict(has_dilation=has_dilation)
+
             except Exception:
-                continue
+                try:
+                    if valid_vit_encoder_for_dpt(name):
+                        supported_models[name] = dict(supported_only_for_dpt = True)
+                except:
+                    continue
+            
+
 
     table = make_table(supported_models)
     print(table)