Custom Segmentation Model for Transfer Learning

I’ve trained a 1 class UNet model that can be detect car shapes for segmentation properly, however I need a 10 class classification process with the same dataset. I’m getting good results from only segmentation, so I’m trying to implement this model allows to make classification too.

I labeled my first dataset as I need to, so is there any way to merge this labeled dataset and 1 class segmentation model together and train them properly ?

Some models such as MaskRCNN output segmentation masks as well as classification results for the objects, so you might take a look at these implementations and check, if you could use a similar workflow for your use case.

You can get two outputs by insert an output from your encoder of UNet, and fully connect them , add two or three dense layers, and output 10 classes as you want.

Thanks for your replies, I really appreciate that.

You can get two outputs by insert an output from your encoder of UNet, and fully connect them , add two or three dense layers, and output 10 classes as you want.

Freeze the whole network and train the dense layers which is added to end of UNet ?

Some models such as MaskRCNN output segmentation masks as well as classification results for the objects, so you might take a look at these implementations and check, if you could use a similar workflow for your use case.

My labels and UNet class parameter is 1, because of that I get only background segmentation. I wonder which one is the more logical approach then the other one; add a layer to end of the network, freeze except these dense layers and train them only to get more classes or training same model with more classes with multi-labeled dataset ?