How to efficiently subsample from large images

username: sakkosjo
cooked: <p>Hi,</p>
<p>I’m new to Pytorch and deep learning in general. I’m developing a bacterial cell segmentation tool for microscopy with Pytorch/Unet. Since bacterial cells are very small (~1 micron wide x 3 microns long), they are only 20 or so pixels wide and I can’t simply load my images (1460 x 1936 pixels ) and scale them down without losing critical information. Instead, we’ve been subsampling regions small enough that can be run on a gpu in batches of 16 or so (160 x 160 pixels) and stitching the model predictions together at the end to get the full image mask. I’m wondering if there is a better/more established method for this kind of subsampling.</p>
<p>Any help would be greatly appreciated.</p>
<p>Here’s some example code:</p>
<pre><code class="lang-auto">#import raw image datasets and resize if necessary
if len(training_data) == len(training_mask): #check list of training and label images
    train_set_size = len(training_data)
    X_data = np.zeros([train_set_size,1460,1936])
    Y_data = np.zeros([train_set_size,1460,1936])
n = 0    
for X_file, y_file, in zip(training_data,training_mask):
    #Load raw images .tiff .tif 
    if X_file[-4:] == '.tif' or X[-4:] == 'tiff':
        X_image_load = np.array(Image.open(join(folder_path,X_file)))
    if X_image_load.shape == (728, 968):
        X_image_for_stack = Image.fromarray(X_image_load).resize((1936, 1460))
    elif X_image_load.shape == (1460, 1936):
        X_image_for_stack = X_image_load
    else:
        print(X_file + 'has wrong size')
    X_data[n,:,:] = np.array(X_image_for_stack) #stack raw images
    
    y_image_load = np.array(Image.open(join(folder_path,y_file)))
    if y_image_load.shape == (1460, 1936):
        y_image_for_stack = np.array(y_image_load)
    elif y_image_load.shape == (728, 968):
        y_image_for_stack = Image.fromarray(y_image_load).resize((1936, 1460))
    elif y_image_load.shape == (1456, 1936):
        y_image_for_stack = Image.fromarray(y_image_load).resize((1936, 1460))
        
    Y_data[n,:,:] = np.array(y_image_for_stack) #stack raw label images (binary masks)
    n += 1

X = []
Y = []
Coverage = 10
image_width = 1936
image_height = 1460
sample_px = 160
n_samples = int(Coverage * image_width * image_height  / (sample_px**2))
#randomly select points to subsample from
H = np.random.randint(int(sample_px/2) + 1,image_height - int(sample_px/2), n_samples) #height sampling
W = np.random.randint(int(sample_px/2) + 1,image_width -int(sample_px/2), n_samples) #width sampling
Stack_idx = np.random.randint(0, train_set_size, n_samples) #image stack sampling
Transp = np.random.randint(0, 2, n_samples) #randomly transpose images
Blur = np.random.randint(0, 2, n_samples) #randomly blur images
#loop over all n_samples for random subsampling and append 160 x 160 array to stacked set of data and labels
for i in range(n_samples):
    tmp_X = X_data[Stack_idx[i],H[i]-int(sample_px/2):H[i]+int(sample_px/2),W[i]-sample_px/2:W[i]+sample_px/2]
    tmp_X = (tmp_X-np.min(tmp_X))/np.max(tmp_X) #normalize data between 0 and 1
    if Transp[i]:
        tmp_X = tmp_X.T #transpose
    if Blur[i] and i &gt; 5000 and (i%20 == 1):
        tmp_X = gaussian_filter(tmp_X,np.random.randint(2, 5)) #blur for later in training epoch
    X.append(tmp_X)
    tmp_Y = Y_data[Stack_idx[i],H[i]-int(sample_px/2):H[i]+int(sample_px/2),W[i]-sample_px/2:W[i]+sample_px/2]
    tmp_Y = tmp_Y/255 #normalize data between 0 and 1
    if Transp[i]:
        tmp_Y = tmp_Y.T #transpose
    Y.append(tmp_Y)
    i += 1

