Update predict.py

utils/predict.py

@@ -23,12 +23,15 @@ def get_model():
 
     return model
 
+
 model = get_model()
 labels = ['liver', 'kidney-right', 'kidney-left', 'femur-right', 'femur-left', 'bladder', 'heart', 'lung-right', 'lung-left', 'spleen', 'pancreas']
 
+
 def predict_label(path):
     frames = load_video(path)
-    prediction = model.predict(tf.expand_dims(frames, axis=0))[0]
+    dataloader = prepare_dataloader(frames)
+    prediction = model.predict(dataloader)[0]
     label = np.argmax(prediction, axis=0)
     label = labels[label]
 
@@ -52,3 +55,30 @@ def load_video(path):
     finally:
         cap.release()
     return np.array(frames)
+
+
+def prepare_dataloader(video):
+    video = tf.expand_dims(video, axis=0)
+    dataset = tf.data.Dataset.from_tensor_slices((video, np.array([0])))
+
+    dataloader = (
+        dataset.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
+        .batch(1)
+        .prefetch(tf.data.AUTOTUNE)
+    )
+    return dataloader
+
+
+@tf.function
+def preprocess(frames: tf.Tensor, label: tf.Tensor):
+    """Preprocess the frames tensors and parse the labels."""
+    # Preprocess images
+    frames = tf.image.convert_image_dtype(
+        frames[
+            ..., tf.newaxis
+        ],  # The new axis is to help for further processing with Conv3D layers
+        tf.float32,
+    )
+    # Parse label
+    label = tf.cast(label, tf.float32)
+    return frames, label