├── README.md ├── anchor-cluster.py ├── check_voc.py ├── coco2voc.py ├── dota1.5tococo.py ├── gen_tfrecord.py ├── gen_voc_trainval_test.py ├── gen_yolo_trainval_test.py ├── rename_files.py ├── roboflow_voc2coco.py ├── visdrone2yolo.py ├── voc2coco.py ├── voc2csv.py ├── voc2txt.py ├── voc2yolo.py ├── voc_augument.py ├── voc_dataset_information.py ├── yolo2voc.py ├── yolo_augument.py └── yolov5tococo.py /README.md: -------------------------------------------------------------------------------- 1 | # OD_dataset_conversion_scripts 2 | Object detection dataset conversion scripts 3 | 4 | 1. PASCAL VOC => YOLO: voc2yolo.py 5 | 2. YOLO => PASCAL VOC: yolo2voc.py 6 | 3. PASCAL VOC => COCO: voc2coco.py 7 | 4. COCO => PASCAL VOC 8 | 9 | Use `utils_cv.detection.data.coco2voc` to complete this conversion. The process is listed below: 10 | - Install **Microsoft utils_cv** package: `pip install git+https://github.com/microsoft/ComputerVision.git@master#egg=utils_cv` 11 | - Import fumction: `from utils_cv.detection.data import coco2voc` 12 | - Function Signature: 13 | ``` 14 | Signature: 15 | coco2voc( 16 | anno_path: str, 17 | output_dir: str, 18 | anno_type: str = 'instance', 19 | download_images: bool = False, 20 | ) -> None 21 | Docstring: 22 | Convert COCO annotation (single .json file) to Pascal VOC annotations 23 | (multiple .xml files). 24 | 25 | Args: 26 | anno_path: path to coco-formated .json annotation file 27 | output_dir: root output directory 28 | anno_type: "instance" for rectangle annotation, or "keypoint" for keypoint annotation. 29 | download_images: if true then download images from their urls. 30 | ``` 31 | 5. PASCAL VOC => CSV: voc2csv.py 32 | 6. PASCAL VOC => TXT: voc2txt.py 33 | 7. PASCAL VOC dataset information: voc_dataset_information.py 34 | 8. PASCAL VOC Augmentation: voc_augument.py 35 | 9. YOLO Augmentation: yolo_augument.py 36 | 10. Rename file names: rename_files.py 37 | 11. Generate VOC/ImageSets/Main/trainval.txt(train.txt,val.txt,test.txt): voc_gen_trainval_test.py 38 | 12. Cluster anchors used in YOLO series: anchor-cluster.py 39 | -------------------------------------------------------------------------------- /anchor-cluster.py: -------------------------------------------------------------------------------- 1 | #coding=utf-8 2 | import xml.etree.ElementTree as ET 3 | import numpy as np 4 | import glob 5 | 6 | def iou(box, clusters): 7 | """ 8 | 计算一个ground truth边界盒和k个先验框(Anchor)的交并比(IOU)值。 9 | 参数box: 元组或者数据，代表ground truth的长宽。 10 | 参数clusters: 形如(k,2)的numpy数组，其中k是聚类Anchor框的个数 11 | 返回：ground truth和每个Anchor框的交并比。 12 | """ 13 | x = np.minimum(clusters[:, 0], box[0]) 14 | y = np.minimum(clusters[:, 1], box[1]) 15 | if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0: 16 | raise ValueError("Box has no area") 17 | intersection = x * y 18 | box_area = box[0] * box[1] 19 | cluster_area = clusters[:, 0] * clusters[:, 1] 20 | iou_ = intersection / (box_area + cluster_area - intersection) 21 | return iou_ 22 | 23 | 24 | def avg_iou(boxes, clusters): 25 | """ 26 | 计算一个ground truth和k个Anchor的交并比的均值。 27 | """ 28 | return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])]) 29 | 30 | def kmeans(boxes, k, dist=np.median): 31 | """ 32 | 利用IOU值进行K-means聚类 33 | 参数boxes: 形状为(r, 2)的ground truth框，其中r是ground truth的个数 34 | 参数k: Anchor的个数 35 | 参数dist: 距离函数 36 | 返回值：形状为(k, 2)的k个Anchor框 37 | """ 38 | # 即是上面提到的r 39 | rows = boxes.shape[0] 40 | # 距离数组，计算每个ground truth和k个Anchor的距离 41 | distances = np.empty((rows, k)) 42 | # 上一次每个ground truth"距离"最近的Anchor索引 43 | last_clusters = np.zeros((rows,)) 44 | # 设置随机数种子 45 | np.random.seed() 46 | 47 | # 初始化聚类中心，k个簇，从r个ground truth随机选k个 48 | clusters = boxes[np.random.choice(rows, k, replace=False)] 49 | # 开始聚类 50 | while True: 51 | # 计算每个ground truth和k个Anchor的距离，用1-IOU(box,anchor)来计算 52 | for row in range(rows): 53 | distances[row] = 1 - iou(boxes[row], clusters) 54 | # 对每个ground truth，选取距离最小的那个Anchor，并存下索引 55 | nearest_clusters = np.argmin(distances, axis=1) 56 | # 如果当前每个ground truth"距离"最近的Anchor索引和上一次一样，聚类结束 57 | if (last_clusters == nearest_clusters).all(): 58 | break 59 | # 更新簇中心为簇里面所有的ground truth框的均值 60 | for cluster in range(k): 61 | clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0) 62 | # 更新每个ground truth"距离"最近的Anchor索引 63 | last_clusters = nearest_clusters 64 | 65 | return clusters 66 | 67 | # 加载自己的数据集，只需要所有labelimg标注出来的xml文件即可 68 | def load_dataset(path): 69 | dataset = [] 70 | for xml_file in glob.glob("{}/*xml".format(path)): 71 | tree = ET.parse(xml_file) 72 | # 图片高度 73 | height = int(tree.findtext("./size/height")) 74 | # 图片宽度 75 | width = int(tree.findtext("./size/width")) 76 | 77 | for obj in tree.iter("object"): 78 | # 偏移量 79 | xmin = int(obj.findtext("bndbox/xmin")) / width 80 | ymin = int(obj.findtext("bndbox/ymin")) / height 81 | xmax = int(obj.findtext("bndbox/xmax")) / width 82 | ymax = int(obj.findtext("bndbox/ymax")) / height 83 | xmin = np.float64(xmin) 84 | ymin = np.float64(ymin) 85 | xmax = np.float64(xmax) 86 | ymax = np.float64(ymax) 87 | if xmax == xmin or ymax == ymin: 88 | print(xml_file) 89 | # 将Anchor的长宽放入dateset，运行kmeans获得Anchor 90 | dataset.append([xmax - xmin, ymax - ymin]) 91 | return np.array(dataset) 92 | 93 | if __name__ == '__main__': 94 | import argparse 95 | import os 96 | parser = argparse.ArgumentParser() 97 | parser.add_argument('--voc-root', help="VOC格式数据集路径", type=str) 98 | parser.add_argument('--clusters', help="anchor数量", type=int, default=9) 99 | parser.add_argument('--input-size', help="输入网络大小", type=str, default=416) 100 | 101 | args = parser.parse_args() 102 | 103 | ANNOTATIONS_PATH = os.path.join(args.voc_root,'Annotations') # xml文件所在文件夹 104 | CLUSTERS = args.clusters #聚类数量，anchor数量 105 | INPUTDIM = args.input_size #输入网络大小 106 | 107 | data = load_dataset(ANNOTATIONS_PATH) 108 | out = kmeans(data, k=CLUSTERS) 109 | print('Boxes:') 110 | print(np.array(out)*INPUTDIM) 111 | print("Accuracy: {:.2f}%".format(avg_iou(data, out) * 100)) 112 | final_anchors = np.around(out[:, 0] / out[:, 1], decimals=2).tolist() 113 | print("Before Sort Ratios:\n {}".format(final_anchors)) 114 | print("After Sort Ratios:\n {}".format(sorted(final_anchors))) 115 | -------------------------------------------------------------------------------- /check_voc.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 检查voc数据集 3 | ''' 4 | from pathlib import Path 5 | import os 6 | import argparse 7 | import numpy as np 8 | 9 | def check_files(ann_root, img_root): 10 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀''' 11 | if os.path.exists(ann_root): 12 | ann = Path(ann_root) 13 | else: 14 | raise Exception("标注文件路径错误") 15 | if os.path.exists(img_root): 16 | img = Path(img_root) 17 | else: 18 | raise Exception("图像文件路径错误") 19 | ann_files = [] 20 | img_files = [] 21 | img_exts = [] 22 | for an, im in zip(ann.iterdir(),img.iterdir()): 23 | ann_files.append(an.stem) 24 | img_files.append(im.stem) 25 | img_exts.append(im.suffix) 26 | 27 | print('图像后缀列表：', np.unique(img_exts)) 28 | if len(np.unique(img_exts)) > 1: 29 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 30 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 31 | if set(ann_files)==set(img_files): 32 | print('标注文件和图像文件匹配') 33 | else: 34 | print('标注文件和图像文件不匹配') 35 | 36 | return np.unique(img_exts)[0] 37 | 38 | if __name__ == "__main__": 39 | 40 | parser = argparse.ArgumentParser() 41 | parser.add_argument('--voc-root', type=str, required=True, 42 | help='VOC格式数据集根目录，该目录下必须包含JPEGImages和Annotations这两个文件夹') 43 | parser.add_argument('--img_dir', type=str, required=False, 44 | help='VOC格式数据集图像存储路径，如果不指定，默认为JPEGImages') 45 | parser.add_argument('--anno_dir', type=str, required=False, 46 | help='VOC格式数据集标注文件存储路径，如果不指定，默认为Annotations') 47 | opt = parser.parse_args() 48 | 49 | print('Pascal VOC格式数据集路径：', opt.voc_root) 50 | 51 | if opt.img_dir is None: 52 | img_dir = 'JPEGImages' 53 | else: 54 | img_dir = opt.img_dir 55 | IMG_DIR = os.path.join(opt.voc_root, img_dir) 56 | if not os.path.exists(IMG_DIR): 57 | raise Exception(f'数据集图像路径{IMG_DIR}不存在！') 58 | 59 | if opt.anno_dir is None: 60 | anno_dir = 'Annotations' 61 | else: 62 | anno_dir = opt.anno_dir 63 | XML_DIR = os.path.join(opt.voc_root, anno_dir) 64 | if not os.path.exists(XML_DIR): 65 | raise Exception(f'数据集图像路径{XML_DIR}不存在！') 66 | 67 | check_files(XML_DIR, IMG_DIR) -------------------------------------------------------------------------------- /coco2voc.py: -------------------------------------------------------------------------------- 1 | ''' 2 | COCO数据集转VOC 3 | `pip install git+https://github.com/microsoft/computervision-recipes.git@master#egg=utils_cv` 4 | 5 | 参考：https://github.com/microsoft/computervision-recipes/blob/master/utils_cv/detection/references/anno_coco2voc.py 6 | 7 | coco instance标注格式转换为voc格式 8 | ''' 9 | 10 | import argparse, json 11 | import cytoolz 12 | from lxml import etree, objectify 13 | import os, re 14 | # from utils_cv.detection.data import coco2voc 15 | from pathlib import Path 16 | import argparse 17 | 18 | def instance2xml_base(anno): 19 | # anno: a dict type containing annotation infomation 20 | E = objectify.ElementMaker(annotate=False) 21 | anno_tree = E.annotation( 22 | E.folder('VOC2014_instance/{}'.format(anno['category_id'])), 23 | E.filename(anno['file_name']), 24 | E.size( 25 | E.width(anno['width']), 26 | E.height(anno['height']), 27 | E.depth(3) 28 | ), 29 | E.segmented(0), 30 | ) 31 | return anno_tree 32 | 33 | 34 | def instance2xml_bbox(anno, bbox_type='xyxy'): 35 | """bbox_type: xyxy (xmin, ymin, xmax, ymax); xywh (xmin, ymin, width, height)""" 36 | assert bbox_type in ['xyxy', 'xywh'] 37 | if bbox_type == 'xyxy': 38 | xmin, ymin, w, h = anno['bbox'] 39 | xmax = xmin+w 40 | ymax = ymin+h 41 | else: 42 | xmin, ymin, xmax, ymax = anno['bbox'] 43 | E = objectify.ElementMaker(annotate=False) 44 | anno_tree = E.object( 45 | E.name(anno['category_id']), 46 | E.bndbox( 47 | E.xmin(xmin), 48 | E.ymin(ymin), 49 | E.xmax(xmax), 50 | E.ymax(ymax) 51 | ), 52 | E.difficult(anno['iscrowd']) 53 | ) 54 | return anno_tree 55 | 56 | 57 | def parse_instance(content, outdir): 58 | categories = {d['id']: d['name'] for d in content['categories']} 59 | 60 | # EDITED - make sure image_id is of type int (and not of type string) 61 | for i in range(len(content['annotations'])): 62 | content['annotations'][i]['image_id'] = int(content['annotations'][i]['image_id']) 63 | 64 | # EDITED - save all annotation .xml files into same sub-directory 65 | anno_dir = os.path.join(outdir, "annotations") 66 | if not os.path.exists(anno_dir): 67 | os.makedirs(anno_dir) 68 | 69 | # merge images and annotations: id in images vs image_id in annotations 70 | merged_info_list = list(map(cytoolz.merge, cytoolz.join('id', content['images'], 'image_id', content['annotations']))) 71 | 72 | # convert category id to name 73 | for instance in merged_info_list: 74 | assert 'category_id' in instance, f"WARNING: annotation error: image {instance['file_name']} has a rectangle without a 'category_id' field." 75 | instance['category_id'] = categories[instance['category_id']] 76 | 77 | # group by filename to pool all bbox in same file 78 | img_filenames = {} 79 | names_groups = cytoolz.groupby('file_name', merged_info_list).items() 80 | for index, (name, groups) in enumerate(names_groups): 81 | print(f"Converting annotations for image {index} of {len(names_groups)}: {name}") 82 | assert not name.lower().startswith(("http:","https:")), "Image seems to be a url rather than local. Need to set 'download_images' = False" 83 | 84 | anno_tree = instance2xml_base(groups[0]) 85 | # if one file have multiple different objects, save it in each category sub-directory 86 | filenames = [] 87 | for group in groups: 88 | filename = os.path.splitext(name)[0] + ".xml" 89 | 90 | # EDITED - save all annotations in single folder, rather than separate folders for each object 91 | #filenames.append(os.path.join(outdir, re.sub(" ", "_", group['category_id']), filename)) 92 | filenames.append(os.path.join(anno_dir, filename)) 93 | 94 | anno_tree.append(instance2xml_bbox(group, bbox_type='xyxy')) 95 | 96 | for filename in filenames: 97 | etree.ElementTree(anno_tree).write(filename, pretty_print=True) 98 | 99 | def coco2voc(anno_file, output_dir, anno_type): 100 | '''对原本代码进行裁剪，只支持instance标注格式''' 101 | if not os.path.exists(output_dir): 102 | os.makedirs(output_dir) 103 | content = json.load(open(anno_file, 'r')) 104 | 105 | if anno_type == 'instance': 106 | # EDITED - save all annotations in single folder, rather than separate folders for each object 107 | # make subdirectories 108 | # sub_dirs = [re.sub(" ", "_", cate['name']) for cate in content['categories']] #EDITED 109 | # for sub_dir in sub_dirs: 110 | # sub_dir = os.path.join(output_dir, str(sub_dir)) 111 | # if not os.path.exists(sub_dir): 112 | # os.makedirs(sub_dir) 113 | parse_instance(content, output_dir) 114 | else: 115 | raise Exception('格式不符合，请检查！') 116 | 117 | if __name__ == '__main__': 118 | 119 | parser = argparse.ArgumentParser() 120 | parser.add_argument('--anno-path', type=str, required=True, 121 | help='COCO .json 格式文件路径') 122 | parser.add_argument('--out-dir', type=str, 123 | help='VCO格式标注文件存储根路径，例如-out-dir=VOC,该路径下会自动生成VOC/annotations文件夹，该文件夹下存储.xml格式标注文件') 124 | parser.add_argument('--anno-type',type=str, default="instance", 125 | help='"instance" for rectangle annotation, or "keypoint" for keypoint annotation.') 126 | 127 | opt = parser.parse_args() 128 | 129 | if opt.out_dir is None: 130 | out_dir = Path(opt.anno_path).parent.parent / 'VOCAnnotations' 131 | else: 132 | out_dir = opt.out_dir 133 | coco2voc(anno_file=opt.anno_path, output_dir=out_dir, anno_type=opt.anno_type) -------------------------------------------------------------------------------- /dota1.5tococo.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from PIL import Image 4 | from tqdm import tqdm 5 | import numpy as np 6 | info = {"description": "DOTA dataset from WHU", "url": "http://caption.whu.edu.cn", "year": 2018, "version": "1.0"} 7 | licenses = {"url": "http://creativecommons.org/licenses/by-nc/2.0/", "id": 1, "name": "Attribution-NonCommercial License"} 8 | categories = [] 9 | cat_names = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 'tennis-court','basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter','container-crane'] 10 | for i, catName in enumerate(cat_names, start=1): 11 | categories.append({"id": i, "name": "%s" % catName, "supercategory": "%s" % catName}) 12 | 13 | images = [] 14 | annotations = [] 15 | aug = "/home/lxy/dota/data/aug" 16 | augmented = "/home/lxy/dota/data/augmented" 17 | train_small = "/home/lxy/dota/data/train_small" 18 | trainsplit_HBB = "/home/lxy/dota/data/trainsplit_HBB" 19 | val_small = "/home/lxy/dota/data/val_small" 20 | valsplit_HBB = r"D:\BaiduNetdiskDownload" 21 | # dataset_path = [augmented, train_small, trainsplit_HBB, val_small, valsplit_HBB] 22 | dataset_path = [valsplit_HBB] 23 | imgid = 0 24 | annid = 0 25 | for path in dataset_path: 26 | img_path = os.path.join(path, "JPEGImages") 27 | label_path = os.path.join(path, "DOTA-v1.5_val_hbb") 28 | for file in tqdm(os.listdir(label_path)): 29 | img_name = file.replace("txt", "png") 30 | im = Image.open(os.path.join(img_path, img_name)) 31 | w, h = im.size 32 | imgid += 1 33 | images.append({"license": 1, "file_name": "%s" % img_name, \ 34 | "height": h, "width": w, "id": imgid}) 35 | 36 | f = open(os.path.join(label_path, file)) 37 | for line in f.readlines(): 38 | line = "".join(line).strip("\n").split(" ") 39 | # a bbox has 4 points, a category name and a difficulty 40 | if len(line) != 10: 41 | print(path, file) 42 | else: 43 | annid += 1 44 | catid = cat_names.index(line[-2]) + 1 45 | w_bbox = int(line[4][:-2]) - int(line[0][:-2]) 46 | h_bbox = int(line[5][:-2]) - int(line[1][:-2]) 47 | # bbox = [line[0], line[1], str(w_bbox)+'.0', str(h_bbox)+'.0'] 48 | bbox = [np.double(line[0]),np.double(line[1]), np.double(w_bbox), np.double(h_bbox)] 49 | seg = [np.double(line[0]), np.double(line[1]), np.double(line[2]), np.double(line[3]),np.double(line[4]), np.double(line[5]), np.double(line[6]), np.double(line[7])] 50 | annotations.append({"id": annid, "image_id": imgid, "category_id": catid, \ 51 | "segmentation": [seg], \ 52 | "area": float(w_bbox*h_bbox), \ 53 | "bbox": bbox, "iscrowd": 0}) 54 | 55 | f.close() 56 | 57 | my_json = {"info": info, "licenses": licenses, "images": images, "annotations": annotations, "categories": categories} 58 | 59 | json_path = os.path.join(valsplit_HBB,'val1.json') 60 | with open(json_path, "w+") as f: 61 | json.dump(my_json, f) 62 | print("writing json file done!") -------------------------------------------------------------------------------- /gen_tfrecord.py: -------------------------------------------------------------------------------- 1 | ''' 2 | tf2.x 版本转换PASCAL VOC至tfrecord格式 3 | 1. 使用labelimg等标注工具制作pascal voc格式数据集，注意：图像存储在JPEGImages文件夹，xml标注文件存储在Annotations文件夹 4 | 2. 将xml格式转换成csv格式，本脚本使用xml_to_csv函数已经在内部实现 5 | 3. 将csv转成TFrecord格式，注意tf1.x版本和tf2.x版本接口是不一样的 6 | 7 | 参考链接：https://www.pythonf.cn/read/109620 8 | 9 | 注意事项：对于自定义数据集，需要指定labels列表 10 | ''' 11 | from __future__ import division 12 | from __future__ import print_function 13 | from __future__ import absolute_import 14 | 15 | import os 16 | import io 17 | import pandas as pd 18 | import tensorflow as tf 19 | 20 | from PIL import Image 21 | # from object_detection.utils import dataset_util 22 | from collections import namedtuple, OrderedDict 23 | from tqdm import tqdm 24 | import argparse 25 | import glob 26 | import xml.etree.ElementTree as ET 27 | from pathlib import Path 28 | # flags = tf.app.flags 29 | # flags.DEFINE_string('csv_input', '', 'Path to the CSV input') 30 | # flags.DEFINE_string('output_path', '', 'Path to output TFRecord') 31 | # FLAGS = flags.FLAGS 32 | # TO-DO replace this with label map 33 | # labels = ['cow', 'tvmonitor', 'car', 'aeroplane', 'sheep', 34 | # 'motorbike', 'train', 'chair', 'person', 'sofa', 35 | # 'pottedplant', 'diningtable', 'horse', 'bottle', 36 | # 'boat', 'bus', 'bird', 'bicycle', 'cat', 'dog'] 37 | 38 | # 根据自定义数据集修改该列表 39 | labels = ['raccoon'] 40 | 41 | def class_text_to_int(row_label): 42 | return labels.index(row_label)+1 43 | 44 | def split(df, group): 45 | data = namedtuple('data', ['filename', 'object']) 46 | gb = df.groupby(group) 47 | return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)] 48 | 49 | def xml_to_csv(xml_anno, data_type): 50 | ''' 51 | xml_anno: pascal voc标准文件路径 52 | data_type:['trainvaltest','train','val','trainval','test'] 53 | ''' 54 | xml_list = [] 55 | # xml_files = [] 56 | txt_file = str(Path(xml_anno).parent/'ImageSets/Main'/f'{data_type}.txt') 57 | xml_files = [os.path.join(xml_anno, k.strip()+'.xml') for k in open(txt_file,'r').readlines()] 58 | # for xml_file in glob.glob(xml_anno + '/*.xml'): 59 | for xml_file in xml_files: 60 | tree = ET.parse(xml_file) 61 | root = tree.getroot() 62 | for member in root.findall('object'): 63 | value = (root.find('filename').text, 64 | int(root.find('size')[0].text), 65 | int(root.find('size')[1].text), 66 | member[0].text, 67 | int(member[4][0].text), 68 | int(member[4][1].text), 69 | int(member[4][2].text), 70 | int(member[4][3].text) 71 | ) 72 | xml_list.append(value) 73 | column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax'] 74 | xml_df = pd.DataFrame(xml_list, columns=column_name) 75 | return xml_df 76 | 77 | def create_tf_example(group, path): 78 | with tf.io.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid: 79 | encoded_jpg = fid.read() 80 | encoded_jpg_io = io.BytesIO(encoded_jpg) 81 | image = Image.open(encoded_jpg_io) 82 | width, height = image.size 83 | 84 | filename = group.filename.encode('utf8') 85 | image_format = opt.format.encode('utf8') 86 | xmins = [] 87 | xmaxs = [] 88 | ymins = [] 89 | ymaxs = [] 90 | classes_text = [] 91 | classes = [] 92 | 93 | for index, row in group.object.iterrows(): 94 | xmins.append(row['xmin'] / width) 95 | xmaxs.append(row['xmax'] / width) 96 | ymins.append(row['ymin'] / height) 97 | ymaxs.append(row['ymax'] / height) 98 | classes_text.append(row['class'].encode('utf8')) 99 | classes.append(class_text_to_int(row['class'])) 100 | 101 | tf_example = tf.train.Example(features=tf.train.Features(feature={ 102 | 'image/height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height])), 103 | 'image/width': tf.train.Feature(int64_list=tf.train.Int64List(value=[width])), 104 | 'image/filename':tf.train.Feature(bytes_list=tf.train.BytesList(value=[filename])), 105 | 'image/source_id': tf.train.Feature(bytes_list=tf.train.BytesList(value=[filename])), 106 | 'image/encoded': tf.train.Feature(bytes_list=tf.train.BytesList(value=[encoded_jpg])), 107 | 'image/format': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_format])), 108 | 'image/object/bbox/xmin': tf.train.Feature(float_list=tf.train.FloatList(value=xmins)), 109 | 'image/object/bbox/xmax': tf.train.Feature(float_list=tf.train.FloatList(value=xmaxs)), 110 | 'image/object/bbox/ymin': tf.train.Feature(float_list=tf.train.FloatList(value=ymins)), 111 | 'image/object/bbox/ymax':tf.train.Feature(float_list=tf.train.FloatList(value=ymaxs)), 112 | 'image/object/class/text': tf.train.Feature(bytes_list=tf.train.BytesList(value=classes_text)), 113 | 'image/object/class/label': tf.train.Feature(int64_list=tf.train.Int64List(value=classes)), 114 | })) 115 | return tf_example 116 | 117 | 118 | def main(voc_root, output_name): 119 | img_path = os.path.join(voc_root, 'JPEGImages') 120 | # examples = pd.read_csv(csv_input) 121 | imgset_path = os.path.join(voc_root, 'ImageSets/Main') 122 | if not os.path.exists(imgset_path): 123 | raise Exception('ImageSets/Main文件夹不存在，请通过脚本生成相应的文件！') 124 | txt_files = ['trainvaltest.txt','train.txt','val.txt','trainval.txt','test.txt'] 125 | 126 | valid_txt = [] 127 | for k in txt_files: 128 | txt = os.path.join(imgset_path, k) 129 | if os.path.exists(txt): 130 | valid_txt.append(k[:-4]) 131 | 132 | if valid_txt: 133 | print(valid_txt) 134 | else: 135 | raise Exception('ImageSets/Main文件夹下不存在train.txt等文件，请检查数据集！') 136 | 137 | for data_type in valid_txt: 138 | output_path = output_name + f'_{data_type}.tfrecord' 139 | output_path = os.path.join(voc_root, output_path) 140 | writer = tf.io.TFRecordWriter(output_path) 141 | examples = xml_to_csv(os.path.join(voc_root, 'Annotations'), data_type) 142 | grouped = split(examples, 'filename') 143 | 144 | for group in tqdm(grouped): 145 | tf_example = create_tf_example(group, img_path) 146 | writer.write(tf_example.SerializeToString()) 147 | 148 | writer.close() 149 | print('Successfully created the TFRecords: {}'.format(output_path)) 150 | 151 | if __name__ == '__main__': 152 | # tf.app.run() 153 | parser = argparse.ArgumentParser() 154 | parser.add_argument("--voc-root", type=str, required=True, help="PASCAL VOC 数据集路径，包含JPEGImages和Annotations两个文件夹") 155 | parser.add_argument("--output_name", type=str, default="voc2020", help="tfrecord文件名称,默认保存在VOC根路径") 156 | parser.add_argument("--format", type=str, default="jpg", help="图像格式") 157 | opt = parser.parse_args() 158 | main(opt.voc_root, opt.output_name) -------------------------------------------------------------------------------- /gen_voc_trainval_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Pascal VOC格式数据集生成ImageSets/Main/train.txt,val.txt,trainval.ttx和test.txt 3 | ''' 4 | from pathlib import Path 5 | import os 6 | import sys 7 | import xml.etree.ElementTree as ET 8 | import random 9 | import argparse 10 | from sklearn.model_selection import train_test_split 11 | from sklearn.utils import shuffle 12 | import shutil 13 | 14 | def mkdir(path): 15 | # 去除首位空格 16 | path = path.strip() 17 | # 去除尾部 \ 符号 18 | path = path.rstrip("\\") 19 | # 判断路径是否存在 20 | # 存在 True 21 | # 不存在 False 22 | isExists = os.path.exists(path) 23 | # 判断结果 24 | if not isExists: 25 | # 如果不存在则创建目录 26 | # 创建目录操作函数 27 | os.makedirs(path) 28 | print(path + ' 创建成功') 29 | return True 30 | else: 31 | # 如果目录存在则不创建，并提示目录已存在 32 | print(path + ' 目录已存在') 33 | return False 34 | 35 | if __name__ == '__main__': 36 | 37 | parser = argparse.ArgumentParser() 38 | parser.add_argument('--voc-root', type=str, required=True, 39 | help='VOC格式数据集根目录，该目录下必须包含JPEGImages和Annotations这两个文件夹') 40 | parser.add_argument('--test-ratio',type=float, default=0.2, 41 | help='验证集比例，默认为0.3') 42 | opt = parser.parse_args() 43 | 44 | voc_root = opt.voc_root 45 | print('Pascal VOC格式数据集路径：', voc_root) 46 | 47 | xml_file = [] 48 | img_files = [] 49 | voc_anno = os.path.join(voc_root, 'Annotations') 50 | 51 | voc_jpeg = os.path.join(voc_root, 'JPEGImages') 52 | 53 | voc_img_set = os.path.join(voc_root, 'ImageSets') 54 | try: 55 | shutil.rmtree(voc_img_set) 56 | except FileNotFoundError as e: 57 | a = 1 58 | mkdir(voc_img_set) 59 | 60 | ImgSetsMain = os.path.join(voc_img_set, 'Main') 61 | try: 62 | shutil.rmtree(ImgSetsMain) 63 | except FileNotFoundError as e: 64 | a = 1 65 | mkdir(ImgSetsMain) 66 | 67 | files = [x.stem for x in Path(voc_jpeg).iterdir() if not x.stem.startswith('.')] 68 | print(files[:10]) 69 | print('>>>随机划分VOC数据集') 70 | print('数据集长度:',len(files)) 71 | files = shuffle(files) 72 | ratio = opt.test_ratio 73 | trainval, test = train_test_split(files, test_size=ratio) 74 | train, val = train_test_split(trainval,test_size=0.2) 75 | print('训练集数量: ',len(train)) 76 | print('验证集数量: ',len(val)) 77 | print('测试集数量: ',len(test)) 78 | 79 | def write_txt(txt_path, data): 80 | '''写入txt文件''' 81 | with open(txt_path,'w') as f: 82 | for d in data: 83 | f.write(str(d)) 84 | f.write('\n') 85 | # 写入各个txt文件 86 | trainvaltest_txt = os.path.join(ImgSetsMain,'trainvaltest.txt') 87 | write_txt(trainvaltest_txt, files) 88 | 89 | trainval_txt = os.path.join(ImgSetsMain,'trainval.txt') 90 | write_txt(trainval_txt, trainval) 91 | 92 | train_txt = os.path.join(ImgSetsMain,'train.txt') 93 | write_txt(train_txt, train) 94 | 95 | val_txt = os.path.join(ImgSetsMain,'val.txt') 96 | write_txt(val_txt, val) 97 | 98 | test_txt = os.path.join(ImgSetsMain,'test.txt') 99 | write_txt(test_txt, test) 100 | -------------------------------------------------------------------------------- /gen_yolo_trainval_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | YOLO格式数据集生成train.txt,val.txt,trainval.ttx和test.txt 3 | ''' 4 | from pathlib import Path 5 | import os 6 | import sys 7 | # from voc2coco import voc_root 8 | import xml.etree.ElementTree as ET 9 | import random 10 | import argparse 11 | from sklearn.model_selection import train_test_split 12 | from sklearn.utils import shuffle 13 | import shutil 14 | 15 | def mkdir(path): 16 | # 去除首位空格 17 | path = path.strip() 18 | # 去除尾部 \ 符号 19 | path = path.rstrip("\\") 20 | # 判断路径是否存在 21 | # 存在 True 22 | # 不存在 False 23 | isExists = os.path.exists(path) 24 | # 判断结果 25 | if not isExists: 26 | # 如果不存在则创建目录 27 | # 创建目录操作函数 28 | os.makedirs(path) 29 | print(path + ' 创建成功') 30 | return True 31 | else: 32 | # 如果目录存在则不创建，并提示目录已存在 33 | print(path + ' 目录已存在') 34 | return False 35 | 36 | def write_txt(txt_path, data): 37 | '''写入txt文件''' 38 | with open(txt_path,'w') as f: 39 | for d in data: 40 | f.write(str(d)) 41 | f.write('\n') 42 | 43 | if __name__ == '__main__': 44 | 45 | parser = argparse.ArgumentParser() 46 | parser.add_argument('--yolo-root', type=str, required=True, 47 | help='YOLO格式数据集根目录，该目录下必须包含images和labels这两个文件夹') 48 | parser.add_argument('--from_voc',type=bool, default=False, 49 | help='从VOC数据集中的ImageSets/Main文件夹下提取') 50 | parser.add_argument('--voc-root',type=str, 51 | help='VOC数据集路径，需要包含ImageSets/Main文件夹') 52 | parser.add_argument('--test-ratio',type=float, default=0.2, 53 | help='测试集比例，默认为0.2') 54 | parser.add_argument('--ext', type=str, default='.png', 55 | help='YOLO图像数据后缀，注意带"." ' ) 56 | opt = parser.parse_args() 57 | 58 | yolo_root = opt.yolo_root 59 | print('YOLO格式数据集路径：', yolo_root) 60 | 61 | yolo_anno_root = os.path.join(yolo_root, 'labels') 62 | assert Path(yolo_anno_root).exists(), '{}不存在'.format(yolo_anno_root) 63 | yolo_img_root = os.path.join(yolo_root, 'images') 64 | assert Path(yolo_img_root).exists(), '{}不存在'.format(yolo_img_root) 65 | 66 | if opt.from_voc: 67 | print('从VOC数据集中分割数据集') 68 | if not opt.voc_root: 69 | raise Exception('需要提供VOC格式路径') 70 | voc_root = opt.voc_root 71 | voc_sets = os.path.join(voc_root,'ImageSets/Main') 72 | voc_img_root = os.path.join(voc_root,'JPEGImages') 73 | if not os.path.exists(voc_img_root): 74 | raise Exception('VOC数据集中没有JPEGImages文件夹') 75 | 76 | img_suffix = set([x.suffix for x in Path(voc_img_root).iterdir()]) 77 | if len(img_suffix) != 1: 78 | raise Exception('VOC数据集中JPEGImages文件夹中的图片格式不一致') 79 | img_suffix = img_suffix.pop() 80 | print('VOC数据集中图片后缀：', img_suffix) 81 | if not os.path.exists(voc_sets): 82 | raise Exception('VOC数据集不存在ImageSets/Main路径') 83 | else: 84 | file_lists = list(Path(voc_sets).iterdir()) 85 | for file in file_lists: 86 | img_ids = [x.strip() for x in open(file,'r').readlines()] 87 | img_full_path = [os.path.join(yolo_img_root, img_id+img_suffix) for img_id in img_ids] 88 | file_to_write = os.path.join(yolo_root,file.name) 89 | write_txt(file_to_write, img_full_path) 90 | else: 91 | print('从YOLO数据集中按比例随机分割数据集') 92 | 93 | files = [str(x) for x in Path(yolo_img_root).iterdir()] 94 | print('数据集长度:',len(files)) 95 | files = shuffle(files) 96 | ratio = opt.test_ratio 97 | trainval, test = train_test_split(files, test_size=ratio) 98 | train, val = train_test_split(trainval, test_size=0.2) 99 | print('训练集数量: ',len(train)) 100 | print('验证集数量: ',len(val)) 101 | print('测试集数量: ',len(test)) 102 | 103 | # 写入各个txt文件 104 | trainval_txt = os.path.join(yolo_root,'trainval.txt') 105 | write_txt(trainval_txt, trainval) 106 | 107 | train_txt = os.path.join(yolo_root,'train.txt') 108 | write_txt(train_txt, train) 109 | 110 | val_txt = os.path.join(yolo_root,'val.txt') 111 | write_txt(val_txt, val) 112 | 113 | test_txt = os.path.join(yolo_root,'test.txt') 114 | write_txt(test_txt, test) 115 | -------------------------------------------------------------------------------- /rename_files.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import os 3 | import numpy as np 4 | from tqdm import tqdm 5 | import shutil 6 | import argparse 7 | 8 | def create_dir(ROOT:str): 9 | if not os.path.exists(ROOT): 10 | os.mkdir(ROOT) 11 | else: 12 | shutil.rmtree(ROOT) # 先删除，再创建 13 | os.mkdir(ROOT) 14 | 15 | def rename_files(img_root, ann_root, img_ext='.png', ann_ext='.txt'): 16 | '''对图片和对应的标签进行重命名，支持voc格式和yolo格式(注意ann_ext后缀)''' 17 | p1 = Path(img_root) 18 | p2 = Path(ann_root) 19 | renamed_img = os.path.join(p1.parent, 'RenamedImgs') 20 | create_dir(renamed_img) 21 | renamed_anno = os.path.join(p2.parent, 'RenamedAnnos') 22 | create_dir(renamed_anno) 23 | 24 | imgs, annos = [], [] 25 | for img, anno in zip(p1.iterdir(),p2.iterdir()): 26 | imgs.append(img.name.split('.')[0]) # 这里用'.'进行分割，因此要保证文件名中只有区分后缀的一个小数点 27 | annos.append(anno.name.split('.')[0]) 28 | imgs= sorted(imgs) 29 | annos = sorted(annos) 30 | # print(imgs[:10], annos[:10]) 31 | # print((set(imgs)|set(annos))-set(imgs)&set(annos)) 32 | print(len(imgs), len(annos)) 33 | print(set(imgs)-set(annos)) 34 | print(set(annos)-set(imgs)) 35 | assert set(imgs)==set(annos) # 检查图片文件名和标签文件名是否一致 36 | 37 | LENGTH = len(imgs) 38 | print('图像数量：', LENGTH) 39 | for new_num, id in tqdm(zip(range(1,LENGTH+1), imgs), total=LENGTH): 40 | src_img_path = os.path.join(img_root, id+img_ext) # 原始Pascal格式数据集的图像全路径 41 | dst_img_path = os.path.join(renamed_img, str(new_num)+img_ext) # coco格式下的图像存储路径 42 | shutil.copy(src_img_path, dst_img_path) 43 | 44 | src_xml_path = os.path.join(ann_root, id+ann_ext) # 原始Pascal格式数据集的图像全路径 45 | dst_xml_path = os.path.join(renamed_anno, str(new_num)+ann_ext) # coco格式下的图像存储路径 46 | shutil.copy(src_xml_path, dst_xml_path) 47 | 48 | def check_files(ann_root, img_root): 49 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀''' 50 | if os.path.exists(ann_root): 51 | ann = Path(ann_root) 52 | else: 53 | raise Exception("标注文件路径错误") 54 | if os.path.exists(img_root): 55 | img = Path(img_root) 56 | else: 57 | raise Exception("图像文件路径错误") 58 | ann_files = [] 59 | img_files = [] 60 | img_exts = [] 61 | anno_exts = [] 62 | for an, im in zip(ann.iterdir(),img.iterdir()): 63 | ann_files.append(an.stem) 64 | img_files.append(im.stem) 65 | img_exts.append(im.suffix) 66 | anno_exts.append(an.suffix) 67 | 68 | print('图像后缀列表：', np.unique(img_exts)) 69 | print('标注文件后缀列表：', np.unique(anno_exts)) 70 | if len(np.unique(img_exts)) > 1: 71 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 72 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 73 | if set(ann_files)==set(img_files): 74 | print('标注文件和图像文件匹配') 75 | else: 76 | print('标注文件和图像文件不匹配') 77 | 78 | return np.unique(img_exts)[0], np.unique(anno_exts)[0] 79 | 80 | if __name__ == '__main__': 81 | 82 | parser = argparse.ArgumentParser() 83 | parser.add_argument('--img', type=str, required=True, 84 | help='数据集图像存储路径') 85 | parser.add_argument('--anno',type=str, required=True, 86 | help='数据集标注文件存储路径') 87 | 88 | opt = parser.parse_args() 89 | 90 | img_root = opt.img 91 | ann_root = opt.anno 92 | img_ext, anno_ext = check_files(ann_root, img_root) 93 | rename_files(img_root, ann_root, img_ext=img_ext, ann_ext=anno_ext) -------------------------------------------------------------------------------- /roboflow_voc2coco.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Roboflow支持导出voc,yolo,coco等多种数据格式，但是导出的图像名称以及组织形式发生了改变，导致image_id和id名称不一致，在yolov5验证的时候会出现一些问题 3 | 本脚本将roboflow导出的Pascal VOC格式数据集转成满足yolov5验证时使用的coco格式数据集（yolov5在训练的时候用yolo格式，但是在验证的时候为了保证指标的 4 | 一致性，yolov5还提供了pycocotools接口，这就需要对应的coco格式数据集。 5 | 6 | 适用项目： 7 | 1. yolov5 8 | 2. mmdetection2 9 | 10 | roboflow导出的voc数据集按照如下方式进行组织（train，test和valid的文件命名不重叠）： 11 | roboflow_vocdata/ 12 | -train/ # 训练集 13 | -1_jpg.rf.679262612f32c8ad16ce6546f276a1c1.jpg 14 | -1_jpg.rf.679262612f32c8ad16ce6546f276a1c1.xml 15 | -2_jpg.rf.1c8ad7446f202205729e6ba1164ee310.jpg 16 | -2_jpg.rf.1c8ad7446f202205729e6ba1164ee310.xml 17 | -... 18 | -test/ # VOC数据集ImageSets 19 | -3_jpg.rf.0a2534752963e87fe7c0cf4a8de33a9c.jpg 20 | -3_jpg.rf.0a2534752963e87fe7c0cf4a8de33a9c.xml 21 | -4_jpg.rf.47e001dc6cbce6aa26b8d12726453f38.jpg 22 | -4_jpg.rf.47e001dc6cbce6aa26b8d12726453f38.xml 23 | -... 24 | -valid/ # VOC数据集图像存储路径 25 | -5_jpg.rf.7911892930670d44d21ae8db3c525171.jpg 26 | -5_jpg.rf.7911892930670d44d21ae8db3c525171.xml 27 | -6_jpg.rf.bf0a49d44a8e3c631c7d0050dfd406ac.jpg 28 | -6_jpg.rf.bf0a49d44a8e3c631c7d0050dfd406ac.xml 29 | -... 30 | -README.roboflow.txt 31 | 32 | 可以看出文件命名已经被改变，pycocotools对coco格式的标注信息限制比较多，这种命名形式下image_id和id名称不一致，在调用相关api的时候会报错。 33 | 34 | 35 | 转换后的coco数据集组织形式为： 36 | ├───annotations 37 | ├──────instances_train.json 38 | ├──────instances_test.json 39 | ├──────instances_val.json 40 | ├───test 41 | ├──────3_jpg.rf.0a2534752963e87fe7c0cf4a8de33a9c.jpg 42 | ├──────... 43 | ├───train 44 | ├──────1_jpg.rf.679262612f32c8ad16ce6546f276a1c1.jpg 45 | ├──────... 46 | └───val 47 | ├──────5_jpg.rf.7911892930670d44d21ae8db3c525171.jpg 48 | ├──────... 49 | 50 | ''' 51 | # -*- coding: utf-8 -*- 52 | """ 53 | Created on Mon Nov 22 10:19:29 2021 54 | 55 | @author: gaoya 56 | """ 57 | import argparse 58 | import json 59 | import os 60 | import sys 61 | from pathlib import Path 62 | from threading import Thread 63 | 64 | import numpy as np 65 | import torch 66 | from tqdm import tqdm 67 | from pycocotools.coco import COCO 68 | # check_requirements(['pycocotools']) 69 | from pycocotools.coco import COCO 70 | from pycocotools.cocoeval import COCOeval 71 | import shutil 72 | 73 | def create_dir(ROOT:str): 74 | if not os.path.exists(ROOT): 75 | os.mkdir(ROOT) 76 | else: 77 | shutil.rmtree(ROOT) # 先删除，再创建 78 | os.mkdir(ROOT) 79 | 80 | def ToPascalFormat(roboflow_voc_root): 81 | # roboflow_voc_root = r"D:\Datasets\cotterpins\cotterpin.v1-640_mosaic_3x.voc" 82 | root = Path(roboflow_voc_root) 83 | 84 | create_dir(os.path.join(roboflow_voc_root, '')) 85 | imgs_dict = {'train':[], 'test':[], 'valid':[]} 86 | xmls_dict = {'train':[], 'test':[], 'valid':[]} 87 | 88 | for subdir in root.iterdir(): 89 | if subdir.is_dir(): 90 | for file in tqdm(subdir.iterdir()): 91 | if file.suffix == '.xml': 92 | print(file) 93 | xmls_dict[subdir.stem].append(str(file)) 94 | 95 | if file.suffix == '.jpg': 96 | print(file) 97 | imgs_dict[subdir.stem].append(str(file)) 98 | #%% 99 | imgs = imgs_dict['train'] + imgs_dict['test'] + imgs_dict['valid'] 100 | xmls = xmls_dict['train'] + xmls_dict['test'] + xmls_dict['valid'] 101 | 102 | img_names = [ Path(k).stem for k in imgs] 103 | xml_names = [ Path(k).stem for k in xmls] 104 | assert (img_names==xml_names) 105 | # 映射字典 106 | map_dict = {k:v for (v,k) in enumerate(img_names)} 107 | 108 | -------------------------------------------------------------------------------- /visdrone2yolo.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 将visdrone数据集转换为yolo格式，visdrone标注数据的格式为： 3 | ,,,,,,, 4 | 该数据集的类别总共有11类 5 | ignored regions(0), pedestrian(1), people(2), bicycle(3), car(4), van(5), truck(6), tricycle(7), 6 | awning-tricycle(8), bus(9), motor(10), others(11) 7 | 8 | yolo格式为： 9 | class x_center y_center width height（归一化数值） 10 | ''' 11 | import os 12 | from pathlib import Path 13 | from PIL import Image 14 | import csv 15 | from tqdm import tqdm 16 | import argparse 17 | import numpy as np 18 | import shutil 19 | def convert(size, box): 20 | dw = 1. / size[0] 21 | dh = 1. / size[1] 22 | x = (box[0] + box[2] / 2) * dw 23 | y = (box[1] + box[3] / 2) * dh 24 | w = box[2] * dw 25 | h = box[3] * dh 26 | return (x, y, w, h) 27 | 28 | def check_files(ann_root, img_root): 29 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀 30 | return： 31 | 返回图像后缀 32 | ''' 33 | if os.path.exists(ann_root): 34 | ann = Path(ann_root) 35 | else: 36 | raise Exception("标注文件路径错误") 37 | if os.path.exists(img_root): 38 | img = Path(img_root) 39 | else: 40 | raise Exception("图像文件路径错误") 41 | ann_files = [] 42 | img_files = [] 43 | img_exts = [] 44 | for an, im in zip(ann.iterdir(),img.iterdir()): 45 | ann_files.append(an.stem) 46 | img_files.append(im.stem) 47 | img_exts.append(im.suffix) 48 | 49 | print('图像后缀列表：', np.unique(img_exts)) 50 | if len(np.unique(img_exts)) > 1: 51 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 52 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 53 | if set(ann_files)==set(img_files): 54 | print('标注文件和图像文件匹配') 55 | else: 56 | print('标注文件和图像文件不匹配') 57 | 58 | return np.unique(img_exts)[0] 59 | 60 | 61 | def create_dir(ROOT:str): 62 | if not os.path.exists(ROOT): 63 | os.mkdir(ROOT) 64 | else: 65 | shutil.rmtree(ROOT) # 先删除，再创建 66 | os.mkdir(ROOT) 67 | 68 | if __name__ == '__main__': 69 | 70 | parser = argparse.ArgumentParser() 71 | parser.add_argument('--visdrone-root', type=str, required=True, 72 | help='visDrone数据集根目录，该目录下必须包含annotations和images的两个文件夹') 73 | 74 | parser.add_argument('--yolo-label-dir', type=str, default=None, 75 | help='directory to save yolo label.') 76 | 77 | opt = parser.parse_args() 78 | 79 | vis_dir = opt.visdrone_root 80 | vis_img_dir = os.path.join(vis_dir, 'images') # visdrone图像存储路径 81 | vis_anno_dir = os.path.join(vis_dir, 'annotations') # visdrone标注文件存储路径 82 | # 检查数据集 83 | img_suffix = check_files(vis_anno_dir, vis_img_dir) 84 | 85 | if opt.yolo_label_dir is None: 86 | yolo_label_dir = os.path.join(vis_dir,'labels') 87 | if not os.path.exists(yolo_label_dir): 88 | os.makedirs(yolo_label_dir) 89 | else: 90 | yolo_label_dir = opt.yolo_label_dir 91 | print('YOLO标签存储路径：', yolo_label_dir) 92 | 93 | total_imgs = len(os.listdir(vis_anno_dir)) 94 | annos = Path(vis_anno_dir).iterdir() 95 | 96 | for anno in tqdm(annos, total=total_imgs): 97 | ans = '' 98 | # print(anno) 99 | if anno.suffix != '.txt': 100 | continue 101 | # load image 102 | with Image.open(os.path.join(vis_img_dir,anno.stem+img_suffix)) as Img: 103 | img_size = Img.size 104 | # read annotation file 105 | # print(img_size) 106 | with open(os.path.join(vis_anno_dir, str(anno)),) as f: 107 | lines = f.readlines() 108 | save_path = os.path.join(yolo_label_dir,anno.stem+'.txt') # path to save yolo format annotation 109 | for line in lines: 110 | row = line.strip().split(',') 111 | if row[4] == '0': 112 | continue 113 | bb = convert(img_size, tuple(map(int, row[:4]))) 114 | ans = ans + str(int(row[5])-1) + ' ' + ' '.join(str(a) for a in bb) + '\n' 115 | with open(save_path, 'w') as outfile: 116 | outfile.write(ans) 117 | # outfile.close() 118 | -------------------------------------------------------------------------------- /voc2coco.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Pascal VOC格式数据集转COCO格式数据集 3 | 适用项目： 4 | 1. https://github.com/zylo117/Yet-Another-EfficientDet-Pytorch 5 | 2. mmdetection2 6 | 7 | 数据集按照如下方式进行组织： 8 | datasets/ 9 | -Annotations/ # VOC格式标注存储路径 10 | -*.xml 11 | -ImageSets/ # VOC数据集ImageSets 12 | -Main/ 13 | -train.txt 14 | -trainval.txt 15 | -val.txt 16 | -test.txt 17 | -JPEGImages/ # VOC数据集图像存储路径 18 | -*.jpg 19 | -CocoFormat/ # 本脚本生成的coco格式数据集默认存储位置 20 | -trainval/ # COCO trainval图像路径 21 | -*.jpg 22 | -train/ # COCO train图像路径 23 | -*.jpg 24 | -val/ # COCO val图像路径 25 | -*.jpg 26 | -test/ # COCO test图像路径 27 | -*.jpg 28 | -annotations # COCO json标注文件路径 29 | -instances_trainval.json 30 | -instances_train.json 31 | -instances_val.json 32 | -instances_test.json 33 | ##============== 重要通告 ===============## 34 | 笔者在使用一些COCO格式目标检测模型的时候，发现如果image_id不为int型的话会有很多问题，例如在使用torchvison中的 35 | COCODetection时会遇到错误： 36 | 37 | path = coco.loadImgs(img_id)[0]['file_name'] 38 | 39 | File "python\lib\site-packages\pycocotools\coco.py", line 230, in loadImgs 40 | return [self.imgs[id] for id in ids] 41 | 42 | File "python\lib\site-packages\pycocotools\coco.py", line 230, in 43 | return [self.imgs[id] for id in ids] 44 | 45 | KeyError: '0' 46 | 47 | 原因是image_id值是[]，直接报错，因此需要考虑将VOC格式下的文件名全部重命名为数字后再进行转换，使用参数选项--rename即可 48 | 49 | ''' 50 | from pathlib import Path 51 | import os 52 | import sys 53 | import xml.etree.ElementTree as ET 54 | import numpy as np 55 | import argparse 56 | from sklearn.model_selection import train_test_split 57 | from sklearn.utils import shuffle 58 | import shutil 59 | import json 60 | from typing import Dict, List 61 | from tqdm import tqdm 62 | import re 63 | from collections import Counter 64 | from imageio import imread 65 | def get_label2id(labels_path: str) -> Dict[str, int]: 66 | ''' 67 | id is 1 start 68 | ''' 69 | with open(labels_path, 'r') as f: 70 | labels_str = f.read().strip().split('\n') 71 | labels_ids = list(range(1, len(labels_str)+1)) 72 | return dict(zip(labels_str, labels_ids)) 73 | 74 | 75 | def get_image_info(ann_path, annotation_root, extract_num_from_imgid=True): 76 | ''' 77 | ann_path：标注文件全路径 78 | annotation_root：xml对根内容进行解析后的内容 79 | extract_num_from_imgid：是否从imageid中提取数字，对于COCO格式数据集最好使用True选项，将image_id转换为整型 80 | ''' 81 | img_name = os.path.basename(ann_path) 82 | img_id = os.path.splitext(img_name)[0] 83 | filename = img_id+ext 84 | 85 | if extract_num_from_imgid and isinstance(img_id, str): 86 | # 采用正则表达式，支持转换的文件命名：0001.png, cls_0021.png, cls0123.jpg, 00123abc.png等 87 | img_id = int(re.findall(r'\d+', img_id)[0]) 88 | 89 | try: 90 | size = annotation_root.find('size') 91 | width = int(size.findtext('width')) 92 | height = int(size.findtext('height')) 93 | except: 94 | img_path = Path(ann_path).parent.parent.joinpath('JPEGImages', filename) 95 | width, height = imread(str(img_path)).shape[:2] 96 | 97 | image_info = { 98 | 'file_name': filename, 99 | 'height': height, 100 | 'width': width, 101 | 'id': img_id 102 | } 103 | return image_info 104 | 105 | 106 | def counting_labels(anno_root: str): 107 | ''' 108 | 获取pascal voc格式数据集中的所有标签名 109 | anno_root: pascal标注文件路径，一般为Annotations 110 | ''' 111 | all_classes = [] 112 | 113 | for xml_file in os.listdir(anno_root): 114 | xml_file = os.path.join(anno_root, xml_file) 115 | # print(xml_file) 116 | xml = open(xml_file,) # encoding='utf-8' 117 | tree=ET.parse(xml) 118 | root = tree.getroot() 119 | for obj in root.iter('object'): 120 | 121 | class_ = obj.find('name').text.strip() 122 | all_classes.append(class_) 123 | 124 | print(Counter(all_classes)) 125 | 126 | labels = sorted(list(set(all_classes))) 127 | print('标签数据：', labels) 128 | print('标签长度：', len(labels)) 129 | print('写入标签信息...{}'.format(os.path.join(opt.voc_root,'labels.txt'))) 130 | with open( os.path.join(opt.voc_root,'labels.txt') , 'w') as f: 131 | for k in labels: 132 | f.write(k) 133 | f.write('\n') 134 | 135 | def get_coco_annotation_from_obj(obj, label2id): 136 | label = obj.findtext('name').strip() 137 | assert label in label2id, f"Error: {label} is not in label2id !" 138 | category_id = label2id[label] 139 | bndbox = obj.find('bndbox') 140 | xmin = int(bndbox.findtext('xmin')) - 1 141 | ymin = int(bndbox.findtext('ymin')) - 1 142 | xmax = int(bndbox.findtext('xmax')) 143 | ymax = int(bndbox.findtext('ymax')) 144 | assert xmax > xmin and ymax > ymin, f"Box size error !: (xmin, ymin, xmax, ymax): {xmin, ymin, xmax, ymax}" 145 | o_width = xmax - xmin 146 | o_height = ymax - ymin 147 | ann = { 148 | 'area': o_width * o_height, 149 | 'iscrowd': 0, 150 | 'bbox': [xmin, ymin, o_width, o_height], 151 | 'category_id': category_id, 152 | 'ignore': 0, 153 | # 起始点是左上角，按照顺时针方向 154 | 'segmentation': [[xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax]] 155 | } 156 | return ann 157 | 158 | def convert_xmls_to_cocojson(annotation_paths: List[str], 159 | label2id: Dict[str, int], 160 | output_jsonpath: str, 161 | extract_num_from_imgid: bool = True): 162 | output_json_dict = { 163 | "images": [], 164 | "type": "instances", 165 | "annotations": [], 166 | "categories": [] 167 | } 168 | bnd_id = 1 # START_BOUNDING_BOX_ID, TODO input as args ? 169 | print('Start converting !') 170 | for a_path in tqdm(annotation_paths): 171 | # Read annotation xml 172 | ann_tree = ET.parse(a_path) 173 | ann_root = ann_tree.getroot() 174 | # print(a_path) 175 | img_info = get_image_info(ann_path=a_path, 176 | annotation_root=ann_root, 177 | extract_num_from_imgid=extract_num_from_imgid) 178 | img_id = img_info['id'] 179 | output_json_dict['images'].append(img_info) 180 | 181 | for obj in ann_root.findall('object'): 182 | ann = get_coco_annotation_from_obj(obj=obj, label2id=label2id) 183 | ann.update({'image_id': img_id, 'id': bnd_id}) 184 | output_json_dict['annotations'].append(ann) 185 | bnd_id = bnd_id + 1 186 | 187 | for label, label_id in label2id.items(): 188 | category_info = {'supercategory': 'electrical_fittings', 'id': label_id, 'name': label} 189 | output_json_dict['categories'].append(category_info) 190 | 191 | with open(output_jsonpath, 'w') as f: 192 | output_json = json.dumps(output_json_dict) 193 | f.write(output_json) 194 | 195 | def create_dir(ROOT:str): 196 | if not os.path.exists(ROOT): 197 | os.makedirs(ROOT) 198 | 199 | def check_files(ann_root, img_root): 200 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀''' 201 | if os.path.exists(ann_root): 202 | ann = Path(ann_root) 203 | else: 204 | raise Exception("标注文件路径错误") 205 | if os.path.exists(img_root): 206 | img = Path(img_root) 207 | else: 208 | raise Exception("图像文件路径错误") 209 | 210 | ann_files = [] 211 | img_files = [] 212 | img_exts = [] 213 | for an in ann.iterdir(): 214 | ann_files.append(an.stem) 215 | 216 | for im in img.iterdir(): 217 | img_files.append(im.stem) 218 | img_exts.append(im.suffix) 219 | 220 | if not len(ann_files)==len(img_files): 221 | raise Exception("图像数据和标注数据数量不一致！") 222 | 223 | print('图像后缀列表：', np.unique(img_exts)) 224 | if len(np.unique(img_exts)) > 1: 225 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 226 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 227 | if set(ann_files)==set(img_files): 228 | print('标注文件和图像文件匹配') 229 | else: 230 | print('标注文件和图像文件不匹配') 231 | 232 | return np.unique(img_exts)[0] 233 | 234 | if __name__ == '__main__': 235 | 236 | parser = argparse.ArgumentParser() 237 | parser.add_argument('--voc-root', type=str, required=True, 238 | help='VOC格式数据集根目录，该目录下必须包含存储图像和标注文件的两个文件夹，例如官方格式下有JPEGImages和Annotations两个文件夹') 239 | parser.add_argument('--img_dir', type=str, required=False, 240 | help='VOC格式数据集图像存储路径，如果不指定，默认为JPEGImages') 241 | parser.add_argument('--anno_dir', type=str, required=False, 242 | help='VOC格式数据集标注文件存储路径，如果不指定，默认为Annotations') 243 | parser.add_argument('--coco-dir', type=str, default='CocoFormatData', 244 | help='COCO数据集存储路径，默认为VOC数据集相同路径下新建文件夹CocoDataset') 245 | parser.add_argument('--test-ratio',type=float, default=0.2, 246 | help='验证集比例，默认为0.3') 247 | parser.add_argument('--rename',type=bool, default=False, 248 | help='是否对VOC数据集进行数字化重命名') 249 | parser.add_argument('--label-file', type=str, required=False, 250 | help='path to label list.') 251 | 252 | opt = parser.parse_args() 253 | 254 | voc_root = opt.voc_root 255 | print('Pascal VOC格式数据集路径：', voc_root) 256 | 257 | xml_file = [] 258 | img_files = [] 259 | 260 | if opt.img_dir is None: 261 | img_dir = 'JPEGImages' 262 | else: 263 | img_dir = opt.img_dir 264 | voc_jpeg = os.path.join(voc_root, img_dir) 265 | if not os.path.exists(voc_jpeg): 266 | raise Exception(f'数据集图像路径{voc_jpeg}不存在！') 267 | 268 | if opt.anno_dir is None: 269 | anno_dir = 'Annotations' 270 | else: 271 | anno_dir = opt.anno_dir 272 | voc_anno = os.path.join(voc_root, anno_dir) # 273 | if not os.path.exists(voc_anno): 274 | raise Exception(f'数据集图像路径{voc_anno}不存在！') 275 | 276 | ext = check_files(voc_anno, voc_jpeg) # 检查图像后缀 277 | assert ext is not None, "请检查图像后缀是否正确！" 278 | print() 279 | ##============================## 280 | ## 对文件进行数字化重命名，需要对ImageSets/Main下的分割数据做相应修改 281 | ##============================## 282 | if opt.rename==True: 283 | renamed_jpeg = os.path.join(voc_root,'RenamedJPEGImages') 284 | create_dir(renamed_jpeg) 285 | renamed_xml = os.path.join(voc_root,'RenamedAnnotations') 286 | create_dir(renamed_xml) 287 | 288 | p1 = Path(voc_jpeg) 289 | p2 = Path(voc_anno) 290 | imgs = sorted([x.stem for x in p1.iterdir() if not x.stem.startswith('.')]) 291 | annos = sorted([x.stem for x in p2.iterdir() if not x.stem.startswith('.')]) 292 | 293 | assert imgs==annos 294 | 295 | # 非规则名称与数字id的映射字典{'a'：0, 'b':1, ...} 296 | names_to_id_dict = {k:v for (v,k) in enumerate(imgs)} 297 | 298 | print('图像数量：', len(imgs)) 299 | for name, id in tqdm(names_to_id_dict.items()): 300 | src_img_path = os.path.join(voc_jpeg, name+ext) # 原始Pascal格式数据集的图像全路径 301 | # print(src_img_path) 302 | dst_img_path = os.path.join(renamed_jpeg, str(id)+ext) # coco格式下的图像存储路径 303 | # print(dst_img_path) 304 | shutil.copy2(src_img_path, dst_img_path) 305 | 306 | src_xml_path = os.path.join(voc_anno, name+'.xml') # 原始Pascal格式数据集的图像全路径 307 | dst_xml_path = os.path.join(renamed_xml, str(id)+'.xml') # coco格式下的图像存储路径 308 | shutil.copy2(src_xml_path, dst_xml_path) 309 | 310 | voc_jpeg = renamed_jpeg # 将重命名后的图像路径赋值给JPEG 311 | voc_anno = renamed_xml # 将重命名后的标注路径赋值给ANNO 312 | 313 | ImgSets = os.path.join(voc_root, 'ImageSets') 314 | if not os.path.exists(ImgSets): 315 | os.mkdir(ImgSets) 316 | ImgSetsMain = os.path.join(ImgSets,'Main') 317 | 318 | create_dir(ImgSetsMain) 319 | 320 | #== COCO 数据集路径 321 | coco_root = os.path.join(str(Path(voc_root).parent), Path(voc_root).stem + opt.coco_dir) # pascal voc转coco格式的存储路径 322 | if os.path.exists(coco_root): 323 | shutil.rmtree(coco_root) 324 | create_dir(coco_root) 325 | 326 | txt_files = ['trainvaltest','train','val','trainval','test'] 327 | 328 | coco_dirs = [] 329 | for dir_ in txt_files: 330 | DIR = os.path.join(coco_root, dir_) 331 | coco_dirs.append(DIR) 332 | create_dir(DIR) 333 | 334 | coco_anno = os.path.join(coco_root, 'annotations') # coco标注文件存放路径 335 | create_dir(coco_anno) 336 | 337 | # 利用VOC ImageSets数据划分信息，注意ImageSets/main/train.txt文件只记录图片名称，没有后缀 338 | # 所有图片名称 339 | files = [x.stem for x in Path(voc_jpeg).iterdir() if not x.stem.startswith('.')] 340 | # files = list(Path(voc_jpeg).iterdir()) 341 | # files = [str(x).replace('.jpg','') for x in files] 342 | print('数据集长度:',len(files)) 343 | assert os.path.exists(os.path.join(voc_root, 'ImageSets/Main/trainval.txt')) 344 | if os.path.exists(os.path.join(voc_root, 'ImageSets/Main/trainval.txt')): 345 | 346 | print('>>>使用ImageSet信息分割数据集') 347 | trainval_file = os.path.join(voc_root, 'ImageSets/Main/trainval.txt') 348 | if opt.rename: 349 | trainval = [names_to_id_dict[i.strip()] for i in open(trainval_file,'r').readlines()] 350 | else: 351 | # trainval = [i.strip() for i in open(trainval_file,'r').readlines()] 352 | trainval = [i.replace('\n','') for i in open(trainval_file,'r').readlines()] 353 | # trainval = [os.path.join(os.path.join(coco_root,'trainval'),name) for name in trainval_name] 354 | 355 | train_file = os.path.join(voc_root, 'ImageSets/Main/train.txt') 356 | if opt.rename: 357 | train = [names_to_id_dict[i.strip()] for i in open(train_file,'r').readlines()] 358 | else: 359 | # train = [i.strip() for i in open(train_file,'r').readlines()] 360 | train = [i.replace('\n','') for i in open(train_file,'r').readlines()] 361 | # train = [os.path.join(os.path.join(coco_root,'train'),name) for name in train_name] 362 | 363 | val_file = os.path.join(voc_root, 'ImageSets/Main/val.txt') 364 | if opt.rename: 365 | val = [names_to_id_dict[i.strip()] for i in open(val_file,'r').readlines()] 366 | else: 367 | val = [i.replace('\n','') for i in open(val_file,'r').readlines()] 368 | # val = [os.path.join(os.path.join(coco_root,'val'),name) for name in val_name] 369 | 370 | test_file = os.path.join(voc_root, 'ImageSets/Main/test.txt') 371 | if opt.rename: 372 | test = [names_to_id_dict[i.strip()] for i in open(test_file,'r').readlines()] 373 | else: 374 | test = [i.replace('\n','') for i in open(test_file,'r').readlines()] 375 | # test = [os.path.join(os.path.join(coco_root,'test'),name) for name in test_name] 376 | 377 | print('>>>训练集数量: ',len(train)) 378 | print('>>>训练集验证集数量: ',len(trainval)) 379 | print('>>>验证集数量: ',len(val)) 380 | print('>>>测试集数量: ',len(test)) 381 | 382 | else: 383 | print('>>>随机划分COCO数据集') 384 | files = shuffle(files) 385 | ratio = opt.test_ratio 386 | trainval, test = train_test_split(files, test_size=ratio) 387 | train, val = train_test_split(trainval,test_size=0.2) 388 | print('训练集数量: ',len(train)) 389 | print('验证集数量: ',len(val)) 390 | print('测试集数量: ',len(test)) 391 | 392 | def write_txt(txt_path, data): 393 | with open(txt_path,'w') as f: 394 | for d in data: 395 | f.write(str(d)) 396 | f.write('\n') 397 | 398 | # 写入各个txt文件 399 | datas = [files, train, val, trainval, test] 400 | 401 | for txt, data in zip(txt_files, datas): 402 | txt_path = os.path.join(ImgSetsMain, txt+'.txt') 403 | write_txt(txt_path, data) 404 | 405 | # 遍历xml文件，得到所有标签值，并且保存为labels.txt 406 | if opt.label_file: 407 | print('从自定义标签文件读取！') 408 | labels = opt.label_file 409 | else: 410 | print('从xml文件自动处理标签！') 411 | counting_labels(voc_anno) 412 | labels = os.path.join(voc_root, 'labels.txt') 413 | 414 | if not os.path.isfile(labels): 415 | raise Exception('需要提供数据集标签文件路径，用于按顺序转换数值id，如果没有，需要手动创建！') 416 | 417 | label2id = get_label2id(labels_path=labels) 418 | print('标签值及其对应的编码值：',label2id) 419 | 420 | for name,imgs,coco_dir in tqdm(zip(txt_files,datas,coco_dirs)): 421 | 422 | annotation_paths = [] 423 | # [1] copy image files 424 | for img in imgs: 425 | # print(img) 426 | annotation_paths.append(os.path.join(voc_anno, str(img)+'.xml')) 427 | src_img_path = os.path.join(voc_jpeg, str(img)+ext) # 原始Pascal格式数据集的图像全路径 428 | # print(src_img_path) 429 | dst_img_path = os.path.join(coco_dir, str(img)+ext) # coco格式下的图像存储路径 430 | # print(dst_img_path) 431 | 432 | shutil.copy2(src_img_path, dst_img_path) 433 | 434 | # [2] convert xml to coco json format files 435 | convert_xmls_to_cocojson( 436 | annotation_paths=annotation_paths, 437 | label2id=label2id, 438 | output_jsonpath=os.path.join(coco_anno, f'instances_{name}.json'), 439 | # img_ids = imgs 440 | extract_num_from_imgid=True # 一定注意这里，COCO格式数据集image_id需要整型，可以从图片名称中抽取id号 441 | ) 442 | -------------------------------------------------------------------------------- /voc2csv.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | ''' 3 | pascal voc格式数据集转csv格式，适用项目： 4 | 1. https://github.com/fizyr/keras-retinanet 5 | 2. https://github.com/yhenon/pytorch-retinanet 6 | 7 | 这种csv格式的数据集更灵活，只要利用不同的机器学习框架提供的数据接口编译一个数据加载类即可，而且csv文本方便处理。 8 | 本程序会生成3个文件：train_csv_annotations.csv,val_csv_annotations.csv和csv_classes.csv， 9 | 其中train_csv_annotations.csv和val_csv_annotations.csv的内容格式为： 10 | path/to/image.jpg,x1,y1,x2,y2,class_name，例如： 11 | 12 | /data/imgs/img_001.jpg,837,346,981,456,cow 13 | /data/imgs/img_002.jpg,215,312,279,391,cat 14 | /data/imgs/img_002.jpg,22,5,89,84,bird 15 | /data/imgs/img_003.jpg,,,,, 16 | ... 17 | 注意是每行一个标注信息，上面例子中img_003表示不含需要训练识别的区域(region of interest,ROI) 18 | 19 | 注意需要绝对路径。 20 | csv_classes.csv的内容格式为： 21 | class_name,id 22 | 例如： 23 | cow,0 24 | cat,1 25 | bird,2 26 | ... 27 | 28 | ''' 29 | import csv 30 | import shutil 31 | import os 32 | from pathlib import Path 33 | import sys 34 | import xml.etree.ElementTree as ET 35 | import numpy as np 36 | import argparse 37 | from sklearn.model_selection import train_test_split 38 | # 获取当前文件所在文件夹 39 | dirname = os.path.dirname(os.path.abspath(__file__)) 40 | print('当前工作路径：',dirname) 41 | 42 | def create_dir(ROOT:str): 43 | if not os.path.exists(ROOT): 44 | os.mkdir(ROOT) 45 | else: 46 | shutil.rmtree(ROOT) # 先删除，再创建 47 | os.mkdir(ROOT) 48 | 49 | class PascalVOC2CSV(object): 50 | def __init__(self, voc_root, xml, imgs, ratio, 51 | trainvaltest_ann='trainvaltest_csv_annotations.csv', 52 | trainval_ann='trainval_csv_annotations.csv', 53 | train_ann='train_csv_annotations.csv', 54 | val_ann='val_csv_annotations.csv', 55 | test_ann='test_csv_annotations.csv', 56 | classes_path='csv_classes.csv', ): 57 | ''' 58 | :param voc_root: VOC数据集根目录 59 | :param xml: 所有Pascal VOC的xml文件路径组成的列表 60 | :param jpgs: 所有图像的文件路径 61 | :param train_ann_path: 训练集标注信息 62 | :param val_ann_path: 验证集标注信息 63 | :param classes_path: classes_path 64 | 65 | 返回值： 66 | 在voc_root根目录生成三个文件：train_csv_annotations.csv,val_csv_annotations.csv和csv_classes.csv 67 | ''' 68 | self.xml = xml 69 | self.imgs = imgs 70 | csv_root = os.path.join(voc_root,'CSVDataset') 71 | create_dir(csv_root) 72 | self.trainvaltest_ann = os.path.join(csv_root, trainvaltest_ann) 73 | self.trainval_ann = os.path.join(csv_root, trainval_ann) 74 | self.train_ann = os.path.join(csv_root, train_ann) 75 | self.val_ann = os.path.join(csv_root, val_ann) 76 | self.test_ann = os.path.join(csv_root, test_ann) 77 | 78 | 79 | self.classes_path = os.path.join(csv_root, classes_path) 80 | self.label=[] 81 | self.annotations=[] 82 | self.ratio = ratio 83 | self.data_transfer() 84 | self.write_file() 85 | self.valid=None 86 | self.train=None 87 | 88 | def data_transfer(self): 89 | for num, (xml_file, img_file) in enumerate( zip(self.xml, self.imgs)): 90 | try: 91 | # print(xml_file) 92 | # 进度输出 93 | sys.stdout.write('\r>> Converting image %d/%d' % ( 94 | num + 1, len(self.xml))) 95 | sys.stdout.flush() 96 | 97 | xml = open(xml_file,encoding='utf-8') 98 | tree=ET.parse(xml) 99 | root = tree.getroot() 100 | self.filename = img_file 101 | 102 | for obj in root.iter('object'): 103 | 104 | self.supercategory = obj.find('name').text.strip() 105 | if self.supercategory not in self.label: 106 | self.label.append(self.supercategory) 107 | 108 | xmlbox = obj.find('bndbox') # 进一步在bndbox寻找 109 | x1 = int(xmlbox.find('xmin').text) 110 | y1 = int(xmlbox.find('ymin').text) 111 | x2 = int(xmlbox.find('xmax').text) 112 | y2 = int(xmlbox.find('ymax').text) 113 | assert x1 < x2 and y1 < y2, 'x1 must be less than x2 and y1 must be less than y2' 114 | self.annotations.append( 115 | [os.path.join(os.path.join(dirname, 'JPEGImages'),self.filename), 116 | x1,y1,x2,y2, 117 | self.supercategory]) 118 | 119 | except: 120 | continue 121 | # print(self.annotations[:10]) 122 | # k = int(len(self.annotations) * self.ratio) # ratio是验证集比例 123 | print('\n按照比例：{:.2f}:{:.2f} 划分训练集和测试集...'.format(1-self.ratio, self.ratio)) 124 | 125 | self.trainval, self.test = train_test_split(self.annotations, test_size=self.ratio) 126 | self.train, self.val = train_test_split(self.trainval, test_size=0.2) 127 | print('训练集数量：', len(self.train)) 128 | print('验证集数量：', len(self.val)) 129 | print('测试集数量：', len(self.test)) 130 | sys.stdout.write('\n') 131 | sys.stdout.flush() 132 | 133 | def write_file(self,): 134 | print(f'写入全部数据集:{self.trainvaltest_ann}') 135 | with open(self.trainvaltest_ann, 'w', newline='') as fp: 136 | csv_writer = csv.writer(fp, dialect='excel') 137 | csv_writer.writerows(self.annotations) 138 | print(f'写入训练集:{self.trainval_ann}') 139 | with open(self.trainval_ann, 'w', newline='') as fp: 140 | csv_writer = csv.writer(fp, dialect='excel') 141 | csv_writer.writerows(self.trainval) 142 | print(f'写入训练集:{self.train_ann}') 143 | with open(self.val_ann, 'w', newline='') as fp: 144 | csv_writer = csv.writer(fp, dialect='excel') 145 | csv_writer.writerows(self.train) 146 | print(f'写入验证集:{self.train_ann}') 147 | with open(self.val_ann, 'w', newline='') as fp: 148 | csv_writer = csv.writer(fp, dialect='excel') 149 | csv_writer.writerows(self.val) 150 | print(f'写入测试集:{self.test_ann}') 151 | with open(self.test_ann, 'w', newline='') as fp: 152 | csv_writer = csv.writer(fp, dialect='excel') 153 | csv_writer.writerows(self.test) 154 | 155 | class_name=sorted(self.label) 156 | print('标签名称：', class_name) 157 | 158 | print('标签长度：', len(class_name)) 159 | class_=[] 160 | for num,name in enumerate(class_name): 161 | class_.append([name,num]) 162 | print(f'写入标签文件:{self.classes_path}...') 163 | with open(self.classes_path, 'w', newline='') as fp: 164 | csv_writer = csv.writer(fp, dialect='excel') 165 | csv_writer.writerows(class_) 166 | 167 | def check_files(ann_root, img_root): 168 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀''' 169 | if os.path.exists(ann_root): 170 | ann = Path(ann_root) 171 | else: 172 | raise Exception("标注文件路径错误") 173 | if os.path.exists(img_root): 174 | img = Path(img_root) 175 | else: 176 | raise Exception("图像文件路径错误") 177 | ann_files = [] 178 | img_files = [] 179 | img_exts = [] 180 | for an, im in zip(ann.iterdir(),img.iterdir()): 181 | ann_files.append(an.stem) 182 | img_files.append(im.stem) 183 | img_exts.append(im.suffix) 184 | 185 | print('图像后缀列表：', np.unique(img_exts)) 186 | if len(np.unique(img_exts)) > 1: 187 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 188 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 189 | if set(ann_files)==set(img_files): 190 | print('标注文件和图像文件匹配') 191 | else: 192 | print('标注文件和图像文件不匹配') 193 | 194 | return np.unique(img_exts)[0] 195 | 196 | if __name__ == '__main__': 197 | 198 | parser = argparse.ArgumentParser() 199 | parser.add_argument('--voc-root', type=str, required=True, 200 | help='VOC格式数据集根目录，该目录下必须包含JPEGImages和Annotations这两个文件夹') 201 | parser.add_argument('--img_dir', type=str, required=False, 202 | help='VOC格式数据集图像存储路径，如果不指定，默认为JPEGImages') 203 | parser.add_argument('--anno_dir', type=str, required=False, 204 | help='VOC格式数据集标注文件存储路径，如果不指定，默认为Annotations') 205 | parser.add_argument('--valid-ratio',type=float, default=0.3, 206 | help='验证集比例，默认为0.3') 207 | opt = parser.parse_args() 208 | 209 | voc_root = opt.voc_root 210 | print('Pascal VOC格式数据集路径：', voc_root) 211 | 212 | xml_file = [] 213 | img_files = [] 214 | 215 | if opt.img_dir is None: 216 | img_dir = 'JPEGImages' 217 | else: 218 | img_dir = opt.img_dir 219 | JPEG = os.path.join(voc_root, img_dir) 220 | 221 | if opt.anno_dir is None: 222 | anno_dir = 'Annotations' 223 | else: 224 | anno_dir = opt.anno_dir 225 | ANNO = os.path.join(voc_root, anno_dir) 226 | 227 | check_files(ANNO, JPEG) 228 | 229 | for k in os.listdir(JPEG): 230 | ''' 231 | 以图片所在路径进行遍历 232 | ''' 233 | img_files.append( os.path.join(JPEG, k)) 234 | xml_file.append( os.path.join(ANNO, k[:-4]+'.xml')) 235 | 236 | PascalVOC2CSV(voc_root, xml_file, img_files, ratio=opt.valid_ratio) 237 | -------------------------------------------------------------------------------- /voc2txt.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 本脚本处理的数据集格式适用程序项目： 3 | 1. https://github.com/Tianxiaomo/pytorch-YOLOv4 4 | 2. https://github.com/YunYang1994/tensorflow-yolov3 5 | 3. https://github.com/YunYang1994/TensorFlow2.0-Examples/tree/master/4-Object_Detection/YOLOV3 6 | 7 | 数据集格式： 8 | # train.txt 9 | xxx/xxx.jpg 18.19,6.32,424.13,421.83,20 323.86,2.65,640.0,421.94,20 10 | xxx/xxx.jpg 48,240,195,371,11 8,12,352,498,14 11 | # image_path x_min, y_min, x_max, y_max, class_id x_min, y_min ,..., class_id 12 | # make sure that x_max < width and y_max < height 13 | ... 14 | ... 15 | ''' 16 | 17 | import os 18 | import argparse 19 | import xml.etree.ElementTree as ET 20 | 21 | def convert_voc_annotation(data_path, data_type, anno_path, use_difficult_bbox=True): 22 | 23 | classes = ['fault'] 24 | img_inds_file = os.path.join(data_path, 'ImageSets', 'Main', data_type + '.txt') 25 | with open(img_inds_file, 'r') as f: 26 | txt = f.readlines() 27 | image_inds = [line.split('/')[-1].strip().split('.')[0].replace(' ','') for line in txt] 28 | 29 | with open(os.path.join(data_path,anno_path), 'a') as f: 30 | for image_ind in image_inds: 31 | image_path = os.path.join(data_path, 'JPEGImages', image_ind + '.png') 32 | annotation = image_path 33 | label_path = os.path.join(data_path, 'Annotations', image_ind + '.xml') 34 | root = ET.parse(label_path).getroot() 35 | objects = root.findall('object') 36 | for obj in objects: 37 | difficult = obj.find('difficult').text.strip() 38 | if (not use_difficult_bbox) and(int(difficult) == 1): 39 | continue 40 | bbox = obj.find('bndbox') 41 | class_ind = classes.index(obj.find('name').text.lower().strip()) 42 | xmin = bbox.find('xmin').text.strip() 43 | xmax = bbox.find('xmax').text.strip() 44 | ymin = bbox.find('ymin').text.strip() 45 | ymax = bbox.find('ymax').text.strip() 46 | annotation += ' ' + ','.join([xmin, ymin, xmax, ymax, str(class_ind)]) 47 | print(annotation) 48 | f.write(annotation + "\n") 49 | return len(image_inds) 50 | 51 | 52 | if __name__ == '__main__': 53 | parser = argparse.ArgumentParser() 54 | parser.add_argument("--voc-root", type=str, required=True, 55 | help='VOC格式数据集根目录，该目录下必须包含JPEGImages，Annotations和ImageSets这三个文件夹') 56 | parser.add_argument("--train_annotation", default="voc_train.txt") 57 | parser.add_argument("--test_annotation", default="voc_test.txt") 58 | opt = parser.parse_args() 59 | 60 | if os.path.exists(os.path.join( opt.voc_root, opt.train_annotation)): 61 | os.remove(os.path.join(opt.voc_root, opt.train_annotation)) 62 | if os.path.exists(os.path.join( opt.voc_root, opt.test_annotation)): 63 | os.remove(os.path.join(opt.voc_root, opt.test_annotation)) 64 | 65 | # trainval包括训练和验证，在此全部当作训练集使用 66 | num1 = convert_voc_annotation(opt.voc_root, 'trainval', opt.train_annotation, False) 67 | 68 | num2 = convert_voc_annotation(opt.voc_root, 'test', opt.test_annotation, False) 69 | print('=> The number of image for train is: %d\nThe number of image for test is:%d' %(num1, num2)) 70 | 71 | 72 | -------------------------------------------------------------------------------- /voc2yolo.py: -------------------------------------------------------------------------------- 1 | ''' 2 | PASCAL VOC格式数据集转YOLO格式数据集 3 | 适合项目地址： 4 | 1. https://github.com/eriklindernoren/PyTorch-YOLOv3 5 | 2. https://github.com/ultralytics/yolov3/ 6 | 3. https://github.com/AlexeyAB 7 | 4. https://github.com/ultralytics/yolov5/ 8 | 9 | 该项目对自定义的数据集格式要求图片要有对应的txt格式标注文件，要求图片存放在images文件夹，标签存放在labels文件夹，例如： 10 | 11 | data/custom/images/train.jpg 12 | data/custom/labels/train.txt 13 | yolo_classes.names 14 | yolo_classes_ssd.names 15 | trainval.txt 16 | train.txt 17 | val.txt 18 | 19 | 当然，images文件夹和labels这两个文件夹名称可以更改，但相应的也要在代码中做修改（PyTorch-YOLOV3项目）： 20 | ```utils/datasets.py: line 65 21 | class ListDataset(Dataset): 22 | def __init__(self, list_path, img_size=416, augment=True, multiscale=True, normalized_labels=True): 23 | with open(list_path, "r") as file: 24 | self.img_files = file.readlines() 25 | 26 | self.label_files = [ 27 | path.replace("images", "labels").replace(".png", ".txt").replace(".jpg", ".txt") 28 | ## ^^^^^^ and ^^^^^^ 修改这两处的值 29 | for path in self.img_files 30 | ] 31 | self.img_size = img_size 32 | self.max_objects = 100 33 | self.augment = augment 34 | ... 35 | ``` 36 | labels/train.txt的标注信息格式为： 37 | 38 | label_idx x_center y_center width height（归一化数值） 39 | label_idx x_center y_center width height（归一化数值） 40 | ... 41 | 42 | trainval.txt,val.txt,test.txt文件每一行记录了图像数据所在的全路径，这几个文件和yolo_classes.names 43 | 会在U版和A版的YOLOv3/v4系列的*.data配置文件中使用。在U版的yolov5模型中，数据配置文件保存在data/*.yaml文件中，其示例内容如下： 44 | ``` 45 | # train and val data as 46 | # 1) directory: path/images/, 47 | # 2) file: path/images.txt, or 48 | # 3) list: [path1/images/, path2/images/] 49 | 50 | train: /data/custom_yolo/trainval.txt 51 | val: /data/custom_yolo/test.txt 52 | 53 | # number of classes 54 | nc: 2 55 | 56 | # class names 57 | names: ['person', 'bicycle'] 58 | ``` 59 | ''' 60 | import xml.etree.ElementTree as ET 61 | import pickle 62 | import os 63 | from os import listdir, getcwd 64 | from os.path import join 65 | import pandas as pd 66 | import numpy as np 67 | from collections import Counter 68 | import argparse 69 | from tqdm import tqdm 70 | from sklearn.model_selection import train_test_split 71 | import sys 72 | import shutil 73 | from pathlib import Path 74 | from imageio import imread 75 | 76 | def counting_labels(anno_root,yolo_root): 77 | ''' 78 | 获取pascal voc格式数据集中的所有标签名 79 | anno_root: pascal标注文件路径，一般为Annotations 80 | ''' 81 | all_classes = [] 82 | 83 | for xml_file in os.listdir(anno_root): 84 | xml_file = os.path.join(anno_root, xml_file) 85 | # print(xml_file) 86 | xml = open(xml_file,encoding='utf-8') 87 | tree=ET.parse(xml) 88 | root = tree.getroot() 89 | for obj in root.iter('object'): 90 | 91 | class_ = obj.find('name').text.strip() 92 | all_classes.append(class_) 93 | 94 | print(Counter(all_classes)) 95 | 96 | labels = list(set(all_classes)) 97 | print('标签数据：', labels) 98 | print('标签长度：', len(labels)) 99 | print('写入标签信息...{}'.format(os.path.join(yolo_root,'yolo_classes.names'))) 100 | with open( os.path.join(yolo_root,'yolo_classes.names') , 'w') as f: 101 | for k in labels: 102 | f.write(k) 103 | f.write('\n') 104 | with open( os.path.join(yolo_root,'yolo_classes_ssd.names') , 'w') as f: 105 | for k in labels: 106 | f.write("\'"+k+"\'"+',') 107 | f.write('\n') 108 | return labels 109 | 110 | 111 | def convert(size, box): 112 | dw = 1./(size[0]) # 宽度缩放比例, size[0]为图像宽度width 113 | dh = 1./(size[1]) 114 | x = (box[0] + box[1])/2.0 - 1 115 | y = (box[2] + box[3])/2.0 - 1 116 | w = box[1] - box[0] 117 | h = box[3] - box[2] 118 | x = x*dw 119 | w = w*dw 120 | y = y*dh 121 | h = h*dh 122 | return (x,y,w,h) # 123 | 124 | def convert_annotation(anno_root:str, image_id, classes, dest_yolo_dir='YOLOLabels'): 125 | ''' 126 | anno_root:pascal格式标注文件路径，一般为Annotations 127 | image_id：文件名（图片名和对应的pascal voc格式标注文件名是一致的） 128 | dest_yolo_dir：yolo格式标注信息目标保存路径，默认为opt.yolo_dir 129 | ''' 130 | in_file = open( os.path.join(anno_root, image_id+'.xml'), encoding='utf-8') 131 | out_file = open(os.path.join(dest_yolo_dir, image_id+'.txt'), 'w') 132 | tree=ET.parse(in_file) 133 | root = tree.getroot() 134 | try: 135 | size = root.find('size') 136 | w = int(size.find('width').text) 137 | h = int(size.find('height').text) 138 | except: 139 | img_path = Path(anno_root).parent.joinpath('JPEGImages', image_id+img_suffix) 140 | w,h = imread(img_path).shape[:2] 141 | 142 | for obj in root.iter('object'): 143 | try: 144 | difficult = obj.find('difficult').text 145 | except: 146 | difficult = 0 147 | cls = obj.find('name').text 148 | if cls not in classes or int(difficult)==1: 149 | continue 150 | cls_id = classes.index(cls) 151 | xmlbox = obj.find('bndbox') 152 | xmin,xmax,ymin,ymax = float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text) 153 | assert xmin=0 and ymin>=0, f"Box size error !: (xmin, ymin, xmax, ymax): {xmin, ymin, xmax, ymax}" 154 | b = (xmin,xmax,ymin,ymax) 155 | bb = convert((w,h), b) 156 | out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n') 157 | 158 | def gen_image_ids(jpeg_root): 159 | ''' 160 | jpeg_root: JPEGImages文件夹路径 161 | ''' 162 | img_ids = [] 163 | 164 | for k in os.listdir(jpeg_root): 165 | img_ids.append(k) # 图片名，含后缀 166 | 167 | return img_ids 168 | 169 | def create_dir(ROOT:str): 170 | if not os.path.exists(ROOT): 171 | os.mkdir(ROOT) 172 | else: 173 | shutil.rmtree(ROOT) # 先删除，再创建 174 | os.mkdir(ROOT) 175 | 176 | def check_files(ann_root, img_root): 177 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀''' 178 | if os.path.exists(ann_root): 179 | ann = Path(ann_root) 180 | else: 181 | raise Exception("标注文件路径错误") 182 | if os.path.exists(img_root): 183 | img = Path(img_root) 184 | else: 185 | raise Exception("图像文件路径错误") 186 | ann_files = [] 187 | img_files = [] 188 | img_exts = [] 189 | for an, im in zip(ann.iterdir(),img.iterdir()): 190 | ann_files.append(an.stem) 191 | img_files.append(im.stem) 192 | img_exts.append(im.suffix) 193 | 194 | print('图像后缀列表：', np.unique(img_exts)) 195 | if len(np.unique(img_exts)) > 1: 196 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 197 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 198 | if set(ann_files)==set(img_files): 199 | print('标注文件和图像文件匹配') 200 | else: 201 | print('标注文件和图像文件不匹配') 202 | 203 | return np.unique(img_exts)[0] 204 | 205 | 206 | if __name__ == '__main__': 207 | 208 | parser = argparse.ArgumentParser() 209 | parser.add_argument('--voc-root', type=str, required=True, 210 | help='VOC格式数据集根目录，该目录下必须包含存储图像和标注文件的两个文件夹') 211 | parser.add_argument('--img_dir', type=str, required=False, 212 | help='VOC格式数据集图像存储路径，如果不指定，默认为JPEGImages') 213 | parser.add_argument('--anno_dir', type=str, required=False, 214 | help='VOC格式数据集标注文件存储路径，如果不指定，默认为Annotations') 215 | parser.add_argument('--yolo-dir',type=str, default='YOLOFormatData', 216 | help='yolo格式数据集保存路径，默认为VOC数据集相同路径下新建文件夹YOLODataset') 217 | parser.add_argument('--valid-ratio',type=float, default=0.3, 218 | help='验证集比例，默认为0.3') 219 | 220 | opt = parser.parse_args() 221 | 222 | voc_root = opt.voc_root 223 | 224 | print('Pascal VOC格式数据集路径：', voc_root) 225 | if opt.img_dir is None: 226 | img_dir = 'JPEGImages' 227 | else: 228 | img_dir = opt.img_dir 229 | jpeg_root = os.path.join(voc_root, img_dir) 230 | if not os.path.exists(jpeg_root): 231 | raise Exception(f'数据集图像路径{jpeg_root}不存在！') 232 | 233 | if opt.anno_dir is None: 234 | anno_dir = 'Annotations' 235 | else: 236 | anno_dir = opt.anno_dir 237 | anno_root = os.path.join(voc_root,anno_dir) 238 | if not os.path.exists(anno_root): 239 | raise Exception(f'数据集图像路径{anno_root}不存在！') 240 | 241 | # 确定图像后缀 242 | img_suffix = check_files(anno_root, jpeg_root) 243 | assert img_suffix is not None, "请检查图像后缀是否正确！" 244 | print('图像后缀：', img_suffix) 245 | # YOLO数据集存储路径，YOLOFormat 246 | dest_yolo_dir = os.path.join(str(Path(voc_root).parent), Path(voc_root).stem+opt.yolo_dir) 247 | # 248 | image_ids = [x.name for x in Path(jpeg_root).iterdir()] 249 | 250 | print('数据集长度：', len(image_ids)) 251 | 252 | if not os.path.exists(dest_yolo_dir): 253 | os.makedirs(dest_yolo_dir) # 创建labels文件夹,存储yolo格式标注文件 254 | 255 | yolo_labels = os.path.join(dest_yolo_dir,'labels') 256 | create_dir(yolo_labels) 257 | yolo_images = os.path.join(dest_yolo_dir,'images') 258 | create_dir(yolo_images) 259 | 260 | classes = counting_labels(anno_root,dest_yolo_dir) 261 | print('数据类别：', classes) 262 | length = len(image_ids) 263 | 264 | for idx, img in enumerate(image_ids): 265 | sys.stdout.write('\r>> Converting image %d/%d' % ( 266 | idx + 1, length)) 267 | sys.stdout.flush() 268 | image_id = img.split('.')[0] 269 | # print(image_id) 270 | # print('图像名称：', image_id) 271 | # 转换标签 272 | convert_annotation(anno_root, image_id, classes, dest_yolo_dir=yolo_labels) 273 | 274 | shutil.copy(os.path.join(voc_root, 'JPEGImages', img), yolo_images) 275 | 276 | ## 生成用于config/custom.data指定的训练训练集和验证集文件yolo_train.txt和yolo_valid.txt 277 | # 该文件的内容就是每行为图片数据在文件系统中的绝对路径 278 | 279 | ratio = opt.valid_ratio # 验证集比例 280 | def write_txt(txt_path, data): 281 | '''写入txt文件''' 282 | with open(txt_path,'w') as f: 283 | for d in data: 284 | f.write(str(d)) 285 | f.write('\n') 286 | 287 | # 所有yolo images名称 288 | files = [x.stem for x in Path(yolo_images).iterdir() if not x.stem.startswith('.')] 289 | 290 | print('数据集长度:',len(files)) 291 | assert os.path.exists(os.path.join(voc_root, 'ImageSets/Main/trainval.txt')) 292 | if os.path.exists(os.path.join(voc_root, 'ImageSets/Main/trainval.txt')): 293 | print('\n使用Pascal VOC ImageSet信息分割数据集') 294 | trainval_file = os.path.join(voc_root, 'ImageSets/Main/trainval.txt') 295 | trainval_name = [i.strip() for i in open(trainval_file,'r').readlines()] 296 | trainval = [os.path.join(yolo_images,name+img_suffix) for name in trainval_name] 297 | 298 | train_file = os.path.join(voc_root, 'ImageSets/Main/train.txt') 299 | train_name = [i.strip() for i in open(train_file,'r').readlines()] 300 | train = [os.path.join(yolo_images,name+img_suffix) for name in train_name] 301 | 302 | val_file = os.path.join(voc_root, 'ImageSets/Main/val.txt') 303 | val_name = [i.strip() for i in open(val_file,'r').readlines()] 304 | val = [os.path.join(yolo_images,name+img_suffix) for name in val_name] 305 | 306 | test_file = os.path.join(voc_root, 'ImageSets/Main/test.txt') 307 | test_name = [i.strip() for i in open(test_file,'r').readlines()] 308 | test = [os.path.join(yolo_images,name+img_suffix) for name in test_name] 309 | 310 | print('训练集数量: ',len(train_name)) 311 | print('训练集验证集数量: ',len(trainval_name)) 312 | print('验证集数量: ',len(val_name)) 313 | print('测试集数量: ',len(test_name)) 314 | 315 | else: 316 | print('\n随即划分YOLO数据集') 317 | 318 | trainval, test = train_test_split(files, test_size=ratio) 319 | train, val = train_test_split(trainval,test_size=0.2) 320 | print('训练集数量: ',len(train)) 321 | print('验证集数量: ',len(val)) 322 | print('测试集数量: ',len(test)) 323 | 324 | # 写入各个txt文件 325 | trainval_txt = os.path.join(dest_yolo_dir,'trainval.txt') 326 | write_txt(trainval_txt, trainval) 327 | 328 | train_txt = os.path.join(dest_yolo_dir,'train.txt') 329 | write_txt(train_txt, train) 330 | 331 | val_txt = os.path.join(dest_yolo_dir,'val.txt') 332 | write_txt(val_txt, val) 333 | 334 | test_txt = os.path.join(dest_yolo_dir,'test.txt') 335 | write_txt(test_txt, test) 336 | -------------------------------------------------------------------------------- /voc_augument.py: -------------------------------------------------------------------------------- 1 | ''' 2 | VOC格式数据集离线扩增 3 | 参考链接： 4 | [1] https://zhuanlan.zhihu.com/p/85292901 5 | [2] 6 | ''' 7 | import xml.etree.ElementTree as ET 8 | import pickle 9 | import os 10 | from os import getcwd 11 | import numpy as np 12 | from PIL import Image 13 | import shutil 14 | import matplotlib.pyplot as plt 15 | from tqdm import tqdm 16 | import imgaug as ia 17 | from imgaug import augmenters as iaa 18 | import argparse 19 | sometimes = lambda aug: iaa.Sometimes(0.5, aug) #建立lambda表达式， 20 | ia.seed(1) 21 | 22 | def read_xml_annotation(root, image_id): 23 | '''从xml标注文件所在路径读取标注框 24 | root: xml文件路径 25 | image_id: xml文件名称（带.xml后缀） 26 | ''' 27 | 28 | in_file = open(os.path.join(root, image_id)) 29 | tree = ET.parse(in_file) 30 | root = tree.getroot() 31 | bndboxlist = [] 32 | 33 | for object in root.findall('object'): # 找到root节点下的所有country节点 34 | bndbox = object.find('bndbox') # 子节点下节点rank的值 35 | 36 | xmin = int(bndbox.find('xmin').text) 37 | xmax = int(bndbox.find('xmax').text) 38 | ymin = int(bndbox.find('ymin').text) 39 | ymax = int(bndbox.find('ymax').text) 40 | # print(xmin,ymin,xmax,ymax) 41 | bndboxlist.append([xmin, ymin, xmax, ymax]) 42 | # print(bndboxlist) 43 | 44 | return bndboxlist 45 | 46 | 47 | def change_xml_list_annotation(root, image_id, new_target, saveroot, id): 48 | ''' 49 | root: 原始voc数据集中xml文件所在路径 50 | image_id：xml文件名，带后缀 51 | new_target: 新的bndbox:[[x1,y1,x2,y2],...[],[]] 52 | saveroot: 扩增数据集后xml文件的保存路径 53 | id: 新的xml标注文件名名称（不含后缀） 54 | ''' 55 | in_file = open(os.path.join(root, str(image_id) + '.xml')) # 这里root分别由两个意思 56 | tree = ET.parse(in_file) 57 | elem = tree.find('filename') 58 | elem.text = (str("%d" % int(id)) + '.jpg') 59 | xmlroot = tree.getroot() 60 | index = 0 61 | 62 | for object in xmlroot.findall('object'): # 找到root节点下的所有country节点 63 | bndbox = object.find('bndbox') # 子节点下节点rank的值 64 | 65 | new_xmin = new_target[index][0] 66 | new_ymin = new_target[index][1] 67 | new_xmax = new_target[index][2] 68 | new_ymax = new_target[index][3] 69 | 70 | xmin = bndbox.find('xmin') 71 | xmin.text = str(new_xmin) # 替换原来的值 72 | ymin = bndbox.find('ymin') 73 | ymin.text = str(new_ymin) 74 | xmax = bndbox.find('xmax') 75 | xmax.text = str(new_xmax) 76 | ymax = bndbox.find('ymax') 77 | ymax.text = str(new_ymax) 78 | 79 | index = index + 1 80 | 81 | tree.write(os.path.join(saveroot, str("%d" % int(id)) + '.xml')) 82 | 83 | 84 | def mkdir(path): 85 | # 去除首位空格 86 | path = path.strip() 87 | # 去除尾部 \ 符号 88 | path = path.rstrip("\\") 89 | # 判断路径是否存在 90 | # 存在 True 91 | # 不存在 False 92 | isExists = os.path.exists(path) 93 | # 判断结果 94 | if not isExists: 95 | # 如果不存在则创建目录 96 | # 创建目录操作函数 97 | os.makedirs(path) 98 | print(path + ' 创建成功') 99 | return True 100 | else: 101 | # 如果目录存在则不创建，并提示目录已存在 102 | print(path + ' 目录已存在') 103 | return False 104 | 105 | 106 | if __name__ == "__main__": 107 | 108 | parser = argparse.ArgumentParser() 109 | parser.add_argument('--voc-root', type=str, required=True, 110 | help='VOC格式数据集根目录，该目录下必须包含JPEGImages和Annotations这两个文件夹') 111 | parser.add_argument('--aug-dir',type=str, default='VOCAugumented', 112 | help='数据增强后保存的路径，默认在voc-root路径下创建一个文件夹Augument进行保存') 113 | parser.add_argument('--aug_num',type=int, default=3, 114 | help='每张图片进行扩增的次数') 115 | parser.add_argument('--ext', type=str, default='.jpg', help='图像后缀，默认为.jpg') 116 | opt = parser.parse_args() 117 | 118 | ext = opt.ext 119 | IMG_DIR = os.path.join(opt.voc_root, "JPEGImages") 120 | XML_DIR = os.path.join(opt.voc_root, "Annotations") 121 | 122 | AUGUMENT = os.path.join(os.path.dirname(opt.voc_root), opt.aug_dir) 123 | if not os.path.exists(AUGUMENT): 124 | os.mkdir(AUGUMENT) 125 | 126 | # 存储增强后的XML文件夹路径 127 | AUG_XML_DIR = os.path.join(AUGUMENT, "Annotations") 128 | try: 129 | shutil.rmtree(AUG_XML_DIR) 130 | except FileNotFoundError as e: 131 | a = 1 132 | mkdir(AUG_XML_DIR) 133 | 134 | # 存储增强后的影像文件夹路径 135 | AUG_IMG_DIR = os.path.join(AUGUMENT, "JPEGImages") 136 | try: 137 | shutil.rmtree(AUG_IMG_DIR) 138 | except FileNotFoundError as e: 139 | a = 1 140 | mkdir(AUG_IMG_DIR) 141 | 142 | AUGLOOP = opt.aug_num # 每张影像增强的数量 143 | 144 | boxes_img_aug_list = [] 145 | new_bndbox = [] 146 | new_bndbox_list = [] 147 | 148 | # 影像增强 149 | ''' 150 | seq = iaa.Sequential([ 151 | iaa.Flipud(0.2), # vertically flip 20% of all images 152 | iaa.Fliplr(0.5), # 镜像 153 | iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs 154 | iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5), 155 | iaa.Affine( 156 | translate_px={"x": 15, "y": 15}, 157 | scale=(0.8, 0.95), 158 | rotate=(-30, 30) 159 | ) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs 160 | ]) 161 | ''' 162 | seq = iaa.Sequential([ 163 | iaa.Sometimes(p=0.5, 164 | #高斯模糊 165 | then_list=[iaa.GaussianBlur(sigma=(0, 0.5))], 166 | #锐化 167 | else_list=[iaa.ContrastNormalization((0.15, 0.75), per_channel=True)] 168 | ), #以p的概率执行then_list的增强方法，以1-p的概率执行else_list的增强方法，其中then_list,else_list默认为None 169 | 170 | iaa.SomeOf(4,[ 171 | # 以下一共10个，随机选7个进行处理，也可以将7改为其他数值，继续对数据集进行扩充 172 | 173 | # 边缘检测，将检测到的赋值0或者255然后叠在原图上 174 | # sometimes(iaa.OneOf([ 175 | # iaa.EdgeDetect(alpha=(0, 0.7)), 176 | # iaa.DirectedEdgeDetect( 177 | # alpha=(0, 0.7), direction=(0.0, 1.0) 178 | # ), 179 | # ])), 180 | 181 | # 将RGB变成灰度图然后乘alpha加在原图上 182 | # iaa.Grayscale(alpha=(0.0, 1.0)), 183 | 184 | # 扭曲图像的局部区域 185 | sometimes(iaa.PiecewiseAffine(scale=(0.001, 0.005))), 186 | 187 | # 每个像素随机加减-10到10之间的数 188 | iaa.Add((-10, 10), per_channel=0.5), 189 | 190 | # 中值模糊 191 | iaa.MedianBlur(k=1, name=None, deterministic=False, random_state=None), 192 | #锐化 193 | iaa.Sharpen(alpha=0, lightness=1, name=None, deterministic=False, random_state=None), 194 | # 从最邻近像素中取均值来扰动。 195 | iaa.AverageBlur(k=1, name=None, deterministic=False, random_state=None), 196 | # 0-0.05的数值，分别乘以图片的宽和高为剪裁的像素个数，保持原尺寸 197 | # iaa.Crop(percent=(0.01, 0.01)), 198 | 199 | # iaa.Affine( 200 | # # 对图片进行仿射变化，scale缩放x,y取值，translate_percent左右上下移动 201 | # # rotate为旋转角度，shear为剪切取值范围0-360 202 | # scale={"x": (0.99, 1), "y": (0.99, 1)}, 203 | # translate_percent={"x": (-0.01, 0.01), "y": (-0.01, 0.01)}, 204 | # rotate=(-1, 1), 205 | # shear=(-1, 1)), 206 | 207 | # 20%的图片像素值乘以0.8-1.2中间的数值,用以增加图片明亮度或改变颜色 208 | iaa.Multiply((0.8, 1.2), per_channel=0.2), 209 | # 随机去掉一些像素点, 即把这些像素点变成0。 210 | iaa.Dropout(p=0, per_channel=False, name=None, deterministic=False, random_state=None), 211 | # 浮雕效果 212 | # iaa.Emboss(alpha=0, strength=2, name=None, deterministic=False, random_state=None), 213 | # loc 噪声均值，scale噪声方差，50%的概率，对图片进行添加白噪声并应用于每个通道 214 | iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.1 * 255), per_channel=0.3)], 215 | ), 216 | ], random_order=True) # 打乱定义图像增强的顺序 217 | 218 | number = 1 # 在原来图像数量基础上叠加 219 | for root, sub_folders, files in os.walk(XML_DIR): 220 | 221 | for name in tqdm(files): 222 | 223 | bndbox = read_xml_annotation(XML_DIR, name) 224 | # 首先将原来的图像和对应的标注文件复制到扩增目录下 225 | shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR) 226 | shutil.copy(os.path.join(IMG_DIR, name[:-4] + ext), AUG_IMG_DIR) 227 | 228 | for epoch in range(AUGLOOP): # 在每个扩增循环里处理 229 | seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机 230 | # 读取图片 231 | img = Image.open(os.path.join(IMG_DIR, name[:-4] + ext)) 232 | # sp = img.size 233 | img = np.asarray(img) 234 | # bndbox 坐标增强 235 | for i in range(len(bndbox)): 236 | bbs = ia.BoundingBoxesOnImage([ 237 | ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]), 238 | ], shape=img.shape) 239 | 240 | bbs_aug = seq_det.augment_bounding_boxes([bbs])[0] # 对bbox扩增 241 | boxes_img_aug_list.append(bbs_aug) 242 | 243 | # new_bndbox_list:[[x1,y1,x2,y2],...[],[]] 244 | n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1))) 245 | n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1))) 246 | n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2))) 247 | n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2))) 248 | if n_x1 == 1 and n_x1 == n_x2: 249 | n_x2 += 1 250 | if n_y1 == 1 and n_y2 == n_y1: 251 | n_y2 += 1 252 | if n_x1 >= n_x2 or n_y1 >= n_y2: 253 | print('error', name) 254 | new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2]) 255 | # 存储变化后的图片 256 | image_aug = seq_det.augment_images([img])[0] # 对图像进行扩增 257 | path = os.path.join(AUG_IMG_DIR, 258 | str("%d" % (len(files) + number)) + ext) 259 | 260 | # image_auged = bbs.draw_on_image(image_aug, size=0) 261 | Image.fromarray(image_aug).save(path) 262 | 263 | # 存储变化后的XML 264 | change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, 265 | len(files) + number) 266 | # print(str("%d" % (len(files) + number)) + ext) 267 | number += 1 268 | new_bndbox_list = [] -------------------------------------------------------------------------------- /voc_dataset_information.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 计算Pascal voc格式数据集的相关信息，包括： 3 | 1）各类别数量，可视化结果 4 | 2）各图片长宽大小，平均大小，可视化结果 5 | 3）锚框长宽大小，平均大小，可视化结果 6 | 4）锚框k-means聚类结果 7 | 5）各锚框占该锚框所在图中的面积比，可视化结果 8 | 6）无参考系的图像质量评估 9 | 10 | ''' 11 | 12 | import xml.etree.ElementTree as ET 13 | import pickle 14 | import os 15 | import sys 16 | from os import listdir, getcwd 17 | from os.path import join 18 | import pandas as pd 19 | import numpy as np 20 | import shutil 21 | import matplotlib.pyplot as plt 22 | plt.rcParams['font.family'] = ['sans-serif'] 23 | plt.rcParams['font.sans-serif'] = ['Times New Roman'] 24 | plt.rcParams['font.size'] = 13 25 | plt.rc('axes', unicode_minus=False) 26 | plt.rc('axes', unicode_minus=False) 27 | # plt.style.use(['science','ieee']) 28 | from collections import Counter 29 | from pathlib import Path 30 | import argparse 31 | from tqdm import tqdm 32 | import seaborn as sns 33 | from matplotlib.transforms import Bbox 34 | from imageio import imread 35 | 36 | def load_dataset(xml_list, anno_root, savefig=True,img_name=''): 37 | '''计算标签 38 | xml_list: xml标注文件名称列表 39 | anno_root: 标注文件路径 40 | ''' 41 | xml_info = [] 42 | length = len(xml_list) 43 | for idx, xml_file in enumerate(xml_list): 44 | sys.stdout.write('\r>> Converting image %d/%d' % ( 45 | idx+1, length)) 46 | sys.stdout.flush() 47 | # print(xml_file) 48 | xml = open(xml_file, encoding='utf-8') 49 | tree=ET.parse(xml) 50 | root = tree.getroot() 51 | try: 52 | filename = root.find('filename').text 53 | except: 54 | filename = root.find('frame').text 55 | 56 | try: 57 | # 图片高度 58 | height = int(root.findtext("./size/height")) 59 | # 图片宽度 60 | width = int(root.findtext("./size/width")) 61 | except: 62 | img_path = Path(anno_root).parent.joinpath('JPEGImages', filename+img_suffix) 63 | width, height = imread(str(img_path)).shape[:2] 64 | 65 | for obj in root.iter('object'): 66 | xmin = int(float(obj.findtext("bndbox/xmin"))) 67 | ymin = int(float(obj.findtext("bndbox/ymin"))) 68 | xmax = int(float(obj.findtext("bndbox/xmax")))+1 69 | ymax = int(float(obj.findtext("bndbox/ymax")))+1 70 | assert xmin < xmax and ymin < ymax, 'xmin, ymin, xmax, ymax must be in ascending order' 71 | value = ( 72 | filename, 73 | width, 74 | height, 75 | obj.find('name').text.strip(), 76 | xmin, 77 | ymin, 78 | xmax, 79 | ymax 80 | ) 81 | if value[3]=='paches': 82 | print(xml_file) 83 | xml_info.append(value) 84 | 85 | column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax'] 86 | 87 | xml_df = pd.DataFrame(xml_info, columns=column_name) 88 | xml_df['box_width'] = xml_df['xmax']-xml_df['xmin'] 89 | xml_df['box_height'] = xml_df['ymax'] - xml_df['ymin'] 90 | 91 | def color(df): 92 | if df['class']=='lost': 93 | return 0 94 | else: 95 | return 1 96 | def area(df): 97 | return (df['box_width']*df['box_height'])/(df['width']*df['height']) 98 | xml_df['box_area'] = xml_df.apply(area, axis=1) 99 | 100 | count = dict(Counter(xml_df['class'])) 101 | print() 102 | print('标签类别：\n', np.unique(xml_df['class'])) 103 | print('标签名称及数量：\n', xml_df['class'].value_counts()) 104 | print('总标签数量：', len(count)) 105 | 106 | ## 图像大小 107 | print('平均图像大小(宽度×高度)：{:.2f}×{:.2f}'.format(np.mean(xml_df['width']), 108 | np.mean(xml_df['height']))) 109 | ## 锚框大小 110 | print('平均锚框大小(宽度×高度)：{:.2f}×{:.2f}'.format(np.mean(xml_df['xmax']-xml_df['xmin']), 111 | np.mean(xml_df['ymax']-xml_df['ymin']))) 112 | 113 | df_group = xml_df.groupby('class') 114 | for cls, df in df_group: 115 | print('类别：', cls) 116 | print('平均锚框大小(宽度×高度)：{:.2f}×{:.2f}'.format(np.mean(df['xmax']-df['xmin']), 117 | np.mean(df['ymax']-df['ymin']))) 118 | 119 | if savefig: 120 | 121 | plt.figure(figsize=(9,9)) 122 | plt.subplot(2,2,1) 123 | count = dict(Counter(xml_df['class'])) 124 | classes = list(count.keys()) 125 | df = pd.Series(list(count.values()), index=count.keys()) 126 | df = df.sort_values(ascending=True) 127 | df.plot(kind='bar',alpha=0.75, rot=0) 128 | # plt.xticks(rotation=90) 129 | plt.ylabel('number of instances') 130 | # plt.title('Distribution of different classes') 131 | 132 | plt.subplot(2,2,2) 133 | # plt.hist(xml_df['box_area']*100, bins=100,) 134 | # x直方图 135 | plt.hist((xml_df['xmax']+xml_df['xmin'])/2.0/xml_df['width'], bins=50,) 136 | # plt.title('Histogram plot of x') 137 | plt.xlabel('x') 138 | plt.ylabel('frequency') 139 | 140 | plt.subplot(2,2,3) 141 | plt.hist((xml_df['ymax']+xml_df['ymin'])/2.0/xml_df['height'], bins=50, ) 142 | # plt.title('Histogram Plot of Box Areas') 143 | plt.xlabel('y') 144 | plt.ylabel('frequency') 145 | 146 | plt.subplot(2,2,4) 147 | for c in classes: 148 | df_ = xml_df[xml_df['class']==c][['box_width','box_height']] 149 | plt.scatter(df_['box_width']/xml_df['width'], df_['box_height']/xml_df['height'],label=c) 150 | 151 | plt.xlabel('w') 152 | plt.ylabel('h') 153 | # plt.title('Scatter Plot of Boxes') 154 | plt.legend(loc='best') 155 | 156 | plt.savefig(os.path.join(voc_stat,f'{img_name}_output.png'), dpi=800,bbox_inches='tight', pad_inches=0.0) 157 | 158 | 159 | return xml_df 160 | 161 | def create_dir(ROOT:str): 162 | if not os.path.exists(ROOT): 163 | os.mkdir(ROOT) 164 | else: 165 | shutil.rmtree(ROOT) # 先删除，再创建 166 | os.mkdir(ROOT) 167 | 168 | if __name__ == '__main__': 169 | 170 | parser = argparse.ArgumentParser() 171 | parser.add_argument('--voc-root', type=str, required=False, 172 | help='VOC格式数据集根目录，该目录下必须包含JPEGImages,Annotationshe ImageSets这3个文件夹，\ 173 | 在ImageSets文件夹下还要有Main/trainval.txt等文件') 174 | parser.add_argument('--img_dir', type=str, required=False, 175 | help='VOC格式数据集图像存储路径，如果不指定，默认为JPEGImages') 176 | parser.add_argument('--anno_dir', type=str, required=False, 177 | help='VOC格式数据集标注文件存储路径，如果不指定，默认为Annotations') 178 | 179 | opt = parser.parse_args() 180 | 181 | voc_root = opt.voc_root 182 | 183 | # check image root and annotation root 184 | if opt.img_dir is None: 185 | img_dir = 'JPEGImages' 186 | else: 187 | img_dir = opt.img_dir 188 | voc_jpeg = os.path.join(voc_root, img_dir) 189 | if not os.path.exists(voc_jpeg): 190 | raise Exception(f'数据集图像路径{voc_jpeg}不存在！') 191 | 192 | if opt.anno_dir is None: 193 | anno_dir = 'Annotations' 194 | else: 195 | anno_dir = opt.anno_dir 196 | voc_anno = os.path.join(voc_root, anno_dir) # 197 | if not os.path.exists(voc_anno): 198 | raise Exception(f'数据集图像路径{voc_anno}不存在！') 199 | # check image set root 200 | if not os.path.exists(os.path.join(voc_root,'ImageSets/Main')): 201 | raise Exception("$VOC_ROOT/ImageSets/Main doesn't exist, please generate them using script voc2coco.py") 202 | 203 | voc_stat = os.path.join(voc_root, 'VOC统计信息') 204 | create_dir(voc_stat) 205 | 206 | anno_root = os.path.join(voc_root,'Annotations') 207 | # image suffix 208 | img_suffix = set([x.suffix for x in Path(voc_jpeg).iterdir()]) 209 | if len(img_suffix) != 1: 210 | raise Exception('VOC数据集中JPEGImages文件夹中的图片格式不一致') 211 | img_suffix = img_suffix.pop() 212 | 213 | print("=========统计所有数据信息============") 214 | all_xml_list = list(Path(anno_root).iterdir()) 215 | df = load_dataset(all_xml_list, anno_root) 216 | df.to_csv(os.path.join(voc_stat,'all_info.csv'), index=False) 217 | 218 | for data_type in ['train', 'trainval', 'val', 'test']: 219 | print(f"\n\n=========统计{data_type}数据信息============") 220 | txt = os.path.join(voc_root, f'ImageSets/Main/{data_type}.txt') 221 | if not os.path.exists(txt): 222 | print(f'文件ImageSets/Main/{data_type}.txt不存在!') 223 | continue 224 | # xml_files = [x.strip() for x in open(txt,'r').readlines()] 225 | xml_files = [x.replace('\n','') for x in open(txt,'r').readlines()] 226 | xml_list = [os.path.join(anno_root, xml_name+'.xml') for xml_name in xml_files] 227 | df = load_dataset(xml_list, anno_root, savefig=True, img_name=data_type) 228 | df.to_csv(os.path.join(voc_stat,f'{data_type}_info.csv'), index=False) 229 | -------------------------------------------------------------------------------- /yolo2voc.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 本脚本将YOLO格式的数据集转换成VOC格式，数据集组织格式为： 3 | 4 | |--data # 数据集根目录 5 | |----images # 存储图像数据 6 | |------0.png 7 | |------1.png 8 | |------... 9 | |----labels # 存储yolo格式的标注数据：class x_center y_center width height，class索引从0开始，xywh的范围是（0，1） 10 | |------0.txt 11 | |------1.txt 12 | |------... 13 | |----classes.txt 14 | 15 | ''' 16 | import xml.etree.ElementTree as ET 17 | import os 18 | import numpy as np 19 | from PIL import Image 20 | import shutil 21 | import matplotlib.pyplot as plt 22 | from tqdm import tqdm 23 | import argparse 24 | from sklearn.model_selection import train_test_split 25 | from sklearn.utils import shuffle 26 | from pathlib import Path 27 | 28 | 29 | def yolo_to_voc_format(x_center, y_center, width, height, img_width, img_height): 30 | '''将yolo格式转换为voc格式''' 31 | 32 | center_x = round(float(str(x_center).strip()) * img_width) 33 | center_y = round(float(str(y_center).strip()) * img_height) 34 | bbox_width = round(float(str(width).strip()) * img_width) 35 | bbox_height = round(float(str(height).strip()) * img_height) 36 | 37 | xmin = int(center_x - bbox_width / 2 ) 38 | ymin = int(center_y - bbox_height / 2) 39 | xmax = int(center_x + bbox_width / 2) 40 | ymax = int(center_y + bbox_height / 2) 41 | 42 | return xmin,ymin,xmax,ymax 43 | 44 | 45 | def voc_to_yolo_format(xmin, ymin, xmax, ymax, img_width, img_height): 46 | '''将voc格式转换为yolo格式''' 47 | dw = 1./(img_width) # 宽度缩放比例, size[0]为图像宽度width 48 | dh = 1./(img_height) 49 | x = (xmin + xmax)/2.0 - 1 50 | y = (ymin + ymax)/2.0 - 1 51 | w = xmax - xmin 52 | h = ymax - ymin 53 | x_center = x*dw 54 | width = w*dw 55 | y_center = y*dh 56 | height = h*dh 57 | 58 | return x_center, y_center, width, height 59 | 60 | def read_yolo_annotation(img_root, label_root, image_id): 61 | '''root：一般是labels, txt文件, 62 | image_id是包含.txt后缀的文件名 63 | 64 | return: 65 | [xmin, ymin, xmax, ymax, label], pascal format 66 | ''' 67 | img_width, img_height = Image.open(os.path.join(img_root, image_id[:-4] + ext)).size 68 | 69 | annos = [x for x in open(os.path.join(label_root, image_id)).readlines()] 70 | bndboxlist = [] 71 | 72 | for anno in annos: # 找到root节点下的所有country节点 73 | lb, x_center, y_center, width, height = anno.split(' ') 74 | 75 | xmin,ymin,xmax,ymax = yolo_to_voc_format(x_center, y_center, width, height, img_width, img_height) 76 | # print(xmin,ymin,xmax,ymax) 77 | bndboxlist.append([xmin, ymin, xmax, ymax, int(lb)]) 78 | # print(bndboxlist) 79 | 80 | return bndboxlist 81 | 82 | def write_xml(img_root, bndbox, save_root, name): 83 | ''' 84 | img_root: yolo images root 85 | bndbox: [[xmin, ymin, xmax, ymax, label],[...],...], pascal format 86 | save_root: directory to save xml file 87 | name: note: with .txt 88 | ''' 89 | img_width, img_height = Image.open(os.path.join(img_root, name[:-4] + ext)).size 90 | 91 | root = ET.Element('annotation') #创建Annotation根节点 92 | ET.SubElement(root, 'filename').text = name[:-4] + ext #创建filename子节点（带后缀） 93 | sizes = ET.SubElement(root,'size') #创建size子节点 94 | ET.SubElement(sizes, 'width').text = str(img_width) #没带脑子直接写了原图片的尺寸...... 95 | ET.SubElement(sizes, 'height').text = str(img_height) 96 | ET.SubElement(sizes, 'depth').text = '3' #图片的通道数：img.shape[2] 97 | for box in bndbox: 98 | xmin, ymin, xmax, ymax, label = box 99 | label_class = classes_idx[label] # 数字标签转换为字符串标签 100 | 101 | objects = ET.SubElement(root, 'object') #创建object子节点 102 | ET.SubElement(objects, 'name').text = label_class 103 | #的类别名 104 | ET.SubElement(objects, 'pose').text = 'Unspecified' 105 | ET.SubElement(objects, 'truncated').text = '0' 106 | ET.SubElement(objects, 'difficult').text = '0' 107 | bndbox = ET.SubElement(objects,'bndbox') 108 | ET.SubElement(bndbox, 'xmin').text = str(xmin) 109 | ET.SubElement(bndbox, 'ymin').text = str(ymin) 110 | ET.SubElement(bndbox, 'xmax').text = str(xmax) 111 | ET.SubElement(bndbox, 'ymax').text = str(ymax) 112 | tree = ET.ElementTree(root) 113 | 114 | filepath = os.path.join( save_root, name[:-4]+'.xml') 115 | tree.write(filepath, encoding='utf-8') 116 | 117 | 118 | def mkdir(path:str): 119 | # 去除首位空格 120 | path = path.strip() 121 | # 去除尾部 \ 符号 122 | path = path.rstrip("\\") 123 | # 判断路径是否存在 124 | # 存在 True 125 | # 不存在 False 126 | isExists = os.path.exists(path) 127 | # 判断结果 128 | if not isExists: 129 | # 如果不存在则创建目录 130 | # 创建目录操作函数 131 | os.makedirs(path) 132 | print(path + ' 创建成功') 133 | return True 134 | else: 135 | # 如果目录存在则不创建，并提示目录已存在 136 | print(path + ' 目录已存在') 137 | return False 138 | 139 | def check_files(img_root): 140 | '''检测图像名称和xml标准文件名称是否一致，检查图像后缀''' 141 | 142 | if os.path.exists(img_root): 143 | img = Path(img_root) 144 | else: 145 | raise Exception("图像文件路径错误") 146 | img_exts = [] 147 | for im in img.iterdir(): 148 | img_exts.append(im.suffix) 149 | 150 | print('图像后缀列表：', np.unique(img_exts)) 151 | if len(np.unique(img_exts)) > 1: 152 | # print('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 153 | raise Exception('数据集包含多种格式图像，请检查！', np.unique(img_exts)) 154 | 155 | return np.unique(img_exts)[0] 156 | 157 | if __name__ == "__main__": 158 | 159 | parser = argparse.ArgumentParser() 160 | parser.add_argument('--yolo-root', type=str, required=True, 161 | help='VOC格式数据集根目录，该目录下必须包含images和labels这两个文件夹，以及classes.txt标签名文件') 162 | parser.add_argument('--voc-outdir',type=str, default='VOCFormatData', 163 | help='Pascal VOC格式数据集存储路径，默认为yolo数据集同级目录下新建文件夹VOCFormatData') 164 | parser.add_argument('--test_ratio', type=float, default=0.2, help='测试集比例，（0,1）之间的浮点数') 165 | 166 | opt = parser.parse_args() 167 | assert os.path.exists(opt.yolo_root) # 确保数据集存在 168 | 169 | yolo_img = os.path.join(opt.yolo_root, "images") 170 | yolo_label = os.path.join(opt.yolo_root, "labels") 171 | ext = check_files(yolo_img) #检查文件后缀 172 | 173 | voc_root = os.path.join(str(Path(opt.yolo_root).parent), Path(opt.yolo_root).stem + opt.voc_outdir) 174 | if not os.path.exists(voc_root): 175 | os.mkdir(voc_root) 176 | # 读取标签名 177 | assert os.path.isfile(os.path.join(opt.yolo_root, 'classes.txt')), "请检查YOLO数据集根目录下的标签文件classes.txt，\ 178 | 若没有，需自行构建（每行一个类别，索引从0开始，和yolo的txt标注文件里的标签匹配！" 179 | classes = [x.strip() for x in open( os.path.join(opt.yolo_root, 'classes.txt'),'r', encoding='utf-8').readlines()] 180 | print('>>>数据集标签：', classes) 181 | classes_dict = {} # 字典{'class1':0, 'class2':1, ...} 182 | classes_idx = {} # 字典：{0:'class0', 1:'class1', ...} 183 | for k, v in enumerate(classes): 184 | classes_dict[v] = k 185 | classes_idx[k] = v 186 | 187 | # 创建存储图像路径 188 | voc_jpeg = os.path.join(voc_root, "JPEGImages") # 存储 189 | try: 190 | shutil.rmtree(voc_jpeg) 191 | except FileNotFoundError as e: 192 | a = 1 193 | mkdir(voc_jpeg) 194 | # 创建存储xml标签路径 195 | voc_anno = os.path.join(voc_root, "Annotations") # 存储XML文件夹路径 196 | try: 197 | shutil.rmtree(voc_anno) 198 | except FileNotFoundError as e: 199 | a = 1 200 | mkdir(voc_anno) 201 | 202 | # 创建存储train.txt, trainval.txt, test.txt标签路径 203 | voc_set_dir = os.path.join(voc_root, "ImageSets/Main") # 204 | if not os.path.exists(voc_set_dir): 205 | os.makedirs(voc_set_dir) 206 | 207 | for root, sub_folders, files in os.walk(yolo_label): 208 | for name in tqdm(files): # name是包含.txt后缀的文件名 209 | # [xmin, ymin, xmax, ymax, int(lb)] 210 | bndbox = read_yolo_annotation(yolo_img, yolo_label, name) # 读取yolo标注文，注意：## 返回voc格式标注框 ## 211 | shutil.copy(os.path.join(yolo_img, name[:-4] + ext), voc_jpeg) # 复制图像至VOC图像存储路径 212 | write_xml(img_root=yolo_img, 213 | bndbox=bndbox, 214 | save_root=voc_anno, 215 | name=name 216 | ) 217 | 218 | # 所有图片名称 219 | files = [x.stem for x in Path(voc_jpeg).iterdir() if not x.stem.startswith('.')] 220 | 221 | # 利用已有的yolo数据划分信息 ${YOLO-ROOT}/trainval.txt, train.txt, val.txt, test.txt 222 | if os.path.exists(os.path.join(opt.yolo_root, 'trainval.txt')) and \ 223 | os.path.exists(os.path.join(opt.yolo_root, 'train.txt')) and \ 224 | os.path.exists(os.path.join(opt.yolo_root, 'val.txt')) and \ 225 | os.path.exists(os.path.join(opt.yolo_root, 'test.txt')): 226 | print('>>> 使用YOLO已有划分数据分割train,val和test') 227 | trainval = [Path(x).stem for x in open(os.path.join(opt.yolo_root, 'trainval.txt')).readlines()] 228 | train = [Path(x).stem for x in open(os.path.join(opt.yolo_root, 'train.txt')).readlines()] 229 | val = [Path(x).stem for x in open(os.path.join(opt.yolo_root, 'val.txt')).readlines()] 230 | test = [Path(x).stem for x in open(os.path.join(opt.yolo_root, 'test.txt')).readlines()] 231 | else: 232 | 233 | print('>>>随即划分train,val和test') 234 | files = shuffle(files) 235 | ratio = opt.test_ratio 236 | trainval, test = train_test_split(files, test_size=ratio) 237 | train, val = train_test_split(trainval,test_size=0.2) 238 | print('训练集数量: ',len(train)) 239 | print('验证集数量: ',len(val)) 240 | print('测试集数量: ',len(test)) 241 | 242 | def write_txt(txt_path, data): 243 | '''写入txt文件''' 244 | with open(txt_path,'w') as f: 245 | for d in data: 246 | f.write(str(d)) 247 | f.write('\n') 248 | # 写入各个txt文件 249 | trainval_txt = os.path.join(voc_set_dir,'trainval.txt') 250 | write_txt(trainval_txt, trainval) 251 | 252 | train_txt = os.path.join(voc_set_dir,'train.txt') 253 | write_txt(train_txt, train) 254 | 255 | val_txt = os.path.join(voc_set_dir,'val.txt') 256 | write_txt(val_txt, val) 257 | 258 | test_txt = os.path.join(voc_set_dir,'test.txt') 259 | write_txt(test_txt, test) 260 | -------------------------------------------------------------------------------- /yolo_augument.py: -------------------------------------------------------------------------------- 1 | ''' 2 | YOLO数据集扩增 3 | ''' 4 | import xml.etree.ElementTree as ET 5 | import pickle 6 | import os 7 | from os import getcwd 8 | import numpy as np 9 | from PIL import Image 10 | import shutil 11 | import matplotlib.pyplot as plt 12 | from tqdm import tqdm 13 | import imgaug as ia 14 | from imgaug import augmenters as iaa 15 | import argparse 16 | 17 | ia.seed(1) 18 | 19 | 20 | def yolo_to_voc_format(x_center, y_center, width, height, img_width, img_height): 21 | '''将yolo格式转换为voc格式，注意输入的x_center, y_center, width, height是str类型''' 22 | 23 | center_x = round(float(str(x_center).strip()) * img_width) 24 | center_y = round(float(str(y_center).strip()) * img_height) 25 | bbox_width = round(float(str(width).strip()) * img_width) 26 | bbox_height = round(float(str(height).strip()) * img_height) 27 | 28 | xmin = int(center_x - bbox_width / 2 ) 29 | ymin = int(center_y - bbox_height / 2) 30 | xmax = int(center_x + bbox_width / 2) 31 | ymax = int(center_y + bbox_height / 2) 32 | 33 | return xmin,ymin,xmax,ymax 34 | 35 | 36 | def voc_to_yolo_format(xmin, ymin, xmax, ymax, img_width, img_height): 37 | '''将voc格式转换为yolo格式''' 38 | dw = 1./(img_width) # 宽度缩放比例, size[0]为图像宽度width 39 | dh = 1./(img_height) 40 | x = (xmin + xmax)/2.0 - 1 41 | y = (ymin + ymax)/2.0 - 1 42 | w = xmax - xmin 43 | h = ymax - ymin 44 | x_center = x*dw 45 | width = w*dw 46 | y_center = y*dh 47 | height = h*dh 48 | 49 | return x_center, y_center, width, height 50 | 51 | def read_yolo_annotation(img_root, label_root, image_id): 52 | ''' 53 | root：一般是labels, txt文件, 54 | image_id是包含.txt后缀的文件名 55 | 56 | return: 57 | [xmin, ymin, xmax, ymax, label] 58 | ''' 59 | img_width, img_height = Image.open(os.path.join(img_root, image_id[:-4] + ext)).size 60 | 61 | annos = [x for x in open(os.path.join(label_root, image_id)).readlines()] 62 | bndboxlist = [] # 存储标注框信息 63 | 64 | for anno in annos: # 找到root节点下的所有country节点 65 | lb, x_center, y_center, width, height = anno.split(' ') # 注意得到的是str类型 66 | 67 | xmin,ymin,xmax,ymax = yolo_to_voc_format(x_center, y_center, width, height, img_width, img_height) 68 | # print(xmin,ymin,xmax,ymax) 69 | bndboxlist.append([xmin, ymin, xmax, ymax, int(lb)]) 70 | # print(bndboxlist) 71 | 72 | return bndboxlist 73 | 74 | def change_yolo_annotation(img_root, label_root, image_id, new_target, saveroot, id): 75 | '''保存新的yolo标注 76 | img_root: 原数据集图像路径 77 | label_root: 原数据标注文件路径 78 | image_id：文件名，带.txt后缀 79 | new_target: new_bndbox_list:[[x1,y1,x2,y2,label],...[],[]] 80 | saveroot: 扩增后标注文件保存路径 81 | id: 扩增后保存标注文件名 82 | ''' 83 | img_width, img_height = Image.open(os.path.join(img_root, image_id[:-4] + ext)).size 84 | 85 | new_annos = [] 86 | for anno in new_target: 87 | xmin, ymin, xmax, ymax, label = anno 88 | x_center, y_center, width, height = voc_to_yolo_format(xmin, ymin, xmax, ymax, img_width, img_height) 89 | label = str(label) 90 | x_center = str(x_center) 91 | y_center = str(y_center) 92 | width = str(width) 93 | height = str(height)+'\n' 94 | 95 | new_annos.append(' '.join((label,x_center,y_center, width, height))) 96 | 97 | # 存储新标注 98 | with open(os.path.join(saveroot, str("%d" % int(id)) + '.txt'),'w') as f: 99 | f.writelines(new_annos) 100 | 101 | 102 | def mkdir(path): 103 | # 去除首位空格 104 | path = path.strip() 105 | # 去除尾部 \ 符号 106 | path = path.rstrip("\\") 107 | # 判断路径是否存在 108 | # 存在 True 109 | # 不存在 False 110 | isExists = os.path.exists(path) 111 | # 判断结果 112 | if not isExists: 113 | # 如果不存在则创建目录 114 | # 创建目录操作函数 115 | os.makedirs(path) 116 | print(path + ' 创建成功') 117 | return True 118 | else: 119 | # 如果目录存在则不创建，并提示目录已存在 120 | print(path + ' 目录已存在') 121 | return False 122 | 123 | 124 | if __name__ == "__main__": 125 | 126 | parser = argparse.ArgumentParser() 127 | parser.add_argument('--yolo-root', type=str, required=True, 128 | help='YOLO格式数据集根目录，该目录下必须包含images和labels这两个文件夹') 129 | parser.add_argument('--aug-dir',type=str, default='YOLOAugumented', 130 | help='数据增强后保存的路径，默认在voc-root路径下创建一个文件夹YOLOAugumented进行保存') 131 | parser.add_argument('--aug_num',type=int, default=5, 132 | help='每张图片进行扩增的次数') 133 | parser.add_argument('--ext', type=str, default='.png', help='图像后缀，默认为.png') 134 | 135 | opt = parser.parse_args() 136 | ext = opt.ext 137 | 138 | IMG_DIR = os.path.join(opt.yolo_root, "images") 139 | LABEL_DIR = os.path.join(opt.yolo_root, "labels") 140 | 141 | AUGUMENT = os.path.join(opt.yolo_root, opt.aug_dir) 142 | if not os.path.exists(AUGUMENT): 143 | os.mkdir(AUGUMENT) 144 | 145 | AUG_IMG_DIR = os.path.join(AUGUMENT, "images") # 存储增强后的影像文件夹路径 146 | try: 147 | shutil.rmtree(AUG_IMG_DIR) 148 | except FileNotFoundError as e: 149 | a = 1 150 | mkdir(AUG_IMG_DIR) 151 | 152 | AUG_LABEL_DIR = os.path.join(AUGUMENT, "labels") # 存储增强后的XML文件夹路径 153 | try: 154 | shutil.rmtree(AUG_LABEL_DIR) 155 | except FileNotFoundError as e: 156 | a = 1 157 | mkdir(AUG_LABEL_DIR) 158 | 159 | AUGLOOP = opt.aug_num # 每张影像增强的数量 160 | 161 | boxes_img_aug_list = [] 162 | new_bndbox = [] 163 | new_bndbox_list = [] 164 | 165 | # 影像增强 166 | seq = iaa.Sequential([ 167 | iaa.Flipud(0.5), # vertically flip 20% of all images 168 | iaa.Fliplr(0.5), # 镜像 169 | iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs 170 | iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5), 171 | iaa.Affine( 172 | translate_px={"x": 15, "y": 15}, 173 | scale=(0.8, 0.95), 174 | rotate=(-30, 30) 175 | ) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs 176 | ]) 177 | 178 | number = 1 179 | for root, sub_folders, files in os.walk(LABEL_DIR): 180 | for name in tqdm(files): 181 | # [xmin, ymin, xmax, ymax, int(lb)] 182 | bndbox = read_yolo_annotation(IMG_DIR, LABEL_DIR, name) # 读取yolo标注文件，name是包含.txt后缀的文件名，注意：## 返回voc格式标注框 ## 183 | shutil.copy(os.path.join(LABEL_DIR, name), AUG_LABEL_DIR) # 复制标注文件 184 | shutil.copy(os.path.join(IMG_DIR, name[:-4] + ext), AUG_IMG_DIR) # 复制图像 185 | 186 | for epoch in range(AUGLOOP): 187 | seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机 188 | # 读取图片 189 | img = Image.open(os.path.join(IMG_DIR, name[:-4] + ext)) 190 | # sp = img.size 191 | img = np.asarray(img) 192 | # bndbox 坐标增强 193 | for i in range(len(bndbox)): 194 | bbs = ia.BoundingBoxesOnImage([ 195 | ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]), 196 | ], shape=img.shape) 197 | 198 | bbs_aug = seq_det.augment_bounding_boxes([bbs])[0] 199 | boxes_img_aug_list.append(bbs_aug) 200 | 201 | # new_bndbox_list:[[x1,y1,x2,y2],...[],[]] 202 | n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1))) 203 | n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1))) 204 | n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2))) 205 | n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2))) 206 | if n_x1 == 1 and n_x1 == n_x2: 207 | n_x2 += 1 208 | if n_y1 == 1 and n_y2 == n_y1: 209 | n_y2 += 1 210 | if n_x1 >= n_x2 or n_y1 >= n_y2: 211 | print('error', name) 212 | new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2, bndbox[i][4]]) # 注意保存标签 213 | 214 | # 存储变化后的图片 215 | image_aug = seq_det.augment_images([img])[0] 216 | path = os.path.join(AUG_IMG_DIR, 217 | str("%d" % (len(files) + number)) + ext) 218 | # image_auged = bbs.draw_on_image(image_aug, size=0) 219 | Image.fromarray(image_aug).save(path) 220 | 221 | # 存储变化后的yolo标注文件 222 | change_yolo_annotation(img_root=IMG_DIR, 223 | label_root=LABEL_DIR, 224 | image_id=name, # 带后缀 225 | new_target=new_bndbox_list, 226 | saveroot=AUG_LABEL_DIR, 227 | id=len(files) + number) 228 | 229 | # print(str("%d" % (len(files) + number)) + '.png') 230 | number = number + 1 231 | new_bndbox_list = [] -------------------------------------------------------------------------------- /yolov5tococo.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Adapt from: 3 | https://github.com/RapidAI/YOLO2COCO/blob/main/yolov5_2_coco.py 4 | ''' 5 | 6 | import argparse 7 | import json 8 | import shutil 9 | from pathlib import Path 10 | import time 11 | 12 | import cv2 13 | from tqdm import tqdm 14 | 15 | 16 | def read_txt(txt_path): 17 | with open(str(txt_path), 'r', encoding='utf-8') as f: 18 | data = list(map(lambda x: x.rstrip('\n'), f)) 19 | return data 20 | 21 | 22 | def mkdir(dir_path): 23 | Path(dir_path).mkdir(parents=True, exist_ok=True) 24 | 25 | 26 | def verify_exists(file_path): 27 | file_path = Path(file_path) 28 | if not file_path.exists(): 29 | raise FileNotFoundError(f'The {file_path} is not exists!!!') 30 | 31 | 32 | class YOLOV5ToCOCO(object): 33 | def __init__(self, dir_path): 34 | self.src_data = Path(dir_path) 35 | self.src = self.src_data.parent 36 | self.train_txt_path = self.src_data / 'train.txt' 37 | self.val_txt_path = self.src_data / 'val.txt' 38 | self.test_txt_path = self.src_data / 'test.txt' 39 | self.classes_path = self.src_data / 'classes.txt' 40 | 41 | # 文件存在性校验 42 | verify_exists(self.src_data / 'images') 43 | verify_exists(self.src_data / 'labels') 44 | verify_exists(self.train_txt_path) 45 | verify_exists(self.val_txt_path) 46 | verify_exists(self.test_txt_path) 47 | verify_exists(self.classes_path) 48 | 49 | # 构建COCO格式目录 50 | self.dst = Path(self.src) / f"{Path(self.src_data).name}_COCOFormat" 51 | self.coco_train = "train" 52 | self.coco_val = "val" 53 | self.coco_test = "test" 54 | self.coco_annotation = "annotations" 55 | self.coco_train_json = self.dst / self.coco_annotation / \ 56 | f'instances_{self.coco_train}.json' 57 | self.coco_val_json = self.dst / self.coco_annotation / \ 58 | f'instances_{self.coco_val}.json' 59 | self.coco_test_json = self.dst / self.coco_annotation / \ 60 | f'instances_{self.coco_test}.json' 61 | 62 | mkdir(self.dst) 63 | mkdir(self.dst / self.coco_train) 64 | mkdir(self.dst / self.coco_val) 65 | mkdir(self.dst / self.coco_test) 66 | mkdir(self.dst / self.coco_annotation) 67 | 68 | # 构建json内容结构 69 | self.type = 'instances' 70 | self.categories = [] 71 | self.annotation_id = 1 72 | 73 | # 读取类别数 74 | self._get_category() 75 | 76 | cur_year = time.strftime('%Y', time.localtime(time.time())) 77 | self.info = { 78 | 'year': int(cur_year), 79 | 'version': '1.0', 80 | 'description': 'For object detection', 81 | 'date_created': cur_year, 82 | } 83 | 84 | self.licenses = [{ 85 | 'id': 1, 86 | 'name': 'Apache License v2.0', 87 | 'url': 'https://github.com/RapidAI/YOLO2COCO/LICENSE', 88 | }] 89 | 90 | def _get_category(self): 91 | class_list = read_txt(self.classes_path) 92 | for i, category in enumerate(class_list, 1): 93 | self.categories.append({ 94 | 'supercategory': category, 95 | 'id': i, 96 | 'name': category, 97 | }) 98 | 99 | def generate(self): 100 | self.train_files = read_txt(self.train_txt_path) 101 | self.valid_files = read_txt(self.val_txt_path) 102 | self.test_files = read_txt(self.test_txt_path) 103 | 104 | train_dest_dir = Path(self.dst) / self.coco_train 105 | self.gen_dataset(self.train_files, train_dest_dir, 106 | self.coco_train_json, mode='train') 107 | 108 | val_dest_dir = Path(self.dst) / self.coco_val 109 | self.gen_dataset(self.valid_files, val_dest_dir, 110 | self.coco_val_json, mode='val') 111 | 112 | test_test_dir = Path(self.dst) / self.coco_test 113 | self.gen_dataset(self.test_files, test_test_dir, 114 | self.coco_test_json, mode='test') 115 | print(f"The output directory is: {str(self.dst)}") 116 | 117 | def gen_dataset(self, img_paths, target_img_path, target_json, mode): 118 | """ 119 | https://cocodataset.org/#format-data 120 | 121 | """ 122 | images = [] 123 | annotations = [] 124 | for img_id, img_path in enumerate(tqdm(img_paths, desc=mode), 1): 125 | img_path = Path(img_path) 126 | 127 | verify_exists(img_path) 128 | 129 | label_path = str(img_path.parent.parent 130 | / 'labels' / f'{img_path.stem}.txt') 131 | 132 | imgsrc = cv2.imread(str(img_path)) 133 | height, width = imgsrc.shape[:2] 134 | 135 | dest_file_name = f'{img_id:012d}.jpg' 136 | save_img_path = target_img_path / dest_file_name 137 | 138 | if img_path.suffix.lower() == ".jpg": 139 | shutil.copyfile(img_path, save_img_path) 140 | else: 141 | cv2.imwrite(str(save_img_path), imgsrc) 142 | 143 | images.append({ 144 | 'date_captured': '2021', 145 | 'file_name': dest_file_name, 146 | 'id': img_id, 147 | 'height': height, 148 | 'width': width, 149 | }) 150 | 151 | if Path(label_path).exists(): 152 | new_anno = self.read_annotation(label_path, img_id, 153 | height, width) 154 | if len(new_anno) > 0: 155 | annotations.extend(new_anno) 156 | else: 157 | raise ValueError(f'{label_path} is empty') 158 | else: 159 | raise FileNotFoundError(f'{label_path} not exists') 160 | 161 | json_data = { 162 | 'info': self.info, 163 | 'images': images, 164 | 'licenses': self.licenses, 165 | 'type': self.type, 166 | 'annotations': annotations, 167 | 'categories': self.categories, 168 | } 169 | with open(target_json, 'w', encoding='utf-8') as f: 170 | json.dump(json_data, f, ensure_ascii=False) 171 | 172 | def read_annotation(self, txt_file, img_id, height, width): 173 | annotation = [] 174 | all_info = read_txt(txt_file) 175 | for label_info in all_info: 176 | # 遍历一张图中不同标注对象 177 | label_info = label_info.split(" ") 178 | if len(label_info) < 5: 179 | continue 180 | 181 | category_id, vertex_info = label_info[0], label_info[1:] 182 | segmentation, bbox, area = self._get_annotation(vertex_info, 183 | height, width) 184 | annotation.append({ 185 | 'segmentation': segmentation, 186 | 'area': area, 187 | 'iscrowd': 0, 188 | 'image_id': img_id, 189 | 'bbox': bbox, 190 | 'category_id': int(category_id)+1, 191 | 'id': self.annotation_id, 192 | }) 193 | self.annotation_id += 1 194 | return annotation 195 | 196 | @staticmethod 197 | def _get_annotation(vertex_info, height, width): 198 | cx, cy, w, h = [float(i) for i in vertex_info] 199 | 200 | cx = cx * width 201 | cy = cy * height 202 | box_w = w * width 203 | box_h = h * height 204 | 205 | # left top 206 | x0 = max(cx - box_w / 2, 0) 207 | y0 = max(cy - box_h / 2, 0) 208 | 209 | # right bottomt 210 | x1 = min(x0 + box_w, width) 211 | y1 = min(y0 + box_h, height) 212 | 213 | segmentation = [[x0, y0, x1, y0, x1, y1, x0, y1]] 214 | bbox = [x0, y0, box_w, box_h] 215 | area = box_w * box_h 216 | return segmentation, bbox, area 217 | 218 | 219 | if __name__ == "__main__": 220 | parser = argparse.ArgumentParser('Datasets converter from YOLOV5 to COCO') 221 | parser.add_argument('--yolov5-root', type=str, 222 | default='datasets/YOLOV5', 223 | help='Dataset root path') 224 | args = parser.parse_args() 225 | 226 | converter = YOLOV5ToCOCO(args.yolov5_root) 227 | converter.generate() 228 | --------------------------------------------------------------------------------