Search code examples
pythondeep-learningpytorchobject-detectiontorch

How can I save the path of an image (string) from my dataloader (PyTorch)?


I've created a dataloader for my object detection task.

However, I cannot place the image/path name to a tensor. Instead I have it indexed, where in the last portion of the dataloader class, I have this:

target = {}
        target['boxes'] = boxes
        target['labels'] = labels
        target['image_id'] = torch.tensor([index])
        target['area'] = area
        target['iscrowd'] = iscrowd
        target['image_name'] = torch.tensor(index)
        

        return image, target

where atm image_id and image_name are the same thing.

When I print out the image_name from the dataloader, I of course get this:

for image, target in valid_data_loader:
  print(target[0]['image_name'])


Output:

tensor(0)
tensor(1)
tensor(2)
tensor(3)
tensor(4)
tensor(5)
tensor(6)
tensor(7)

I'm aware that strings can't be saved into torch tensors, so is there any way I can refer back to the original image name rather than the index of the tensor? Or would I just have to use the number that comes out and refer back to the dataset class (not dataloader)?

I ultimately want to save the image name, and attributes such as bounding box info to a separate numpy dataframe.


Solution

  • Ok, so this is a bit ad-hoc and not exactly what I was thinking but here is one method I have used to retrieve the paths/image names. I basically find the id from the dataloader by removing it from the tensor. I then use the tensor_id to find the corresponding id in the original dataframe:

    for image, target in valid_data_loader:
      tensor_id = target[0]['image_name'].item()
      print(valid_df.iloc[tensor_id]['image_id'])
    

    I don't know if this is efficient though but it got what I wanted...