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)