x_train_tensor = torch.Tensor(X)
y_train_tensor = torch.Tensor(Y)
dataset = TensorDataset(x_train_tensor, y_train_tensor)
</code></pre>

username: ptrblck
cooked: <p>You could probably use <code>unfold</code> directly to create these patches.<br>
Have a look at <a href="https://discuss.pytorch.org/t/patch-making-does-pytorch-have-anything-to-offer/33850/11">this post</a> for a simple example.</p>

username: sakkosjo
cooked: <p>Thanks, that post was very helpful. I’m still a little confused about how I should be feeding the patches into my model for prediction and putting the predictions back together.</p>
<pre><code class="lang-auto">x = torch.randn(1, 500, 500)  # batch, c, h, w
k = 160 # kernel size
d = 160# stride
patches = x.unfold(1, k, d).unfold(2, k, d)
unfold_shape = patches.size()
patches = patches.contiguous().view(-1, 1,k, k) #my model takes tensors of shape (1,1,160,160)
print(patches.shape)
data = torch.utils.data.DataLoader(patches,batch_size=2,num_workers=0)
for batch in data:
   output = model(batch)
#collect outputs
#fold

</code></pre>

username: ptrblck
cooked: <p>You won’t be able to create the same input shape using the specified shapes, as the input shape of 500 divided by the kernel size of 160 will have a remainder.<br>
If that’s important, have a look at <a href="https://discuss.pytorch.org/t/creating-nonoverlapping-patches-from-3d-data-and-reshape-them-back-to-the-image/51210/6">this code</a>, which adds padding to the input.</p>

username: sakkosjo
cooked: <p>Thanks for the clarification. The input shape above was arbitrary, but padding will still be useful for my actual images. My question was a bit more basic, relating to the use of the trained model on images for segmentation. If I’m starting with an image, unfolding it, getting model predictions, and then piecing the original shape back together, what is the best way to keep track of each patch, if we assume they’re not overlapping (for now)?</p>
<p>My original code started with an empty numpy array with the shape of the original image and filled it with the patches, but it seems like this isn’t the most efficient way to process the images. I ultimately want to use the <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader" rel="nofollow noopener"> <code>torch.utils.data.DataLoader</code> </a>to feed all of the patches into the model, but I can’t quite understand how to keep track of their location in the original image.</p>
<p>My old code for getting patches:</p>
<pre><code class="lang-auto">if image.shape == (1460, 1936): #actual image shape
        x_points = np.uint16(np.linspace(80, 1460-80, 10)) #split x into 10 slices
        y_points = np.uint16(np.linspace(80, 1936-80, 12)) #split y into 12 slices
        output_mask = np.zeros([1460, 1936]) #create placeholder
        for x_coords in tqdm(x_points): #loop over x and y coordinates and fill in the placeholder
            for y_coords in y_points:
                tmp_data = image[x_coords-80:x_coords+80,y_coords-80:y_coords+80]
                tmp_data = (tmp_data-np.min(tmp_data))/np.max(tmp_data)
                X = torch.from_numpy(tmp_data)
                X = X.view(1,1,160,160).to(device, dtype = torch.float)
                prediction = model(X)
                tmp_out = prediction.cpu().detach()[0][1]&gt;threshold
                output_mask[x_coords-80:x_coords+80,y_coords-80:y_coords+80] = tmp_out
