Wrong number of classes in visualize_segmentation

[gaetanmuck]

trafficstars

Hello there, I think I have spotted a small mistake in the code.

TL;DR:

n_classes should be np.max(seg_arr) + 1 and not np.max(seg_arr)

How to see the bug:

Visualize an image without specifying the class number: Observe a color on the image that is not in the legend (black). It is 'normal' because in get_colored_segmentation_image, not found classes number are set to [0, 0, 0]. Once the class number is correctly set (does not use the default code) then it works as expected.

Where is it?

File: https://github.com/divamgupta/image-segmentation-keras/blob/master/keras_segmentation/predict.py Function: visualize_segmentation Line: 104 Found: n_classes = np.max(seg_arr) Should: n_classes = np.max(seg_arr) + 1

Hope it will help Cheers!