读取文件夹内图像--方法2

GTX_AI

读取文件夹内图像：
供深度学习算法调用：Pytorch_Unet图像分割

1. import torch.utils.data as data
   
2. import PIL.Image as Image
   
3. import os
   
4. 
   
5. def make_dataset_txt(root, txt_file):
   
6.     imgs=[]
   
7.     with open(txt_file) as f:
   
8.         indexs=f.readlines()
   
9.     for i in range( len(indexs) ):
   
10.         index = str( indexs[i] ).split('.')[0]
    
11.         img=os.path.join(root, index+".jpg")
    
12.         mask=os.path.join(root,index+".png")
    
13.         imgs.append((img,mask))
    
14.     return imgs
    
15. 
    
16. class LiverDataset(data.Dataset):
    
17.     def __init__(self, root, txt_file, transform=None, target_transform=None):
    
18.         imgs = make_dataset_txt(root, txt_file)
    
19.         self.imgs = imgs
    
20.         self.transform = transform
    
21.         self.target_transform = target_transform
    
22. 
    
23.     def __getitem__(self, index):
    
24.         x_path, y_path = self.imgs[index]
    
25.         img_x = Image.open(x_path)
    
26.         img_y = Image.open(y_path)
    
27.         if self.transform is not None:
    
28.             img_x = self.transform(img_x)
    
29.         if self.target_transform is not None:
    
30.             img_y = self.target_transform(img_y)
    
31.         return img_x, img_y
    
32. 
    
33.     def __len__(self):
    
34.         return len(self.imgs)

复制代码

txt_file:

1. 000_label.png
   
2. 001_label.png
   
3. 002_label.png
   
4. 003_label.png
   
5. 004_label.png
   
6. 005_label.png
   
7. 006_label.png
   
8. 007_label.png
   
9. 008_label.png
   
10. 009_label.png
    
11. 010_label.png
    
12. 011_label.png
    
13. ……

复制代码