如何解决如何获取我放入Pytorch的Dataloader中的图像的文件名
我使用pytorch加载这样的图像:
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)
然后:
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)
解决方法
DataLoader
基本上无法获得文件名。但是在Dataset
(即上述问题中的InfDataloader
)中,您可以从张量中获取文件名。
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,1))
img_tor = torch.from_numpy(img_tor).float()
img_tor = self.normalize(img_tor)
return img_np,img_tor,name
我在这里添加行
name = self.img_paths[idx]
并退回。
所以
with torch.no_grad():
for batch_idx,(img_np,name) in enumerate(inf_dataloader,start=1):
img_tor = img_tor.to(device)
pred_masks,_ = model(img_tor)
我能知道这个名字。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。