笔者在数据预处理时需要把一组文件夹下图像读取为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

以上

参考链接

评论可见,请评论后查看内容,谢谢!!!
 您阅读本篇文章共花了: