PascalVocSource = PascalVocSource()

khandy/label/detect.py

@@ -13,15 +13,21 @@ import lxml.builder
 import numpy as np
 
 
-__all__ = ['DetectIrObject', 'DetectIrRecord', 'load_detect',
-           'save_detect', 'convert_detect', 'replace_detect_label',
-           'load_coco_class_names']
+__all__ = [
+    "DetectIrObject",
+    "DetectIrRecord",
+    "load_detect",
+    "save_detect",
+    "convert_detect",
+    "replace_detect_label",
+    "load_coco_class_names",
+]
 
 
 @dataclass
 class DetectIrObject:
-    """Intermediate Representation Format of Object
-    """
+    """Intermediate Representation Format of Object"""
+
     label: str
     x_min: float
     y_min: float
@@ -31,8 +37,8 @@ class DetectIrObject:
 
 @dataclass
 class DetectIrRecord:
-    """Intermediate Representation Format of Record
-    """
+    """Intermediate Representation Format of Record"""
+
     filename: str
     width: int
     height: int
@@ -41,9 +47,9 @@ class DetectIrRecord:
 
 @dataclass
 class PascalVocSource:
-    database: str = ''
-    annotation: str = ''
-    image: str = ''
+    database: str = ""
+    annotation: str = ""
+    image: str = ""
 
 
 @dataclass
@@ -64,7 +70,7 @@ class PascalVocBndbox:
 @dataclass
 class PascalVocObject:
     name: str
-    pose: str = 'Unspecified'
+    pose: str = "Unspecified"
     truncated: int = 0
     difficult: int = 0
     bndbox: Optional[PascalVocBndbox] = None
@@ -72,10 +78,10 @@ class PascalVocObject:
 
 @dataclass
 class PascalVocRecord:
-    folder: str = ''
-    filename: str = ''
-    path: str = ''
-    source: PascalVocSource = PascalVocSource()
+    folder: str = ""
+    filename: str = ""
+    path: str = ""
+    source: PascalVocSource = field(default_factory=PascalVocSource)
     size: Optional[PascalVocSize] = None
     segmented: int = 0
     objects: List[PascalVocObject] = field(default_factory=list)
@@ -87,40 +93,40 @@ class PascalVocHandler:
         pascal_voc_record = PascalVocRecord()
 
         xml_tree = ET.parse(filename)
-        pascal_voc_record.folder = xml_tree.find('folder').text
-        pascal_voc_record.filename = xml_tree.find('filename').text
-        pascal_voc_record.path = xml_tree.find('path').text
-        pascal_voc_record.segmented = xml_tree.find('segmented').text
+        pascal_voc_record.folder = xml_tree.find("folder").text
+        pascal_voc_record.filename = xml_tree.find("filename").text
+        pascal_voc_record.path = xml_tree.find("path").text
+        pascal_voc_record.segmented = xml_tree.find("segmented").text
 
-        source_tag = xml_tree.find('source')
+        source_tag = xml_tree.find("source")
         pascal_voc_record.source = PascalVocSource(
-            database=source_tag.find('database').text,
+            database=source_tag.find("database").text,
             # annotation=source_tag.find('annotation').text,
             # image=source_tag.find('image').text
         )
 
-        size_tag = xml_tree.find('size')
+        size_tag = xml_tree.find("size")
         pascal_voc_record.size = PascalVocSize(
-            width=int(size_tag.find('width').text),
-            height=int(size_tag.find('height').text),
-            depth=int(size_tag.find('depth').text)
+            width=int(size_tag.find("width").text),
+            height=int(size_tag.find("height").text),
+            depth=int(size_tag.find("depth").text),
         )
 
-        object_tags = xml_tree.findall('object')
+        object_tags = xml_tree.findall("object")
         for index, object_tag in enumerate(object_tags):
-            bndbox_tag = object_tag.find('bndbox')
+            bndbox_tag = object_tag.find("bndbox")
             bndbox = PascalVocBndbox(
-                xmin=float(bndbox_tag.find('xmin').text) - 1,
-                ymin=float(bndbox_tag.find('ymin').text) - 1,
-                xmax=float(bndbox_tag.find('xmax').text) - 1,
-                ymax=float(bndbox_tag.find('ymax').text) - 1
+                xmin=float(bndbox_tag.find("xmin").text) - 1,
+                ymin=float(bndbox_tag.find("ymin").text) - 1,
+                xmax=float(bndbox_tag.find("xmax").text) - 1,
+                ymax=float(bndbox_tag.find("ymax").text) - 1,
             )
             pascal_voc_object = PascalVocObject(
-                name=object_tag.find('name').text,
-                pose=object_tag.find('pose').text,
-                truncated=object_tag.find('truncated').text,
-                difficult=object_tag.find('difficult').text,
-                bndbox=bndbox
+                name=object_tag.find("name").text,
+                pose=object_tag.find("pose").text,
+                truncated=object_tag.find("truncated").text,
+                difficult=object_tag.find("difficult").text,
+                bndbox=bndbox,
             )
             pascal_voc_record.objects.append(pascal_voc_object)
         return pascal_voc_record
