Extract only mask from the output

[SahilChachra]

First of all, thanks for the code and the model.
Question - How can I extract/save only the mask (RGB, without any label text on it) during inference?
Which function shall I modify to get the RGB mask?

Thanks!