File size: 4,230 Bytes
31607dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
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 |