|
|
import gradio as gr |
|
|
import numpy as np |
|
|
from tensorflow.keras.models import load_model |
|
|
from huggingface_hub import hf_hub_download |
|
|
from PIL import Image |
|
|
|
|
|
repo_id = "KevinJuanCarlos23/VehicleClassification" |
|
|
filename = "vehicle_classification_model.keras" |
|
|
|
|
|
model_path = hf_hub_download(repo_id=repo_id, filename=filename) |
|
|
|
|
|
|
|
|
try: |
|
|
model = load_model(model_path) |
|
|
print("Model loaded successfully!") |
|
|
except OSError as e: |
|
|
print("Error loading model:", e) |
|
|
|
|
|
class_labels = ['Auto Rickshaw', 'Bike', 'Car', 'Motorcycle', 'Plane', 'Ship', 'Train'] |
|
|
|
|
|
def predict_vehicle(img): |
|
|
img = img.resize((224, 224)) |
|
|
img = np.array(img) / 255.0 |
|
|
img = np.expand_dims(img, axis=0) |
|
|
|
|
|
predictions = model.predict(img) |
|
|
predicted_class = np.argmax(predictions, axis=1)[0] |
|
|
confidence = np.max(predictions, axis=1)[0] |
|
|
predicted_label = class_labels[predicted_class] |
|
|
|
|
|
return {"vehicle_type": predicted_label, "confidence": float(confidence)} |
|
|
|
|
|
gr.Interface( |
|
|
fn=predict_vehicle, |
|
|
inputs=gr.Image(type="pil", label="Upload Image"), |
|
|
outputs=gr.JSON(label="Prediction Result"), |
|
|
title="Vehicle Classification", |
|
|
description="Upload an image of a vehicle and the model will predict its type.", |
|
|
live=True |
|
|
).launch() |
