Where can I find the resulting image of semantic segmentation?

[Kenta-b]

Where can I find the code to display the results of semantic segmentation? I want to save the predicted images together.

[ethanwharris]

Hi @Kenta-b thanks for the question! The conversion from labels to image happens here: https://github.com/PyTorchLightning/lightning-flash/blob/master/flash/image/segmentation/serialization.py You could probably convert the predicted labels to an image like this:

from flash.image.segmentation.serialization import SegmentationLabels

preds = model.predict(...)
pred_tensor = torch.tensor(preds[0][0])
image = SegmentationLabels.labels_to_image(pred_tensor, SegmentationLabels.create_random_labels_map(num_classes))

This is definitely something we can do better. Ideally there would be a serializer which just gives back segmentation masks as tensors or PIL images, would you be interested in making a PR that implements this?

Thanks 😃

[Kenta-b]

Thank you!
I was able to automate it safely.

I'm still a student, so I don't have the skills to make PR.