@@ -158,18 +164,17 @@ class PascalVocHandler:
             )
             xml.append(object_tag)
 
-        if not filename.endswith('.xml'):
-            filename = filename + '.xml'
-        with open(filename, 'wb') as f:
-            f.write(lxml.etree.tostring(
-                xml, pretty_print=True, encoding='utf-8'))
+        if not filename.endswith(".xml"):
+            filename = filename + ".xml"
+        with open(filename, "wb") as f:
+            f.write(lxml.etree.tostring(xml, pretty_print=True, encoding="utf-8"))
 
     @staticmethod
     def to_ir(pascal_voc_record: PascalVocRecord) -> DetectIrRecord:
         ir_record = DetectIrRecord(
             filename=pascal_voc_record.filename,
             width=pascal_voc_record.size.width,
-            height=pascal_voc_record.size.height
+            height=pascal_voc_record.size.height,
         )
         for pascal_voc_object in pascal_voc_record.objects:
             ir_object = DetectIrObject(
@@ -177,7 +182,7 @@ class PascalVocHandler:
                 x_min=pascal_voc_object.bndbox.xmin,
                 y_min=pascal_voc_object.bndbox.ymin,
                 x_max=pascal_voc_object.bndbox.xmax,
-                y_max=pascal_voc_object.bndbox.ymax
+                y_max=pascal_voc_object.bndbox.ymax,
             )
             ir_record.objects.append(ir_object)
         return ir_record
@@ -186,11 +191,7 @@ class PascalVocHandler:
     def from_ir(ir_record: DetectIrRecord) -> PascalVocRecord:
         pascal_voc_record = PascalVocRecord(
             filename=ir_record.filename,
-            size=PascalVocSize(
-                width=ir_record.width,
-                height=ir_record.height,
-                depth=3
-            )
+            size=PascalVocSize(width=ir_record.width, height=ir_record.height, depth=3),
         )
         for ir_object in ir_record.objects:
             pascal_voc_object = PascalVocObject(
@@ -200,24 +201,36 @@ class PascalVocHandler:
                     ymin=ir_object.y_min,
                     xmax=ir_object.x_max,
                     ymax=ir_object.y_max,
-                )
+                ),
             )
             pascal_voc_record.objects.append(pascal_voc_object)
         return pascal_voc_record
 
 
 class _NumpyEncoder(json.JSONEncoder):
-    """ Special json encoder for numpy types """
+    """Special json encoder for numpy types"""
 
     def default(self, obj):
         if isinstance(obj, (np.bool_,)):
             return bool(obj)
-        elif isinstance(obj, (np.int_, np.intc, np.intp, np.int8,
-                              np.int16, np.int32, np.int64, np.uint8,
-                              np.uint16, np.uint32, np.uint64)):
+        elif isinstance(
+            obj,
+            (
+                np.int_,
+                np.intc,
+                np.intp,
+                np.int8,
+                np.int16,
+                np.int32,
+                np.int64,
+                np.uint8,
+                np.uint16,
+                np.uint32,
+                np.uint64,
+            ),
+        ):
             return int(obj)
-        elif isinstance(obj, (np.float_, np.float16, np.float32,
-                              np.float64)):
+        elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)):
             return float(obj)
         elif isinstance(obj, (np.ndarray,)):
             return obj.tolist()
@@ -238,7 +251,7 @@ class LabelmeShape:
 
 @dataclass
 class LabelmeRecord:
-    version: str = '4.5.6'
+    version: str = "4.5.6"
     flags: dict = field(default_factory=dict)
     shapes: List[LabelmeShape] = field(default_factory=list)
     imagePath: Optional[str] = None
@@ -267,10 +280,10 @@ class LabelmeHandler:
         ir_record = DetectIrRecord(
             filename=labelme_record.imagePath,
             width=labelme_record.imageWidth,
-            height=labelme_record.imageHeight
+            height=labelme_record.imageHeight,
         )
         for labelme_shape in labelme_record.shapes:
