new version

app.py

@@ -5,9 +5,9 @@ import numpy as np
 import gradio as gr
 
 tf.keras.optimizers.AdamW = AdamW
-model = from_pretrained_keras("keras-io/vit-small-ds")
+model = from_pretrained_keras("keras-io/vit_small_ds_v2")
+
 
-IMAGE_SIZE = 72
 
 def softmax(x):
     f_x = np.exp(x) / np.sum(np.exp(x))
@@ -16,10 +16,8 @@ def softmax(x):
 labels = ["apple", "aquarium_fish", "baby", "bear", "beaver", "bed", "bee", "beetle", "bicycle", "bottle", "bowl", "boy", "bridge", "bus", "butterfly", "camel", "can", "castle", "caterpillar", "cattle", "chair", "chimpanzee", "clock", "cloud", "cockroach", "couch", "cra", "crocodile", "cup", "dinosaur", "dolphin", "elephant", "flatfish", "forest", "fox", "girl", "hamster", "house", "kangaroo", "keyboard", "lamp", "lawn_mower", "leopard", "lion", "lizard", "lobster", "man", "maple_tree", "motorcycle", "mountain", "mouse", "mushroom", "oak_tree", "orange", "orchid", "otter", "palm_tree", "pear", "pickup_truck", "pine_tree", "plain", "plate", "poppy", "porcupine", "possum", "rabbit", "raccoon", "ray", "road", "rocket", "rose", "sea", "seal", "shark", "shrew", "skunk", "skyscraper", "snail", "snake", "spider", "squirrel", "streetcar", "sunflower", "sweet_pepper", "table", "tank", "telephone", "television", "tiger", "tractor", "train", "trout", "tulip", "turtle", "wardrobe", "whale", "willow_tree", "wolf", "woman", "worm"]
 
 def classify_image(image):
-  #inp = inp.reshape((-1, 299, 299, 3))
-  resized_image = tf.image.resize(
-    tf.convert_to_tensor([image]), size=(IMAGE_SIZE, IMAGE_SIZE))
-  pred = model.predict(resized_image)
+  image = image.reshape((1, 32, 32, 3))
+  pred = model.predict(image)
   prediction = softmax(pred)[0]
   return {labels[i]: float(prediction[i]) for i in range(100)}