Force single blob segmentation in U-Net model

AviZ: <p>I have a vanilla U-Net model for 2 classes image segmentation.<br>
The classes are background and a single type of object.</p>
<p>The object is always a <a href="https://en.wikipedia.org/wiki/Connected_space" rel="noopener nofollow ugc">connected set</a> of pixels.<br>
As a binary mask it is a single blob of pixels with no holes in it.</p>
<p>I wonder if there is a loss function to force such shape on the model output.<br>
As for now the output I get is sometimes several blobs.</p>
<p>I currently use the <a href="https://arxiv.org/abs/1706.05721" rel="noopener nofollow ugc">Tversky loss</a>.</p>
<p>I thought about adding a loss based on the maximum and minimum coordinates of the target blob and the output blob. Enforcing the rectangle. Is there a gradient friendly to do so?</p>

KFrank: <p>Hi Avi!</p>
<aside class="quote no-group" data-username="AviZ" data-post="1" data-topic="207673" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/e0b2c6/48.png" class="avatar"> AviZ:</div>
<blockquote>
<p>The object is always a <a href="https://en.wikipedia.org/wiki/Connected_space" rel="noopener nofollow ugc">connected set</a> of pixels.<br>
As a binary mask it is a single blob of pixels with no holes in it.</p>
</blockquote>
</aside>
<p>As you describe it, you might choose to treat your problem as one<br>
of <em>instance segmentation,</em> rather than just semantic segmentation.</p>
<aside class="quote no-group">
<blockquote>
<p>I wonder if there is a loss function to force such shape on the model output.</p>
</blockquote>
</aside>
<p>I don’t think that there is a natural loss function* for this kind of thing.</p>
<p>However, you might be able to contrive something.  For example, you<br>
might say that if a pixel is predicted to be a background pixel, but some<br>
of the pixels on either side of it are foreground pixels, it should be given<br>
some additional penalty for being classified as background.  But this is<br>
exactly the kind of thing – the status of adjacent pixels –  that should be<br>
able to be learned by a semantic-segmentation model (such as U-Net).</p>
<aside class="quote no-group">
<blockquote>
<p>As for now the output I get is sometimes several blobs.</p>
</blockquote>
</aside>
<p>Is your use case such that each image has exactly one blob in it?<br>
If so, predicting a single blob might be as simple as <a href="https://en.wikipedia.org/wiki/Closing_(morphology)" rel="noopener nofollow ugc">closing</a> (dilating<br>
and then eroding) the prediction of your U-Net.  (Note, that although<br>
quite simple, this is a version of semantic segmentation followed by<br>
post-processing.)</p>
<p>If a single image can contain multiple separate blobs, you’re really<br>
now looking at an <em>instance segmentation</em> problem and you should<br>
consider instance-segmentation algorithms.</p>
<p>If the individual blobs are well-separated, then post-processing with<br>
a closing operation could still work well.</p>
<p>Note that there is a class of instance-segmentation algorithms<br>
that rely on “enhanced” semantic segmentation followed by<br>
post-processing (for example, <a href="https://arxiv.org/pdf/1806.03535" rel="noopener nofollow ugc">StarDist,</a> which uses a U-Net for<br>
its initial semantic-segmentation).</p>
<aside class="quote no-group">
<blockquote>
<p>I currently use the <a href="https://arxiv.org/abs/1706.05721" rel="noopener nofollow ugc">Tversky loss</a>.</p>
</blockquote>
</aside>
<p>I would recommend using <code>BCEWithLogitsLoss</code> (possibly with<br>
<code>pos_weight</code>, if you have a significant foreground-background<br>
imbalance) and only add something like Tversky loss to it if you<br>
can show that doing so improves performance.</p>
<p>*)  You could look a adding a <a href="https://proceedings.neurips.cc/paper_files/paper/2019/file/2d95666e2649fcfc6e3af75e09f5adb9-Paper.pdf" rel="noopener nofollow ugc">topological-loss</a> term to your overall<br>
loss, but this might be overkill, and is something of a nuisance to<br>
implement.</p>
<p>Best.</p>
<p>K. Frank</p>

AviZ: <p>Indeed, in my case it is assured there is a single blob.<br>
Are there models specialized in a single instance segmentation?</p>
<p>Basically I’d like o to have a bounding box.<br>
Yet I don’t want a fully featured object detection model with anchors and grid.<br>
Just a simple loss that will also punish for the bounding box being different than the bounding box of the reference mask.</p>

