Updated Data Loader, Model Architecture

Author: animikhaich

utils/data_loader.py

@@ -117,11 +117,9 @@ def augment_batch(self, image, label) -> tuple:
         if tf.random.normal([1]) < 0:
             image = tf.image.random_brightness(image, 0.2)
         if self.NUM_CHANNELS == 3 and tf.random.normal([1]) < 0:
-            image = tf.image.random_hue(image, 0.3)
+            image = tf.image.random_hue(image, 0.1)
         if self.NUM_CHANNELS == 3 and tf.random.normal([1]) < 0:
             image = tf.image.random_saturation(image, 0, 15)
-        if tf.random.normal([1]) < 0:
-            image = tf.image.random_flip_up_down(image)
 
         image = tf.image.random_flip_left_right(image)
         image = tf.image.random_jpeg_quality(image, 10, 100)
@@ -197,9 +195,7 @@ def visualize_batch(self, augment=True) -> None:
 
 
 if __name__ == "__main__":
-    dataset_root_dir = (
-        "/home/ani/Documents/pycodes/Dataset/mit-indoor-scenes/indoorCVPR_09/Images"
-    )
+    dataset_root_dir = "/home/ani/Documents/pycodes/Dataset/gender/Training/"
     data_loader = ImageClassificationDataLoader(
         data_dir=dataset_root_dir,
         image_dims=(512, 512),

utils/model.py

@@ -180,8 +180,8 @@ def init_network(self):
             backbone=self.backbone, input_shape=self.input_shape, classes=self.classes
         )
         x = base.output
+        x = tf.keras.layers.Conv2D(64, (1, 1), activation="relu")(x)
         x = tf.keras.layers.Dropout(0.5)(x)
-        x = tf.keras.layers.AveragePooling2D(pool_size=(5, 5))(x)
         x = tf.keras.layers.Flatten()(x)
         x = tf.keras.layers.Dense(units=128, activation="relu")(x)
         x = tf.keras.layers.Dropout(0.5)(x)
@@ -191,10 +191,10 @@ def init_network(self):
         self.model = tf.keras.models.Model(inputs=[base.input], outputs=[x])
         return self.model
 
-    def get_model_summary(self):
+    def print_model_summary(self):
         if self.model is None:
             return f"Please use the `init_network` Method to  the model first"
-        return self.model.summary()
+        self.model.summary()
 
     def get_model(self):
         if self.history is None:
@@ -248,6 +248,7 @@ def train(
         val_generator=None,
         val_steps=None,
         epochs=500,
+        print_summary=False,
     ):
 
         if self.metrics is None:
@@ -264,6 +265,10 @@ def train(
             optimizer=self.optimizer,
             metrics=self.metrics,
         )
+
+        if print_summary:
+            self.model.summary()
+
         self.history = self.model.fit(
             x=train_generator,
             epochs=epochs,