Generate Semantic Segmentation Output mask

[shreshtashetty]

Hi
I was trying to carry out inference from the pre-trained model following the ros2 deployment instructions. I was able to execute that successfully.
I want to enable the semantic segmentation output as well, so to enable that I followed below steps.

1. Generated .onnx file using this command with --enable-semantic flag.

python3 onnx_conversion.py -p <checkpoint> -o <onnx_file_path> --enable-semantic

and converted it to trt.engine file

2. Then checked the dimensions for the segmentation mask using the output of ONNX file

for output in onnx_model.graph.output:
    print(f"Output Name: {output.name}, Shape: {output.type.tensor_type.shape}")

3. After that modified the x_mobility_navigator.py file and made below modifications
   a. Initialized the variable here self.segmentation_output = np.zeros((1, 1, 7, 320, 512), dtype=np.float32)
   b. Then in this function added cuda.mem_alloc, cuda.memcpy_htod, cuda.memcpy_dtoh for the same

4. When I ran the inference I got the masks but the values are zero

• I am getting the right shape (1, 1, 7, 320, 512)
• but all the values are zero

Segmentation Output: [[[[[0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    ...
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]
[x_mobility_navigator-1]     [0. 0. 0. ... 0. 0. 0.]]]]]

Any suggestions and swift response will be highly appreciated.
Thanks in advance, and anything to be kept in mind as I am planning on generating the dataset and train it for my use-case with 5 classes.

[shreshtashetty]

#Update
I got the semantic segmentation to work but the output is not what I expected, but the code had no documentation so wasn't able to make out how to interpret the output.

Segmentation Output: [[[[[ 7.4267483   9.66694    10.827346   ... 14.840271   12.94972
[x_mobility_navigator-1]      10.131485  ]
[x_mobility_navigator-1]     [ 8.763442   10.688378   11.996544   ... 17.369646   14.8487425
[x_mobility_navigator-1]      11.381178  ]
[x_mobility_navigator-1]     [ 9.512807   11.547724   12.91613    ... 17.478561   14.99352
[x_mobility_navigator-1]      12.375072  ]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [-4.6494937  -4.6010637  -4.9401145  ...  0.6400943   0.7787968
[x_mobility_navigator-1]       1.1041331 ]
[x_mobility_navigator-1]     [-3.663425   -3.9732387  -4.097291   ...  0.4910759   0.9742001
[x_mobility_navigator-1]       0.88267493]
[x_mobility_navigator-1]     [-0.07677049  1.0375022   1.20701    ...  3.3003602   2.8207648
[x_mobility_navigator-1]       2.8552713 ]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[-0.8575109  -0.664371   -1.8897895  ... -3.33995    -2.7081063
[x_mobility_navigator-1]      -1.2290901 ]
[x_mobility_navigator-1]     [-3.1589923  -2.970861   -3.55552    ... -5.522795   -4.6392198
[x_mobility_navigator-1]      -1.867524  ]
[x_mobility_navigator-1]     [-3.614818   -3.4836454  -4.3968067  ... -8.575841   -7.8700027
[x_mobility_navigator-1]      -4.3531966 ]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [10.573848   12.645357   15.283156   ... 15.7275305  14.076453
[x_mobility_navigator-1]      11.784009  ]
[x_mobility_navigator-1]     [ 9.887658   12.159548   13.885411   ... 13.1086855  11.520699
[x_mobility_navigator-1]       9.84858   ]
[x_mobility_navigator-1]     [ 9.3639555  11.172118   11.986088   ... 10.955655   10.139017
[x_mobility_navigator-1]       9.041935  ]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[-0.68744195 -0.29546118  0.3909457  ...  3.0548515   2.8289714
[x_mobility_navigator-1]       1.1503994 ]
[x_mobility_navigator-1]     [ 2.027917    2.4642344   2.604083   ...  5.81062     5.136679
[x_mobility_navigator-1]       2.856278  ]
[x_mobility_navigator-1]     [ 1.2630867   1.2173995   1.7031786  ...  4.13748     3.2807117
[x_mobility_navigator-1]       0.9366235 ]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [-4.189419   -3.5965734  -3.8184323  ... -6.0579195  -5.2113295
[x_mobility_navigator-1]      -3.168498  ]
[x_mobility_navigator-1]     [-4.043027   -3.1924691  -3.2702775  ... -4.238928   -3.3972902
[x_mobility_navigator-1]      -2.4845042 ]
[x_mobility_navigator-1]     [-6.076063   -6.2295537  -5.957752   ... -7.3533473  -6.4057684
[x_mobility_navigator-1]      -4.3637056 ]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    ...
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[-1.7212461  -2.7670915  -2.6434405  ... -3.854366   -4.32727
[x_mobility_navigator-1]      -3.9652631 ]
[x_mobility_navigator-1]     [-2.15209    -3.6468601  -3.4497328  ... -5.7337217  -6.2256575
[x_mobility_navigator-1]      -5.5341763 ]
[x_mobility_navigator-1]     [-2.1236742  -4.1555095  -4.257398   ... -5.676658   -5.8456745
[x_mobility_navigator-1]      -4.9922748 ]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [-0.3886564   0.9593225  -0.40121654 ... -1.9532456  -1.9088852
[x_mobility_navigator-1]      -1.9303128 ]
[x_mobility_navigator-1]     [-1.0288689   1.4890827   0.65898454 ... -0.33653173 -0.35149306
[x_mobility_navigator-1]      -1.2426891 ]
[x_mobility_navigator-1]     [-2.9272106  -0.9635485  -1.3262808  ... -2.028891   -2.147013
[x_mobility_navigator-1]      -2.153096  ]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[-4.2822385  -4.522312   -4.928849   ... -3.927567   -4.0011864
[x_mobility_navigator-1]      -4.5269127 ]
[x_mobility_navigator-1]     [-4.038172   -5.0027504  -5.746505   ... -5.381513   -5.5103507
[x_mobility_navigator-1]      -5.252159  ]
[x_mobility_navigator-1]     [-3.7158968  -5.022142   -5.6779113  ... -3.8676746  -3.9631667
[x_mobility_navigator-1]      -4.30433   ]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [-2.7755098  -1.6228032  -1.2014444  ... -2.1661494  -2.0145369
[x_mobility_navigator-1]      -2.9842207 ]
[x_mobility_navigator-1]     [-2.885616   -2.1845543  -1.9621587  ... -2.4616888  -2.6872191
[x_mobility_navigator-1]      -3.2258584 ]
[x_mobility_navigator-1]     [-4.3623447  -4.44032    -4.466794   ... -4.350885   -4.721442
[x_mobility_navigator-1]      -4.6268597 ]]
[x_mobility_navigator-1] 
[x_mobility_navigator-1]    [[-0.7530664  -1.0236669  -1.6409763  ...  1.7984413   1.5678668
[x_mobility_navigator-1]       2.541337  ]
[x_mobility_navigator-1]     [-1.5141847  -1.1958674  -2.2503934  ...  2.1389194   1.6713094
[x_mobility_navigator-1]       0.87952703]
[x_mobility_navigator-1]     [-1.8632065  -1.3333427  -1.7088321  ...  1.166529    0.789874
[x_mobility_navigator-1]       0.27983984]
[x_mobility_navigator-1]     ...
[x_mobility_navigator-1]     [-3.020556   -3.851541   -3.7923927  ... -3.96487    -2.237349
[x_mobility_navigator-1]      -1.7071831 ]
[x_mobility_navigator-1]     [-2.3813386  -3.120698   -3.213468   ... -3.0247722  -1.2806818
[x_mobility_navigator-1]      -1.2975016 ]
[x_mobility_navigator-1]     [-1.2004359  -1.5602779  -1.7289892  ... -1.7851927  -0.08681303
[x_mobility_navigator-1]      -0.47794083]]]]]

