Merge pull request #12 from GreatV/dev

add output_format option

Author: GreatV

.gitignore

@@ -1,3 +1,5 @@
+dataset/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -14,7 +14,7 @@ If you've already marked your segmentation dataset by LabelMe, it's easy to use
 
 ## New
 - export data as yolo polygon annotation (for YOLOv5 v7.0 segmentation)
-
+- Now you can choose the output format of the label text. The available options are `plygon` and `bbox`.
 ## Installation
 
 ```console
@@ -30,6 +30,8 @@ pip install labelme2yolo
 
 **--json_name (Optional)** Convert single LabelMe JSON file.
 
+**--output_format (Optional)** The output format of label.
+
 ## How to Use
 
 ### 1. Convert JSON files, split training, validation and test dataset by --val_size and --test_size

src/labelme2yolo/__about__.py

@@ -1,5 +1,5 @@
-# SPDX-FileCopyrightText: 2022-present Wang Xin <[email protected]>
+# SPDX-FileCopyrightText: 2023-present Wang Xin <[email protected]>
 #
 # SPDX-License-Identifier: MIT
 
-__version__ = '0.0.8'
+__version__ = '0.0.9'

src/labelme2yolo/cli/__init__.py

@@ -16,14 +16,14 @@ def run():
         type=float,
         nargs="?",
         default=None,
-        help="Please input the validation dataset size, for example 0.1 ",
+        help="Please input the validation dataset size, for example 0.1.",
     )
     parser.add_argument(
         "--test_size",
         type=float,
         nargs="?",
         default=None,
-        help="Please input the validation dataset size, for example 0.1 ",
+        help="Please input the test dataset size, for example 0.1.",
     )
     parser.add_argument(
         "--json_name",
@@ -32,13 +32,21 @@ def run():
         default=None,
         help="If you put json name, it would convert only one json file to YOLO.",
     )
+    parser.add_argument(
+        "--output_format",
+        type=str,
+        default="polygon",
+        help='The default output format for labelme2yolo is "polygon".'
+        ' However, you can choose to output in bbox format by specifying the "bbox" option.',
+    )
+
     args = parser.parse_args()
 
     if not args.json_dir:
         parser.print_help()
         return 0
 
-    convertor = Labelme2YOLO(args.json_dir)
+    convertor = Labelme2YOLO(args.json_dir, args.output_format)
 
     if args.json_name is None:
         convertor.convert(val_size=args.val_size, test_size=args.test_size)

src/labelme2yolo/l2y.py

@@ -93,27 +93,28 @@ def get_label_id_map(json_dir):
     return OrderedDict([(label, label_id) for label_id, label in enumerate(label_set)])
 
 
-def extend_point_list(point_list):
+def extend_point_list(point_list, format="polygon"):
     xmin = min([float(point) for point in point_list[::2]])
     xmax = max([float(point) for point in point_list[::2]])
     ymin = min([float(point) for point in point_list[1::2]])
     ymax = max([float(point) for point in point_list[1::2]])
 
-    return np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
+    if (format == "polygon"):
+        return np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
+    if (format == "bbox"):
+        return np.array([xmin, ymin, xmax - xmin, ymax - ymin])
 
 
 def save_yolo_label(json_name, label_dir_path, target_dir, yolo_obj_list):
-    txt_path = os.path.join(
-        label_dir_path, target_dir, json_name.replace(".json", ".txt")
-    )
+    txt_path = os.path.join(label_dir_path,
+                            target_dir,
+                            json_name.replace(".json", ".txt"))
 
     with open(txt_path, "w+") as f:
-        for yolo_obj_idx, yolo_obj in enumerate(yolo_obj_list):
-            yolo_obj_line = (
-                "%s %s %s %s %s\n" % yolo_obj
-                if yolo_obj_idx + 1 != len(yolo_obj_list)
-                else "%s %s %s %s %s" % yolo_obj
-            )
+        for yolo_obj in yolo_obj_list:
+            label, points = yolo_obj
+            points = [str(item) for item in points]
+            yolo_obj_line = f"{label} {' '.join(points)}\n"
             f.write(yolo_obj_line)
 
 
@@ -130,10 +131,11 @@ def save_yolo_image(json_data, json_name, image_dir_path, target_dir):
 
 class Labelme2YOLO(object):
 
-    def __init__(self, json_dir):
+    def __init__(self, json_dir, output_format):
         self._json_dir = json_dir
+        self._output_format = output_format
 
-        self._label_id_map = self._get_label_id_map(self._json_dir)
+        self._label_id_map = get_label_id_map(self._json_dir)
 
     def _make_train_val_dir(self):
         self._label_dir_path = os.path.join(self._json_dir,
@@ -152,20 +154,6 @@ def _make_train_val_dir(self):
 
             os.makedirs(yolo_path)
 
-    @staticmethod
-    def _get_label_id_map(self, json_dir):
-        label_set = set()
-
-        for file_name in os.listdir(json_dir):
-            if file_name.endswith('json'):
-                json_path = os.path.join(json_dir, file_name)
-                data = json.load(open(json_path))
-                for shape in data['shapes']:
-                    label_set.add(shape['label'])
-
-        return OrderedDict([(label, label_id)
-                            for label_id, label in enumerate(label_set)])
-
     def _train_test_split(self, folders, json_names, val_size, test_size):
         if len(folders) > 0 and 'train' in folders and 'val' in folders and 'test' in folders:
             train_folder = os.path.join(self._json_dir, 'train/')
@@ -233,28 +221,28 @@ def covert_json_to_text(self, target_dir, json_name):
         print('Converting %s for %s ...' %
               (json_name, target_dir.replace('/', '')))
 
-        img_path = self._save_yolo_image(json_data,
-                                         json_name,
-                                         self._image_dir_path,
-                                         target_dir)
+        img_path = save_yolo_image(json_data,
+                                   json_name,
+                                   self._image_dir_path,
+                                   target_dir)
         yolo_obj_list = self._get_yolo_object_list(json_data, img_path)
-        self._save_yolo_label(json_name,
-                              self._label_dir_path,
-                              target_dir,
-                              yolo_obj_list)
+        save_yolo_label(json_name,
+                        self._label_dir_path,
+                        target_dir,
+                        yolo_obj_list)
 
     def convert_one(self, json_name):
         json_path = os.path.join(self._json_dir, json_name)
         json_data = json.load(open(json_path))
 
         print('Converting %s ...' % json_name)
 
-        img_path = self._save_yolo_image(json_data, json_name,
-                                         self._json_dir, '')
+        img_path = save_yolo_image(json_data, json_name,
+                                   self._json_dir, '')
 
         yolo_obj_list = self._get_yolo_object_list(json_data, img_path)
-        self._save_yolo_label(json_name, self._json_dir,
-                              '', yolo_obj_list)
+        save_yolo_label(json_name, self._json_dir,
+                        '', yolo_obj_list)
 
     def _get_yolo_object_list(self, json_data, img_path):
         yolo_obj_list = []
@@ -298,35 +286,14 @@ def _get_other_shape_yolo_object(self, shape, img_h, img_w):
         points[::2] = [float(point[0]) / img_w for point in point_list]
         points[1::2] = [float(point[1]) / img_h for point in point_list]
         if len(points) == 4:
-            points = extend_point_list(points)
+            if self._output_format == "polygon":
+                points = extend_point_list(points)
+            if self._output_format == "bbox":
+                points = extend_point_list(points, "bbox")
         label_id = self._label_id_map[shape['label']]
 
         return label_id, points.tolist()
 
-    @staticmethod
-    def _save_yolo_label(self, json_name, label_dir_path, target_dir, yolo_obj_list):
-        txt_path = os.path.join(label_dir_path,
-                                target_dir,
-                                json_name.replace('.json', '.txt'))
-
-        with open(txt_path, 'w+') as f:
-            for yolo_obj in yolo_obj_list:
-                label, points = yolo_obj
-                points = [str(item) for item in points]
-                yolo_obj_line = f"{label} {' '.join(points)}\n"
-                f.write(yolo_obj_line)
-
-    @staticmethod
-    def _save_yolo_image(self, json_data, json_name, image_dir_path, target_dir):
-        img_name = json_name.replace('.json', '.png')
-        img_path = os.path.join(image_dir_path, target_dir, img_name)
-
-        if not os.path.exists(img_path):
-            img = img_b64_to_arr(json_data['imageData'])
-            PIL.Image.fromarray(img).save(img_path)
-
-        return img_path
-
     def _save_dataset_yaml(self):
         yaml_path = os.path.join(
             self._json_dir, 'YOLODataset/', 'dataset.yaml')