笔者在数据预处理时需要把一组文件夹下图像读取为tensor格式并分组保存,最后存在一个list下,以下给出代码
import torch
from PIL import Image
import numpy as np
import os
import pandas as pd
def split_images_into_groups(folder_path):
# 设置每组的大小
# 可按需自行修改
group_sizes = [190, 190, 190, 190, 190, 240]
# 获取文件夹中的所有图像
images = os.listdir(folder_path)
# 对图像进行排序以确保正确的顺序
images.sort()
# 将图像分成6组
image_groups = []
start = 0
for size in group_sizes:
end = start + size
image_groups.append(images[start:end])
start = end
# 将每个图像读取为numpy数组并添加到一个名为“image”的list中
image = []
for group in image_groups:
group_images = []
for img_name in group:
img = Image.open(os.path.join(folder_path, img_name))
# 图片转成numpy数组
img_array = np.array(img)
# numpy数组转tensor
tensor = torch.tensor(img_array)
tensor = tensor.permute(2, 0, 1)
tensor = tensor.float()
group_images.append(tensor)
stacked_tensor = torch.stack(group_images, dim=0)
image.append(stacked_tensor)
# 现在“image”列表包含6个张量组,分别具有大小[190, 图像高度, 图像宽度, 通道数]或[240, 图像高度, 图像宽度, 通道数]
return image
以上
参考链接
发表评论