-            if labelme_shape.shape_type != 'rectangle':
+            if labelme_shape.shape_type != "rectangle":
                 continue
             ir_object = DetectIrObject(
                 label=labelme_shape.label,
@@ -287,14 +300,16 @@ class LabelmeHandler:
         labelme_record = LabelmeRecord(
             imagePath=ir_record.filename,
             imageWidth=ir_record.width,
-            imageHeight=ir_record.height
+            imageHeight=ir_record.height,
         )
         for ir_object in ir_record.objects:
             labelme_shape = LabelmeShape(
                 label=ir_object.label,
-                shape_type='rectangle',
-                points=[[ir_object.x_min, ir_object.y_min],
-                        [ir_object.x_max, ir_object.y_max]]
+                shape_type="rectangle",
+                points=[
+                    [ir_object.x_min, ir_object.y_min],
+                    [ir_object.x_max, ir_object.y_max],
+                ],
             )
             labelme_record.shapes.append(labelme_shape)
         return labelme_record
@@ -320,23 +335,26 @@ class YoloRecord:
 class YoloHandler:
     @staticmethod
     def load(filename, **kwargs) -> YoloRecord:
-        assert 'image_filename' in kwargs
-        assert 'width' in kwargs and 'height' in kwargs
+        assert "image_filename" in kwargs
+        assert "width" in kwargs and "height" in kwargs
 
         records = khandy.load_list(filename)
         yolo_record = YoloRecord(
-            filename=kwargs.get('image_filename'),
-            width=kwargs.get('width'),
-            height=kwargs.get('height'))
+            filename=kwargs.get("image_filename"),
+            width=kwargs.get("width"),
+            height=kwargs.get("height"),
+        )
         for record in records:
             record_parts = record.split()
-            yolo_record.objects.append(YoloObject(
-                label=record_parts[0],
-                x_center=float(record_parts[1]),
-                y_center=float(record_parts[2]),
-                width=float(record_parts[3]),
-                height=float(record_parts[4]),
-            ))
+            yolo_record.objects.append(
+                YoloObject(
+                    label=record_parts[0],
+                    x_center=float(record_parts[1]),
+                    y_center=float(record_parts[2]),
+                    width=float(record_parts[3]),
+                    height=float(record_parts[4]),
+                )
+            )
         return yolo_record
 
     @staticmethod
@@ -344,9 +362,10 @@ class YoloHandler:
         records = []
         for object in yolo_record.objects:
             records.append(
-                f'{object.label} {object.x_center} {object.y_center} {object.width} {object.height}')
-        if not filename.endswith('.txt'):
-            filename = filename + '.txt'
+                f"{object.label} {object.x_center} {object.y_center} {object.width} {object.height}"
+            )
+        if not filename.endswith(".txt"):
+            filename = filename + ".txt"
         khandy.save_list(filename, records)
 
     @staticmethod
@@ -354,23 +373,23 @@ class YoloHandler:
         ir_record = DetectIrRecord(
             filename=yolo_record.filename,
             width=yolo_record.width,
-            height=yolo_record.height
+            height=yolo_record.height,
         )
         for yolo_object in yolo_record.objects:
-            x_min = (yolo_object.x_center - 0.5 *
-                     yolo_object.width) * yolo_record.width
-            y_min = (yolo_object.y_center - 0.5 *
-                     yolo_object.height) * yolo_record.height
-            x_max = (yolo_object.x_center + 0.5 *
-                     yolo_object.width) * yolo_record.width
-            y_max = (yolo_object.y_center + 0.5 *
-                     yolo_object.height) * yolo_record.height
+            x_min = (yolo_object.x_center - 0.5 * yolo_object.width) * yolo_record.width
+            y_min = (
+                yolo_object.y_center - 0.5 * yolo_object.height
+            ) * yolo_record.height
+            x_max = (yolo_object.x_center + 0.5 * yolo_object.width) * yolo_record.width
+            y_max = (
+                yolo_object.y_center + 0.5 * yolo_object.height
+            ) * yolo_record.height
             ir_object = DetectIrObject(
                 label=yolo_object.label,
                 x_min=x_min,
                 y_min=y_min,
                 x_max=x_max,
-                y_max=y_max
+                y_max=y_max,
             )
             ir_record.objects.append(ir_object)
         return ir_record
