Update model.py

model.py

@@ -1,38 +1,27 @@
-from tensorflow.keras.models import Sequential
-from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
-from tensorflow.keras.layers import BatchNormalization
-from tensorflow.keras.regularizers import l2
+from tensorflow.keras.applications import VGG16
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout
+from tensorflow.keras.optimizers import Adam
 
-def build_model(input_shape, num_classes):
-    model = Sequential([
-        Conv2D(32, (3, 3), activation='relu',padding='same', input_shape=input_shape,kernel_regularizer='l2'),
-        BatchNormalization(),
-        MaxPooling2D((2, 2)),
-        
-        Conv2D(64, (3, 3), activation='relu',padding='same',kernel_regularizer='l2'),
-        BatchNormalization(),
-        MaxPooling2D((2, 2)),
-        
-        Conv2D(128, (3, 3), activation='relu',padding='same',kernel_regularizer='l2'),
-        BatchNormalization(),
-        MaxPooling2D((2, 2)),
-        
-        Conv2D(256, (3, 3), activation='relu',padding='same',kernel_regularizer='l2'),
-        BatchNormalization(),
-        MaxPooling2D((2, 2)),
-        
-        Conv2D(256, (3, 3), activation='relu',padding='same',kernel_regularizer='l2'),
-        BatchNormalization(),
-        MaxPooling2D((2, 2)),
-        
-        Flatten(),
-        Dense(512, activation='relu',kernel_regularizer='l2'),
-        #BatchNormalization(),
-        Dropout(0.5),
-        Dense(256, activation='relu',kernel_regularizer='l2'),
-        #BatchNormalization(),
-        Dropout(0.5),
-        Dense(3, activation='softmax')  # Assuming 3 classes
-    ])
+def build_model(num_classes):
+    # Load VGG16 model without the top layers
+    base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
 
-    return model
+    # Adding additional layers
+    x = base_model.output
+    x = GlobalAveragePooling2D()(x)
+    x = Dense(1024, activation='relu')(x)
+    x = Dropout(0.5)(x)
+    predictions = Dense(num_classes, activation='softmax')(x)
+
+    # Creating the final model
+    model = Model(inputs=base_model.input, outputs=predictions)
+
+    # Freezing the layers except the last layers
+    for layer in base_model.layers:
+        layer.trainable = False
+
+    # Compile the model
+    model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
+
+    return model