train_model.py

import os
import numpy as np
from Mask_RCNN.mrcnn.model import MaskRCNN
from Mask_RCNN.mrcnn.config import Config
from train_dataset import MaskDataset

# Specify constants
DIRECTORY = '/kaggle/input/face-mask-detection'
SIZE = None
TRAIN_EPOCHS = 5

class MaskConfig(Config):
    NAME = 'mask_cfg'
    NUM_CLASSES = 1 + 3
    IMAGES_PER_GPU = 1
    STEPS_PER_EPOCH = 100
    VALIDATION_STEPS = 10

# Prepare datasets
train_set = MaskDataset()
train_set.load_dataset(dataset_directory=DIRECTORY, train=True, size=SIZE)
train_set.prepare()

val_set = MaskDataset()
val_set.load_dataset(dataset_directory=DIRECTORY, train=False, size=SIZE)
val_set.prepare()

# Create model configuration
config = MaskConfig()
config.display()

# Create model
model = MaskRCNN(mode='training', model_dir='./', config=config)

# Load pretrained weights
model.load_weights(
    "./mask_rcnn_coco.h5",
    by_name=True,
    exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"]
)

# Train the model
model.train(
    train_set,
    val_set,
    learning_rate=config.LEARNING_RATE,
    epochs=TRAIN_EPOCHS,
    layers='heads'
)

print("Training completed.")