用Python写一段随机旋转数据集图片的代码,在旋转的同时修改对应标签文件
时间: 2023-05-24 14:06:41 浏览: 128
以下是使用Python和OpenCV库对数据集图片进行随机旋转并修改对应标签文件的代码:
```python
import os
import cv2
import numpy as np
# 定义函数,实现随机旋转图片和修改对应标签文件
def rotate_dataset_images(src_path, dst_path, angle_range, label_path):
"""
src_path: 原始数据集路径
dst_path: 旋转后输出的数据集路径
angle_range: 旋转角度范围,如(0, 360)
label_path: 对应标签文件路径
"""
# 如果输出目录不存在则创建
if not os.path.exists(dst_path):
os.makedirs(dst_path)
# 读取标签文件,获取每张图片对应的标签
labels = {}
with open(label_path, "r") as f:
for line in f:
img_path, label = line.strip().split() # 获取图片路径和标签
labels[img_path] = label
# 遍历原始数据集,随机旋转并输出到指定目录
for img_path in os.listdir(src_path):
img = cv2.imread(os.path.join(src_path, img_path)) # 读取图片
h, w = img.shape[:2] # 获取图片尺寸
angle = np.random.uniform(angle_range[0], angle_range[1]) # 随机旋转角度
rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0) # 生成旋转矩阵
rotated_img = cv2.warpAffine(img, rotation_matrix, (w, h)) # 应用旋转矩阵
rotated_img_path = os.path.join(dst_path, f"rotated_{angle:.2f}_{img_path}") # 构造输出路径
cv2.imwrite(rotated_img_path, rotated_img) # 输出旋转后的图片
# 更新对应标签文件中的标签
with open(label_path, "r") as f:
data = f.readlines()
with open(label_path, "w") as f:
for line in data:
if img_path in line:
f.write(f"{rotated_img_path} {labels[img_path]}\n")
else:
f.write(line)
# 测试代码
if __name__ == "__main__":
src_path = "/path/to/dataset"
dst_path = "/path/to/rotated_dataset"
angle_range = (0, 360)
label_path = "/path/to/label.txt"
rotate_dataset_images(src_path, dst_path, angle_range, label_path)
```
上述代码中,我们定义了一个名为`rotate_dataset_images()`的函数,用于随机旋转数据集图片和修改对应标签文件。该函数的参数如下:
- `src_path`:原始数据集路径。
- `dst_path`:旋转后输出的数据集路径。
- `angle_range`:旋转角度范围,由一个二元组表示,如`(0, 360)`表示旋转角度在$0^{\circ}$到$360^{\circ}$之间随机选择。
- `label_path`:对应标签文件路径。
首先,我们在函数中读取标签文件,将图片路径和对应标签存储在一个字典中。然后,我们遍历原始数据集中的每张图片,并针对每张图片进行随机旋转。在旋转过程中,我们使用`cv2.getRotationMatrix2D()`函数生成旋转矩阵,再使用`cv2.warpAffine()`函数应用旋转矩阵,得到旋转后的图片。我们通过构造输出路径将旋转后的图片输出到指定目录中。在输出旋转后的图片之后,我们需要更新对应标签文件中的标签。我们读取标签文件的全部内容到一个列表中,在更新目标图片对应的标签的同时将其它行的内容原封不动地写回到标签文件中。这里我们采用了读写标签文件分别打开的方式,确保更新图片标签不影响原有标签内容。
在测试代码中,我们设置了以下参数:
- 原始数据集路径为`/path/to/dataset`。
- 旋转后输出的数据集路径为`/path/to/rotated_dataset`,如果该目录不存在则会自动创建。
- 旋转角度范围为$(0^{\circ}, 360^{\circ})$。
- 对应标签文件路径为`/path/to/label.txt`。
运行函数即可实现对数据集图片的随机旋转和对应标签的修改。
阅读全文