Classify Pixels Using Deep Learning - Probability Rasters

[MadelineHayes1]

I have a trained U-Net model that I am using to classify wetland pixels. I used the API to train the model, and it is saved as an .emd and .dlpk. I am using the "Classify Pixels Using Deep Learning" tool in ArcGIS Pro, version 3.4.2. I would like probability rasters as well as hard classifications as the output when running this tool. Based on the documentation, there is an argument called "return_probability_raster" that would provide the probabilities instead of the binary classifications. However, it looks like this argument is only valid for MultiTaskRoadExtractor, ConnectNet, BDCNEdgeDetector, HEDEdgeDetector, and MMSegmentation. 

Can I get probability rasters from a U-Net architecture? I don't understand why it would be limited to the architectures mentioned above, but I can't find a way around that. 

[PatrikOlsson]

I have exactly the same problem. Chatgpt and I really tried yesterday to solve this. I changed .dlpk to .zip so I could access the "baked in" emd file. Adjusted emd ("SupportsProbability": true) and some other stuff, replaced the old emd file with the new one. re-zipped and changed the file extension back to .dlpk. But still no progress. I also passed the argument "return_probability_raster True". The classifying 1 or 0 works just fine with excellent result. But I want a probability for each pixel and not just an binary result

[ShivaniPathak]

Hi @MadelineHayes1  & @PatrikOlsson,

Thank you for your question. You're correct that the return_probability_raster parameter is currently supported only by a limited set of models such as MultiTaskRoadExtractor, ConnectNet, BDCNEdgeDetector, and HEDEdgeDetector. These models are primarily designed for tasks where a probability raster makes more intuitive sense—especially when they generate binary outputs like edge detection or presence/absence classification.

At this time, U-Net models trained using the ArcGIS API do not support generating probability rasters through the “Classify Pixels Using Deep Learning” tool. This limitation exists because models that support multiclass classification, such as U-Net, are currently designed to output hard class labels rather than class-wise probabilities. In contrast, binary models like HEDEdgeDetector or ConnectNet can output a single-channel probability raster indicating the likelihood of a pixel belonging to the positive class.

That said, if you could share more about how the probability raster would benefit your workflow—for instance, if you’re planning to apply a custom threshold, perform uncertainty analysis, or post-process soft predictions—we’d be happy to take that feedback back to our development team for consideration in future updates.

Let me know how you’re planning to use the probabilities, and I’d be glad to follow up!