Spaces:
Runtime error
Runtime error
import numpy as np | |
import tensorflow as tf | |
from tensorflow import keras | |
from huggingface_hub import from_pretrained_keras | |
from .lr_schedule import WarmUpCosine | |
from .constants import Config, class_vocab | |
from keras.utils import load_img, img_to_array | |
from tensorflow_addons.optimizers import AdamW | |
import matplotlib.pyplot as plt | |
import pandas as pd | |
import random | |
config = Config() | |
##Load Model | |
model = from_pretrained_keras("keras-io/shiftvit", custom_objects={"WarmUpCosine":WarmUpCosine, "AdamW": AdamW}) | |
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data() | |
AUTO = tf.data.AUTOTUNE | |
def predict(image_path): | |
""" | |
This function is used for fetching predictions corresponding to input_image. | |
It outputs confidence scores corresponding to each class on which the model was trained | |
""" | |
test_image1 = load_img(image_path,target_size =(32,32)) | |
test_image = img_to_array(test_image1) | |
test_image = np.expand_dims(test_image, axis =0) | |
test_image = test_image.astype('uint8') | |
predict_ds = tf.data.Dataset.from_tensor_slices(test_image) | |
predict_ds = predict_ds.shuffle(config.buffer_size).batch(config.batch_size).prefetch(AUTO) | |
logits = model.predict(predict_ds) | |
prob = tf.nn.softmax(logits) | |
confidences = {} | |
prob_list = prob.numpy().flatten().tolist() | |
sorted_prob = np.argsort(prob)[::-1].flatten() | |
for i in sorted_prob: | |
confidences[class_vocab[i]] = float(prob_list[i]) | |
return confidences | |
def predict_batch(image_path): | |
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)) | |
test_ds = test_ds.batch(config.batch_size).prefetch(AUTO) | |
slice = test_ds.take(1) | |
slice_pred = model.predict(slice) | |
slice_pred = tf.nn.softmax(slice_pred) | |
saved_plot = "plot.jpg" | |
fig = plt.figure() | |
predictions_df = pd.DataFrame() | |
num = random.randint(0,50) | |
for images, labels in slice: | |
for i, j in zip(range(num,num+6), range(6)): | |
ax = plt.subplot(3, 3, j + 1) | |
plt.imshow(images[i].numpy().astype("uint8")) | |
output = np.argmax(slice_pred[i]) | |
prob_list = slice_pred[i].numpy().flatten().tolist() | |
sorted_prob = np.argsort(slice_pred[i])[::-1].flatten() | |
prob_scores = {"image": "image "+ str(j), "first": f"predicted {class_vocab[sorted_prob[0]]} with {round(prob_list[sorted_prob[0]] * 100,2)}% confidence", | |
"second": f"predicted {class_vocab[sorted_prob[1]]} is {round(prob_list[sorted_prob[1]] * 100,2)}% confidence", | |
"third": f"predicted {class_vocab[sorted_prob[2]]} is {round(prob_list[sorted_prob[2]] * 100,2)}% confidence"} | |
predictions_df = predictions_df.append(prob_scores,ignore_index=True) | |
plt.title(f"image {j} : {class_vocab[output]}") | |
plt.axis("off") | |
plt.savefig(saved_plot,bbox_inches='tight') | |
return saved_plot, predictions_df | |