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)

# Load model
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()