读取图像文件及其对应标签信息问题

黎明丿晓小

代码如下：

1. import torch
   
2. import torch.nn as nn
   
3. from torch.utils.data import DataLoader
   
4. from torchvision import transforms
   
5. from model import VGG16
   
6. import matplotlib.pyplot as plt
   
7. import os
   
8. from PIL import Image
   
9. 
   
10. # Set random seed
    
11. torch.manual_seed(0)
    
12. 
    
13. # Set parameters and data paths
    
14. lr = 0.0001
    
15. num_epochs = 20
    
16. batch_size = 8
    
17. data_path = 'D:/Backup/datasets/insect_multi-label'
    
18. 
    
19. # Define image loader function
    
20. def load_image_information(path):
    
21.     # Image root directory
    
22.     image_root_dir = data_path
    
23.     # Get image path
    
24.     image_dir = os.path.join(image_root_dir, path)
    
25.     # Open image in RGB format
    
26.     # PyTorch DataLoader expects images to be opened using PIL
    
27.     # It's recommended to use this method to read images
    
28.     # When reading grayscale images, use convert('L')
    
29.     return Image.open(image_dir).convert('RGB')
    
30. 
    
31. # Define custom dataset class
    
32. class CustomDataset(nn.Module):
    
33.     def __init__(self, txt_file, transform=None, loader=None):
    
34.         super(CustomDataset, self).__init__()
    
35.         self.transform = transform
    
36.         self.loader = loader
    
37. 
    
38.         self.images = []
    
39.         self.labels = []
    
40.         with open(txt_file, 'r') as file:
    
41.             for line in file:
    
42.                 #line = line.strip().split()
    
43.                 line.strip('\n')
    
44.                 line.rstrip()
    
45.                 line = line.split()
    
46.                 image_path = line[0]
    
47.                 label = [float(x) for x in line[1:len(line)]]
    
48.                 self.images.append(image_path)
    
49.                 self.labels.append(label)
    
50. 
    
51.     def __len__(self):
    
52.         return len(self.images)
    
53. 
    
54.     def __getitem__(self, index):
    
55.         image_path = self.images[index]
    
56.         label = self.labels[index]
    
57. 
    
58.         image = self.loader(image_path)
    
59. 
    
60.         if self.transform is not None:
    
61.             image = self.transform(image)
    
62. 
    
63.         return image, torch.FloatTensor(label)
    
64. 
    
65. def main():
    
66.     # Data processing
    
67.     transform_train = transforms.Compose([transforms.RandomResizedCrop(224),
    
68.                                           transforms.RandomHorizontalFlip(),
    
69.                                           transforms.ToTensor()])
    
70.     transform_val = transforms.Compose([transforms.Resize(256),
    
71.                                         transforms.CenterCrop(224),
    
72.                                         transforms.ToTensor()])
    
73. 
    
74.     train_dataset = CustomDataset(r'D:/Backup/datasets/insect_multi-label/train/train.txt',
    
75.                                   transform=transform_train, loader=load_image_information)
    
76.     train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
    
77. 
    
78.     val_dataset = CustomDataset(r'D:/Backup/datasets/insect_multi-label/val/val.txt',
    
79.                                 transform=transform_val, loader=load_image_information)
    
80.     val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)
    
81. 
    
82.     # Initialize model
    
83.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
84.     model = VGG16(num_classes=len(train_dataset.labels[0])).to(device)
    
85. 
    
86.     # Define loss function and optimizer
    
87.     criterion = nn.BCELoss()
    
88.     optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    
89. 
    
90.     # Train model
    
91.     train_losses = []
    
92.     val_losses = []
    
93.     for epoch in range(num_epochs):
    
94.         # Training
    
95.         train_loss = 0.0
    
96.         model.train()
    
97.         for images, labels in train_loader:
    
98.             images, labels = images.to(device), labels.to(device)
    
99.             optimizer.zero_grad()
    
100. 
     
101.             # Forward pass
     
102.             outputs = model(images)
     
103.             loss = criterion(outputs, labels)
     
104. 
     
105.             # Backward pass and optimize
     
106.             loss.backward()
     
107.             optimizer.step()
     
108. 
     
109.             train_loss += loss.item() * images.size(0)
     
110.         train_loss = train_loss / len(train_loader.dataset)
     
111.         train_losses.append(train_loss)
     
112. 
     
113.         # Validation
     
114.         val_loss = 0.0
     
115.         model.eval()
     
116.         with torch.no_grad():
     
117.             for images, labels in val_loader:
     
118.                 images, labels = images.to(device), labels.to(device)
     
119.                 outputs = model(images)
     
120.                 loss = criterion(outputs, labels)
     
121.                 val_loss += loss.item() * images.size(0)
     
122.             val_loss = val_loss / len(val_loader.dataset)
     
123.             val_losses.append(val_loss)
     
124. 
     
125.         print('Epoch [{}/{}], Train Loss: {:.4f}, Val Loss: {:.4f}'.
     
126.               format(epoch + 1, num_epochs, train_loss, val_loss))
     
127. 
     
128.     # Save the trained model
     
129.     torch.save(model.state_dict(), 'vgg16_multi_label_classification.pth')
     
130. 
     
131.     # Visualize the training process
     
132.     plt.plot(range(num_epochs), train_losses, '-b', label='train')
     
133.     plt.plot(range(num_epochs), val_losses, '-r', label='validation')
     
134.     plt.legend(loc='lower right')
     
135.     plt.xlabel('epoch')
     
136.     plt.ylabel('loss')
     
137.     plt.show()
     
138. 
     
139. 
     
140. if __name__ == '__main__':
     
141.     main()

复制代码

运行此代码时报错：
Traceback (most recent call last):
  File "E:/PycharmProjects/Multi-label classification/VGGNet_ML_test/train.py", line 238, in <module>
    main()
  File "E:/PycharmProjects/Multi-label classification/VGGNet_ML_test/train.py", line 172, in main
    transform=transform_train, loader=load_image_information)
  File "E:/PycharmProjects/Multi-label classification/VGGNet_ML_test/train.py", line 144, in __init__
    label = [float(x) for x in line[1:len(line)]]
  File "E:/PycharmProjects/Multi-label classification/VGGNet_ML_test/train.py", line 144, in <listcomp>
    label = [float(x) for x in line[1:len(line)]]
ValueError: could not convert string to float: 'photos.jpg'

当选择断点为label时，debug后结果如图：可以看到label的输出结果为[1.0 0.0 1.0]，转换float形式的原因是BCELoss计算时用float形式。
但是，当选择断点为transform=transform_train时，debug结果表明label = [float(x) for x in line[1:len(line)]]这行代码中的x读取为'photo.jpg'，如图，看代码x从1.0开始读入，但是代码把图像文件名也读进去了，所以转换不了。
比较奇怪，不知道怎么解决，求大神指点 ，孩子快急死了

附txt文件内容，如图所示，其中以.jpg结尾的为图像文件，后续数字表示：1代表含此标签，0表示不含此标签；一共三类标签。
附图像数据集格式如下：
dataset
----train
     ----ants
          ----1.jpg
          ----2.jpg
          ...
     ----bees
          ----1.jpg
          ----2.jpg
          ...
     ----disease
          ----1.jpg
          ----2.jpg
          ...
----train.txt
----val
     ----ants
          ----1.jpg
          ----2.jpg
          ...
     ----bees
          ----1.jpg
          ----2.jpg
          ...
     ----disease
          ----1.jpg
          ----2.jpg
          ...
----val.txt