Problem loading images

vision
1

Hi! I’m following the code provided in https://github.com/yunlongdong/FCN-pytorch-easiest to train an FCN, and in each epoch, I plot each mask and the prediction of the model to see it’s progress. However, I’ve noticed it’s just showing me one image, and it’s corresponding mask. Even seeing the test results, only that one image is loaded. Is there anything wrong with my dataloader?

class MyDataset(Dataset):

    def __init__(self, root_dir, image_dir, mask_dir, label, transform):
      self.dataset_path = root_dir
      self.image_dir = image_dir
      self.mask_dir = mask_dir
      mask_full_path = os.path.join(self.dataset_path, self.mask_dir)
      self.mask_file_list = [f for f in listdir(root_dir+mask_dir) if isfile(root_dir+join(mask_dir, f))]
      random.shuffle(self.mask_file_list)
      self.mapping = {
      0.007782101167315175: 0,
      0.5019455252918288: 0,
      0.9961089494163424: 1,
      1.0: 1}
      self.transform = transform
    def __len__(self):
        return len(self.mask_file_list)

    def mask_to_class(self, mask):
      for k in self.mapping:
          mask[mask == k] = self.mapping[k]
      return mask
    def __getitem__(self, idx):

        file_name = self.mask_file_list[index]
        img_name = os.path.join(self.dataset_path, self.image_dir, file_name.replace('.png', '.dcm'))
        mask_name = os.path.join(self.dataset_path, self.mask_dir, file_name)
        imgA = dicom_to_numpy(img_name)
        mask_p = img_to_numpy(mask_name)
        mask_p = cv2.resize(mask_p, (256, 256))
        imgB = self.mask_to_class(mask_p)
        imgA = cv2.resize(imgA, (256, 256))

        imgB = imgB.astype('uint8')
        imgB = onehot(imgB, 2)
        imgB = imgB.swapaxes(0, 2).swapaxes(1, 2)
        imgB = torch.FloatTensor(imgB)
        if self.transform:
            imgA = self.transform(imgA)    
        item = {'A':imgA, 'B':imgB}
        return item

image_dir = 'mass/'
mask_dir = 'mask/'
label = 'mass'
fulldir = '/'
full_data = MyDataset(fulldir, image_dir, mask_dir, label, transform=transform)

train_size = int(0.8 * len(full_data))
test_size = len(full_data) - train_size
train_data, test_data = torch.utils.data.random_split(full_data, [train_size, test_size])

train_loader = DataLoader(train_data, batch_size=4, shuffle=True, num_workers=4)

test_loader = DataLoader(test_data, batch_size=4, shuffle=True, num_workers=4)

Thanks for your help!

2

Could you post the code showing, how you draw the image and mask, and how you are plotting it in your training?

3

So this is my training code:

vgg_model = VGGNet(requires_grad=True)
fcn_model = FCNs(pretrained_net=vgg_model, n_class=2)
fcn_model = fcn_model.cuda()
criterion = nn.BCELoss().cuda()
optimizer = optim.SGD(fcn_model.parameters(), lr=1e-2, momentum=0.7)
saving_index = 0
i = 0
iou = 0
iou_bg = 0
for epo in range(30):
  saving_index +=1
  index = 0
  epo_loss = 0
  for item in train_loader:
      index += 1
      start = time.time()
      input = item['A']
      y = item['B']
      input = torch.autograd.Variable(input)
      y = torch.autograd.Variable(y)

      input = input.float()
      input = input.cuda()
      y = y.cuda()

      optimizer.zero_grad()
      output = fcn_model(input)
      output = nn.functional.sigmoid(output)
      loss = criterion(output, y)
      loss.backward()
      iter_loss = loss.item()
      epo_loss += iter_loss
      optimizer.step()

      output_np = output.cpu().data.numpy().copy()
      output_bg_np = np.argmax(output_np, axis=1)
      output_np = np.argmin(output_np, axis=1)
      
      y_np = y.cpu().data.numpy().copy()
      y_np = np.argmin(y_np, axis=1)
      input = input.cpu().data.numpy().copy().squeeze(1)
      y_not = np.where(y_np[0, :, :] == 1, 0, 1)


      intersect = (output_np[0, :, :]*y_np[0, :, :]).sum()
      union = (output_np[0, :, :] + y_np[0, :, :]).sum()
      intersect_bg = (output_bg_np*y_not).sum()
      union_bg = (output_bg_np + y_not).sum()
      iu = intersect / (union-intersect+0.00001)
      iu_bg = intersect_bg / (union_bg-intersect_bg+0.00001)

      fig, ax = plt.subplots(1,2, figsize=(10,3))  # 1 row, 2 columns
      ax[0].imshow(output_np[0, :, :])
      ax[1].imshow(y_np[0, :, :])
      iou = iu + iou
      iou_bg = iu_bg + iou_bg

      i = i + 1
      print("IOU is:", iou/i, "bg is: ", iou_bg/i)
      print("mean IOU is:", (iou/i+iou_bg/i)/2)

      if np.mod(index, 5) ==1:
          print('epoch {}, {}/{}, loss is {}'.format(epo, index, len(input), iter_loss))
          fig, ax = plt.subplots(1,2, figsize=(10,3))  # 1 row, 2 columns
          ax[0].imshow(output_np[0, :, :])
          ax[1].imshow(y_np[0, :, :])
  print('epoch loss = %f'%(epo_loss/len(train_loader)))
  
  if np.mod(saving_index, 5)==1:
      torch.save(fcn_model, './fcn_model_{}.pt'.format(epo))
      print('saveing checkpoints/fcn_model_{}.pt'.format(epo))
print("IOU is:", iou/i, "bg is: ", iou_bg/i)
print("mean IOU is:", (iou/i+iou_bg/i)/2)
4

The code should plot the first output of each batch.
If you are shuffling the data, it should be random, so I’m not sure what’s going on.

Also, why are you using the lowest probabilities via output_np = np.argmin(output_np, axis=1)?

1 Like
5

With a great amount of shame, I declare that in MyDataset, I’ve passed the parameter “idx” as index in getitem(), but I’ve used a new parameter “index” in the function! I’m sorry for wasting your time!

6

Haha, that’s a tricky bug and I completely missed it. :wink:
Glad to hear it’s working now.

1 Like