from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
import numpy as np
from google.colab import files

# Step 1: Upload the model to Colab (run this in a Colab cell)
uploaded = files.upload()  # Upload the .h5 model file

# Step 2: Load the trained model
MODEL_PATH = "setosys_dogs_model.h5"  # Adjust if needed
model = load_model(MODEL_PATH)

# Step 3: Define the class labels manually (as per your model's training setup)
# You need to know the classes that were used during model training
class_labels = ["Labrador Retriever", "German Shepherd", "Golden Retriever", "Bulldog", "Poodle"]  # Example, update this list

# Step 4: Define image preprocessing function for EfficientNetV2
def preprocess_image(image_path):
    """Preprocess the image to match the model's input requirements."""
    img = load_img(image_path, target_size=(224, 224))  # Resize image to model input size
    img_array = img_to_array(img)
    img_array = preprocess_input(img_array)  # EfficientNetV2 preprocessing
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    return img_array

# Step 5: Define prediction function
def predict_dog_breed(image_path):
    """Predict the breed of the dog in the uploaded image."""
    img_array = preprocess_image(image_path)
    predictions = model.predict(img_array)
    
    # Check the shape of the predictions to make sure the output is correct
    print("Predictions Shape:", predictions.shape)
    
    class_idx = np.argmax(predictions)  # Index of the highest prediction probability
    confidence = float(np.max(predictions))  # Confidence score
    
    # Get predicted breed and its confidence score
    predicted_breed = class_labels[class_idx] if class_idx < len(class_labels) else "Unknown"
    
    return predicted_breed, confidence

# Step 6: Upload and test with an image
uploaded_image = files.upload()  # Upload a test image
image_path = list(uploaded_image.keys())[0]  # Get the filename of the uploaded image

# Step 7: Run prediction
breed, confidence = predict_dog_breed(image_path)
print(f"Predicted Breed: {breed}, Confidence: {confidence:.2f}")