1

I use pytorch to load images like this:

inf_data = InfDataloader(img_folder=args.imgs_folder, target_size=args.img_size)
inf_dataloader = DataLoader(inf_data, batch_size=1, shuffle=True, num_workers=2)

And then:

    with torch.no_grad():
        for batch_idx, (img_np, img_tor) in enumerate(inf_dataloader, start=1):

            img_tor = img_tor.to(device)
            pred_masks, _ = model(img_tor)

But I want to get the file name of the image. Can anyone help me with this? Thanks a lot!

Nice
  • 97
  • 1
  • 6
  • Have you checked the documentation? – AMC Aug 21 '20 at 21:12
  • I have already find the solution yesterday. I'm sorry that I didn't realize this yesterday. I'm pretty new to Pytorch. You actually could not get file name from `DataLoader` directly. Here I posted my answer. I hope it could help others. – Nice Aug 22 '20 at 15:30

1 Answers1

2

The DataLoader basically can not get the name of the file. But in Dataset, which is the InfDataloader in the question mentioned above, you can get the name of file from the tensor.

class InfDataloader(Dataset):
    """
    Dataloader for Inference.
    """
    def __init__(self, img_folder, target_size=256):
        self.imgs_folder = img_folder

        self.img_paths = []

        img_path = self.imgs_folder + '/'
        img_list = os.listdir(img_path)
        img_list.sort()
        img_list.sort(key=lambda x: int(x[:-4]))  ##文件名按数字排序
        img_nums = len(img_list)
        for i in range(img_nums):
            img_name = img_path + img_list[i]
            self.img_paths.append(img_name)

        # self.img_paths = sorted(glob.glob(self.imgs_folder + '/*'))

        print(self.img_paths)


        self.target_size = target_size
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])

    def __getitem__(self, idx):
        """
        __getitem__ for inference
        :param idx: Index of the image
        :return: img_np is a numpy RGB-image of shape H x W x C with pixel values in range 0-255.
        And img_tor is a torch tensor, RGB, C x H x W in shape and normalized.
        """
        img = cv2.imread(self.img_paths[idx])
        name = self.img_paths[idx]

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Pad images to target size
        img_np = pad_resize_image(img, None, self.target_size)
        img_tor = img_np.astype(np.float32)
        img_tor = img_tor / 255.0
        img_tor = np.transpose(img_tor, axes=(2, 0, 1))
        img_tor = torch.from_numpy(img_tor).float()
        img_tor = self.normalize(img_tor)

        return img_np, img_tor, name

Here I add the line name = self.img_paths[idx] and return it.

So,

 with torch.no_grad():
        for batch_idx, (img_np, img_tor, name) in enumerate(inf_dataloader, start=1):
            img_tor = img_tor.to(device)
            pred_masks, _ = model(img_tor)

I could get the name.

Nice
  • 97
  • 1
  • 6