Fixed issue #6

label_convert/vis_coco.py

@@ -4,10 +4,11 @@
 import argparse
 import json
 import platform
-from pathlib import Path
 import random
+from pathlib import Path
 
 import cv2
+import numpy as np
 
 
 class VisCOCO:
@@ -56,7 +57,14 @@ def __call__(self, img_id: int, json_path, img_path):
             class_name = categories_dict[class_id]
             class_color = color[class_id - 1]
 
-            x, y, w, h = list(map(int, anno["bbox"]))
+            # plot sgmentations
+            segs = anno.get("segmentation", None)
+            if segs is not None:
+                segs = np.array(segs).reshape(-1, 2)
+                cv2.polylines(image, np.int32([segs]), 2, class_color)
+
+            # plot rectangle
+            x, y, w, h = [round(v) for v in anno["bbox"]]
             cv2.rectangle(
                 image, (int(x), int(y)), (int(x + w), int(y + h)), class_color, 2
             )
@@ -101,10 +109,10 @@ def main():
     parser.add_argument(
         "--json_path",
         type=str,
-        default="tests/test_files/COCO_dataset/annotations/instances_train2017.json",
+        default="tests/test_files/yolov5_dataset_coco/annotations/instances_train2017.json",
     )
     parser.add_argument(
-        "--img_dir", type=str, default="tests/test_files/COCO_dataset/train2017"
+        "--img_dir", type=str, default="tests/test_files/yolov5_dataset_coco/train2017"
     )
     args = parser.parse_args()
 

label_convert/yolov5_to_coco.py

@@ -10,6 +10,7 @@
 from typing import List, Tuple, Union
 
 import cv2
+import numpy as np
 from tqdm import tqdm
 
 ValueType = Union[str, Path, None]
@@ -104,7 +105,11 @@ def convert(self, img_list, save_img_dir, mode):
             images.append(image_dict)
 
             label_path = self.data_dir / "labels" / f"{Path(img_path).stem}.txt"
-            annotation = self.get_annotation(
+            if not label_path.exists():
+                warnings.warn(f"{label_path} not exists. Skip")
+                continue
+
+            annotation = self.read_annotation(
                 label_path, img_id, image_dict["height"], image_dict["width"]
             )
             annotations.extend(annotation)
@@ -122,8 +127,6 @@ def convert(self, img_list, save_img_dir, mode):
     def get_image_info(self, img_path, img_id, save_img_dir):
         img_path = Path(img_path)
         if self.data_dir.as_posix() not in img_path.as_posix():
-            # relative path (relative to the data_dir)
-            # e.g. images/images(3).jpg
             img_path = self.data_dir / img_path
 
         self.verify_exists(img_path)
@@ -146,43 +149,7 @@ def get_image_info(self, img_path, img_id, save_img_dir):
         }
         return image_info
 
-    def get_annotation(self, label_path: Path, img_id, height, width):
-        def get_box_info(vertex_info, height, width):
-            cx, cy, w, h = [float(i) for i in vertex_info]
-
-            cx = cx * width
-            cy = cy * height
-            box_w = w * width
-            box_h = h * height
-
-            # left top
-            x0 = max(cx - box_w / 2, 0)
-            y0 = max(cy - box_h / 2, 0)
-
-            # right bottom
-            x1 = min(x0 + box_w, width)
-            y1 = min(y0 + box_h, height)
-
-            segmentation = [[x0, y0, x1, y0, x1, y1, x0, y1]]
-            bbox = [x0, y0, box_w, box_h]
-            area = box_w * box_h
-            return segmentation, bbox, area
-
-        if not label_path.exists():
-            annotation = [
-                {
-                    "segmentation": [],
-                    "area": 0,
-                    "iscrowd": 0,
-                    "image_id": img_id,
-                    "bbox": [],
-                    "category_id": -1,
-                    "id": self.annotation_id,
-                }
-            ]
-            self.annotation_id += 1
-            return annotation
-
+    def read_annotation(self, label_path: Path, img_id, height, width):
         annotation = []
         label_list = self.read_txt(str(label_path))
         for i, one_line in enumerate(label_list):
@@ -192,7 +159,19 @@ def get_box_info(vertex_info, height, width):
                 continue
 
             category_id, vertex_info = label_info[0], label_info[1:]
