Spaces:
Build error
Build error
| 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() | |