@@ -378,15 +397,11 @@ class YoloHandler:
     @staticmethod
     def from_ir(ir_record: DetectIrRecord) -> YoloRecord:
         yolo_record = YoloRecord(
-            filename=ir_record.filename,
-            width=ir_record.width,
-            height=ir_record.height
+            filename=ir_record.filename, width=ir_record.width, height=ir_record.height
         )
         for ir_object in ir_record.objects:
-            x_center = (ir_object.x_max + ir_object.x_min) / \
-                (2 * ir_record.width)
-            y_center = (ir_object.y_max + ir_object.y_min) / \
-                (2 * ir_record.height)
+            x_center = (ir_object.x_max + ir_object.x_min) / (2 * ir_record.width)
+            y_center = (ir_object.y_max + ir_object.y_min) / (2 * ir_record.height)
             width = abs(ir_object.x_max - ir_object.x_min) / ir_record.width
             height = abs(ir_object.y_max - ir_object.y_min) / ir_record.height
             yolo_object = YoloObject(
@@ -422,31 +437,32 @@ class CocoHandler:
     def load(filename, **kwargs) -> List[CocoRecord]:
         json_data = khandy.load_json(filename)
 
-        images = json_data['images']
-        annotations = json_data['annotations']
-        categories = json_data['categories']
+        images = json_data["images"]
+        annotations = json_data["annotations"]
+        categories = json_data["categories"]
 
         label_map = {}
         for cat_item in categories:
-            label_map[cat_item['id']] = cat_item['name']
+            label_map[cat_item["id"]] = cat_item["name"]
 
         coco_records = OrderedDict()
         for image_item in images:
-            coco_records[image_item['id']] = CocoRecord(
-                filename=image_item['file_name'],
-                width=image_item['width'],
-                height=image_item['height'],
-                objects=[])
+            coco_records[image_item["id"]] = CocoRecord(
+                filename=image_item["file_name"],
+                width=image_item["width"],
+                height=image_item["height"],
+                objects=[],
+            )
 
         for annotation_item in annotations:
             coco_object = CocoObject(
-                label=label_map[annotation_item['category_id']],
-                x_min=annotation_item['bbox'][0],
-                y_min=annotation_item['bbox'][1],
-                width=annotation_item['bbox'][2],
-                height=annotation_item['bbox'][3])
-            coco_records[annotation_item['image_id']
-                         ].objects.append(coco_object)
+                label=label_map[annotation_item["category_id"]],
+                x_min=annotation_item["bbox"][0],
+                y_min=annotation_item["bbox"][1],
+                width=annotation_item["bbox"][2],
+                height=annotation_item["bbox"][3],
+            )
+            coco_records[annotation_item["image_id"]].objects.append(coco_object)
         return list(coco_records.values())
 
     @staticmethod
@@ -462,7 +478,7 @@ class CocoHandler:
                 x_min=coco_object.x_min,
                 y_min=coco_object.y_min,
                 x_max=coco_object.x_min + coco_object.width,
-                y_max=coco_object.y_min + coco_object.height
+                y_max=coco_object.y_min + coco_object.height,
             )
             ir_record.objects.append(ir_object)
         return ir_record
@@ -470,9 +486,7 @@ class CocoHandler:
     @staticmethod
     def from_ir(ir_record: DetectIrRecord) -> CocoRecord:
         coco_record = CocoRecord(
-            filename=ir_record.filename,
-            width=ir_record.width,
-            height=ir_record.height
+            filename=ir_record.filename, width=ir_record.width, height=ir_record.height
         )
         for ir_object in ir_record.objects:
             coco_object = CocoObject(
@@ -480,26 +494,25 @@ class CocoHandler:
                 x_min=ir_object.x_min,
                 y_min=ir_object.y_min,
                 width=ir_object.x_max - ir_object.x_min,
-                height=ir_object.y_max - ir_object.y_min
+                height=ir_object.y_max - ir_object.y_min,
             )
             coco_record.objects.append(coco_object)
         return coco_record
 
 
 def load_detect(filename, fmt, **kwargs) -> DetectIrRecord:
-    if fmt == 'labelme':
+    if fmt == "labelme":
         labelme_record = LabelmeHandler.load(filename, **kwargs)
         ir_record = LabelmeHandler.to_ir(labelme_record)
-    elif fmt == 'yolo':
+    elif fmt == "yolo":
         yolo_record = YoloHandler.load(filename, **kwargs)
         ir_record = YoloHandler.to_ir(yolo_record)
-    elif fmt in ('voc', 'pascal', 'pascal_voc'):
+    elif fmt in ("voc", "pascal", "pascal_voc"):
         pascal_voc_record = PascalVocHandler.load(filename, **kwargs)
         ir_record = PascalVocHandler.to_ir(pascal_voc_record)
