Package: seg-torch
Semantic-Segmentation-PytorchPytorch implementation of FCN, UNet, PSPNet and various encoder models.These are the reference implementation of the models.FCN (Fully Convolutional Networks for Sementic Segmentation)[Paper]UNet (Convolutional Networks for Biomedical Image Segmentation)[Paper]PSPNet (Pyramid Scene Parsing Network)[Paper]ModelsThis project supports models as follow:model nameencoder modeldecoder modelfcn8_vgg11VGG 11FCN8fcn8_vgg13VGG 13FCN8fcn8_vgg16VGG 16FCN8fcn8_vgg19VGG 19FCN8fcn16_vgg11VGG 11FCN16fcn16_vgg13VGG 13FCN16fcn16_vgg16VGG 16FCN16fcn16_vgg19VGG 19FCN16fcn32_vgg11VGG 11FCN32fcn32_vgg13VGG 13FCN32fcn32_vgg16VGG 16FCN32fcn32_vgg19VGG 19FCN32fcn8_resnet18Resnet-18FCN8fcn8_resnet34Resnet-34FCN8fcn8_resnet50Resnet-50FCN8fcn8_resnet101Resnet-101FCN8fcn8_resnet152Resnet-152FCN8fcn16_resnet18Resnet-18FCN16fcn16_resnet34Resnet-34FCN16fcn16_resnet50Resnet-50FCN16fcn16_resnet101Resnet-101FCN16fcn16_resnet152Resnet-152FCN16fcn32_resnet18Resnet-18FCN32fcn32_resnet34Resnet-34FCN32fcn32_resnet50Resnet-50FCN32fcn32_resnet101Resnet-101FCN32fcn32_resnet152Resnet-152FCN32fcn8_mobilenet_v2MobileNet-v2FCN8fcn16_mobilenet_v2MobileNet-v2FCN16fcn32_mobilenet_v2MobileNet-v2FCN32unetUnetUnetunet_vgg11VGG11Unetunet_vgg13VGG13Unetunet_vgg16VGG16Unetunet_vgg19VGG19Unetunet_resnet18Resnet-18Unetunet_resnet34Resnet-34Unetunet_resnet50Resnet-50Unetunet_resnet101Resnet-101Unetunet_resnet152Resnet-152Unetunet_mobilenet_v2MobileNet-v2Unetpspnet_vgg11VGG11PSPNetpspnet_vgg13VGG13PSPNetpspnet_vgg16VGG16PSPNetpspnet_vgg19VGG19PSPNetpspnet_resnet18Resnet-18PSPNetpspnet_resnet34Resnet-34PSPNetpspnet_resnet50Resnet-50PSPNetpspnet_resnet101Resnet-101PSPNetpspnet_resnet152Resnet-152PSPNetpspnet_mobilenet_v2MobileNet-v2PSPNetExample results of the pspnet_mobilenet_v2 model:Input ImageGround Truth ImageResult ImageGetting StartedRequirementspytorch>= 1.5.0torchvision>= 0.5.0TensorboardX>= 0.2.0opencv-pythontqdmInstallationpipinstallseg-torchorgitclonehttps://github.com/IanTaehoonYoo/semantic-segmentation-pytorch/cdsemantic-segmentation-pytorch
pythonsetup.pyinstallPreparing the data for trainingIn this project, the data for training is the[Cityspaces]. You can run this project using the sample dataset in the segmentation/test/dataset/cityspaces folder. If you want to run this project using another dataset, please refer to the dataset format as bellow.There are two folders which are the training images folder and the groundtruth labeled images folder.The training image and the labeled image must have the same file name and size.The training image must be the RGB image, and the labeled image should have the class value, the range [0, n_classes].Example code to use this project with pythonfromtorchvisionimporttransformsfromsegmentation.data_loader.segmentation_datasetimportSegmentationDatasetfromsegmentation.data_loader.transformimportRescale,ToTensorfromsegmentation.trainerimportTrainerfromsegmentation.predictimport*fromsegmentation.modelsimportall_modelsfromutil.loggerimportLoggertrain_images=r'dataset/cityspaces/images/train'test_images=r'dataset/cityspaces/images/test'train_labled=r'dataset/cityspaces/labeled/train'test_labeled=r'dataset/cityspaces/labeled/test'if__name__=='__main__':model_name="fcn8_vgg16"device='cuda'batch_size=4n_classes=34num_epochs=10image_axis_minimum_size=200pretrained=Truefixed_feature=Falselogger=Logger(model_name=model_name,data_name='example')### Loadercompose=transforms.Compose([Rescale(image_axis_minimum_size),ToTensor()])train_datasets=SegmentationDataset(train_images,train_labled,n_classes,compose)train_loader=torch.utils.data.DataLoader(train_datasets,batch_size=batch_size,shuffle=True,drop_last=True)test_datasets=SegmentationDataset(test_images,test_labeled,n_classes,compose)test_loader=torch.utils.data.DataLoader(test_datasets,batch_size=batch_size,shuffle=True,drop_last=True)### Modelmodel=all_models.model_from_name[model_name](n_classes,batch_size,pretrained=pretrained,fixed_feature=fixed_feature)model.to(device)###Load model###please check the foloder: (.segmentation/test/runs/models)#logger.load_model(model, 'epoch_15')### Optimizersifpretrainedandfixed_feature:#fine tunningparams_to_update=model.parameters()print("Params to learn:")params_to_update=[]forname,paraminmodel.named_parameters():ifparam.requires_grad==True:params_to_update.append(param)print("\t",name)optimizer=torch.optim.Adadelta(params_to_update)else:optimizer=torch.optim.Adadelta(model.parameters())### Train#scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)trainer=Trainer(model,optimizer,logger,num_epochs,train_loader,test_loader)trainer.train()#### Writing the predict result.predict(model,r'dataset/cityspaces/input.png',r'dataset/cityspaces/output.png')Pre-trained models (Encoder models)This project uses pre-trained models such as VGG, ResNet, and MobileNet from the torchvision library. If you want the fine-tunning model, you can change the input parameters which are 'pretrained' and 'fixed_feature' when calling a model. And then, you should set the optimizer to freeze the model like as follow.model=all_models.model_from_name[model_name](n_classes,batch_size,pretrained=pretrained,fixed_feature=fixed_feature)# Optimizersifpretrainedandfixed_feature:#fine-tunningparams_to_update=model.parameters()print("Params to learn:")params_to_update=[]forname,paraminmodel.named_parameters():ifparam.requires_grad==True:params_to_update.append(param)print("\t",name)optimizer=torch.optim.Adadelta(params_to_update)else:optimizer=torch.optim.Adadelta(model.parameters())Getting the learning results on TensorboardThe Logger class is to write the result such as mean IoU, accuracy, loss, and predict labeled images. The logger class gets the model name and the data name. So, it can generate the tensorboard files automatically in the runs folder, .\segmentation\runs\Here is example command to see the resulttensorboard--logdir=%project_path\segmentation\runs--hostlocalhostIf you don't know about Tensorboard, please refer to[Tensorboard]Saving and loading check pointsThe trainer class can save the check point automatically depends on argument is called 'check_point_epoch_stride'. So check points will be saved for every epoch stride in the runs folder, ./segmentation/runs/models.Also, you can load the check point using the logger class. Here are example codes, please refer to as bellow."""Save check point.Please check the runs folder, ./segmentation/runs/models"""check_point_stride=30# the checkpoint is saved for every 30 epochs.#'model_name' and 'data_name' are to set a path to save the check point.# So you should set the same the Logger's arguemnts when you load the check point.logger=Logger(model_name="pspnet_mobilenet_v2",data_name='example')trainer=Trainer(model,optimizer,logger,num_epochs,train_loader,test_loader,check_point_epoch_stride=check_point_stride)"""Load check point."""n_classes=33batch_size=4# The Logger's arguemnts should be the same as when you train the model.logger=Logger(model_name="pspnet_mobilenet_v2",data_name='example')model=all_models.model_from_name[model_name](n_classes,batch_size)logger.load_model(model,'epoch_253')Cite This ProjectIf you find this code useful, please consider the following BibTeX entry.@misc{seg-pytorch,author={IanYoo},title={{sementic-segmentation-pytorch:PytorchimplementationofFCN,UNet,PSPNetandvariousencodermodels}},howpublished={\url{https://github.com/IanTaehoonYoo/semantic-segmentation-pytorch}},year={2020}}