Update model_tensorflow.py

model_tensorflow.py

@@ -1,48 +1,52 @@
 import tensorflow as tf
 from tensorflow.keras import layers, Model
 
-
-def separable_block(x, filters: int, dropout_rate: float = 0.25):
-    x = layers.SeparableConv2D(
-        filters, kernel_size=3, padding="same", use_bias=False)(x)
+# ==============================
+# Separable Convolution Block
+# ==============================
+def separable_block(x, filters, dropout_rate=0.25):
+    x = layers.SeparableConv2D(filters, 3, padding="same", use_bias=False)(x)
     x = layers.BatchNormalization()(x)
     x = layers.Activation("relu")(x)
 
-    x = layers.SeparableConv2D(
-        filters, kernel_size=3, padding="same", use_bias=False)(x)
+    x = layers.SeparableConv2D(filters, 3, padding="same", use_bias=False)(x)
     x = layers.BatchNormalization()(x)
     x = layers.Activation("relu")(x)
 
-    x = layers.MaxPooling2D(pool_size=2)(x)
-    x = layers.SpatialDropout2D(rate=dropout_rate)(x)
+    x = layers.MaxPooling2D()(x)
+    x = layers.SpatialDropout2D(dropout_rate)(x)
 
     return x
 
 
-def build_sara_tf_model(input_shape=(150, 150, 3), num_classes: int = 6) -> Model:
+# ==============================
+# Model Definition
+# ==============================
+def build_sara_tf_model(input_shape=(150, 150, 3), num_classes=6):
 
-    inputs = tf.keras.Input(shape=input_shape, name="image_input")
+    inputs = tf.keras.Input(shape=input_shape)
 
-    x = layers.Conv2D(32, kernel_size=3, padding="same",
-                      use_bias=False, name="stem_conv")(inputs)
-    x = layers.BatchNormalization(name="stem_bn")(x)
-    x = layers.Activation("relu", name="stem_relu")(x)
+    # Stem
+    x = layers.Conv2D(32, 3, padding="same", use_bias=False)(inputs)
+    x = layers.BatchNormalization()(x)
+    x = layers.Activation("relu")(x)
 
-    x = separable_block(x, filters=64,  dropout_rate=0.25)
-    x = separable_block(x, filters=128, dropout_rate=0.25)
-    x = separable_block(x, filters=256, dropout_rate=0.30)
+    # Feature extractor
+    x = separable_block(x, 64, 0.25)
+    x = separable_block(x, 128, 0.25)
+    x = separable_block(x, 256, 0.30)
 
-    x = layers.GlobalAveragePooling2D(name="gap")(x)
+    # Head
+    x = layers.GlobalAveragePooling2D()(x)
 
-    x = layers.Dense(128, use_bias=False, name="fc1")(x)
-    x = layers.BatchNormalization(name="fc1_bn")(x)
-    x = layers.Activation("relu", name="fc1_relu")(x)
-    x = layers.Dropout(0.5, name="fc1_drop")(x)
+    x = layers.Dense(128, use_bias=False)(x)
+    x = layers.BatchNormalization()(x)
+    x = layers.Activation("relu")(x)
+    x = layers.Dropout(0.5)(x)
 
-    outputs = layers.Dense(num_classes, activation="softmax",
-                           name="predictions")(x)
+    outputs = layers.Dense(num_classes, activation="softmax")(x)
 
-    model = Model(inputs=inputs, outputs=outputs, name="SaraCNN_TF")
+    model = Model(inputs, outputs, name="SaraCNN_TF")
 
     model.compile(
         optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
@@ -52,12 +56,3 @@ def build_sara_tf_model(input_shape=(150, 150, 3), num_classes: int = 6) -> Mode
 
     return model
 
-
-if __name__ == "__main__":
-    model = build_sara_tf_model()
-
-    model.build((None, 150, 150, 3))
-
-    model.load_weights("ton_modele_weights.h5")
-
-    print("Modèle chargé correctement !")