Spaces:
Runtime error
Runtime error
File size: 2,511 Bytes
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 94 95 96 97 98 99 100 101 102 |
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 PIL import Image
from model import Classifier
# 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 sample images
sample_images = {
"dog": "./test_images/dog.jpeg",
"cat": "./test_images/cat.jpeg",
"butterfly": "./test_images/butterfly.jpeg",
"elephant": "./test_images/elephant.jpg",
"horse": "./test_images/horse.jpeg",
}
# 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(torch.exp(output[0]), dim=0)
topk_prob, topk_label = torch.topk(probabilities, 3)
# Return the top 3 predictions
return {labels[i]: float(probabilities[i]) for i in range(3)}
except Exception as e:
print(f"Error predicting image: {e}")
return []
# Define the interface
def app():
title = "Animal-10 Image Classification"
description = "Classify images using a custom CNN model and deploy using Gradio."
gr.Interface(
title=title,
description=description,
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/elephant.jpg",
"./test_images/horse.jpeg",
],
).launch()
# Run the app
if __name__ == "__main__":
app()
|