每个样本文件夹中图片数目不一,如何把一个文件夹中的图片加载到一起,例如五张图片变成(15,256,256)的张量形式加载在dataloader上
时间: 2024-09-19 08:14:34 浏览: 23
如果每个样本文件夹内的图片数量不一致,但你想统一成相同的维度(比如(15, 256, 256)),你可以在加载数据时采取以下策略:
1. **填充或截断**:对于少于15张图片的文件夹,可以选择填充额外的黑色或者其他固定颜色的图片来达到15张。对于多余的图片,则可以选择忽略或只使用前15张。你可以使用PIL库创建占位图。
2. **分块读取**:根据每个文件夹实际的图片数量,将其分为几个部分,每一部分保证有15张图片。然后,把这些部分合并成一个大批次,再按需切分成(15, 256, 256)的张量。
3. **动态数据集**:定义一个可变长度的数据集,根据每个文件夹的实际内容动态调整。每次读取时,根据剩余文件夹的图片数量生成相应的张量。
以下是一个简单的例子,假设我们使用`os.listdir`获取文件夹内图片路径,并用`Image.open`读取图片:
```python
import os
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
class CustomDataset(Dataset):
def __init__(self, folder_path):
self.img_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg')] # 假设都是.jpg格式
self.max_images_per_folder = 15
self.image_shape = (256, 256)
def __len__(self):
return len(self.img_files) // self.max_images_per_folder
def __getitem__(self, idx):
img_files_in_batch = self.img_files[idx * self.max_images_per_folder : (idx + 1) * self.max_images_per_folder]
images = [Image.open(os.path.join(folder_path, img)) for img in img_files_in_batch]
# 根据实际图片数量进行填充或截断
padded_images = pad_or_truncate(images, self.image_shape, fill_value=0)
return torch.stack(padded_images)
def pad_or_truncate(images, target_shape, fill_value=0):
# 实现填充或截断功能...
pass
# 使用Dataloader加载数据
folder_path = 'path_to_your_samples'
dataset = CustomDataset(folder_path)
dataloader = DataLoader(dataset, batch_size=self.max_images_per_folder, shuffle=True)
```
阅读全文