-    elif fmt == 'coco':
+    elif fmt == "coco":
         coco_records = CocoHandler.load(filename, **kwargs)
-        ir_record = [CocoHandler.to_ir(coco_record)
-                     for coco_record in coco_records]
+        ir_record = [CocoHandler.to_ir(coco_record) for coco_record in coco_records]
     else:
         raise ValueError(f"Unsupported detect label fmt. Got {fmt}")
     return ir_record
@@ -507,16 +520,16 @@ def load_detect(filename, fmt, **kwargs) -> DetectIrRecord:
 
 def save_detect(filename, ir_record: DetectIrRecord, out_fmt):
     os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True)
-    if out_fmt == 'labelme':
+    if out_fmt == "labelme":
         labelme_record = LabelmeHandler.from_ir(ir_record)
         LabelmeHandler.save(filename, labelme_record)
-    elif out_fmt == 'yolo':
+    elif out_fmt == "yolo":
         yolo_record = YoloHandler.from_ir(ir_record)
         YoloHandler.save(filename, yolo_record)
-    elif out_fmt in ('voc', 'pascal', 'pascal_voc'):
+    elif out_fmt in ("voc", "pascal", "pascal_voc"):
         pascal_voc_record = PascalVocHandler.from_ir(ir_record)
         PascalVocHandler.save(filename, pascal_voc_record)
-    elif out_fmt == 'coco':
+    elif out_fmt == "coco":
         raise ValueError("Unsupported for `coco` now!")
     else:
         raise ValueError(f"Unsupported detect label fmt. Got {out_fmt}")
@@ -524,25 +537,32 @@ def save_detect(filename, ir_record: DetectIrRecord, out_fmt):
 
 def _get_format(record):
     if isinstance(record, LabelmeRecord):
-        return ('labelme',)
+        return ("labelme",)
     elif isinstance(record, YoloRecord):
-        return ('yolo',)
+        return ("yolo",)
     elif isinstance(record, PascalVocRecord):
-        return ('voc', 'pascal', 'pascal_voc')
+        return ("voc", "pascal", "pascal_voc")
     elif isinstance(record, CocoRecord):
-        return ('coco',)
+        return ("coco",)
     elif isinstance(record, DetectIrRecord):
-        return ('ir', 'detect_ir')
+        return ("ir", "detect_ir")
     else:
         return ()
 
 
 def convert_detect(record, out_fmt):
-    allowed_fmts = ('labelme', 'yolo', 'voc', 'coco',
-                    'pascal', 'pascal_voc', 'ir', 'detect_ir')
+    allowed_fmts = (
+        "labelme",
+        "yolo",
+        "voc",
+        "coco",
+        "pascal",
+        "pascal_voc",
+        "ir",
+        "detect_ir",
+    )
     if out_fmt not in allowed_fmts:
-        raise ValueError(
-            "Unsupported label format conversions for given out_fmt")
+        raise ValueError("Unsupported label format conversions for given out_fmt")
     if out_fmt in _get_format(record):
         return record
 
@@ -557,17 +577,17 @@ def convert_detect(record, out_fmt):
     elif isinstance(record, DetectIrRecord):
         ir_record = record
     else:
-        raise TypeError('Unsupported type for record')
+        raise TypeError("Unsupported type for record")
 
-    if out_fmt in ('ir', 'detect_ir'):
+    if out_fmt in ("ir", "detect_ir"):
         dst_record = ir_record
-    elif out_fmt == 'labelme':
+    elif out_fmt == "labelme":
         dst_record = LabelmeHandler.from_ir(ir_record)
-    elif out_fmt == 'yolo':
+    elif out_fmt == "yolo":
         dst_record = YoloHandler.from_ir(ir_record)
-    elif out_fmt in ('voc', 'pascal', 'pascal_voc'):
+    elif out_fmt in ("voc", "pascal", "pascal_voc"):
         dst_record = PascalVocHandler.from_ir(ir_record)
-    elif out_fmt == 'coco':
+    elif out_fmt == "coco":
         dst_record = CocoHandler.from_ir(ir_record)
     return dst_record
 
@@ -590,5 +610,5 @@ def replace_detect_label(record: DetectIrRecord, label_map, ignore=True):
 
 def load_coco_class_names(filename):
     json_data = khandy.load_json(filename)
-    categories = json_data['categories']
-    return [cat_item['name'] for cat_item in categories]
+    categories = json_data["categories"]
+    return [cat_item["name"] for cat_item in categories]