Custom UNet segmentation model uses only 2 colors for output map

[caiusdebucean]

Hello,
I have a segmentation UNET model that I have trained with TAO on 3 classes: "building", "floor", "other". I used:

• toolkit_version: 3.22.05
• deepstream_version: 6.0
• jetpack_version: 4.6.1-b110

I have already modified the model to train on 3 classes instead of 2, the evaluation worked and gave satisfactory results on the validation and test data, I have modified the config file and the labels.txt file for this repo, but the output is always binary (e.g. using 2 colors).

This is the specs file used for training:

random_seed: 42
model_config {
  model_input_width: 320
  model_input_height: 320
  model_input_channels: 3
  num_layers: 18
  all_projections: true
  arch: "resnet"
  use_batch_norm: False
  training_precision {
    backend_floatx: FLOAT32
  }
}

training_config {
  batch_size: 32
  epochs: 50
  log_summary_steps: 10
  checkpoint_interval: 5
  loss: "cross_dice_sum"
  learning_rate:0.0001
  regularizer {
    type: L2
    weight: 2e-5
  }
  optimizer {
    adam {
      epsilon: 9.99999993923e-09
      beta1: 0.899999976158
      beta2: 0.999000012875
    }
  }
}

dataset_config {
  dataset: "custom"
  augment: False
  augmentation_config {
    spatial_augmentation {
    hflip_probability : 0.5
    vflip_probability : 0.5
    crop_and_resize_prob : 0.5
  }
  brightness_augmentation {
    delta: 0.2
  }
}
input_image_type: "color"
train_images_path:"/home/jetson/Downloads/test_img/images/train"
train_masks_path:"/home/jetson/Downloads/test_img/masks/train"

val_images_path:"/home/jetson/Downloads/test_img/images/val"
val_masks_path:"/home/jetson/Downloads/test_img/masks/val"

test_images_path:"/home/jetson/Downloads/test_img/images/test"

data_class_config {
  target_classes {
    name: "building"
    mapping_class: "building"
    label_id: 1
  }
  target_classes {
    name: "floor"
    mapping_class: "floor"
    label_id: 2
  }
    target_classes {
    name: "other"
    mapping_class: "other"
    label_id: 0
  }
}
}

The evaluation process went good, providing a Segmentation accuracy of 87% on validation data.
I have converted the model using the following command for the tlt-converter:

./tao-converter -d 3,320,320 -k nvidia_tlt -e /home/jetson/Downloads/segmentator.engine -t fp16 -p input_1:0,1x3x320x320,4x3x320x320,16x3x320x320 /home/jetson/Downloads/model_segmentator.etlt

Also, I am using this config file for the ./ds-tao-segmentation:

[property]
gpu-id=0
net-scale-factor=0.007843
model-color-format=0
offsets=127.5;127.5;127.5
labelfile-path=./unet_labels.txt
##Replace following path to your model file
model-engine-file=/home/jetson/Downloads/segmentator.engine
#current DS cannot parse onnx etlt model, so you need to
#convert the etlt model to TensoRT engine first use tao-convert
tlt-encoded-model=/home/jetson/Downloads/model_segmentator.etlt
tlt-model-key=tlt_encode
infer-dims=3;320;320
batch-size=1
## 0=FP32, 1=INT8, 2=FP16 mode
network-mode=2
num-detected-classes=3
interval=0
gie-unique-id=1
network-type=2
output-blob-names=softmax_1
segmentation-threshold=0.0
##specify the output tensor order, 0(default value) for CHW and 1 for HWC
segmentation-output-order=1

[class-attrs-all]
roi-top-offset=0
roi-bottom-offset=0
detected-min-w=0
detected-min-h=0
detected-max-w=0
detected-max-h=0

with the labels.txt file as:

other
building
floor

I am running inference using this command:

/ds-tao-segmentation -c ../../configs/unet_tao/pgie_unet_tao_config.txt  -i ../../../dashcam.h264 -b 2 -d

For visualization, however, I get a binary segmentation mask instead of a segmentation maps with 3 colors.:

Could you kindly help me how I can:

1. make the segmentation mask centered and cropped only to its dimensions
2. have a segmentation map that uses more than 2 colors
3. what does the model-color-format parameter from the config file do? When changing from 0 to 1 I get a slightly more complete segmentation map, with the same colors, as seen below:
   

All the best!

[gitmhg]

yes i get same answer,did you have good idea