-            segmentation, bbox, area = get_box_info(vertex_info, height, width)
+            point_nums = len(vertex_info)
+            if point_nums == 4:
+                segmentation, bbox, area = self.get_annotation_from_rectangle(
+                    vertex_info, height, width
+                )
+            elif point_nums > 4:
+                segmentation, bbox, area = self.get_annotation_from_poly(
+                    vertex_info, height, width
+                )
+            else:
+                warnings.warn("The nums of points are less than 4. Skip")
+                continue
+
             annotation.append(
                 {
                     "segmentation": segmentation,
@@ -207,6 +186,44 @@ def get_box_info(vertex_info, height, width):
             self.annotation_id += 1
         return annotation
 
+    @staticmethod
+    def get_annotation_from_rectangle(vertex_info, height, width):
+        cx, cy, w, h = [float(i) for i in vertex_info]
+
+        cx = cx * width
+        cy = cy * height
+        box_w = w * width
+        box_h = h * height
+
+        x0 = max(cx - box_w / 2, 0)
+        y0 = max(cy - box_h / 2, 0)
+        x1 = min(x0 + box_w, width)
+        y1 = min(y0 + box_h, height)
+
+        segmentation = [[x0, y0, x1, y0, x1, y1, x0, y1]]
+        bbox = [x0, y0, box_w, box_h]
+        area = box_w * box_h
+        return segmentation, bbox, area
+
+    @staticmethod
+    def get_annotation_from_poly(vertex_info: List[str], height, width):
+        points = np.array(vertex_info).astype(np.float64).reshape(-1, 2)
+
+        new_points = np.copy(points)
+        new_points[:, 0] = points[:, 0] * width
+        new_points[:, 1] = points[:, 1] * height
+
+        segmentation = new_points.tolist()
+
+        x0, y0 = np.min(new_points, axis=0)
+        x1, y1 = np.max(new_points, axis=0)
+        box_w = x1 - x0
+        box_h = y1 - y0
+        bbox = [x0, y0, box_w, box_h]
+
+        area = box_w * box_h
+        return segmentation, bbox, area
+
     @staticmethod
     def read_txt(txt_path):
         with open(str(txt_path), "r", encoding="utf-8") as f:
@@ -231,7 +248,10 @@ def write_json(json_path, content: dict):
 def main():
     parser = argparse.ArgumentParser("Datasets converter from YOLOV5 to COCO")
     parser.add_argument(
-        "--data_dir", type=str, default="dataset/YOLOV5", help="Dataset root path"
+        "--data_dir",
+        type=str,
+        default="tests/test_files/yolov5_dataset",
+        help="Dataset root path",
     )
     parser.add_argument(
         "--mode_list", type=str, default="train,val", help="generate which mode"

label_convert/yolov5_yaml_to_coco.py

@@ -5,10 +5,12 @@
 import json
 import shutil
 import time
+import warnings
 from pathlib import Path
-from typing import Union
+from typing import List, Union
 
 import cv2
+import numpy as np
 import yaml
 from tqdm import tqdm
 
@@ -26,9 +28,13 @@ def __init__(
         with open(yaml_path, "r", encoding="utf-8") as f:
             self.cfg = yaml.safe_load(f)
 
-        self.data_dir = Path(self.cfg.get("path"))
-        self.train_dir = self.data_dir / self.cfg.get("train")[0]
-        self.val_dir = self.data_dir / self.cfg.get("val")[0]
+        data_dir = self.cfg.get("path", None)
+        if data_dir is None:
+            data_dir = ""
+
+        self.data_dir = Path(data_dir)
+        self.train_dir = self.data_dir / self.cfg.get("train")
+        self.val_dir = self.data_dir / self.cfg.get("val")
 
         nc = self.cfg.get("nc")
         self.names = self.cfg.get("names")
@@ -96,19 +102,6 @@ def _init_json(self):
     def get_img_list(self, data_dir: Path):
         return list(data_dir.rglob("*.*"))
 
-    def _get_data_dir(self, mode):
-        data_dir = self.cfg.get(mode)
-        if data_dir:
-            if isinstance(data_dir, str):
-                full_path = [str(self.data_dir / data_dir)]
-            elif isinstance(data_dir, list):
-                full_path = [str(self.data_dir / one_dir) for one_dir in data_dir]
-            else:
-                raise TypeError(f"{data_dir} is not str or list.")
-        else:
-            raise ValueError(f"{mode} dir is not in the yaml.")
-        return full_path
-
     def _get_category(self):
         categories = []
         for i, category in enumerate(self.names, start=1):
@@ -149,11 +142,13 @@ def gen_dataset(self, img_paths, target_img_path, target_json, mode):
             label_name = f"{img_path.stem}.txt"
             label_path = self.data_dir / "labels" / img_path.parent.name / label_name
             if not label_path.exists():
-                raise FileNotFoundError(f"{label_path} not exists")
+                warnings.warn(f"{label_path} not exists. Skip")
+                continue
 
             new_anno = self.read_annotation(label_path, img_id, height, width)
             if len(new_anno) <= 0:
-                raise ValueError(f"{label_path} is empty")
+                warnings.warn(f"{label_path} is empty. Skip")
+                continue
             annotations.extend(new_anno)
 
         json_data = {
@@ -167,16 +162,30 @@ def gen_dataset(self, img_paths, target_img_path, target_json, mode):
         with open(target_json, "w", encoding="utf-8") as f:
             json.dump(json_data, f, ensure_ascii=False)
 
-    def read_annotation(self, txt_file, img_id, height, width):
+    def read_annotation(self, label_path, img_id, height, width):
         annotation = []
-        all_info = self.read_txt(txt_file)
+        all_info = self.read_txt(label_path)
         for label_info in all_info:
             label_info = label_info.split(" ")
-            if len(label_info) < 5:
+            label_len = len(label_info)
+            if label_len < 5:
                 continue
 
             category_id, vertex_info = label_info[0], label_info[1:]
-            segmentation, bbox, area = self._get_annotation(vertex_info, height, width)
+
+            point_nums = len(vertex_info)
+            if point_nums == 4:
+                segmentation, bbox, area = self.get_annotation_from_rectangle(
+                    vertex_info, height, width
+                )
+            elif point_nums > 4:
+                segmentation, bbox, area = self.get_annotation_from_poly(
+                    vertex_info, height, width
+                )
+            else:
+                warnings.warn("The nums of points are less than 4. Skip")
+                continue
+
             annotation.append(
                 {
                     "segmentation": segmentation,
@@ -192,7 +201,7 @@ def read_annotation(self, txt_file, img_id, height, width):
         return annotation
 
     @staticmethod
-    def _get_annotation(vertex_info, height, width):
+    def get_annotation_from_rectangle(vertex_info, height, width):
         cx, cy, w, h = [float(i) for i in vertex_info]
 
         cx = cx * width
@@ -210,6 +219,25 @@ def _get_annotation(vertex_info, height, width):
         area = box_w * box_h
         return segmentation, bbox, area
 
+    @staticmethod
+    def get_annotation_from_poly(vertex_info: List[str], height, width):
+        points = np.array(vertex_info).astype(np.float64).reshape(-1, 2)
+
+        new_points = np.copy(points)
+        new_points[:, 0] = points[:, 0] * width
+        new_points[:, 1] = points[:, 1] * height
+
+        segmentation = new_points.tolist()
+
+        x0, y0 = np.min(new_points, axis=0)
+        x1, y1 = np.max(new_points, axis=0)
+        box_w = x1 - x0
+        box_h = y1 - y0
+        bbox = [x0, y0, box_w, box_h]
+
+        area = box_w * box_h
+        return segmentation, bbox, area
+
     @staticmethod
     def read_txt(txt_path):
         with open(str(txt_path), "r", encoding="utf-8") as f:
@@ -232,7 +260,7 @@ def main():
     parser.add_argument(
         "--yaml_path",
         type=str,
-        default="dataset/YOLOV5_yaml/sample.yaml",
+        default="tests/test_files/crack.v1i.yolov7pytorch/data.yaml",
         help="Dataset cfg file",
     )
     args = parser.parse_args()

tests/test_files/yolov5_dataset/labels/0dcddf72-0cd7-4577-b59b-ed191863a13d_png.rf.7ed0e14fcc5dde85d883c6d41187908b.txt

@@ -0,0 +1,4 @@
+0 0.3716715984375 0.41050295937500003 0.401257396875 0.3272928984375 0.41605029531250004 0.22744082968749998 0.35872780937500004 0.1756656796875 0.2921597640625 0.17011834374999998 0.238535503125 0.14977810625000001 0.20894970468750002 0.12019230781250001 0.1923076921875 0.11279585781250001 0.15162721875000001 0.11279585781250001 0.12019230781250001 0.123890534375 0.0924556203125 0.1386834328125 0.05917159843749999 0.17751479375 0.0480769234375 0.21819526718749999 0.0573224859375 0.3254437859375 0.1109467453125 0.390162721875 0.21449704062500002 0.43269230781249995 0.3716715984375 0.41050295937500003
+0 0.06841715937500001 0.8265532546875001 0.15902366875 0.867233728125 0.2625739640625 0.8468934906250001 0.3291420125 0.7618343203125 0.35318047343750003 0.669378696875 0.31989645 0.5898668625 0.2644230765625 0.571375740625 0.22928994062500002 0.5824704140625 0.1830621296875 0.552884615625 0.12019230781250001 0.560281065625 0.0314349109375 0.6065088750000001 0 0.7100591703125 0.0388313609375 0.78772189375 0.06841715937500001 0.8265532546875001
+0 0.5732248515625 0.519600590625 0.6508875749999999 0.528846153125 0.680473371875 0.501109465625 0.6841715984375 0.48076923125000004 0.7193047328125 0.4419378703125 0.7451923078125 0.40495562031250004 0.77847633125 0.3476331375 0.7766272187500001 0.330991125 0.76553254375 0.2736686390625 0.7451923078125 0.21449704062500002 0.713757396875 0.1756656796875 0.687869821875 0.1553254453125 0.6675295859375 0.13313609375 0.6471893484375 0.1183431953125 0.5954142015625 0.1183431953125 0.560281065625 0.123890534375 0.5159023671875 0.1664201171875 0.5094897671874999 0.2353555578125 0.5214497046875 0.286612425 0.48076923125000004 0.3254437859375 0.45673076874999996 0.377218934375 0.4678254453125 0.44563609375 0.510355028125 0.5122041421875 0.5732248515625 0.519600590625
+0 0.5954142015625 0.8949704140625 0.6268491125 0.8857248515625 0.689718934375 0.8487426031249999 0.7636834328125001 0.7377958578125 0.7525887578125 0.63239645 0.6915680484375 0.5510355031249999 0.6157544374999999 0.538091715625 0.545488165625 0.5436390531249999 0.4974112421875 0.560281065625 0.43823964374999996 0.6527366875 0.44748520625 0.7193047328125 0.4881656796875 0.78402366875 0.5232988171875 0.841346153125 0.5954142015625 0.8949704140625

tests/test_files/yolov5_dataset/labels/8ae4af51-4376-4d42-b6bd-54d17c8c5e02_png.rf.2db43284642a5e92dfe56830435549a1.txt

@@ -0,0 +1,6 @@
+0 0.0166420125 0.6231508875 0.040680473437499996 0.6582840234375 0.07211538437500001 0.6860207093749999 0.12204142031250001 0.6860207093749999 0.16457100468750002 0.6471893484375 0.12573964375 0.5898668625 0.079511834375 0.5399408296875 0.055473371875 0.5491863890625 0.012943785937499999 0.5843195265625 0.0166420125 0.6231508875
+0 0.46597633125 0.7433431953125 0.48076923125000004 0.739644971875 0.49371301718750005 0.693417159375 0.4585798828125 0.6490384609375 0.399408284375 0.6120562140625 0.3698224859375 0.6009615390625 0.3254437859375 0.5861686390624999 0.26627219062499996 0.6009615390625 0.2644230765625 0.6213017734374999 0.2644230765625 0.625 0.2699704140625 0.6749260359375 0.3291420125 0.713757396875 0.393860946875 0.7359467453125 0.46597633125 0.7433431953125
+0 0.704511834375 0.6860207093749999 0.7451923078125 0.636094675 0.7267011828125 0.588017753125 0.6860207093749999 0.526997040625 0.63239645 0.5510355031249999 0.571375740625 0.5972633125 0.5732248515625 0.6434911249999999 0.5935650890625 0.687869821875 0.661982246875 0.721153846875 0.704511834375 0.6860207093749999
+0 0.27181952656249997 0.41420118281249996 0.31989645 0.393860946875 0.3365384609375 0.3439349109375 0.3346893484375 0.301405325 0.3439349109375 0.28476331250000003 0.380917159375 0.24963017812500002 0.393860946875 0.19600591875 0.366124259375 0.1312869828125 0.30880177343750004 0.1017011828125 0.26627219062499996 0.1109467453125 0.20525147812500003 0.1090976328125 0.16272189374999999 0.13498520625 0.12019230781250001 0.16087278124999999 0.099852071875 0.257026628125 0.1368343203125 0.36427514843750003 0.17936390625 0.403106509375 0.22559171562500002 0.41605029531250004 0.27181952656249997 0.41420118281249996
+0 0.6009615390625 0.2810650890625 0.6268491125 0.24963017812500002 0.6342455625000001 0.22559171562500002 0.6545857984375 0.20525147812500003 0.6397928984375 0.1571745546875 0.61760355 0.11649408125 0.5639792906250001 0.105399409375 0.5048076921875 0.1553254453125 0.49926035625000004 0.188609465625 0.501109465625 0.247781065625 0.5343934906250001 0.2773668625 0.6009615390625 0.2810650890625
+0 0.9837278093749999 0.1997041421875 0.9948224859374999 0.1553254453125 0.9615384609375001 0.107248521875 0.9079142015624999 0.073964496875 0.8653846156249999 0.070266271875 0.79326923125 0.1183431953125 0.80251479375 0.16087278124999999 0.8450443796875 0.197855028125 0.919008875 0.20894970468750002 0.9837278093749999 0.1997041421875

tests/test_files/yolov5_dataset/train.txt

@@ -1 +1,2 @@
-images/images(3).jpg
+images/images(3).jpg
+images/0dcddf72-0cd7-4577-b59b-ed191863a13d_png.rf.7ed0e14fcc5dde85d883c6d41187908b.jpg

tests/test_files/yolov5_dataset/val.txt

@@ -1 +1,2 @@
 images/images(13).jpg
+images/8ae4af51-4376-4d42-b6bd-54d17c8c5e02_png.rf.2db43284642a5e92dfe56830435549a1.jpg

tests/test_files/yolov5_yaml_dataset/labels/train/0dcddf72-0cd7-4577-b59b-ed191863a13d_png.rf.7ed0e14fcc5dde85d883c6d41187908b.txt

@@ -0,0 +1,4 @@
+0 0.3716715984375 0.41050295937500003 0.401257396875 0.3272928984375 0.41605029531250004 0.22744082968749998 0.35872780937500004 0.1756656796875 0.2921597640625 0.17011834374999998 0.238535503125 0.14977810625000001 0.20894970468750002 0.12019230781250001 0.1923076921875 0.11279585781250001 0.15162721875000001 0.11279585781250001 0.12019230781250001 0.123890534375 0.0924556203125 0.1386834328125 0.05917159843749999 0.17751479375 0.0480769234375 0.21819526718749999 0.0573224859375 0.3254437859375 0.1109467453125 0.390162721875 0.21449704062500002 0.43269230781249995 0.3716715984375 0.41050295937500003
+0 0.06841715937500001 0.8265532546875001 0.15902366875 0.867233728125 0.2625739640625 0.8468934906250001 0.3291420125 0.7618343203125 0.35318047343750003 0.669378696875 0.31989645 0.5898668625 0.2644230765625 0.571375740625 0.22928994062500002 0.5824704140625 0.1830621296875 0.552884615625 0.12019230781250001 0.560281065625 0.0314349109375 0.6065088750000001 0 0.7100591703125 0.0388313609375 0.78772189375 0.06841715937500001 0.8265532546875001
+0 0.5732248515625 0.519600590625 0.6508875749999999 0.528846153125 0.680473371875 0.501109465625 0.6841715984375 0.48076923125000004 0.7193047328125 0.4419378703125 0.7451923078125 0.40495562031250004 0.77847633125 0.3476331375 0.7766272187500001 0.330991125 0.76553254375 0.2736686390625 0.7451923078125 0.21449704062500002 0.713757396875 0.1756656796875 0.687869821875 0.1553254453125 0.6675295859375 0.13313609375 0.6471893484375 0.1183431953125 0.5954142015625 0.1183431953125 0.560281065625 0.123890534375 0.5159023671875 0.1664201171875 0.5094897671874999 0.2353555578125 0.5214497046875 0.286612425 0.48076923125000004 0.3254437859375 0.45673076874999996 0.377218934375 0.4678254453125 0.44563609375 0.510355028125 0.5122041421875 0.5732248515625 0.519600590625
+0 0.5954142015625 0.8949704140625 0.6268491125 0.8857248515625 0.689718934375 0.8487426031249999 0.7636834328125001 0.7377958578125 0.7525887578125 0.63239645 0.6915680484375 0.5510355031249999 0.6157544374999999 0.538091715625 0.545488165625 0.5436390531249999 0.4974112421875 0.560281065625 0.43823964374999996 0.6527366875 0.44748520625 0.7193047328125 0.4881656796875 0.78402366875 0.5232988171875 0.841346153125 0.5954142015625 0.8949704140625