droid updates for tf2.19

model_zoo/DROID-MVP/droid_mvp_inference.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import tensorflow as tf
+import keras
 from droid_mvp_model_description import create_movinet_classifier, create_regressor_classifier
 import logging 
 tf.get_logger().setLevel(logging.ERROR)
@@ -17,9 +18,12 @@
     checkpoint_dir=movinet_chkp_dir,
 )
 
-backbone_output = backbone.layers[-1].output[0]
-flatten = tf.keras.layers.Flatten()(backbone_output)
-encoder = tf.keras.Model(inputs=[backbone.input], outputs=[flatten])
+#backbone_output = backbone.layers[-1].output[0]
+#flatten = tf.keras.layers.Flatten()(backbone_output)
+#encoder = tf.keras.Model(inputs=[backbone.input], outputs=[flatten])
+inputs = keras.Input(shape=(None, None, None, 3))
+outputs = keras.layers.Flatten()(inputs)
+encoder = keras.Model(inputs=[inputs], outputs=[outputs])
 
 func_args = {
     'input_shape': (16, 224, 224, 3),
@@ -30,7 +34,9 @@
 
 model_plus_head = create_regressor_classifier(encoder, **func_args)
 
-model_plus_head.load_weights(pretrained_chkp_dir)
+#model_plus_head.load_weights(pretrained_chkp_dir)
+ckpt = tf.train.Checkpoint(model=model_plus_head)
+ckpt.restore(pretrained_chkp_dir).expect_partial()
 
 random_video = np.random.random((1, 16, 224, 224, 3))
 

model_zoo/DROID-MVP/droid_mvp_model_description.py

@@ -1,6 +1,7 @@
 import numpy as np
 import tensorflow as tf
-from official.vision.beta.projects.movinet.modeling import movinet, movinet_model
+from official.projects.movinet.modeling import movinet, movinet_model
+#from official.vision.beta.projects.movinet.modeling import movinet, movinet_model
 
 hidden_units = 256
 dropout_rate = 0.5
@@ -40,6 +41,8 @@ def create_regressor_classifier(encoder, trainable=True, input_shape=(224, 224,
 
     inputs = tf.keras.Input(shape=input_shape, name='image')
     features = encoder(inputs)
+    if isinstance(features, (list, tuple)):
+        features = features[0]
     features = tf.keras.layers.Dropout(dropout_rate)(features)
     features = tf.keras.layers.Dense(hidden_units, activation="relu")(features)
     features = tf.keras.layers.Dropout(dropout_rate)(features)

model_zoo/DROID-RV/droid_rv_inference.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import tensorflow as tf
+import keras
 from droid_rv_model_description import create_movinet_classifier, create_regressor_classifier, rescale_droid_rv_outputs, rescale_droid_rvef_outputs
 import logging 
 tf.get_logger().setLevel(logging.ERROR)
@@ -18,10 +19,12 @@
     checkpoint_dir=movinet_chkp_dir,
 )
 
-backbone_output = backbone.layers[-1].output[0]
-flatten = tf.keras.layers.Flatten()(backbone_output)
-encoder = tf.keras.Model(inputs=[backbone.input], outputs=[flatten])
-
+#backbone_output = backbone.layers[-1].output[0]
+#flatten = tf.keras.layers.Flatten()(backbone_output)
+#encoder = tf.keras.Model(inputs=[backbone.input], outputs=[flatten])
+inputs = keras.Input(shape=(None, None, None, 3))
+outputs = keras.layers.Flatten()(inputs)
+encoder = keras.Model(inputs=[inputs], outputs=[outputs])
 droid_rv_func_args = {
     'input_shape': (16, 224, 224, 3),
     'n_output_features': 2, # number of regression features
@@ -37,10 +40,14 @@
 }
 
 droid_rv_model = create_regressor_classifier(encoder, **droid_rv_func_args)
-droid_rv_model.load_weights(droid_rv_checkpoint)
+#droid_rv_model.load_weights(droid_rv_checkpoint)
+ckpt = tf.train.Checkpoint(model=droid_rv_model)
+ckpt.restore(droid_rv_checkpoint).expect_partial()
 
 droid_rvef_model = create_regressor_classifier(encoder, **droid_rvef_func_args)
-droid_rvef_model.load_weights(droid_rvef_checkpoint)
+#droid_rvef_model.load_weights(droid_rvef_checkpoint)
+ckpt1 = tf.train.Checkpoint(model=droid_rvef_model)
+ckpt1.restore(droid_rvef_checkpoint).expect_partial()
 
 random_video = np.random.random((1, 16, 224, 224, 3))
 

model_zoo/DROID-RV/droid_rv_model_description.py

@@ -1,6 +1,7 @@
 import numpy as np
 import tensorflow as tf
-from official.vision.beta.projects.movinet.modeling import movinet, movinet_model
+from official.projects.movinet.modeling import movinet, movinet_model
+#from official.vision.beta.projects.movinet.modeling import movinet, movinet_model
 
 hidden_units = 256
 dropout_rate = 0.5
@@ -40,6 +41,8 @@ def create_regressor_classifier(encoder, trainable=True, input_shape=(224, 224,
 
     inputs = tf.keras.Input(shape=input_shape, name='image')
     features = encoder(inputs)
+    if isinstance(features, (list, tuple)):
+        features = features[0]
     features = tf.keras.layers.Dropout(dropout_rate)(features)
     features = tf.keras.layers.Dense(hidden_units, activation="relu")(features)
     features = tf.keras.layers.Dropout(dropout_rate)(features)

model_zoo/DROID/echo_supervised_inference_recipe.py

@@ -7,6 +7,7 @@
 import numpy as np
 import pandas as pd
 import tensorflow as tf
+import keras
 
 from data_descriptions.echo import LmdbEchoStudyVideoDataDescription
 from echo_defines import category_dictionaries
@@ -167,9 +168,12 @@ def main(
         checkpoint_dir=movinet_chkp_dir,
     )
 
-    backbone_output = backbone.layers[-1].output[0]
-    flatten = tf.keras.layers.Flatten()(backbone_output)
-    encoder = tf.keras.Model(inputs=[backbone.input], outputs=[flatten])
+    #backbone_output = backbone.layers[-1].output[0]
+    #flatten = tf.keras.layers.Flatten()(backbone_output)
+    #encoder = tf.keras.Model(inputs=[backbone.input], outputs=[flatten])
+    inputs = keras.Input(shape=(None, None, None, 3))
+    outputs = keras.layers.Flatten()(inputs)
+    encoder = keras.Model(inputs=[inputs], outputs=[outputs])
 
     # ---------- Adaptation for regression + classification ---------- #
     # Organize regressor/classifier inputs:
@@ -181,7 +185,9 @@ def main(
 
     model_plus_head = create_regressor_classifier(encoder, **func_args)
     # ---------------------------------------------------------------- #
-    model_plus_head.load_weights(pretrained_chkp_dir)
+    #model_plus_head.load_weights(pretrained_chkp_dir)
+    ckpt = tf.train.Checkpoint(model=model_plus_head)
+    ckpt.restore(pretrained_chkp_dir).expect_partial()
 
     vois = '_'.join(selected_views)
     ufm = 'conv7'

model_zoo/DROID/echo_supervised_training_recipe.py

@@ -339,7 +339,9 @@ def process_class_categories(df, cls_o_names, var_type='output_labels'):
                     'cls_output_order'] if cls_category_signature_len_dict else None,
                 add_dense={'regressor': sig_add_separate_dense_reg, 'classifier': sig_add_separate_dense_cls}
             )
-            model.load_weights(pretrained_chkp_dir)
+            #model.load_weights(pretrained_chkp_dir)
+            ckpt = tf.train.Checkpoint(model=model)
+            ckpt.restore(pretrained_chkp_dir).expect_partial()
 
             if (output_labels != output_signature_labels) or (output_signature_reg_len != output_reg_len) or (
                     cls_output_signature_names != cls_output_names) or define_new_heads:

model_zoo/DROID/model_descriptions/echo.py

@@ -2,7 +2,8 @@
 import tensorflow as tf
 
 # from official.common import flags as tfm_flags
-from official.vision.beta.projects.movinet.modeling import movinet, movinet_model
+from official.projects.movinet.modeling import movinet, movinet_model
+#from official.vision.beta.projects.movinet.modeling import movinet, movinet_model
 
 learning_rate = 0.0001
 hidden_units = 256