pytorch dataloader(dataset.ImageFolder) 내 파일 경로 가져오기

1. Dataloader shuffle=False일 때,

dataset = torchvision.datasets.ImageFolder(<folder_path>, transforms=<transforms>)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=<batch_size>, shuffle=False)


# 이미 기존 DataLoader에서 shuffle을 사용하지 않을 경우, "dataloader.dataset.samples"안에 
# (파일경로, 라벨) 튜플로 이루어진 list형태로 저장되어 있다.
allFiles, _ = map(list, zip(*dataloader.dataset.samples))

for i, (inputs, labels) in enumerate(dataloader):
	inputs = inputs.to(device)
    labels = labels.to(device)
    
    for j in range(inputs.size()[0]):
    	print("현재 파일에 대한 파일경로")
        print(allFiles[ i * <batch_size> + j ])

 

2. Dataloader shuffle=True일 때,

# 별도의 Dataset을 기존 ImageFolder dataset과 동일한 역할을 수행하도록 만들어준다.
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, root_dir, transforms=None):
    	self.root_dir = root_dir
        self.classes = os.listdir(self.root_dir)
        self.transforms = transforms
        self.data = []
        self.labels = []
        
        for idx, cls in enumerate(self.classes):
        	cls_dir = os.path.join(self.root_dir, cls)
            for img in glob(os.path.join(cls_dir, '*.jpg')):
            	self.data.append(img)
                self.labels.append(idx)
                
    def __getitem__(self, idx):
    	img_path, label = self.data[idx], self.labels[idx]
        img = PIL.Image.open(img_path)
        
        if self.transforms:
        	img = self.transforms(img)
        
        sample = {'image':img, 'label':label, 'filename':img_path}
        
        return sample
    
    def __len__(self):
    	return len(self.data)

custom_dataset = CustomDataset(root_dir=<folder_path>, transforms=<transforms>)
dataloader = torch.utils.data.DataLoader(custom_dataset, batch_size=<batch_size>,shuffle = True)

for i, (sample) in enumerate(dataloader):
	inputs = sample['image'].to(device)
    labels = sample['label'].to(device)
    names = sample['filename']
    
    for j in range(inputs.size()[0]):
        print(names[j])

 

 

출처, 도움 주신 블로그 : https://hayunjong83.tistory.com/46

PyTorch Dataloader에서 원본 파일명을 알고 싶다.