上下左右翻转照片以及标注信息扩充数据集

目录

前言：

示例项目数据结构：

源代码：

运行代码后生成的项目结构：

效果：

---

前言：

使用yolo训练模型时，遇到数据集很小的情况（一两百张），训练出来的模型效果不好，可以选择扩充数据集重新训练，这篇文章提供将照片左右翻转，上下翻转，以及将标注信息进行对应的翻转，这里的标注信息是yolo五列格式的：

分别代表：类别、归一化的中心点的x坐标、归一化的中心点的y坐标、归一化的目标狂的宽、归一化的目标框的高

示例项目数据结构：

源代码：

from PIL import Image
import os

root_dir = 'D:\python\PycharmProjects\photo-deal\\transpose-image\data'
images_dir = "images"
labels_dir = "labels"
for file in os.listdir(os.path.join(root_dir, images_dir)):
    name, ext = os.path.splitext(file)
    image = Image.open(os.path.join(root_dir, images_dir, file))
    # 获取图像的长和宽
    width, height = image.size
    image.transpose(Image.FLIP_TOP_BOTTOM).save(os.path.join(root_dir, images_dir, name + '_tb' + ext))  # 上下翻转
    # 读取对应的txt
    tbtext_list = []
    with open(os.path.join(root_dir, labels_dir, name + '.txt'), "r") as annofile:
        for line in annofile:
            line = line.strip()
            rects = line.split(" ")
            tbtext_list.append(
                rects[0] + " " + rects[1] + " " + "{:.6f}".format((1.0 - float(rects[2]))) + " " + rects[3] + " " +
                rects[4])
    with open(os.path.join(root_dir, labels_dir, name + '_tb.txt'), 'a') as classify_annofile:
        for item in tbtext_list:
            classify_annofile.write(item + "\n")  # 写入内容并换行
    # 翻转图像并保存
    image.transpose(Image.FLIP_LEFT_RIGHT).save(os.path.join(root_dir, images_dir, name + '_lr' + ext))  # 左右翻转
    # 翻转txt文件
    # 读取对应的txt
    text_list = []
    with open(os.path.join(root_dir, labels_dir, name + '.txt'), "r") as annofile:
        for line in annofile:
            line = line.strip()
            rects = line.split(" ")
            text_list.append(
                rects[0] + " " + "{:.6f}".format((1.0 - float(rects[1]))) + " " + rects[2] + " " + rects[3] + " " +
                rects[4])
    with open(os.path.join(root_dir, labels_dir, name + '_lr.txt'), 'a') as classify_annofile:
        for item in text_list:
            classify_annofile.write(item + "\n")  # 写入内容并换行
    print(file + "图像已处理")

运行代码后生成的项目结构：

效果：

原图：

翻转后：