Update app.py

app.py

@@ -1,3 +1,6 @@
+import tarfile
+import wandb
+
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -42,7 +45,16 @@ def preprocess_input(image):
         "pixel_values": tf.expand_dims(image, 0)
     }
 
-def get_predictions(image):
+def get_predictions(wb_token, image):
+    if MODEL is None:
+        wandb.login(key=wb_token)
+        path = wandb.use_artifact('tfx-vit-pipeline/final_model:latest', type='model').download()
+
+        tar = tarfile.open(f"{path}/model.tar.gz")
+        tar.extractall(path=".")
+
+        MODEL = tf.keras.models.load_model("./model")
+    
     preprocessed_image = preprocess_input(image)
     prediction = MODEL.predict(preprocessed_image)
     probs = tf.nn.softmax(prediction['logits'], axis=1)
@@ -57,10 +69,16 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         image_if = gr.Image()
-        label_if = gr.Label()
+        label_if = gr.Label(num_top_classes=3)
 
     classify_if = gr.Button()
 
+    classify_if.click(
+        get_predictions,
+        [wb_token_if, image_if],
+        label_if
+    )
+
 # demo = gr.Interface(
 #     get_predictions,
 #     gr.inputs.Image(),