Spaces:
Build error
Build error
| import tensorflow as tf | |
| from tensorflow import keras | |
| from tensorflow.keras import layers | |
| import numpy as np | |
| import cv2 | |
| class ImageClassifier: | |
| def __init__(self): | |
| self.model = None | |
| def preprocess_image(self, image): | |
| # Resize the image to (32, 32) | |
| resized_image = cv2.resize(image, (32, 32)) | |
| # # Convert the image to grayscale | |
| # gray_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY) | |
| # # # Normalize the pixel values between 0 and 1 | |
| # normalized_image = gray_image.astype("float32") / 255.0 | |
| # # # Transpose the dimensions to match the model's input shape | |
| # transposed_image = np.transpose(normalized_image, (1, 2, 0)) | |
| # # # Expand dimensions to match model input shape (add batch dimension) | |
| # img_array = np.expand_dims(transposed_image, axis=0) | |
| return resized_image | |
| def load_dataset(self): | |
| # Set up the dataset | |
| (x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data() | |
| # Normalize pixel values between 0 and 1 | |
| x_train = x_train.astype("float32") / 255.0 | |
| x_test = x_test.astype("float32") / 255.0 | |
| return (x_train, y_train), (x_test, y_test) | |
| # def build_model(self, x_train): | |
| # # Define the model architecture | |
| # model = keras.Sequential([ | |
| # # keras.Input(shape=x_train.shape[1]), | |
| # layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'), | |
| # layers.MaxPooling2D(pool_size=(2, 2)), | |
| # layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'), | |
| # layers.MaxPooling2D(pool_size=(2, 2)), | |
| # layers.Flatten(), | |
| # layers.Dropout(0.5), | |
| # layers.Dense(10, activation="softmax") | |
| # ]) | |
| # # Compile the model | |
| # model.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"]) | |
| # self.model = model | |
| def build_model(self, x_train): | |
| # Define the model architecture | |
| model = keras.Sequential([ | |
| layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'), | |
| layers.BatchNormalization(), | |
| layers.MaxPooling2D(pool_size=(2, 2)), | |
| layers.Dropout(0.25), | |
| layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'), | |
| layers.BatchNormalization(), | |
| layers.MaxPooling2D(pool_size=(2, 2)), | |
| layers.Dropout(0.25), | |
| layers.Conv2D(128, kernel_size=(3, 3), activation="relu", padding='same'), | |
| layers.BatchNormalization(), | |
| layers.MaxPooling2D(pool_size=(2, 2)), | |
| layers.Dropout(0.25), | |
| layers.Flatten(), | |
| layers.Dense(256, activation="relu"), | |
| layers.BatchNormalization(), | |
| layers.Dropout(0.5), | |
| layers.Dense(10, activation="softmax") | |
| ]) | |
| # Compile the model | |
| optimizer = keras.optimizers.RMSprop(learning_rate=0.001) | |
| model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"]) | |
| self.model = model | |
| def train_model(self, x_train, y_train, batch_size, epochs, validation_split): | |
| # Train the model | |
| self.model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, verbose=1, validation_split=validation_split) | |
| def evaluate_model(self, x_test, y_test): | |
| # Evaluate the model on the test set | |
| score = self.model.evaluate(x_test, y_test, verbose=0) | |
| print("Test loss:", score[0]) | |
| print("Test accuracy:", score[1]) | |
| def save_model(self, filepath): | |
| # Save the trained model | |
| self.model.save(filepath) | |
| def load_model(self, filepath): | |
| # Load the trained model | |
| self.model = keras.models.load_model(filepath) | |
| def classify_image(self, image, top_k=3): | |
| # Preprocess the image | |
| preprocessed_image = self.preprocess_image(image) | |
| # Perform inference | |
| predicted_probs = self.model.predict(np.array([preprocessed_image])) | |
| top_classes = np.argsort(predicted_probs[0])[-top_k:][::-1] | |
| top_probs = predicted_probs[0][top_classes] | |
| return top_classes, top_probs |