KFrank: <p>Hi Avi!</p>
<aside class="quote no-group" data-username="AviZ" data-post="3" data-topic="207673" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/e0b2c6/48.png" class="avatar"> AviZ:</div>
<blockquote>
<p>Indeed, in my case it is assured there is a single blob.<br>
Are there models specialized in a single instance segmentation?</p>
</blockquote>
</aside>
<p>If you know that each image contains exactly one blob (that is, one<br>
instance), then semantic segmentation is, by definition, instance<br>
segmentation.  So a semantic-segmentation model (such as U-Net)<br>
is just what you want.</p>
<aside class="quote no-group">
<blockquote>
<p>Basically I’d like o to have a bounding box.</p>
</blockquote>
</aside>
<p>If you get your predicted pixels right, the bounding box that you trivially<br>
compute from those pixels will be right.  So, focus on getting your<br>
semantic-segmentation to work well.</p>
<aside class="quote no-group">
<blockquote>
<p>Just a simple loss that will also punish for the bounding box being different than the bounding box of the reference mask.</p>
</blockquote>
</aside>
<p>I don’t understand why you are so focused on the bounding box.  It’s<br>
not really fundamental – it’s more a technical ingredient of certain<br>
detection and instance-segmentation algorithms.</p>
<p>Suppose that the bounding box you compute from your predicted pixels<br>
is, say, two or three pixels too high (but that the predicted pixels are<br>
altogether quite accurate).  Now suppose that you add a bounding-box<br>
loss term to your overall loss and now your predicted bounding box is<br>
exactly right, but, say, 500 pixels (that all happen to lie within the<br>
correct bounding box) are predicted incorrectly.  Is that really a better<br>
outcome for your use case?</p>
<p>When you change your loss to make some aspect of your predictions<br>
better, you will – all else being equal – make some other aspect of your<br>
predictions worse.  After all, you’ve told your model – via your choice of<br>
loss function – that you care more about the precise bounding box,<br>
rather than the specific pixels.</p>
<p>Best.</p>
<p>K. Frank</p>

AviZ: <p><a class="mention" href="/u/kfrank">@KFrank</a> , the problem I face is I have many small blobs outside of the area of the object.<br>
Those are false alarms.</p>
<p>If I make the model limit itself to a rectangle it will mitigate the issue.<br>
I don’t want a fully featured object detector, just a smart and effective way to connect the bounding box of the prediction to the bounding box of the ground truth labels.</p>
<p>Is there a way to define such loss which is gradient friendly?<br>
Given the output mask of the model, how can I build a regression model for the bounding box?</p>

KFrank: <p>Hi Avi!</p>
<aside class="quote no-group" data-username="AviZ" data-post="5" data-topic="207673" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/e0b2c6/48.png" class="avatar"> AviZ:</div>
<blockquote>
<p>the problem I face is I have many small blobs outside of the area of the object.</p>
</blockquote>
</aside>
<p>Well, you haven’t told us anything about your actual use case, so it’s<br>
hard to say …</p>
<aside class="quote no-group">
<blockquote>
<p>Those are false alarms.</p>
</blockquote>
</aside>
<p>It sounds like you have a large “blob” (whatever that means in your<br>
actual use case …) that you wish to segment as foreground and<br>
several smaller blobs, which, other than their sizes, look similar to the<br>
large blob, that you wish to segment as background.</p>
<p>Note that a U-Net model has a “field of view.”  That is, any single pixel<br>
in the output depends <em>only</em> on the input pixels that are in the output<br>
pixel’s field of view.  For example, roughly speaking, if your U-Net<br>
has five layers that downsample by a factor of two, any given output<br>
pixel will only depend on a 32x32 patch of input pixels centered on<br>
the output pixel.</p>
<p>If your field of view is not large enough to more-or-less contain the<br>
largest of your “small” background blobs (and, except for their sizes,<br>
foreground blobs and background blobs look the same), then your<br>
U-Net won’t have enough information <em>in it’s field of view</em> to distinguish<br>
small blobs from large.</p>
<p>You can increase the field of view of your U-Net by adding more<br>
downsampling layers or by downsampling by a larger factor in one or<br>
more of your layers.  For example, if you downsampled by a factor of<br>
four in five layers, your field of view would become, roughly speaking,<br>
1024x1024.</p>
<aside class="quote no-group">
<blockquote>
<p>Is there a way to define such loss which is gradient friendly?<br>
Given the output mask of the model, how can I build a regression model for the bounding box?</p>
</blockquote>
</aside>
<p>The standard approach for regressing bounding boxes is to use fixed<br>
anchors and regress differentiable offsets.  I don’t know of a slick way<br>
of doing anything simpler.</p>
<p>But it sounds like you are trying to reinvent the wheel.  Whether your<br>
problem is technically instance segmentation, with your small blobs<br>
vs. large blobs, it is certainly similar.  So if you want something akin to<br>
instance segmentation and you want to regress bounding boxes, you<br>
ought to look at bounding-box-based instance segmentation models.<br>
Mask R-CNN is one such.</p>
<p>Best.</p>
<p>K. Frank</p>

AviZ: <p><a class="mention" href="/u/kfrank">@KFrank</a> ,</p>
<p>Appreciate the good dialogue.</p>
<p>I am aware of the concept of receptive field in CNN.<br>
I don’t think it is the concept in my case.</p>
<p>Problem wise, you may think I want to segment a single object within a frame.<br>
Namely, I know in prior it appears once in the image.</p>
<p>An example, segment a t-shirt in an image.<br>
The shirt only happens once in the image.</p>
<p>Yet the U-Net segment most of yet also small, false alarm, blobs in the background.</p>
<p>I thought about anchors.<br>
Can I set a single one which is the center of mass of the largest blob of the U-Net output?</p>
<p>The motivation is to take advantage of the strong prior I have:</p>
<ol>
<li>The blob appears only once in each image.</li>
<li>The object to segment is a connected set.</li>
</ol>
<p>Taking advantage of those 2 assumption should make results better.</p>