Car_Vision / app.py
courte's picture
add app
b69eaf0
import tensorflow as tf
import numpy as np
from PIL import Image
import gradio as gr
import os
# Load Keras model
model = tf.keras.models.load_model("car_brand_classifier_final.h5")
# Define the path to the dataset containing car brand folders
DATASET_PATH = "Car_Sales_vision_ai_project"
# Get the list of car brands (folder names) in the dataset
CAR_BRANDS = sorted(os.listdir(DATASET_PATH))
# Ensure the CAR_BRANDS list contains only valid directories
CAR_BRANDS = [brand for brand in CAR_BRANDS if os.path.isdir(os.path.join(DATASET_PATH, brand))]
# Print the list of car brands for debugging
print("Car brands:", CAR_BRANDS)
# Define the preprocessing function
def preprocess_image(image):
"""
Preprocess the input image for the model.
"""
if isinstance(image, np.ndarray): # If Gradio gives a NumPy array, convert it to PIL Image
image = Image.fromarray(image)
# Check if the image is valid
if image is None or image.size == 0:
raise ValueError("Invalid or blank image uploaded.")
image = image.convert("RGB") # Ensure it's in RGB mode
image = image.resize((299, 299)) # Resize to match model input size
image = np.array(image) / 255.0 # Normalize pixel values
image = np.expand_dims(image, axis=0) # Add batch dimension
return image
# Define the prediction function
def predict(image):
"""
Predict the car brand and return the predicted brand and sample images.
"""
try:
# Preprocess the image
processed_image = preprocess_image(image)
# Make a prediction
predictions = model.predict(processed_image)
predicted_class_index = np.argmax(predictions, axis=1)[0]
predicted_brand = CAR_BRANDS[predicted_class_index]
# Debugging: Print the predicted brand
print(f"Predicted brand: {predicted_brand}")
# Get sample images from the predicted brand folder
brand_folder = os.path.join(DATASET_PATH, predicted_brand)
print(f"Predicted brand folder: {brand_folder}")
sample_images = []
if os.path.exists(brand_folder):
print(f"Found {len(os.listdir(brand_folder))} files in the folder.")
for filename in os.listdir(brand_folder)[:5]: # Limit to 5 sample images
img_path = os.path.join(brand_folder, filename)
if os.path.isfile(img_path):
sample_images.append(Image.open(img_path).resize((200, 200))) # Resize for consistency
else:
print("Brand folder does not exist.")
# Return the predicted brand and sample images as separate outputs
return predicted_brand, sample_images or ["No images found for this brand."]
except Exception as e:
print(f"Error during prediction: {e}")
return "Error", ["An error occurred during prediction."]
# Create the Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="numpy"), # Ensure Gradio passes a NumPy array
outputs=[
gr.Textbox(label="Predicted Brand"),
gr.Gallery(label="Sample Images") # Display images in a gallery
],
title="Car Vision",
description="Upload an image of a car to classify its brand and view sample images."
)
# Launch the app
iface.launch(share=True)