对voc2007数据集进行person的单类数据抽取

最新推荐文章于 2023-06-06 19:20:10 发布

图像处理大大大大大牛啊

最新推荐文章于 2023-06-06 19:20:10 发布

阅读量1.1k

点赞数 1

文章标签： voc 数据集 python

本文链接：https://blog.csdn.net/qq_26696715/article/details/123655256

版权

神经网络专栏收录该内容

16 篇文章

订阅专栏

对voc2007数据集进行person的单类数据抽取

在这里插入图片描述

使用方式

ython tovoc.py VOC2007 VOC 14 1

4个参数分别为
VOC2007原始数据集路径
新数据集路径
标签下标
新标签下标

tovoc.py

#!/usr/bin/ python
# -*- encoding: utf-8 -*-
'''
@文件    :tovoc.py
@说明    :python tovoc.py VOC2007 VOC 14 1
@时间    :2022/03/21 16:43:49
@作者    :刘子沫
@邮箱    :spiritai@qq.com
@版本    :1.0
'''


from __future__ import print_function

import argparse
import glob
import os
import os.path as osp
import sys
import shutil
# shutil.copy(源文件, 目标文件) 目标文件可以是目录
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("input_dir", help="input annotated directory")
    parser.add_argument("output_dir", help="output dataset directory")
    parser.add_argument("inclass", help="labels num")
    parser.add_argument("outclass", help="labels num")
    args = parser.parse_args()

    if osp.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    os.makedirs(osp.join(args.output_dir, "JPEGImages"))
    os.makedirs(osp.join(args.output_dir, "SegmentationClassPNG"))
    os.makedirs(osp.join(args.output_dir, "SegmentationObjectPNG"))
    os.makedirs(osp.join(args.output_dir, "labels"))
    print("Creating dataset:", args.output_dir)

    for filename in glob.glob(osp.join(args.input_dir,"labels", "*.txt")):
        print("Generating dataset from:", filename)

        base = osp.splitext(osp.basename(filename))[0]

        in_txt_file = osp.join(
            args.input_dir, "labels", base + ".txt"
        )
        # yololabels
        out_txt_file = osp.join(
            args.output_dir, "labels", base + ".txt"
        )
        datas = []
        nwedatas = []
        flag = True
        with open(in_txt_file, "r") as f1:
            datas = f1.readlines()
        for dt in datas:
            a,b,c,d,e = dt.replace("\n","").split(" ")
            if a==str(args.inclass):
                a = args.outclass
            else:
                flag=False
                break
            nwedatas.append([a,b,c,d,e])

        if not flag:
            continue
        try:
            shutil.copy(osp.join(args.input_dir, "SegmentationClass", base + ".png"), osp.join(args.output_dir, "SegmentationClassPNG", base + ".png"))
            shutil.copy(osp.join(args.input_dir, "SegmentationObject", base + ".png"), osp.join(args.output_dir, "SegmentationObjectPNG", base + ".png"))
        except:
            continue

        with open(out_txt_file, "w+") as f2:
            for cls_id,x, y, w, h in nwedatas:
                f2.write("%s %s %s %s %s\n"%(cls_id,x, y, w, h))


        shutil.copy(osp.join(args.input_dir, "JPEGImages", base + ".jpg"), osp.join(args.output_dir, "JPEGImages", base + ".jpg"))



if __name__ == "__main__":
    main()