* 使用脚本制作数据集

Author: jadehh

dataset_tools/jade_create_object_dection_datasets.py

@@ -14,12 +14,6 @@
 import xml.etree.ElementTree as ET
 from dataset_tools.jade_voc_datasets import GetXmlClassesNames
 
-def CreateSavePath(save_path):
-    if os.path.exists(save_path):
-        return save_path
-    else:
-        os.makedirs(save_path)
-        return save_path
 
 def ProcessXml(xml_path):
     # Read the XML annotation file.
@@ -108,6 +102,85 @@ def CreateYearsDatasets(dir,year=None,save_path=None,rate=0.95):
     CreateLabelList(save_path)
 
 
+# VOC 数据集转换为Darknet数据集
+def CreateYearsDarknetVocDatasets(dir, year=None, save_path=None, rate=0.95,VOC_CLASSES=None):
+    years = os.listdir(dir)
+    if year is None:
+        progressBar1 = ProgressBar(len(years))
+    else:
+        progressBar1 = ProgressBar(1)
+    if os.path.exists(save_path):
+        pass
+    else:
+        os.makedirs(save_path)
+    if year is None:
+        for year in years:
+            if os.path.isdir(os.path.join(dir, year)):
+                if os.path.exists(os.path.join(dir, year, DIRECTORY_IMAGES)) and os.path.exists(
+                        os.path.join(dir, year, DIRECTORY_ANNOTATIONS)):
+                    CreateDarknetVocDatasets(os.path.join(dir, year), save_path, rate,VOC_CLASSES)
+            progressBar1.update()
+    else:
+        if os.path.isdir(os.path.join(dir, year)):
+            if os.path.exists(os.path.join(dir, year, DIRECTORY_IMAGES)) and os.path.exists(os.path.join(dir, year, DIRECTORY_ANNOTATIONS)):
+                CreateDarknetVocDatasets(os.path.join(dir, year), save_path, rate,VOC_CLASSES)
+        progressBar1.update()
+
+    with open(os.path.join(save_path,"classes.txt"),"wb") as f:
+        for class_name in VOC_CLASSES:
+            f.write((class_name+"\n").encode("utf-8"))
+
+
+def convert_voc_to_yolo(xml_dir, output_dir, classes):
+    tree = ET.parse(xml_dir)
+    root = tree.getroot()
+    img_w = int(root.find('size/width').text)
+    img_h = int(root.find('size/height').text)
+
+    with open(os.path.join(output_dir), 'w') as f:
+        for obj in root.findall('object'):
+            cls_name = obj.find('name').text
+            cls_id = classes.index(cls_name)
+            bbox = obj.find('bndbox')
+            x_center = (int(bbox.find('xmin').text) + int(bbox.find('xmax').text)) / 2 / img_w
+            y_center = (int(bbox.find('ymin').text) + int(bbox.find('ymax').text)) / 2 / img_h
+            width = (int(bbox.find('xmax').text) - int(bbox.find('xmin').text)) / img_w
+            height = (int(bbox.find('ymax').text) - int(bbox.find('ymin').text)) / img_h
+            f.write(f"{cls_id} {x_center} {y_center} {width} {height}\n")
+
+def CreateDarknetVocDataset(dir,save_path,image_files,dataset_type,remove_label="None",VOC_CLASSES=None):
+    save_image_path = CreateSavePath(os.path.join(save_path,"images",dataset_type,))
+    save_label_path = CreateSavePath(os.path.join(save_path,"labels",dataset_type))
+    for image_file in image_files:
+        with open(os.path.join(dir, DIRECTORY_IMAGES, image_file), "rb") as f2:
+            if len(f2.read()) == 0:
+                pass
+            else:
+                class_name_list = GetXmlClassesNames(os.path.join(dir, DIRECTORY_ANNOTATIONS, image_file[:-4] + ".xml"))
+                if len(class_name_list) > 0 and remove_label not in class_name_list:
+                    shutil.copy(os.path.join(dir, DIRECTORY_IMAGES, image_file), save_image_path)
+                    convert_voc_to_yolo(os.path.join(dir,DIRECTORY_ANNOTATIONS,image_file[:-4] + ".xml"),os.path.join(save_label_path,image_file[:-4] + ".txt"),VOC_CLASSES)
+                else:
+                    print(os.path.join(dir, DIRECTORY_ANNOTATIONS, image_file[:-4] + ".xml"))
+
+
+
+
+def CreateDarknetVocDatasets(dir,save_path,rate,VOC_CLASSES):
+    """
+    :param dir:
+    """
+    image_files = os.listdir(os.path.join(dir, DIRECTORY_IMAGES))
+    train_image_files = random.sample(image_files, int(len(image_files) *rate))
+    test_image_files = [file for file in image_files if file not in train_image_files]
+    CreateDarknetVocDataset(dir,save_path,train_image_files,"train",VOC_CLASSES=VOC_CLASSES)
+    CreateDarknetVocDataset(dir,save_path,test_image_files,"test",VOC_CLASSES=VOC_CLASSES)
+
+
+
+
+
+
 ##制作VOC数据集
 def CreateVOCDataset(dir, datasetname,save_path=None,rate=0.95):
     """

main.py

@@ -9,12 +9,28 @@
 from dataset_tools.jade_create_paddle_text_detection_datasets import *
 from dataset_tools.jade_create_paddle_ocr_datasets import *
 from dataset_tools.jade_create_object_dection_datasets import CreateYearsDatasets
+
+def test_create_paddle_years_datasets(args):
+    CreateYearsDatasets(args.input_dataset_dir,None,save_path=args.save_dataset_dir,rate=0.95)
+
 if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser(description="制作数据集脚本")
+    parser.add_argument("--dataset_type", default='paddle', help="制作数据集的类型")
+    parser.add_argument("--input_dataset_dir", default='test', help="数据集的地址")
+    parser.add_argument("--save_dataset_dir", default='test/output_seals_01', help="保存数据集的地址")
+    parser.add_argument("--voc_labels",  nargs='+',default="", help="类别")
+    args = parser.parse_args()
+    print(list(args.voc_labels))
+    if args.dataset_type == "paddle_detection":
+        CreateYearsDatasets(args.input_dataset_dir, None, save_path=args.save_dataset_dir, rate=0.95)
+    elif args.dataset_type == "yolo_detection":
+        CreateDarknetVocDatasets(args.input_dataset_dir,  save_path=args.save_dataset_dir, rate=0.95, VOC_CLASSES=args.voc_labels)
+
     #removeNolabelDatasets(r"F:\数据集\关键点检测数据集\定制版箱号关键点数据集\2022-03-09")
     #create_text_detection_datasets(r"F:\数据集\关键点检测数据集\定制版箱号关键点数据集",r"E:\Data\字符检测识别数据集\定制版箱号关键点数据集",0.95)
     #CreatePaddleOCRDatasets(root_path="E:\Data\字符检测识别数据集\镇江大港厂内车牌关键点检测数据集", save_path="E:\Data\OCR\镇江大港厂内车牌识别数据集",dataset_type="镇江厂内车牌数据集")
     #removeNolabelVocDatasets(r"E:\Data\VOC数据集\集装箱残损检测数据集")
     #CreateYearsDatasets(r"E:\Data\VOC数据集\集装箱残损检测数据集")
-    create_text_detection_datasets(r"F:\数据集\关键点检测数据集\箱号关键点数据集",r'E:\Data\字符检测识别数据集\箱号关键点数据集')
     #CreatePaddleOCRDatasets(r'F:\数据集\VOC数据集\箱门检测数据集\ContainVOC', save_path="E:\Data\OCR\箱号识别数据集",dataset_type="箱号数据集")
     #CreateYearsDatasets("F:\数据集\VOC数据集\验残集装箱检测数据集",0.95)

test.py

@@ -15,6 +15,11 @@ def testCreateYearsDatasets():
     """
     CreateYearsDatasets(r"F:\数据集\VOC数据集\验残集装箱检测数据集",  save_path=r"E:\Data\VOC数据集\验残集装箱检测数据集")
 
+def testCreateYearsDarknetVocDatasets():
+    # VOC_CLASSES = ["container"]
+    VOC_CLASSES = ["FRONTEND","DOOREND","UPEND","slide","bromine_tank"]
+    CreateYearsDarknetVocDatasets(r"F:\数据集\VOC数据集\箱门检测数据集\ContainVOC",  save_path=r"E:\Data\VOC数据集\箱门检测数据集\ContainerVOCDarknet",VOC_CLASSES=VOC_CLASSES)
+
 
 if __name__ == '__main__':
-    testCreateYearsDatasets()
+    testCreateYearsDarknetVocDatasets()