Spaces:
Runtime error
Runtime error
File size: 2,983 Bytes
cacf2d0 bafba7e cacf2d0 0170c1c cacf2d0 0170c1c cacf2d0 dfd9dbc cacf2d0 dfd9dbc 252f69c cacf2d0 dfd9dbc cacf2d0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
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
|