[Task] Model Training

[Description] Model Training and saved the model.
[Author]

@alifalhasan

models/football_logo_model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9150fe90396b2098a79532af7278286a7aa8ea39fbc99199982349c33f1ffd0
+size 302849424

requirements.txt

@@ -1 +1,4 @@
-gradio
+gradio
+numpy
+tensorflow
+scipy

src/train/trainer.py

@@ -0,0 +1,98 @@
+"""train_model.py
+
+This module trains a simple image classification model using TensorFlow/Keras.
+
+"""
+
+import os
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.preprocessing import image
+import numpy as np
+from scipy import misc
+
+# Constants
+IMG_SIZE = (224, 224)
+BATCH_SIZE = 32
+EPOCHS = 10
+NUM_CLASSES = 5
+
+def preprocess_image(img):
+    """
+    Preprocess the input image for model prediction.
+
+    Parameters:
+    - img: Input image.
+
+    Returns:
+    - img: Preprocessed image.
+    """
+    img = image.img_to_array(img)
+    img = np.expand_dims(img, axis=0)
+    img /= 255.0
+    return img
+
+def train_model(dataset_path):
+    """
+    Train the image classification model.
+
+    Parameters:
+    - dataset_path (str): Path to the dataset.
+
+    Returns:
+    - model: Trained Keras model.
+    """
+
+    # Ensure the dataset path is correct
+    dataset_path = os.path.abspath(dataset_path)
+
+    # Data preprocessing
+    datagen = ImageDataGenerator(
+        rescale=1./255,
+        validation_split=0.2
+    )
+
+    train_generator = datagen.flow_from_directory(
+        dataset_path,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode='categorical',
+        subset='training'
+    )
+
+    validation_generator = datagen.flow_from_directory(
+        dataset_path,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode='categorical',
+        subset='validation'
+    )
+
+    # Model definition
+    model = Sequential()
+    model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)))
+    model.add(MaxPooling2D((2, 2)))
+    model.add(Flatten())
+    model.add(Dense(64, activation='relu'))
+    model.add(Dense(NUM_CLASSES, activation='softmax'))
+
+    # Model compilation
+    model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
+
+    # Model training
+    history = model.fit(
+        train_generator,
+        epochs=EPOCHS,
+        validation_data=validation_generator
+    )
+
+    return model
+
+if __name__ == "__main__":
+    script_directory = os.path.dirname(os.path.abspath(__file__))
+    dataset_path = os.path.join(script_directory, '../../data')
+    trained_model = train_model(dataset_path)
+    trained_model.save(os.path.join(script_directory, '../../models/football_logo_model.h5'))