将文件夹内的多张png图像转化为tensor并保存
时间: 2024-03-22 14:37:20 浏览: 150
将文件夹内的多张PNG图像转换为张量并保存的步骤如下:
1. 使用Python的Pillow库打开PNG图像文件。
2. 将图像数据转换为NumPy数组。
3. 将NumPy数组转换为PyTorch张量。
4. 将PyTorch张量保存为.pt文件。
5. 遍历文件夹中的所有PNG图像文件,将它们转换为张量并保存。
以下是一个示例代码,假设有一个文件夹'images',其中包含多张PNG图像文件,我们要将它们转换为张量并保存为一个.pt文件。
```python
import os
import torch
import numpy as np
from PIL import Image
# 定义文件夹路径和.pt文件名
folder_path = 'images'
file_name = 'images.pt'
# 定义一个空字典用于保存张量
tensor_dict = {}
# 遍历文件夹中的所有PNG图像文件
for filename in os.listdir(folder_path):
if filename.endswith('.png'):
# 打开PNG图像并转换为NumPy数组
img_path = os.path.join(folder_path, filename)
img_arr = np.array(Image.open(img_path))
# 将NumPy数组转换为PyTorch张量
tensor = torch.from_numpy(img_arr).permute(2, 0, 1).float()
# 将张量保存到字典中
key = os.path.splitext(filename)[0] # 使用文件名作为键值
tensor_dict[key] = tensor
# 保存张量字典为.pt文件
torch.save(tensor_dict, file_name)
```
在这个示例代码中,我们使用了`os.listdir()`函数遍历文件夹中的所有文件,并使用`os.path.join()`函数拼接文件路径。我们还使用了`os.path.splitext()`函数获取文件名和扩展名,并使用文件名作为键值来保存张量。最后,我们使用`torch.save()`函数将张量字典保存为一个.pt文件。
阅读全文