miss import `matplotlib` in `train_augment2.py` and `train_augment1.py`

[vinthony]

Hi,

the script of train_augment2.py and train_augment1.py need to add

import matplotlib

to create the data.