I expected as this is a pixel wise classification I thought the output would be some sort of pixel wise softmax or something. but this is very ambiguous the values are ranging from anywhere from -35 to +35 and seeing model code there is no activation function, so this could easily be raw convolution outputs (ideally there should be pixel wise something like softmax to generate these masks, but looking in code I think that is not the case) so executing the ros examples, if I normalize this output (mask.fromarray(mask.astype(np.uint8) * 255)) I got below outputs. which I am not able to figure out how to infer.

Image

Class - 0

Class - 1

Class -2

Class - 3

Class - 4

Class - 5

Class - 6

Hoping for some clarification and things to be kept in mind to train it for our use case with different number of classes.
Thanks in advance

[nv-liuw]

Hi @shreshtashetty, right, these are raw values from StyleGan. To visualize the semantic segmentation, we need np.argmax with something like below:

SEMANTIC_COLORS = np.array(
    [
        [128, 128, 128],    # Background
        [0, 255, 0],    # NavigableSurface
        [255, 165, 0],    # Forklift
        [0, 0, 255],    # Pallet
        [255, 255, 0],    # Cone
        [255, 0, 255],    # Sign
        [255, 0, 0]    # Fence
    ],
    dtype=np.uint8)


def publish_semantic(self):
      semantic_image = SEMANTIC_COLORS[np.argmax(self.semantic, axis=0)]
      semantic_ros_image = self.cv_bridge.cv2_to_imgmsg(semantic_image, encoding="rgb8")
      self.semantic_publisher.publish(semantic_ros_image)

[rvnandwani]

Hi @nv-liuw
Something on the similar lines, so in literature it is mentioned that the model follows "multi-stage training pipeline" so my understanding was that first we train the model for world understanding i.e. through random actions then using that we train the action policy for generating actions using that embedding. But the train script usage that is highlighted here Steps 3,4 - the model is trained for just world model or both world as well as action policy.

Then how to carry out multi-stage training and if that is actually helpful ?
And is there any provision of transfer learning, or bare minimum learning from a checkpoint, so that if we analyzed the model needs more data or more epochs we are not training it from scratch, we are training it from a trained checkpoint

[nv-liuw]

@rvnandwani Right, in stage one, the random action dataset is used to train only the world model part, and on stage two, both world model and action policy are trained together.

To carry out multi-stage training, we can train the world model first, then use the result checkpoint for stage two.