Solved by ptrblck in post #2 
              
              
                This sounds rather like a segmentation use case, where each pixel will get a class prediction. 
torchvision provides a few segmentation models, which you could try out for your use case. 
Let me know, if I misunderstood your question.