import streamlit as st
import numpy as np
import pandas as pd
import cv2
from tensorflow.keras.models import load_model

# Load the trained model
model = load_model('dog_model.h5')

# List of breeds
breeds = [
    'affenpinscher', 'afghan_hound', 'african_hunting_dog', 'airedale',
    'american_staffordshire_terrier', 'appenzeller', 'australian_terrier',
    'basenji', 'basset', 'beagle', 'bedlington_terrier', 'bernese_mountain_dog',
    'black-and-tan_coonhound', 'blenheim_spaniel', 'bloodhound', 'bluetick',
    'border_collie', 'border_terrier', 'borzoi', 'boston_bull', 'bouvier_des_flandres',
    'boxer', 'brabancon_griffon', 'briard', 'brittany_spaniel', 'bull_mastiff',
    'cairn', 'cardigan', 'chesapeake_bay_retriever', 'chihuahua', 'chow',
    'clumber', 'cocker_spaniel', 'collie', 'curly-coated_retriever', 'dandie_dinmont',
    'dhole', 'dingo', 'doberman', 'english_foxhound', 'english_setter',
    'english_springer', 'entlebucher', 'eskimo_dog', 'flat-coated_retriever',
    'french_bulldog', 'german_shepherd', 'german_short-haired_pointer',
    'giant_schnauzer', 'golden_retriever', 'gordon_setter', 'great_dane',
    'great_pyrenees', 'greater_swiss_mountain_dog', 'groenendael', 'ibizan_hound',
    'irish_setter', 'irish_terrier', 'irish_water_spaniel', 'irish_wolfhound',
    'italian_greyhound', 'japanese_spaniel', 'keeshond', 'kelpie',
    'kerry_blue_terrier', 'komondor', 'kuvasz', 'labrador_retriever',
    'lakeland_terrier', 'leonberg', 'lhasa', 'malamute', 'malinois', 'maltese_dog',
    'mexican_hairless', 'miniature_pinscher', 'miniature_poodle',
    'miniature_schnauzer', 'newfoundland', 'norfolk_terrier', 'norwegian_elkhound',
    'norwich_terrier', 'old_english_sheepdog', 'otterhound', 'papillon', 'pekinese',
    'pembroke', 'pomeranian', 'pug', 'redbone', 'rhodesian_ridgeback', 'rottweiler',
    'saint_bernard', 'saluki', 'samoyed', 'schipperke', 'scotch_terrier',
    'scottish_deerhound', 'sealyham_terrier', 'shetland_sheepdog', 'shih-tzu',
    'siberian_husky', 'silky_terrier', 'soft-coated_wheaten_terrier',
    'staffordshire_bullterrier', 'standard_poodle', 'standard_schnauzer',
    'sussex_spaniel', 'tibetan_mastiff', 'tibetan_terrier', 'toy_poodle', 'toy_terrier',
    'vizsla', 'walker_hound', 'weimaraner', 'welsh_springer_spaniel',
    'west_highland_white_terrier', 'whippet', 'wire-haired_fox_terrier', 'yorkshire_terrier'
]

# Streamlit app
st.title("Dog Breed Classifier 🐶")
images = ["1.jpg", "2.jpg", "3.jpg", "4.jpg", "5.jpg", "6.jpg"]
current_row = 0
for _ in range(2):
  cols = st.columns(3)
  for col, image in zip(cols, images[current_row:current_row+3]):
      col.image(image)
  current_row += 3

st.write("This model is traned on 120 different breeds of dogs using VGG16. Upload an image of a dog to classify its breed. You can use these sample images.")

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

def preprocess_image(image, image_size=(224, 224)):
    image = cv2.resize(image, image_size)
    image = image / 255.0
    return image

if uploaded_file is not None:
    # Read the uploaded image
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    img = cv2.imdecode(file_bytes, 1)
    
    # Convert BGR image to RGB
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # Display the uploaded image
    st.image(img_rgb, caption='Uploaded Image.', use_column_width=True)
    
    # Preprocess the image
    img_processed = preprocess_image(img)
    img_processed = np.expand_dims(img_processed, axis=0)
    
    # Predict the breed
    prediction = model.predict(img_processed)
    predicted_breed = breeds[np.argmax(prediction)]

    # Display probabilities for top 3 breeds
    top_3_indices = prediction[0].argsort()[-3:][::-1]
    top_3_breeds = [(breeds[i], prediction[0][i]) for i in top_3_indices]
    st.write("Top 3 predicted breeds:")
    for breed, prob in top_3_breeds:
        st.header(f"{breed}: {prob:.4f}%")