Image Segmentation performed on GTA 5 Games Dataset using UNet Architecture
Dataset URL: http://download.visinf.tu-darmstadt.de/data/from_games
UNet Paper: https://arxiv.org/abs/1505.04597
Required Libraries: torch, numpy, PIL, glob, torchsummary, argparse, os, cv2
datagenerator.py
: To create custom data generation that we can use in PyTorch Code.
model.py
: Implemented U-Net architecture here.
main.py
: Contains train, validation functions with metrics used for segmentation.
test.py
: Test on given images and save the predicted output as images.
Run main.py
to start training the model.
python main.py -i image_directory
-l label_directory
-lr learning_rate
-e epochs
-b batch_size
-cp checkpoint_saved
For testing:
python test.py -i image_directory
-l label_directory
-s save_predicted_directory
-cp checkpoint_saved
Output:
Here are the sample predictions from my implementation of UNet Model