ColdHearted's picture
Update app.py
f3db019 verified
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import gdown
import os
# Load model
def load_model():
model_path = "resnet50_cifar10_model.h5"
if not os.path.exists(model_path):
url = "https://drive.google.com/uc?id=13KgM2DddlsscFQx4uoYK0lesSE6-DAo3"
gdown.download(url, model_path, quiet=False)
model = tf.keras.models.load_model(model_path)
return model
model = load_model()
class_names = ['Airplane', 'Automobile', 'Bird', 'Cat', 'Deer',
'Dog', 'Frog', 'Horse', 'Ship', 'Truck']
# Prediction function
def predict_cifar10(image):
image = image.convert("RGB")
img = image.resize((32, 32))
img_array = np.array(img) / 255.0
img_array = np.expand_dims(img_array, axis=0)
prediction = model.predict(img_array)
predicted_label = class_names[np.argmax(prediction)]
confidence = float(np.max(prediction)) * 100
return {predicted_label: confidence}
# Gradio Interface
iface = gr.Interface(
fn=predict_cifar10,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=3),
title="πŸš€ CIFAR-10 Image Classifier using ResNet50",
description="Upload an image, and the model will classify it into one of the 10 CIFAR-10 classes."
)
iface.launch()