Training of Object Detection model

This Python script allows you to train a Pytorch Object Detection model with preprocessed HDF5 data and labeled bounding boxes.

The source code of this sample can be found in <install-prefix>/share/metavision/sdk/ml/python_samples/train_detection when installing Metavision SDK from installer or packages. For other deployment methods, check the page Path of Samples.

Expected Output

Object Detection training result:

  • checkpoints (models at different training stages)

  • log files

  • videos on test dataset

A test demo is shown below:



Setup & requirements

To run the script, you need:

  • path to the output folder

  • path to the training dataset:

    • a folder containing 3 sub folders, named train, val, test.

    • each subfolder should contain one or multiple h5 files and their corresponding <h5 filename>_bbox.npy labels.

    • a dictionary file named label_map_dictionary.json, which contains all the detection categories.

Note

To have a quick test of the pipeline, you can simply set the path of training dataset to “toy_problem”. This will train the model on Moving MNIST Dataset.

How to start

To run the script:

python train_detection.py /path/to/logging /path/to/dataset

To find the full list of options, run:

python train_detection.py -h