Adding classification to segmentation model approach

username: aksg87
cooked: <p>Hello all,</p>
<p>If you have a model that segments a structure, like kidneys in this case:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/1/b12ba9f30c0dbf2e50e876040c473b00d22d2b02.png" alt="image" data-base62-sha1="phk3P3khTxpFktgMJiNcbyKuhX4" width="233" height="124"></p>
<p>What is a good approach for eliminating images that have a few positive pixels but no kidneys? A simple approach is thresholding the number of pixels but I figured I would get better results with a classifier on top of the segmentation model.</p>
<p>In this case, I was thinking of training the segmentation models on images with just kidneys and then training the classifier on the normal distribution of images (with and without kidneys) using the first segmentation model. I am just questioning if I should use the segmentation model as an encoder only, or just take the final output of the segmentation model and add a few dense layers on top?</p>
<p>Thoughts or feedback appreciated!</p>
<p>My model setup for reference:</p>
<pre data-code-wrap="python"><code class="lang-python"># https://github.com/qubvel/segmentation_models.pytorch
ENCODER = 'efficientnet-b6'
ENCODER_WEIGHTS = 'imagenet'
ACTIVATION = None 
DEVICE = 'cuda'
n_class = 2

    model = smp.Unet(
        encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, classes=n_class, activation=ACTIVATION,
    )</code></pre>

username: KFrank
cooked: <p>Hello Akshay!</p>
<aside class="quote no-group" data-username="aksg87" data-post="1" data-topic="81212" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aksg87/48/15602_2.png" class="avatar"> aksg87:</div>
<blockquote>
<p>What is a good approach for eliminating images that have a few positive pixels but no kidneys? A simple approach is thresholding the number of pixels but I figured I would get better results with a classifier on top of the segmentation model.</p>
</blockquote>
</aside>
<p>My intuition tells me that <em>if</em> your model does a good job of segmenting<br>
the kidney images, then thresholding the count of predicted kidney<br>
pixels should do an essentially perfect job of classifying kidney vs.<br>
non-kidney.</p>
<p>But, with the following qualification:</p>
<p>If you train your segmenter on non-kidney images (no kidney pixels)<br>
as well as kidney images, and it segments non-kidney images<br>
successfully, predicting no, or just a few, kidney pixels, thresholding<br>
should work very well.</p>
<p>But if you only train your segmenter on kidney images, I could see<br>
the segmenter becoming strongly biased to predict that an image<br>
has many kidney pixels (because all of your training images do),<br>
and therefore predicting a lot of kidney pixels in non-kidney images.</p>
<p>[quote]<br>
In this case, I was thinking of training the segmentation models on images with just kidneys<br>
[/quote</p>
<p>Regardless of which path you take, I would lean against training only<br>
with images of kidneys.  As a general rule, you’re better off training with<br>
the full spectrum of data your model will be seeing in its final application<br>
(if you can).</p>
<aside class="quote no-group">
<blockquote>
<p>and then training the classifier on the normal distribution of images (with and without kidneys) using the first segmentation model. I am just questioning if I should use the segmentation model as an encoder only, or just take the final output of the segmentation model and add a few dense layers on top?</p>
</blockquote>
</aside>
<p>I do think post-training a classifier on top of your pre-trained segmenter<br>
makes sense.  I don’t have intuition about whether you should graft your<br>
classifier onto you segmenter at the “encoder” level or at the final layer.</p>
<p>One thing you might want to experiment with is training just your<br>
classifier weights vs. training all of the weights, including those that<br>
you had pre-trained for the segmenter.  In the latter case, it might<br>
make sense to first freeze your segmenter weights while you partially<br>
train your classifier weights to get them in the right ball park, and then<br>
fine tune your full-model weights, training both the segmenter and<br>
classifier weights jointly.</p>
<p>(Of course if it were me, I would go with thresholding.  Do you have<br>
any examples where thresholding misclassifies?)</p>
<p>Good luck.</p>
<p>K. Frank</p>