Possibility to change the location of the training folder and file names (model/config) (i did a improvement)

[Fatichti]

Current Behaviour?

Currently it is impossible to specify the location of the folders containing the images/annotations and the folder containing the JSON, models, logs... since only the main folder is defined with the method setDataDirectory(data_directory=PATH_FOLDER_ROOT_TRAINING)

Standalone code to reproduce the issue

trainer = DetectionModelTrainer()
trainer.setDataDirectory(anyFolder)
trainer.setModelTypeAsYOLOv3() 
trainer.setTrainConfig(object_names_array=["anyClass"], batch_size=BATCH_SIZE, num_experiments=NUMBER_EXP)
trainer.setDataDirectory(data_directory=PATH_FOLDER_ROOT_TRAINING)
trainer.trainModel()

Relevant log output

Here is the typical tree structure generated automatically :

├───cache
├───json
├───logs
├───models
├───train
│   ├───annotations
│   └───images
└───validation
    ├───annotations
    └───images

Proposal for improvement

1. Move to the init custom detection script and open it

cd $PATH_TO_IMAGEAI/Detection/Custom  
code __init__.py   // because i have visual studio code installed

2. Adding properties to the DetectionModelTrainer class :

class DetectionModelTrainer:
    def __init__(self):
        ...
        self.__train_metafolder = ""
        self.__train_dectection_configJSON_filename = ""
        self.__train_dectection_modelH5_filename = ""

3. Adding methods to the same class

def setTrainMetaFolder(self, path):
    """
    'setTrainMetaFolder' is used to define the location of the folder where the cache, log, models and json files are stored.
    :param path: the path where the training files will be located
    :return:
    """
    self.__train_metafolder = path

def setDetectionConfigJSONFileName(self, filename):
    """
    'setDetectionConfigJSONFileName' is used to set the JSON file name of the training
    :param filename: the name of the file
    :return:
    """
    self.__train_dectection_configJSON_filename = filename

def setDetectionModelH5FileName(self, filename):
    """
    'setDetectionModelH5FileName' is used to set the name of the H5 file of the drive
    :param filename: the name of the file
    :return:
    """
    self.__train_dectection_modelH5_filename = filename

4. Edit setDataDirectory() method

before

os.makedirs(os.path.join(...
...
self.__logs_directory = os.path.join... "logs")

after (rename metadata folder and remove useless json folder)

os.makedirs(os.path.join(self.__train_metafolder, "cache"), exist_ok=True)
self.__train_cache_file = os.path.join(self.__train_metafolder, "cache", "detection_train_data.pkl")
self.__validation_cache_file = os.path.join(self.__train_metafolder, "cache", "detection_test_data.pkl")

os.makedirs(os.path.join(self.__train_metafolder, "models"), exist_ok=True)

os.makedirs(os.path.join(self.__train_metafolder, "logs"), exist_ok=True)

self.__model_directory = os.path.join(self.__train_metafolder, "models")
self.__train_weights_name = os.path.join(self.__model_directory, "detection_model-")
self.__logs_directory = os.path.join(self.__train_metafolder, "logs")

5. Modify 2 lines of the setTrainConfig() method

def setTrainConfig(self,  object_names_array, batch_size=4, num_experiments=100, train_from_pretrained_model=""):
     ...
    with open(os.path.join(self.__json_directory, self.__train_dectection_configJSON_filename), "w+") as json_file:
          json.dump(json_data, json_file, indent=4, separators=(",", " : "),
          ensure_ascii=True)

print("Detection configuration saved in ", os.path.join(self.__json_directory, self.__train_dectection_configJSON_filename))

6. Adding some lines in the _create_callbacks method

def _create_callbacks(self, saved_weights_name, model_to_save):
    # change this line to give the option to default the name automatically
    if self.__train_dectection_modelH5_filename == "":
      customFilePath = saved_weights_name + 'ex-{epoch:03d}--loss-{loss:08.3f}.h5'
    else:
      customFilePath = self.__train_dectection_modelH5_filename
    
    checkpoint = CustomModelCheckpoint(
                model_to_save=model_to_save,
                filepath=customFilePath, 
                monitor='loss',
                verbose=0,
                save_best_only=True,
                mode='min',
                period=1
            )

Implementation of the improvement

trainer = DetectionModelTrainer()
trainer.setDataDirectory(anyFolder)
trainer.setModelTypeAsYOLOv3() 
# START NEW LINES
trainer.setTrainMetaFolder(PATH_FOLDER_TRAINING_METADATA)
trainer.setDetectionModelH5FileName(PATH_FILE_MODEL_TRAINING_PTH + FILENAME_TRAINING_H5)
trainer.setDetectionConfigJSONFileName(PATH_FILE_MODEL_TRAINING_JSON + FILENAME_TRAINING_JSON)
# END NEW LINES
trainer.setTrainConfig(object_names_array=["anyClass"], batch_size=BATCH_SIZE, num_experiments=NUMBER_EXP)
trainer.setDataDirectory(data_directory=PATH_FOLDER_ROOT_TRAINING)
trainer.trainModel()

It's done! 🎉