Package: fastseg
Fast Semantic SegmentationThis respository aims to provide accuratereal-time semantic segmentationcode for mobile devices in PyTorch, with pretrained weights on Cityscapes. This can be used for efficient segmentation on a variety of real-world street images, including datasets like Mapillary Vistas, KITTI, and CamVid.fromfastsegimportMobileV3Largemodel=MobileV3Large.from_pretrained().cuda().eval()model.predict(images)The models are implementations ofMobileNetV3(both large and small variants) with a modified segmentation head based onLR-ASPP. The top model was able to achieve72.3%mIoU accuracy on Cityscapesval, while running at up to37.3 FPSon a GPU. Please see below for detailed benchmarks.Currently, you can do the following:Load pretrained MobileNetV3 semantic segmentation models.Easily generate hard segmentation labels or soft probabilities for street image scenes.Evaluate MobileNetV3 models on Cityscapes, or your own dataset.Export models for production with ONNX.If you have any feature requests or questions, feel free to leave them as GitHub issues!Table of ContentsWhat's New?August 12th, 2020August 11th, 2020OverviewRequirementsPretrained Models and MetricsUsageRunning InferenceExporting to ONNXTraining from ScratchContributionsWhat's New?September 29th, 2020Releasedtraining codefor semantic segmentation modelsAugust 12th, 2020Added pretrained weights forMobileV3Smallwith 256 filtersAugust 11th, 2020Initial releaseImplementations ofMobileV3LargeandMobileV3Smallwith LR-ASPPPretrained weights forMobileV3Largewith 128/256 filters, andMobileV3Smallwith 64/128 filtersInference, ONNX export, and optimization scriptsOverviewHere's an excerpt from theoriginal paperintroducing MobileNetV3:This paper starts the exploration of how automated search algorithms and network design can work together to harness complementary approaches improving the overall state of the art. Through this process we create two new MobileNet models for release: MobileNetV3-Large and MobileNetV3-Small, which are targeted for high and low resource use cases. These models are then adapted and applied to the tasks of object detection and semantic segmentation.For the task of semantic segmentation (or any dense pixel prediction), we propose a new efficient segmentation decoder Lite Reduced Atrous Spatial Pyramid Pooling (LR-ASPP).We achieve new state of the art results for mobile classification, detection and segmentation.MobileNetV3-Large LRASPP is 34% faster than MobileNetV2 R-ASPP at similar accuracy for Cityscapes segmentation.This project tries to faithfully implement MobileNetV3 for real-time semantic segmentation, with the aims of being efficient, easy to use, and extensible.RequirementsThis code requires Python 3.7 or later. It has been tested to work with PyTorch versions 1.5 and 1.6. To install the package, simply runpip install fastseg. Then you can get started with a pretrained model:# Load a pretrained MobileNetV3 segmentation model in inference modefromfastsegimportMobileV3Largemodel=MobileV3Large.from_pretrained().cuda()model.eval()# Open a local image as inputfromPILimportImageimage=Image.open('street_image.png')# Predict numeric labels [0-18] for each pixel of the imagelabels=model.predict_one(image)More detailed examples are given below. As an alternative, instead of installingfastsegfrom pip, you can clone this repository and install thegeffnetpackage(along with other dependencies) by runningpip install -r requirements.txtin the project root.Pretrained Models and MetricsI was able to train a few models close to or exceeding the accuracy described in the originalSearching for MobileNetV3paper. Each was trained only on thegtFinelabels from Cityscapes for around 12 hours on an Nvidia DGX-1 node, with 8 V100 GPUs.ModelSegmentation HeadParametersmIoUInferenceTensorRTWeights?MobileV3LargeLR-ASPP, F=2563.6M72.3%21.1 FPS30.7 FPS✔MobileV3LargeLR-ASPP, F=1283.2M72.3%25.7 FPS37.3 FPS✔MobileV3SmallLR-ASPP, F=2561.4M67.8%30.3 FPS39.4 FPS✔MobileV3SmallLR-ASPP, F=1281.1M67.4%38.2 FPS52.4 FPS✔MobileV3SmallLR-ASPP, F=641.0M66.9%46.5 FPS61.9 FPS✔The accuracy is within0.3%of the original paper, which reported 72.6% mIoU and 3.6M parameters on the Cityscapesvalset. Inference was tested on a single V100 GPU with full-resolution 2MP images (1024 x 2048) as input. It runs roughly 4x faster on half-resolution (512 x 1024) images.The "TensorRT" column shows benchmarks I ran after exporting optimized ONNX models toNvidia TensorRTwith fp16 precision. Performance is measured by taking average GPU latency over 100 iterations.UsageRunning InferenceThe easiest way to get started with inference is to clone this repository and use theinfer.pyscript. For example, if you have street images namedcity_1.pngandcity_2.png, then you can generate segmentation labels for them with the following command.$pythoninfer.pycity_1.pngcity_2.pngOutput:==> Creating PyTorch MobileV3Large model
==> Loading images and running inference
Loading city_1.png
Generated colorized_city_1.png
Generated composited_city_1.png
Loading city_2.png
Generated colorized_city_2.png
Generated composited_city_2.pngOriginalColorizedCompositedTo interact with the models programmatically, first install thefastsegpackage with pip, as described above. Then, you can import and construct models in your own Python code, which are instances of PyTorchnn.Module.fromfastsegimportMobileV3Large,MobileV3Small# Load a pretrained segmentation modelmodel=MobileV3Large.from_pretrained()# Load a segmentation model from a local checkpointmodel=MobileV3Small.from_pretrained('path/to/weights.pt')# Create a custom model with random initializationmodel=MobileV3Large(num_classes=19,use_aspp=False,num_filters=256)To run inference on an image or batch of images, you can use the methodsmodel.predict_one()andmodel.predict(), respectively. These methods take care of the preprocessing and output interpretation for you; they take PIL Images or NumPy arrays as input and return a NumPy array.(You can also run inference directly withmodel.forward(), which will return a tensor containing logits, but be sure to normalize the inputs to have mean 0 and variance 1.)importtorchfromPILimportImagefromfastsegimportMobileV3Large,MobileV3Small# Construct a new model with pretrained weights, in evaluation modemodel=MobileV3Large.from_pretrained().cuda()model.eval()# Run inference on an imageimg=Image.open('city_1.png')labels=model.predict_one(img)# returns a NumPy array containing integer labelsassertlabels.shape==(1024,2048)# Run inference on a batch of imagesimg2=Image.open('city_2.png')batch_labels=model.predict([img,img2])# returns a NumPy array containing integer labelsassertbatch_labels.shape==(2,1024,2048)# Run forward pass directlydummy_input=torch.randn(1,3,1024,2048,device='cuda')withtorch.no_grad():dummy_output=model(dummy_input)assertdummy_output.shape==(1,19,1024,2048)The output labels can be visualized with colorized and composited images.fromfastseg.imageimportcolorize,blendcolorized=colorize(labels)# returns a PIL Imagecolorized.show()composited=blend(img,colorized)# returns a PIL Imagecomposited.show()Exporting to ONNXTheonnx_export.pyscript can be used to convert a pretrained segmentation model to ONNX. You should specify the image input dimensions when exporting. See the usage instructions below:$ python onnx_export.py --help
usage: onnx_export.py [-h] [--model MODEL] [--num_filters NUM_FILTERS]
                      [--size SIZE] [--checkpoint CHECKPOINT]
                      OUTPUT_FILENAME

