How would you train a segmentation model with varying image sizes?

chokevin8: <p>Hi, I’m trying to train a two-step process for image instance segmentation of two classes, and this is multi-label since a pixel can belong in both classes at the same time (a part of the image can be both class A and class B, for example, class A can be a larger area and class B can exist inside it, hence some of the area being both class A and B). However, the multi-label nature isn’t the problem- the problem arises due to the training data image sizes being very different (as small as 10 pixels by 10 pixels, and as big as about 300 by 300 pixels) and also rectangular.</p>
<p>This is because this model is being trained to be a downstream model that takes the inference of a Yolo object detection model’s bounding boxes as an input. The resolution and scale of the images are the same, since they are just different sizes of bounding boxes, but if I were to train a vanilla UNet, I would have to resize and pad them to a uniform size to make sure aspect ratio the same. However, this would not be ideal since the background of my original images are not white, and the backgrounds of each image are always different, so padding might not be the best option.</p>
<p>I just wanted to ask the community about their opinions, of what the best idea would be. Pretty much every model out there requires a fixed image size (which are mostly square and product of 16’s). Maybe trying a scale-invariant model like DeepLabv3+ might be better since it uses ASPP? But DeepLab also requires a uniform input image size.</p>
<p>Would appreciate any inputs from people who has had similar problems, thank you!</p>

KFrank: <p>Hi Choke!</p>
<aside class="quote no-group quote-modified" data-username="chokevin8" data-post="1" data-topic="183280" 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/chokevin8/48/54351_2.png" class="avatar"> chokevin8:</div>
<blockquote>
<p>the problem arises due to the training data image sizes being very different<br>
…<br>
if I were to train a vanilla UNet, I would have to resize and pad them to a uniform size to make sure aspect ratio the same. However, this would not be ideal since the background of my original images are not white, and the backgrounds of each image are always different</p>
</blockquote>
</aside>
<p>The short story:  Use U-Net with  <em>reflection padding</em> and be sure to<br>
understand the padding requirements of a properly-architected U-Net.</p>
<p>As you mention, I would consider using U-Net.  It’s a solid, widely-used<br>
semantic-segmentation model.  As a fully-convolutional model, it can be<br>
trained on, and then used for inference for, images of varying sizes.</p>
<p>The architecture of the original U-Net was designed quite carefully, and,<br>
as a result of that design, while it accepts images of arbitrarily-large size,<br>
it can’t accept input images of <em>arbitrary</em> size. I call the acceptable input<br>
sizes <em>conformable</em> with the specific U-Net (but I don’t think this term is<br>
standard).</p>
<p>Let me assume that your input images and ground-truth targets are the<br>
same size.  You will then need to pad your targets to be conformable with<br>
the output of the U-Net and pad your inputs to be conformable with the<br>
input to the U-Net.  The (padded) input size will be larger than the (padded)<br>
target size, because the image shrinks when it flows through the U-Net as<br>
unpadded convolutions nibble away at its edges.</p>
<p>I strongly recommend that you use <em>reflection padding</em> for the input.  The<br>
padding for the target won’t matter (because you will mask it), so I would<br>
just use zeros.</p>
<p>It sounds like you are performing a multi-label, two-class segmentation.<br>
So your U-Net output should have two channels (per pixel) and you should<br>
use <code>BCEWithLogitsLoss</code> as your loss criterion.</p>
<p>However, only your original unpadded target image contains meaningful<br>
information.  So you should mask your loss computation to only include<br>
that original unpadded region.  I would suggest using<br>
<code>BCEWithLogitsLoss (reduction = 'none')</code> and then perform the<br>
masked reduction yourself.  (You could also use <code>BCEWithLogitsLoss</code>’s<br>
<code>weight</code> constructor argument to perform this masking.)</p>
<p>One last technical point that’s a pytorch rather than U-Net issue.  If you<br>
use a batch size greater than one, then the individual samples within a<br>
batch all have to be the same size (because they’re slices of the a single<br>
pytorch tensor).  So you have to pad your input and target images to the<br>
largest (padded) size of any of the samples in the batch.</p>
<p>Depending on your use case, you may well not need to resize any of<br>
your images or change their aspect ratios.  Within a batch, all of the<br>
images will be padded to have the same size (and hence the same<br>
aspect ratio), but the pixels inside of the padding on the borders will<br>
be the pixels from the original images without any resizing or stretching.</p>
<p>Best.</p>
<p>K. Frank</p>