from keras.models import load_model
import numpy as np

# Load the saved model
# model = load_model('model.h5')

# predicted_image = model.predict(np.expand_dims(input_image, axis=0))


import tensorflow as tf
from PIL import Image
import numpy as np
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Load CIFAR-10 dataset
(x_train, y_train), (x_test, y_test) = cifar10.load_data()

# Normalize pixel values to be between 0 and 1
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# One-hot encode the labels
y_train = to_categorical(y_train, num_classes=10)
y_test = to_categorical(y_test, num_classes=10)






# predicted = model.predict(x_test)

# np.argmax(predicted, axis = 1)


model = Sequential([
    # First convolutional layer
    Conv2D(96, (11, 11), strides=(1, 1), activation='relu', input_shape=(32, 32, 3)),
    MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),
    # Second convolutional layer
    Conv2D(256, (5, 5), padding='same', activation='relu'),
    MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),
    # Third convolutional layer
    Conv2D(384, (3, 3), padding='same', activation='relu'),
    # Fourth convolutional layer
    Conv2D(384, (3, 3), padding='same', activation='relu'),
    # Fifth convolutional layer
    Conv2D(256, (3, 3), padding='same', activation='relu'),
    MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),
    # Flatten the convolutional layers output for fully connected layers
    Flatten(),
    # First fully connected layer
    Dense(4096, activation='relu'),
    Dropout(0.5),
    # Second fully connected layer
    Dense(4096, activation='relu'),
    Dropout(0.5),
    # Output layer
    Dense(10, activation='softmax')
])

# Compile the model with a lower learning rate
optimizer = Adam(learning_rate=0.0001)
model.compile(optimizer=optimizer,
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# Data augmentation
datagen = ImageDataGenerator(
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True,
)

datagen.fit(x_train)

# Train the model with data augmentation
model.fit(datagen.flow(x_train, y_train, batch_size=128), epochs=25, validation_data=(x_test, y_test))


classes = {
    0 : 'Airplane',
    1 : 'Automobile',
    2 : 'Bird',
    3 : 'Cat',
    4 : 'Deer',
    5 : 'Dog',
    6 : 'Frog',
    7 : 'Horse',
    8 : 'Ship',
    9 : 'Truck'
}

def prediction(input_img):
    # Define the transformation
    transform = transforms.Compose([
        transforms.Resize(32),
        transforms.CenterCrop(32),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    pil_image = Image.fromarray(input_img.astype('uint8'))
    # Apply the transformation
    transformed_image = np.array(transform(pil_image).T)
    input_image = np.expand_dims(transformed_image, axis=0)
    output = model.predict(input_image)
    # print(transformed_image.shape)
    # print(transformed_image)
    # plt.imshow(transformed_image)
    # plt.show()
    return classes[np.argmax(output)]
demo = gr.Interface(prediction, gr.Image(), "text")
demo.launch()