made resizing operations compatible with None

convert_to_tf.py

@@ -196,13 +196,13 @@ def main(args):
     dataset_name = args.ckpt_path.split("/")[-2].lower()
     tf_params_path = f"{variant}_{task.lower()}_{dataset_name}.h5"
 
-    tf_model.save_weights(tf_params_path)
-    print(f"Model params serialized to {tf_params_path}.")
+    # tf_model.save_weights(tf_params_path)
+    # print(f"Model params serialized to {tf_params_path}.")
     saved_model_path = tf_params_path.replace(".h5", "")
     tf_model.save(saved_model_path)
     print(f"SavedModel serialized to {saved_model_path}.")
-    push_to_hub_keras(tf_model, repo_path_or_name=f"sayakpaul/{saved_model_path}")
-    print("Model pushed to Hugging Face Hub.")
+    # push_to_hub_keras(tf_model, repo_path_or_name=f"sayakpaul/{saved_model_path}")
+    # print("Model pushed to Hugging Face Hub.")
 
 
 def parse_args():

create_maxim_model.py

@@ -4,7 +4,7 @@
 from maxim.configs import MAXIM_CONFIGS
 
 
-def Model(variant=None, input_resolution=(256, 256), **kw) -> keras.Model:
+def Model(variant=None, input_resolution=(None, None), **kw) -> keras.Model:
     """Factory function to easily create a Model variant like "S".
 
     Args:

maxim/blocks/others.py

@@ -38,17 +38,9 @@ def UpSampleRatio(
     """Upsample features given a ratio > 0."""
 
     def apply(x):
-        n, h, w, c = (
-            K.int_shape(x)[0],
-            K.int_shape(x)[1],
-            K.int_shape(x)[2],
-            K.int_shape(x)[3],
-        )
-
         # Following `jax.image.resize()`
         x = Resizing(
-            height=int(h * ratio),
-            width=int(w * ratio),
+            ratio=ratio,
             method="bilinear",
             antialias=True,
             name=f"{name}_resizing_{K.get_uid('Resizing')}",

maxim/layers.py

@@ -76,28 +76,28 @@ def get_config(self):
 
 
 @tf.keras.utils.register_keras_serializable("maxim")
-class Resizing(layers.Layer):
-    def __init__(self, height, width, antialias=True, method="bilinear", **kwargs):
+class Resizing(tf.keras.layers.Layer):
+    def __init__(self, ratio: int, method="bilinear", antialias=True, **kwargs):
         super().__init__(**kwargs)
-        self.height = height
-        self.width = width
-        self.antialias = antialias
+        self.ratio = ratio
         self.method = method
+        self.antialias = antialias
 
-    def call(self, x):
-        return tf.image.resize(
-            x,
-            size=(self.height, self.width),
-            antialias=self.antialias,
+    def __call__(self, img):
+        n, h, w, c = tf.shape(img)
+        x = tf.image.resize(
+            img,
+            (h // self.ratio, w // self.ratio),
             method=self.method,
+            antialias=self.antialias,
         )
+        return x
 
     def get_config(self):
         config = super().get_config().copy()
         config.update(
             {
-                "height": self.height,
-                "width": self.width,
+                "ratio": self.ratio,
                 "antialias": self.antialias,
                 "method": self.method,
             }

maxim/maxim.py

@@ -103,8 +103,7 @@ def apply(x):
         # Get multi-scale input images
         for i in range(1, num_supervision_scales):
             resizing_layer = Resizing(
-                height=h // (2 ** i),
-                width=w // (2 ** i),
+                ratio=(2 ** i),
                 method="nearest",
                 antialias=True,  # Following `jax.image.resize()`.
                 name=f"initial_resizing_{K.get_uid('Resizing')}",
@@ -222,7 +221,7 @@ def apply(x):
                     [
                         UpSampleRatio(
                             num_channels=(2 ** i) * features,
-                            ratio=2 ** (j - i),
+                            ratio=int(1 / (2 ** (j - i))),
                             use_bias=use_bias,
                             name=f"UpSampleRatio_{K.get_uid('UpSampleRatio')}",
                         )(enc)
@@ -265,7 +264,7 @@ def apply(x):
                     [
                         UpSampleRatio(
                             num_channels=(2 ** i) * features,
-                            ratio=2 ** (depth - j - 1 - i),
+                            ratio=int(1 / (2 ** (depth - j - 1 - i))),
                             use_bias=use_bias,
                             name=f"UpSampleRatio_{K.get_uid('UpSampleRatio')}",
                         )(skip)