Add support for probability map in 'segment'. Only placeholders for now