vaishanthr/Image-Classifier-TensorFlow

1import tensorflow as tf2from tensorflow import keras3from tensorflow.keras import layers4import numpy as np5import cv26 7 8class ImageClassifier:9    def __init__(self):10        self.model = None11 12    def preprocess_image(self, image):13        # Resize the image to (32, 32)14        resized_image = cv2.resize(image, (32, 32))15 16        # # Convert the image to grayscale17        # gray_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)18 19        # # # Normalize the pixel values between 0 and 120        # normalized_image = gray_image.astype("float32") / 255.021 22        # # # Transpose the dimensions to match the model's input shape23        # transposed_image = np.transpose(normalized_image, (1, 2, 0))24 25        # # # Expand dimensions to match model input shape (add batch dimension)26        # img_array = np.expand_dims(transposed_image, axis=0)27        return resized_image28 29    def load_dataset(self):30        # Set up the dataset31        (x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()32 33        # Normalize pixel values between 0 and 134        x_train = x_train.astype("float32") / 255.035        x_test = x_test.astype("float32") / 255.036 37        return (x_train, y_train), (x_test, y_test)38 39    # def build_model(self, x_train):40    #     # Define the model architecture41    #     model = keras.Sequential([42    #         # keras.Input(shape=x_train.shape[1]),43    #         layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'),44    #         layers.MaxPooling2D(pool_size=(2, 2)),45    #         layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'),46    #         layers.MaxPooling2D(pool_size=(2, 2)),47    #         layers.Flatten(),48    #         layers.Dropout(0.5),49    #         layers.Dense(10, activation="softmax")50    #     ])51 52    #     # Compile the model53    #     model.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])54 55    #     self.model = model56    57    def build_model(self, x_train):58      # Define the model architecture59      model = keras.Sequential([60          layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'),61          layers.BatchNormalization(),62          layers.MaxPooling2D(pool_size=(2, 2)),63          layers.Dropout(0.25),64 65          layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'),66          layers.BatchNormalization(),67          layers.MaxPooling2D(pool_size=(2, 2)),68          layers.Dropout(0.25),69 70          layers.Conv2D(128, kernel_size=(3, 3), activation="relu", padding='same'),71          layers.BatchNormalization(),72          layers.MaxPooling2D(pool_size=(2, 2)),73          layers.Dropout(0.25),74 75          layers.Flatten(),76          layers.Dense(256, activation="relu"),77          layers.BatchNormalization(),78          layers.Dropout(0.5),79 80          layers.Dense(10, activation="softmax")81      ])82 83      # Compile the model84      optimizer = keras.optimizers.RMSprop(learning_rate=0.001)85      model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])86 87      self.model = model88 89    def train_model(self, x_train, y_train, batch_size, epochs, validation_split):90        # Train the model91        self.model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, verbose=1, validation_split=validation_split)92 93    def evaluate_model(self, x_test, y_test):94        # Evaluate the model on the test set95        score = self.model.evaluate(x_test, y_test, verbose=0)96        print("Test loss:", score[0])97        print("Test accuracy:", score[1])98 99    def save_model(self, filepath):100        # Save the trained model101        self.model.save(filepath)102 103    def load_model(self, filepath):104        # Load the trained model105        self.model = keras.models.load_model(filepath)106 107    def classify_image(self, image, top_k=3):108        # Preprocess the image109        preprocessed_image = self.preprocess_image(image)110 111        # Perform inference112        predicted_probs = self.model.predict(np.array([preprocessed_image]))113        top_classes = np.argsort(predicted_probs[0])[-top_k:][::-1]114        top_probs = predicted_probs[0][top_classes]115 116        return top_classes, top_probs