UNET Semantic Segmentation: where resize mask shape if input shape != mask shape?

alessandro_mondin: <p>Hello all,</p>
<p>I am trying to code from scratch the UNET where input dimensions in the paper are:<br>
input.shape = (batch_size, 3, 572, 572) and out.shape = (batch_size, 1, 388, 388).<br>
link here: <a href="https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/" rel="noopener nofollow ugc">UNET_PAPER</a><br>
In this case which is the best practice to resize out.shape?<br>
It would be better to resize it in the Dataset class and load targets already of shape==out.shape or I can just resize the target in the forward method of the UNET? Is resizing a cheap operation or this choice might heavily slow down the computations?<br>
Thanks all!</p>

KFrank: <p>Hi Alessandro!</p>
<aside class="quote no-group quote-modified" data-username="alessandro_mondin" data-post="1" data-topic="156191" 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/alessandro_mondin/48/46669_2.png" class="avatar"> alessandro_mondin:</div>
<blockquote>
<p>input.shape = (batch_size, 3, 572, 572) and out.shape = (batch_size, 1, 388, 388).</p>
<p>In this case which is the best practice to resize out.shape?</p>
</blockquote>
</aside>
<p>You shouldn’t resize <code>out.shape</code>.</p>
<p>Leaving aside the “tiling” of large images (as discussed in the paper<br>
you cite), the shapes for the use case you describe would typically be<br>
<code>[batch_size, 1, 388, 388]</code> for both <code>out.shape</code> and <code>target.shape</code><br>
(where <code>target</code> is your ground-truth segmentation mask).  Your initial,<br>
unmodified <code>input.shape</code> would be <code>[batch_size, 3, 388, 388]</code><br>
(having the same spatial extent).</p>
<p>However, you would use “mirroring” to increase the spatial extent of your<br>
input images to the necessary <code>[572, 572]</code>.  Quoting from page 3 of the<br>
<a href="https://arxiv.org/pdf/1505.04597.pdf" rel="noopener nofollow ugc">arxiv version</a> of the original U-Net paper:</p>
<blockquote>
<p>To predict the pixels in the border region of the image, the missing<br>
context is extrapolated by mirroring the input image.</p>
</blockquote>
<p>What’s going on here, is rather than “pollute” the borders by adding<br>
padding to the convolutions, the original U-Net just lets the convolutions<br>
snip off a few border pixels at every convolution stage, so the “image”<br>
shrinks as it passes through the network.  To make up for this, they pad<br>
the borders of the input image with the right, rather large, number of pixels<br>
needed to make up for the shrinkage, and – in my mind, sensibly – they<br>
do this with mirroring.</p>
<p>This is the approach you should take.</p>
<p>(When using tiling, they snip out larger, 572x572 tiles out of the input<br>
image and smaller, 388x388 tiles out of the masks, and only use mirroring<br>
to enlarge tiles that lie on the edge of the input image.)</p>
<p>Best.</p>
<p>K. Frank</p>

alessandro_mondin: <p>Hello Frank,<br>
Thank for helping<br>
I am still confused tbh:<br>
i.e. in this implementation that I found in GitHub <a href="https://github.com/hayashimasa/UNet-PyTorch/blob/main/train.py" rel="noopener nofollow ugc">here</a> at line 126 he is cropping the mask to 388 too.<br>
Furthermore I don’t understand what is the purpose of applying that huge padding that you mention to increase the size of the input from 388 to 572. In this case wouldn’t we be adding too much noise into the model considering that 388<strong>2=150544,<br>
572</strong>2=327184–&gt; 327184-150544/327184= 53% of the pixels of each input image would be zeros<br>
Maybe I did not understand you’re suggestions on padding/mirroring but I am very confused</p>

alessandro_mondin: <p>UPDATE: I understood that you meant using torchvision.transforms.Pad(padding_mode=‘reflect’) to mirror the image but in my case. However still why is him cropping the mask? Is it wrong?</p>

KFrank: <p>Hi Alessandro!</p>
<aside class="quote no-group quote-modified" data-username="alessandro_mondin" data-post="3" data-topic="156191" 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/alessandro_mondin/48/46669_2.png" class="avatar"> alessandro_mondin:</div>
<blockquote>
<p>Hello Frank,<br>
this implementation that I found in GitHub</p>
</blockquote>
</aside>
<p>There are any number of U-Net implementations scattered hither and<br>
yon across the internet.  I can’t tell you how good they are or how closely<br>
they hew to the original U-Net paper.</p>
<aside class="quote no-group">
<blockquote>
<p>Furthermore I don’t understand what is the purpose of applying that huge<br>
padding that you mention to increase the size of the input from 388 to 572.</p>
</blockquote>
</aside>
<p>First, read the original U-Net paper carefully, understand how the image<br>
gets its edges trimmed off (by the unpadded convolutions) as it passes<br>
through the network, and understand how <em>tiling</em> works in the use case<br>
of the original paper.</p>
<p>With tiling, you are only actually padding (adding “noise” to) those tiles<br>
that are along the edge of the (large) image.</p>
<p>Now consider a use case where your labelled training data consists of<br>
images and masks of size 572x572, and where the images you will wish<br>
to perform inference on (roughly speaking, for example, your test set) are<br>
also of size 572x572.</p>
<p>You can either train by cropping the segmentation mask down to 388x388<br>
and, at inference time, only generate a segmentation prediction for the<br>
central 388x388 subimage of your 572x572 input image.  Segmenting only<br>
part of your input image may be good enough for your use case.</p>
<p>Or you can mirror (“reflection-pad”) your input images up to whatever size*<br>
gets trimmed down to (or just larger than) 572x572, train using all of your<br>
labelled training data – that is, using all of the information in your full<br>
572x572 segmentation masks – and then be able to perform inference<br>
on your full 572x572 input images, rather than just their 388x388 central<br>
subimages.</p>
<p>Your segmentation performance near the edges of your full 572x572 input<br>
images may not be as good as near the middle (but you can evaluate any<br>
such degradation in performance using your test set), but might well be<br>
worth more than not performing any segmentation near the edges.</p>
<p>Your choice.</p>
<p>*)  Determining which images sizes flow through U-Net’s convolutions<br>
and 2x2 downsamplings cleanly (so that only even-sized “images” get<br>
downsampled) and what the final output prediction sizes are is something<br>
of a fussy computation, but straightforward.</p>
<p>Best.</p>
<p>K. Frank</p>

alessandro_mondin: <p>Hello Frank,<br>
After spending some more hours I finally got the point!<br>
Thanks a lot for your help <img src="https://discuss.pytorch.org/images/emoji/apple/cowboy_hat_face.png?v=12" title=":cowboy_hat_face:" class="emoji" alt=":cowboy_hat_face:" loading="lazy" width="20" height="20"><br>
Alessandro</p>