Create tensorflow_model.py

projects/piguardian/ai_ml/machine_learning/tensorflow_model.py

@@ -0,0 +1,47 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Dense, Conv2D, MaxPooling2D, Flatten
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
+
+class TensorFlowModel:
+    def __init__(self, input_shape, num_classes, learning_rate=0.001):
+        self.input_shape = input_shape
+        self.num_classes = num_classes
+        self.learning_rate = learning_rate
+        self.model = self.build_model()
+
+    def build_model(self):
+        model = Sequential()
+        model.add(Conv2D(32, (3, 3), activation='relu', input_shape=self.input_shape))
+        model.add(MaxPooling2D((2, 2)))
+        model.add(Conv2D(64, (3, 3), activation='relu'))
+        model.add(MaxPooling2D((2, 2)))
+        model.add(Conv2D(128, (3, 3), activation='relu'))
+        model.add(MaxPooling2D((2, 2)))
+        model.add(Flatten())
+        model.add(Dense(128, activation='relu'))
+        model.add(Dense(self.num_classes, activation='softmax'))
+        model.compile(optimizer=Adam(lr=self.learning_rate), loss='categorical_crossentropy', metrics=['accuracy'])
+        return model
+
+    def train(self, X_train, y_train, X_val, y_val, epochs=10, batch_size=32):
+        early_stopping = EarlyStopping(monitor='val_loss', patience=5, min_delta=0.001)
+        model_checkpoint = ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True, mode='min')
+        self.model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, validation_data=(X_val, y_val), callbacks=[early_stopping, model_checkpoint])
+
+    def evaluate(self, X_test, y_test):
+        y_pred = self.model.predict(X_test)
+        y_pred_class = np.argmax(y_pred, axis=1)
+        y_test_class = np.argmax(y_test, axis=1)
+        accuracy = accuracy_score(y_test_class, y_pred_class)
+        report = classification_report(y_test_class, y_pred_class)
+        matrix = confusion_matrix(y_test_class, y_pred_class)
+        return accuracy, report, matrix
+
+    def save_model(self, filename):
+        self.model.save(filename)
+
+    def load_model(self, filename):
+        self.model = tf.keras.models.load_model(filename)