Przeglądaj źródła

debug ourdataset

yjh0410 1 rok temu
rodzic
commit
49c3fbdbba
3 zmienionych plików z 31 dodań i 22 usunięć
  1. 21 6
      README.md
  2. 1 1
      config/data_config/dataset_config.py
  3. 9 15
      tools/convert_ours_to_coco.py

+ 21 - 6
README.md

@@ -300,7 +300,26 @@ OurDataset
 |  ...
 ```
 
-- Step-2: Convert ourdataset to COCO format.
+- Step-2: Make the configuration for our dataset.
+```Shell
+cd <PyTorch_YOLO_Tutorial_HOME>
+cd config/data_config
+```
+You need to edit the `dataset_cfg` defined in `dataset_config.py`. You can refer to the `ourdataset` defined in `dataset_cfg` to modify the relevant parameters, such as `num_classes`, `classes_names`, to adapt to our dataset.
+
+For example:
+```Shell
+dataset_cfg = {
+    'ourdataset':{
+        'data_name': 'AnimalDataset',
+        'num_classes': 9,
+        'class_indexs': (0, 1, 2, 3, 4, 5, 6, 7, 8),
+        'class_names': ('bird', 'butterfly', 'cat', 'cow', 'dog', 'lion', 'person', 'pig', 'tiger', ),
+    },
+}
+```
+
+- Step-3: Convert ourdataset to COCO format.
 
 ```Shell
 cd <PyTorch_YOLO_Tutorial_HOME>
@@ -336,11 +355,7 @@ OurDataset
 |  ...
 ```
 
-- Step-3 Define our class labels.
-
-Please open `dataset/ourdataset.py` file and change `our_class_labels = ('cat',)` according to our definition of categories.
-
-- Step-4 Check
+- Step-4 Check the data.
 
 ```Shell
 cd <PyTorch_YOLO_Tutorial_HOME>

+ 1 - 1
config/data_config/dataset_config.py

@@ -46,7 +46,7 @@ dataset_cfg = {
     'ourdataset':{
         'data_name': 'AnimalDataset',
         'num_classes': 9,
-        'class_indexs': None,
+        'class_indexs': (0, 1, 2, 3, 4, 5, 6, 7, 8),
         'class_names': ('bird', 'butterfly', 'cat', 'cow', 'dog', 'lion', 'person', 'pig', 'tiger', ),
     },
 

+ 9 - 15
tools/convert_ours_to_coco.py

@@ -3,14 +3,14 @@ import json
 import xml.etree.ElementTree as ET
 import glob
 
+import sys
+sys.path.append("..")
+from config.data_config.dataset_config import dataset_cfg
+data_config = dataset_cfg["ourdataset"]
+num_classes = data_config["num_classes"]
+categories = data_config["class_names"]
 START_BOUNDING_BOX_ID = 1
-PRE_DEFINE_CATEGORIES = None
-# If necessary, pre-define category and its id
-#  PRE_DEFINE_CATEGORIES = {"aeroplane": 1, "bicycle": 2, "bird": 3, "boat": 4,
-#  "bottle":5, "bus": 6, "car": 7, "cat": 8, "chair": 9,
-#  "cow": 10, "diningtable": 11, "dog": 12, "horse": 13,
-#  "motorbike": 14, "person": 15, "pottedplant": 16,
-#  "sheep": 17, "sofa": 18, "train": 19, "tvmonitor": 20}
+PRE_DEFINE_CATEGORIES = {categories[i]: i + 1 for i in range(num_classes)} 
 
 
 def get(root, name):
@@ -73,13 +73,7 @@ def convert(xml_files, json_file):
             print('[{}] / [{}]'.format(i, len(xml_files)))
         tree = ET.parse(xml_file)
         root = tree.getroot()
-        path = get(root, "path")
-        if len(path) == 1:
-            filename = os.path.basename(path[0].text)
-        elif len(path) == 0:
-            filename = get_and_check(root, "filename", 1).text
-        else:
-            raise ValueError("%d paths found in %s" % (len(path), xml_file))
+        filename = get_and_check(root, "filename", 1).text
         ## The filename must be a number
         image_id = get_filename_as_int(filename)
         size = get_and_check(root, "size", 1)
@@ -156,4 +150,4 @@ if __name__ == "__main__":
     print("Number of xml files: {}".format(len(xml_files)))
     print("Converting to COCO format ...")
     convert(xml_files, json_file)
-    print("Success: {}".format(args.json_file))
+    print("Success: {}".format(json_file))