Package: dextr
PyTorch implementation of DEXTRAn implementation ofDEXTR.
The original implementation can be found athttps://github.com/scaelles/DEXTR-PyTorch.This implementation is intended for use as a library.Installation> pip install dextrPython Inference APISeedemo.pyfor an example of using thedextrinference API.We have trained a ResNet-101 based U-Net DEXTR model on the Pascal VOC 2012 training set. You can download ithere.You can load this model -- downloading it automatically -- like so:fromdextr.modelimportDextrModel# Load the model (automatically downloads if necessary)# You can also provide a `map_location` paramter to load it onto a specific devicemodel=DextrModel.pascalvoc_resunet101()Alternatively you can load a model that you have trained yourself from a file:MODEL_PATH='...'dextr_model=torch.load(MODEL_PATH,map_location='cuda:0')Use thepredictmethod to predict a mask for an object in an image, identified by its extreme points:mask=dextr_model.predict([image],[extreme_points])[0]You can perform inference on multiple images with one call.
TheDextrModel.predictmethod takes a list
of images and extreme points as either a list of(4, [y, x])NumPy arrays or
one(N, 4, [y, x])shaped NumPy array.The images that you use as input can take the form of either NumPy arrays or PIL Images. Each image should
have a corresponding list of four extreme points. It returns a list of masks; each
mask is the same size as the corresponding input image:Training using the command linetrain_dextr.pyprogramTrain a DEXTR network using the Pascal VOC datasetThis will train a DEXTR model using aU-Netwith a ResNet-101 based encoder.
It should take several hours on an nVidia 1080Ti GPU.Download the Pascal VOC 2012 datasetdevelopment kitCreate a file calleddextr.cfgwith the following contents:[paths]pascal_voc=<path to VOC2012 diretory>Train the DEXTR model by running:> python train_dextr.py pascal_resunet101 --dataset=pascal_voc --arch=resunet101The namepascal_resunet101is the name of the job; STDOUT will be logged tologs/log_pascal_resunet101.txtand the model
file will be saved tocheckpoints/pascal_resunet101.pth. You can give the job any name you like.Fine tuning a DEXTR network using a custom data setThere are two types of data set you can use:Each input image has a corresponding label image, where label images have an integer pixel type such that each
pixel gives the index of the object that covers it, or 0 for background. The Pascal VOC dataset is arranged in
this way.Each input image has a corresponding set of mask images that form a stack. Each mask image is an 8-bit greyscale
image that corresponds to an object/instance and identifies the pixels covered by it.Please arrange your custom data set so that the image file names (excluding extension) match or are a prefix
to the label/mask image file names. E.g. the imageimg0.jpgwill match the label fileimg0.pngorimg0_labels.png.
For mask stack datasetsimg0.jpgwould match to the mask imagesimg0_mask0.png, ...img0_maskN.png.
The images and labels can live in separate directories; they are matched by filenameonly.In these examples, we assume that you have downloaded the pre-trained DEXTR model linked above.Training using a label image data set> python train_dextr.py my_model_from_labels --dataset=custom_label --train_image_pat=/mydataset/train/input/*.jpg --train_target_pat=/mydataset/train/labels/*.png --arch=resunet101 --load_model=dextr_pascalvoc_resunet101-a2d81727.pthThe input and label images are given to the--train_image_patand--train_target_patoptions.
You can specify validation images using the--val_image_patand--val_target_patoptions in a similar way.--load_model=dextr_pascalvoc_resunet101-a2d81727.pthindicates that we should start by loading the
model trained on Pascal VOC above and fine-tune it, rather than starting from an ImageNet classifier.You can specify that the label index 255 should be ignore by adding--label_ignore_index=255.You could train using the entire (train and validation) Pascal VOC data set using:> python train_dextr.py my_model_from_pascal --dataset=custom_label --train_image_pat=/pascal/VOC2012/JPEGImages/*.jpg --train_target_pat=/pascal/VOC2012/SegmentationObjects/*.png --label_ignore_index=255 --arch=resunet101Training using a mask stack data set> python train_dextr.py my_model_from_masks --dataset=custom_mask --train_image_pat=/mydataset/train/input/*.jpg --train_target_pat=/mydataset/train/masks/*.png --arch=resunet101 --load_model=dextr_pascalvoc_resunet101-a2d81727.pthPython training APIThetraining_loopfunction within thedextr.modelmodule provides a simple training loop that can be used
for training or fine-tuning models. Seetrain_dextr.pyfor usage.