Fix kubric data loading of video shape (None, 256, 256, None) (#327)

google-research · Jun 6, 2024 · 082a020 · 082a020
1 parent 7fa640a
commit 082a020
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/challenges/point_tracking/dataset.py b/challenges/point_tracking/dataset.py
@@ -929,6 +929,7 @@ def add_tracks(data,
   )
   video = tf.slice(video, start, size)
   video = tf.image.resize(tf.cast(video, tf.float32), train_size)
+  video.set_shape([num_frames, train_size[0], train_size[1], 3])
   if vflip:
     video = video[:, ::-1, :, :]
     target_points = target_points * np.array([1, -1])