Trainging code

[asadabbas09]

Thanks for sharing the test code.

Do you plan to share the training code? I would like to train it on my own custom dataset.