Command line script to export a pretrained segmentation model to ONNX.

positional arguments:
  OUTPUT_FILENAME       filename of output model (e.g.,
                        mobilenetv3_large.onnx)

optional arguments:
  -h, --help            show this help message and exit
  --model MODEL, -m MODEL
                        the model to export (default MobileV3Large)
  --num_filters NUM_FILTERS, -F NUM_FILTERS
                        the number of filters in the segmentation head
                        (default 128)
  --size SIZE, -s SIZE  the image dimensions to set as input (default
                        1024,2048)
  --checkpoint CHECKPOINT, -c CHECKPOINT
                        filename of the weights checkpoint .pth file (uses
                        pretrained by default)Theonnx_optimize.pyscript optimizes exported models. If you're looking to deploy a model to TensorRT or a mobile device, you might also want to run it throughonnx-simplifier.Training from ScratchPlease see theekzhang/semantic-segmentationrepository for the training code used in this project, as well as documentation about how to train your own custom models.ContributionsPull requests are always welcome! A big thanks to Andrew Tao and Karan Sapra fromNVIDIA ADLRfor helpful discussions and for lending me their training code, as well as Branislav Kisacanin, without whom this wouldn't be possible.I'm grateful for advice from: Ching Hung, Eric Viscito, Franklyn Wang, Jagadeesh Sankaran, and Zoran Nikolic.Licensed under the MIT License.