Lane segmentation model trained with tensorflow implementation of MobileNetV2 based U-Net
Training environment is contenerized, however you still need some stuff installed on host machine. This includes:
- Docker
- NVIDIA driver package
- nvidia-docker
This will allow to take full advantage of host GPUs from docker container.
If you are running Ubuntu 16.04 then use envsetup.sh which will take care of setting up your host with all necessities.
sudo dos2unix envsetup.sh && \
sudo chmod -x envsetup.sh && \
sudo bash envsetup.sh
First, build docker image with
docker build . --tag mobileunet:v1
then you can start training with
docker run -d --gpus all -v $(pwd)/data:/app/data \
-e --mode="train" \
-e --data-train=${DATA_DIR} \
-e --epochs=${EPOCHS} \
-e --model-name=${MODEL_NAME} mobileunet:v1
Container has volume mapped to $(pwd)/data so make sure such path exists on your host. This directory should also contain training data. Baseline for this model was trained on CU Lane Dataset. Tensorboard files and json logs will be available at $(pwd)/data/logs, model checkpoints and .h5 are saved to $(pwd)/data/models.
Run python test.py -m image -f ${IMG_PATH} -s ${SAVEPOINT} to test the model. You can pass -m image to run model on .jpg or -m video to run it on video feed. FlatBuffer model (.tflite) supported by passing --flatbuff
In order to test model on Coral USB Accelerator, prepare local env and make sure your model has been compiled, then run python tpu_test.py -m ${MODEL_PATH} -v ${VIDEO_PATH}