How to avoid 'blocky' output when stitching 3D patches back together?

username: farazk86
cooked: <p>Hi,</p>
<p>I have volumes of size <code>[120*768*768]</code> which I have broken down into patches of size <code>[32*256*256]</code> to fit on GPU.</p>
<p>for prediction of a full volume the patches are stitched back together to get the full output. I’m using the code shared here (<a href="https://github.com/wolny/pytorch-3dunet/blob/master/predict.py#L97" rel="noopener nofollow ugc">https://github.com/wolny/pytorch-3dunet/blob/master/predict.py#L97</a>) for stitching the patches back into full size.</p>
<p>Basically what I’m doing is initializing a <code>prediction_map</code> of size <code>[120*768*768]</code> and a <code>normalization_mask</code> of the same size.</p>
<pre><code class="lang-auto"># initialize the output prediction arrays
prediction_maps = [np.zeros(prediction_maps_shape, dtype='float32') ]
# initialize normalization mask in order to average out probabilities of overlapping patches
normalization_masks = [np.zeros(prediction_maps_shape, dtype='float32')]
</code></pre>
<p>Then probabilities of each patch by network model is aggregated into this at the positions given by <code>u_index</code>.</p>
<pre><code class="lang-auto"># accumulate probabilities into the output prediction array
prediction_map[u_index] += u_prediction
# count voxel visits for normalization
 normalization_mask[u_index] += 1
</code></pre>
<p><code>index</code> is returned by data loader which keeps track of the patch slices.</p>
<p>The author of the code uses an <code>unpad</code> script to prevent the very thing that is happening with my output here:</p>
<pre><code class="lang-auto"># unpad in order to avoid block artifacts in the output probability maps
u_prediction, u_index = utils.unpad(prediction, index, volume_shape)
</code></pre>
<p><a href="https://github.com/wolny/pytorch-3dunet/blob/master/unet3d/utils.py#L121" rel="noopener nofollow ugc">https://github.com/wolny/pytorch-3dunet/blob/master/unet3d/utils.py#L121</a></p>
<p>In the end the aggregated into <code>prediction_map</code> is divided by the <code>normalization_ask</code> to average the output.</p>
<p>This whole process should not show a blocky output but still I get a result like this for patch size <code>[32, 512, 512]</code> and stride <code>[2, 256, 256]</code>. There is considerable overlap and as per author’s guidelines this should not cause block artefacts.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/e/ecd40827442568e9c87080fe49069431f65a506d.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/ecd40827442568e9c87080fe49069431f65a506d" title="auto_TP100"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/e/ecd40827442568e9c87080fe49069431f65a506d_2_500x500.jpeg" alt="auto_TP100" data-base62-sha1="xN4Xt2w5ekSLPaqkV5fnczFu5DD" width="500" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/e/ecd40827442568e9c87080fe49069431f65a506d_2_500x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/e/ecd40827442568e9c87080fe49069431f65a506d_2_750x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/e/ecd40827442568e9c87080fe49069431f65a506d.jpeg 2x" data-dominant-color="202020"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">auto_TP100</span><span class="informations">768×768 23.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The network is working perfectly for individual patches (Ive checked from TensorBoardX) but its failing at this stitching stage.</p>
<p>And when I use a patch size of <code>[32, 256, 256]</code> and stride <code>[2, 128, 128]</code>, I get the following strange output:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/7/7d39d19797242b62bd8ba5c5fc355ea7c48f4dd2.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7d39d19797242b62bd8ba5c5fc355ea7c48f4dd2" title="bad"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/7/7d39d19797242b62bd8ba5c5fc355ea7c48f4dd2_2_500x500.jpeg" alt="bad" data-base62-sha1="hRNyI1MHulBFviowKHe1i5QONZU" width="500" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/7/7d39d19797242b62bd8ba5c5fc355ea7c48f4dd2_2_500x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/7/7d39d19797242b62bd8ba5c5fc355ea7c48f4dd2_2_750x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/7/7d39d19797242b62bd8ba5c5fc355ea7c48f4dd2.jpeg 2x" data-dominant-color="1E1E1E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">bad</span><span class="informations">768×768 17.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>What am I doing wrong? Is there any better way to stitch these patches back together? Would <code>torch.unfold</code> and <code>torch.fold</code> help me here?</p>
<p>Some thoughts/ help would be appreciated</p>
<p>Thank you</p>

username: farazk86
cooked: <p>Forgot to mention that the patches are sliced in the data loader here: <a href="https://github.com/wolny/pytorch-3dunet/blob/master/datasets/hdf5.py#L13" rel="nofollow noopener">https://github.com/wolny/pytorch-3dunet/blob/master/datasets/hdf5.py#L13</a></p>
<p>Its simple indexing of the volume using the given stride.</p>

username: cwat_tum
cooked: <p>Hi Faraz,</p>
<p>I’m facing a similar issue with 3D UNet segmentaiton of multiple organs, and I also haven’t yet found a solution to this problem. Take a look here for example:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/c/9c77041bfb51fd98cfa60bb866757b3cc3c9d85c.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/9c77041bfb51fd98cfa60bb866757b3cc3c9d85c" title="edges"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/c/9c77041bfb51fd98cfa60bb866757b3cc3c9d85c_2_690x318.jpeg" alt="edges" data-base62-sha1="mk9tw9faP5T2izvna9we8YdQYYQ" width="690" height="318" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/c/9c77041bfb51fd98cfa60bb866757b3cc3c9d85c_2_690x318.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/9/c/9c77041bfb51fd98cfa60bb866757b3cc3c9d85c.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/c/9c77041bfb51fd98cfa60bb866757b3cc3c9d85c.jpeg 2x" data-dominant-color="2E2D30"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">edges</span><span class="informations">748×345 28.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>You can see in the predictions there’s some serious blockiness! Checking the individual outputs also gives good results, so I’m not sure where this is coming from.</p>
<p>Did you ever find a solution to this problem?</p>
<p>-Chris</p>

username: sr1945
cooked: <p>Hi <a class="mention" href="/u/farazk86">@farazk86</a>, <a class="mention" href="/u/cwat_tum">@cwat_tum</a> ,</p>
<p>I am running into a similar issue with patch-based 3D UNet for brain cortex segmentation, and was wondering if either of you found a solution to this problem?</p>
<p>Thanks!</p>