Skip to content

Latest commit

 

History

History
26 lines (18 loc) · 1.16 KB

README.md

File metadata and controls

26 lines (18 loc) · 1.16 KB

UNet for Image Segmentation

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.

Commands to run:

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


Prediction