Upload 2 files

app.py

@@ -7,7 +7,7 @@ from huggingface_hub import from_pretrained_keras
 from predict import predict_label
 
 ##Create list of examples to be loaded
-example_list = glob.glob("examples/*")
+example_list = glob.glob("*.mp4")
 example_list = list(map(lambda el:[el], example_list))
 
 demo = gr.Blocks()
@@ -31,7 +31,7 @@ with demo:
         gr.Markdown("The model is trained to classify videos belonging to the following classes: liver, kidney-right, kidney-left, femur-right, femur-left, bladder, heart, lung-right, lung-left, spleen, pancreas")
 
         with gr.Column():
-            gr.Examples(example_list, [input_video], [output_label], predict_label, cache_examples=True)
+            gr.Examples(example_list, [input_video], [output_label], predict_label, cache_examples=False)
 
     submit_button.click(predict_label, inputs=input_video, outputs=output_label)
 

predict.py

@@ -1,4 +1,5 @@
 import cv2
+# import imageio
 import numpy as np
 import tensorflow as tf
 from huggingface_hub import from_pretrained_keras
@@ -6,12 +7,31 @@ from tensorflow.keras.optimizers import Adam
 
 from constants import LEARNING_RATE
 
+def get_model():
+    """
+    Download the model from the Hugging Face Hub and compile it.
+    """
+    model = from_pretrained_keras("pablorodriper/video-vision-transformer")
+
+    model.compile(
+        optimizer=Adam(learning_rate=LEARNING_RATE),
+        loss="sparse_categorical_crossentropy",
+        # metrics=[
+        #     keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
+        #     keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
+        # ],
+    )
+
+    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]
     label = np.argmax(prediction, axis=0)
+    label = labels[label]
 
     return label
 
@@ -33,21 +53,3 @@ def load_video(path):
     finally:
         cap.release()
     return np.array(frames)
-
-
-def get_model():
-    """
-    Download the model from the Hugging Face Hub and compile it.
-    """
-    model = from_pretrained_keras("pablorodriper/video-vision-transformer")
-
-    model.compile(
-        optimizer=Adam(learning_rate=LEARNING_RATE),
-        loss="sparse_categorical_crossentropy",
-        # metrics=[
-        #     keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
-        #     keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
-        # ],
-    )
-
-    return model