最近一直在复现目标检测和实例分割相关模型,其中训练代码时多数用coco2017数据集,我下载了之后发现尽管数据集只有600多MB,尽管文件不大但是含有12多w图片,传到我租的单卡服务器实在是太久了,Xftp还容易卡顿,因此我找到了开源的规模很小的coco数据集,即coco128,由名字即可知道包含128张训练照片,但是为了我复现模型方便,我决定写一个脚本来将这128张图片按照7:2:1的比例,划分为train、valid、test文件夹。
coco128数据集的文件夹结构如下:
在这里插入图片描述
coco128数据集的资源如下:
链接:https://pan.baidu.com/s/11acYzpRayPnej0tIWyDUKA?pwd=6qu6
提取码:6qu6
–来自百度网盘超级会员V5的分享

1.python脚本如下:

import os
import random
from tqdm import tqdm
# 指定 images 文件夹路径,images文件夹里面是图片,这里的路径换为自己的,一个是图片路径
image_dir = "/root/autodl-tmp/detr-main/datasets/coco128/images/train"
# 指定 labels 文件夹路径,一个是标注文件路径

label_dir = "/root/autodl-tmp/detr-main/datasets/coco128/label"

# 创建一个空列表来存储有效图片的路径

valid_images = []

# 创建一个空列表来存储有效 label 的路径

valid_labels = []

# 遍历 images 文件夹下的所有图片

for image_name in os.listdir(image_dir):

    # 获取图片的完整路径

    image_path = os.path.join(image_dir, image_name)

    # 获取图片文件的扩展名

    # ext = os.path.splitext(image_name)[-1]

    # 根据扩展名替换成对应的 label 文件名
    label_name=image_name.replace(".jpg",".txt")
    # label_name = image_name.replace(ext, ".txt")

    # 获取对应 label 的完整路径

    label_path = os.path.join(label_dir, label_name)

    # 判断 label 是否存在

    if not os.path.exists(label_path):

        # 删除图片
        
        os.remove(image_path)

        print("deleted:", image_path)

    else# 将图片路径添加到列表中

        valid_images.append(image_path)

        # 将label路径添加到列表中

        valid_labels.append(label_path)

        # print("valid:", image_path, label_path)

# 遍历每个有效图片路径

for i in tqdm(range(len(valid_images))):

    image_path = valid_images[i]

    label_path = valid_labels[i]

    # 随机生成一个概率

    r = random.random()

    # 判断图片应该移动到哪个文件夹

    # train:valid:test = 7:3:1

    if r < 0.1:

        # 移动到 test 文件夹,这里的三个路径换为自己的

        destination = "/root/autodl-tmp/detr-main/datasets/test"

    elif r < 0.2:

        # 移动到 valid 文件夹

        destination = "/root/autodl-tmp/detr-main/datasets/valid"

    else:

        # 移动到 train 文件夹

        destination = "/root/autodl-tmp/detr-main/datasets/train"

    # 生成目标文件夹中图片的新路径

    image_destination_path = os.path.join(destination, "images", os.path.basename(image_path))

    # 移动图片到目标文件夹

    os.rename(image_path, image_destination_path)

    # 生成目标文件夹中 label 的新路径

    label_destination_path = os.path.join(destination, "labels", os.path.basename(label_path))

    # 移动 label 到目标文件夹

    os.rename(label_path, label_destination_path)

print("valid images:", valid_images)

#输出有效label路径列表

print("valid labels:", valid_labels)

脚本正常运行后生成文件夹如下:
在这里插入图片描述

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