</code></pre>
<p>Does it make sense to convert from a tensor back to a numpy array for each patch?<br>
The code below gives me a padded tensor of size (108, 1, 160, 160)</p>
<pre><code class="lang-auto">x = torch.randn(1, 1460, 1936)  # batch, h, w
k = 160 # kernel size
d = 160# stride
x = F.pad(x,(x.size(2)%k // 2,x.size(2)%k // 2,
             x.size(1)%k // 2,x.size(1)%k // 2))
patches = x.unfold(1, k, d).unfold(2, k, d) 
unfold_shape = patches.size()
patches = patches.contiguous().view(-1, 1,k, k) #(108, 1, 160, 160)
</code></pre>
<p>Since I can’t run all patches through at once, I need to batch, collect predictions, and then unfold. That’s what I don’t understand conceptually how to do with Pytorch. Thanks again for your help.</p>

username: sakkosjo
cooked: <p>I’ve figured out how to reconstruct the patches after model prediction with the code below. I’m not sure if the padding calculations were correct in the <a href="https://discuss.pytorch.org/t/creating-nonoverlapping-patches-from-3d-data-and-reshape-them-back-to-the-image/51210/6">linked post</a>, but they were giving me half of remainder for padding, rather than the amount to pad to a multiple of 160.</p>
<pre><code class="lang-auto">x0 = torch.randn(1460, 1936)  # h, w
print(x0.shape)
# kernel size
k = 160 
# stride
d = 160
#hpadding
hpad = (k-x0.size(0)%k) // 2 
#wpadding
wpad = (k-x0.size(1)%k) // 2 
#pad x0
x = F.pad(x0,(wpad,wpad,hpad,hpad)) 
print(x.shape)
#unfold
patches = x.unfold(0, k, d).unfold(1, k, d) 
unfold_shape = patches.size()
#reshape to (batch,1,h,w)
patches = patches.contiguous().view(-1, 1,k, k)
print(patches.shape)
#create storage tensor
temp = torch.empty(patches.shape) 

#loop over all patches, feed model predictions back into storage tensor
for i,patch in tqdm(enumerate(patches)):
    temp[i,:,:,:] = model(patch.view(-1,1,k,k).to(device, dtype = torch.float)).cpu().detach()[0][1]

# Reshape back
patches_orig = temp.view(unfold_shape)
output_h = unfold_shape[0] * unfold_shape[2]
output_w = unfold_shape[1] * unfold_shape[3]
patches_orig = patches_orig.permute(0, 2, 1, 3).contiguous()
patches_orig = patches_orig.view(output_h, output_w)
#slice away padding
reconstructed_image = patches_orig[hpad:patches_orig.shape[0]-hpad,wpad:patches_orig.shape[1]-wpad]
print(reconstructed_image.shape)
</code></pre>
<p>The last thing I’m stuck on is getting the index from the batch loader to reconstruct the patched image. <a href="https://discuss.pytorch.org/t/how-does-one-obtain-indicies-from-a-dataloader/16847/3">This post</a> has a lot of discussion on the topic, but not any simple answers.</p>

username: ptrblck
cooked: <p>Wouldn’t the provided code snippet automatically reconstruct the image using the correct patches?<br>
Also, if you put all patches into the batch dimension, you could avoid the for loop, which should be slower than a single forward pass.</p>

username: sakkosjo
cooked: <p>Yes, the snippet does reconstruct the patched image, albeit one patch at a time. So you’re suggesting something like the following for batch processing:</p>
<pre><code class="lang-auto">reshaped = patches.view(-1,k,k) #(batch,h,w)
temp = torch.empty(reshaped.shape)
loader = Dataloader(dataset=reshaped,batch_size=16)
for batch in loader:
   x = batch.view(-1,1,k,k).to(device, dtype = torch.float)
   temp[batch_idx,:,:] = model(x)
</code></pre>
<p>I don’t understand how to get the patch index from each batch to fill the empty tensor (temp). Can you clarify how that should work? Thanks!</p>

username: bugsuse
cooked: <p>I have the same problem, and could you solve it?</p>
<p>I want to patch image using the <code>unfold</code> function, the commands as follows,</p>
<pre data-code-wrap="python"><code class="lang-python">img = io.imread('02.png')
img = img[np.newaxis, :]
img = (np.pad(img, ((0, 0), (26, 26), (26, 26)), mode='constant')/255.).astype(np.float32)
img = torch.unsqueeze(torch.from_numpy(img), dim=0)

kc, kh, kw = 1, 64, 64  # kernel size
dc, dh, dw = 1, 64, 64  # stride

patches = img.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw)
unfold_shape = patches.size()
patches = patches.contiguous().view(patches.size(0), -1, kh, kw)
</code></pre>
<p>and get the patched image as follow using the below commands,</p>
<pre data-code-wrap="python"><code class="lang-python">fig, ax = plt.subplots(figsize=(8, 8), nrows=8, ncols=8)
plt.subplots_adjust(hspace=0.02, wspace=0.005)
for i, axes in enumerate(ax.ravel()):
    axes.imshow(patches[0, i]*70, vmin=0, vmax=70, cmap='pyart_NWSRef')
    axes.set_axis_off()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/7/d7e4a9601f71da23b4f59f4a30580c3497836576.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/d7e4a9601f71da23b4f59f4a30580c3497836576" title="截屏2020-11-28 下午7.21.16"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7e4a9601f71da23b4f59f4a30580c3497836576_2_255x250.jpeg" alt="截屏2020-11-28 下午7.21.16" data-base62-sha1="uNSyO4KRZi3tOS0NpxRn7D55v1Q" width="255" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7e4a9601f71da23b4f59f4a30580c3497836576_2_255x250.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7e4a9601f71da23b4f59f4a30580c3497836576_2_382x375.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7e4a9601f71da23b4f59f4a30580c3497836576_2_510x500.jpeg 2x" data-dominant-color="3BDECB"><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">截屏2020-11-28 下午7.21.16</span><span class="informations">916×898 131 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>

username: bugsuse
cooked: <p>then calculate convolution using nn.Conv2d,</p>
<pre data-code-wrap="python"><code class="lang-python">patches_conv = nn.Conv2d(64, 64, 3, 1, 1)(patches)

fig, ax = plt.subplots(figsize=(8, 8), nrows=8, ncols=8)
plt.subplots_adjust(hspace=0.02, wspace=0.005)
for i, axes in enumerate(ax.ravel()):
    axes.imshow(patches_conv[0, i].detach().numpy()*70, vmin=0, vmax=70, cmap='pyart_NWSRef')
    axes.set_axis_off()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/4/24a8f69afcce674dae427c3b2010cfef27edee0e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/24a8f69afcce674dae427c3b2010cfef27edee0e" title="截屏2020-11-28 下午7.29.38"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24a8f69afcce674dae427c3b2010cfef27edee0e_2_256x249.jpeg" alt="截屏2020-11-28 下午7.29.38" data-base62-sha1="5ejbQKvqFSJD80UjHTLstzTXXvE" width="256" height="249" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24a8f69afcce674dae427c3b2010cfef27edee0e_2_256x249.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24a8f69afcce674dae427c3b2010cfef27edee0e_2_384x373.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24a8f69afcce674dae427c3b2010cfef27edee0e_2_512x498.jpeg 2x" data-dominant-color="36D6EC"><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">截屏2020-11-28 下午7.29.38</span><span class="informations">914×892 148 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>It does not reconstruct the patched model prediction image. It does not works using the <code>F.fold</code> function to reconstruct the image, and get the same image.</p>
<p>So I’m very confused, could you help me?</p>

username: bugsuse
cooked: <p>Sorry for disturb you. I want to unfold the input image to feed into the model and fold the patched model output back. Could you help me? <a class="mention" href="/u/ptrblck">@ptrblck</a> and the original image is below.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/a/cad6c676c266827108b1845e7cf8a289dd08020b.png" alt="02" data-base62-sha1="sWoBntnAJDuCcKKpnvU58ToAWO7" width="230" height="230"></p>
<p>Thanks so much!</p>