Spaces:
Runtime error
Runtime error
File size: 2,152 Bytes
4084d4a dbee91e 74bc43a dbee91e 4084d4a d2c0b63 dbee91e 31104c0 dbee91e 4084d4a 95ce0a4 4084d4a 987e698 4084d4a |
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 89 90 91 92 93 |
import os
import sys
current = os.path.dirname(os.path.realpath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)
import albumentations as A
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import torch
from albumentations.pytorch import ToTensorV2
from model import Classifier
from PIL import Image
# Load the model
model = Classifier.load_from_checkpoint("./models/checkpoint.ckpt")
model.eval()
# Define labels
labels = [
"dog",
"horse",
"elephant",
"butterfly",
"chicken",
"cat",
"cow",
"sheep",
"spider",
"squirrel",
]
# Preprocess function
def preprocess(image):
image = np.array(image)
resize = A.Resize(224, 224)
normalize = A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
to_tensor = ToTensorV2()
transform = A.Compose([resize, normalize, to_tensor])
image = transform(image=image)["image"]
return image
# Define the function to make predictions on an image
def predict(image):
try:
image = preprocess(image).unsqueeze(0)
# Prediction
# Make a prediction on the image
with torch.no_grad():
output = model(image)
# convert to probabilities
probabilities = torch.nn.functional.softmax(output[0])
# get top probabilities
topk_prob, topk_label = torch.topk(probabilities, 3)
# Return the top 3 predictions
return {
labels[label]: float(prob) for label, prob in zip(topk_label, topk_prob)
}
except Exception as e:
print(f"Error predicting image: {e}")
return []
# Define the interface
def app():
title = "Animal-10 Image Classification"
gr.Interface(
title=title,
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(
num_top_classes=3,
),
examples=[
"./test_images/dog.jpeg",
"./test_images/cat.jpeg",
"./test_images/butterfly.jpeg",
"./test_images/horse.jpeg",
],
).launch()
# Run the app
if __name__ == "__main__":
app()
|