- 更新文档

  - 添加api参考
- 修复bug

Author: yatengLG

ISAT/annotation.py

@@ -5,23 +5,24 @@
 from PIL import Image
 import numpy as np
 from json import load, dump
-from typing import List
+from typing import List, Union
 
+__all__ = ['Object', 'Annotation']
 
 class Object:
-    """
-    A class to represent an annotation object.
+    r"""A class to represent an annotation object.
 
-    category: The category of the object.
-    group: The group of the object.
-    segmentation: The vertices of the object.
-    area: The area of the object.
-    layer: The layer of the object.
-    bbox: The bbox of the object.
-    iscrowd: The crowd tag of the object.
-    note: The note of the object.
+    Arguments:
+        category (str): The category of the object.
+        group (int): The group of the object.
+        segmentation (list | tuple): The vertices of the object.[(x1, y1), (x2, y2), ...]
+        area (float): The area of the object.
+        layer (int): The layer of the object.
+        bbox (list | tuple): The bbox of the object. [xmin, ymin, xmax, ymax]
+        iscrowd (bool): The crowd tag of the object.
+        note (str): The note of the object.
     """
-    def __init__(self, category:str, group:int, segmentation, area, layer, bbox, iscrowd=0, note=''):
+    def __init__(self, category: str, group: int, segmentation: Union[list, tuple], area: float, layer: int, bbox: Union[list, tuple], iscrowd: bool=False, note: str=''):
         self.category = category
         self.group = group
         self.segmentation = segmentation
@@ -33,7 +34,23 @@ def __init__(self, category:str, group:int, segmentation, area, layer, bbox, isc
 
 
 class Annotation:
-    def __init__(self, image_path, label_path):
+    r"""A class to represent an annotation containing many objects.
+
+    Arguments:
+        image_path (str): The path to the image.
+        label_path (str): The path to the label file.
+
+    Attributes:
+        description (str): Always 'ISAT'.
+        img_folder (str): The path to the folder where the images are located.
+        img_name (str): The name of the image.
+        label_path (str): The path to the label file.
+        note (str): The note of the image.
+        height (int): The height of the image.
+        width (int): The width of the image.
+        depth (int): The depth of the image.
+    """
+    def __init__(self, image_path:str, label_path:str):
         img_folder, img_name = os.path.split(image_path)
         self.description = 'ISAT'
         self.img_folder = img_folder
@@ -52,9 +69,12 @@ def __init__(self, image_path, label_path):
             print('Warning: Except image has 2 or 3 ndim, but get {}.'.format(image.ndim))
         del image
 
-        self.objects:List[Object,...] = []
+        self.objects:List[Object, ] = []
 
     def load_annotation(self):
+        r"""
+        Load annotation from self.label_path
+        """
         if os.path.exists(self.label_path):
             with open(self.label_path, 'r', encoding='utf-8') as f:
                 dataset = load(f)
@@ -79,7 +99,8 @@ def load_annotation(self):
                         group = obj.get('group', 0)
                         if group is None: group = 0
                         segmentation = obj.get('segmentation', [])
-                        iscrowd = obj.get('iscrowd', 0)
+                        iscrowd = obj.get('iscrowd', False)
+                        iscrowd = iscrowd if isinstance(iscrowd, bool) else bool(iscrowd)
                         note = obj.get('note', '')
                         area = obj.get('area', 0)
                         layer = obj.get('layer', 2)
@@ -92,6 +113,9 @@ def load_annotation(self):
         return self
 
     def save_annotation(self):
+        r"""
+        Save annotation to self.label_path
+        """
         dataset = {}
         dataset['info'] = {}
         dataset['info']['description'] = self.description

ISAT/configs.py

@@ -2,11 +2,16 @@
 from enum import Enum
 import os
 
+__all__ = ['load_config', 'save_config', 'STATUSMode', 'DRAWMode', 'MAPMode', 'CONTOURMode']
+
 ISAT_ROOT = os.path.split(os.path.abspath(__file__))[0]
+"""Project root - ISAT/"""
 SOFTWARE_CONFIG_FILE = os.path.join(ISAT_ROOT, 'software.yaml')
+"""Software config file - ISAT/software.yaml"""
 CONFIG_FILE = os.path.join(ISAT_ROOT, 'isat.yaml')
+"""Category config file - ISAT/isat.yaml"""
 CHECKPOINT_PATH = os.path.join(ISAT_ROOT, 'checkpoints')
-SHORTCUT_FILE = os.path.join(ISAT_ROOT, 'shortcut.yaml')
+"""Checkpoints save root - ISAT/checkpoints"""
 
 os.makedirs(os.path.join(CHECKPOINT_PATH, 'tmp'), exist_ok=True)
 
@@ -18,42 +23,62 @@
     with open(CONFIG_FILE, 'w') as f:
         pass
 
-def load_config(file):
+def load_config(file: str) -> dict:
+    r"""
+    Load config file
+
+    Arguments:
+        file (str): config file path
+
+    Returns:
+        dict: config dict
+    """
     with open(file, 'rb')as f:
         cfg = yaml.load(f.read(), Loader=yaml.FullLoader)
     return cfg
 
-def save_config(cfg, file):
+def save_config(cfg: dict, file: str) -> None:
+    """
+    Save config file
+
+    Arguments:
+        cfg (dict): config dict
+        file (str): config file path
+    """
     s = yaml.dump(cfg, allow_unicode=True)
     with open(file, 'w', encoding='utf-8') as f:
         f.write(s)
-    return True
 
 class STATUSMode(Enum):
+    """The status mode"""
     VIEW = 0
     CREATE = 1
     EDIT = 2
     REPAINT = 3
 
 class DRAWMode(Enum):
+    """The draw mode."""
     POLYGON = 0
+    """Manually draw polygon"""
     SEGMENTANYTHING = 1
+    """Segment anything with point prompt."""
     SEGMENTANYTHING_BOX = 2
-
-class CLICKMode(Enum):
-    POSITIVE = 0
-    NEGATIVE = 1
+    """Segment anything with box prompt."""
 
 class MAPMode(Enum):
+    """Canvas show map mode"""
     LABEL = 0
     SEMANTIC = 1
     INSTANCE = 2
 
 class CONTOURMode(Enum):
-    SAVE_MAX_ONLY = 0       # 只保留最多顶点的mask（一般为最大面积）
-    SAVE_EXTERNAL = 1       # 只保留外轮廓
-    SAVE_ALL = 2            # 保留所有轮廓
-
+    """
+    Contour Mode - ways to convert masks to polygons.
+    """
+    SAVE_MAX_ONLY = 0
+    """Only save max contour."""
+    SAVE_EXTERNAL = 1
+    """Only save external contour."""
+    SAVE_ALL = 2
+    """Only save all contour."""
 
-if __name__ == '__main__':
-    cfg = load_config('/mnt/disk2/PycharmProjects/ISAT_with_segment_anything/ISAT/isat.yaml')

ISAT/formats/coco.py

@@ -9,12 +9,25 @@
 import tqdm
 import numpy as np
 
+__all__ = ['COCO']
 
 class COCO(ISAT):
+    r"""
+    COCO format
+
+    Attributes:
+        keep_crowd (bool): keep the crowded objects
+    """
     def __init__(self):
         self.keep_crowd = True
 
-    def read_from_coco(self, annotation_file):
+    def read_from_coco(self, annotation_file: str) -> bool:
+        r"""
+        Load annotations from a COCO json file.
+
+        Arguments:
+            annotation_file (str): coco json annotations file.
+        """
         self.annos.clear()
         self.cates = ()
 
@@ -48,13 +61,13 @@ def read_from_coco(self, annotation_file):
 
                 segmentations = anno_coco.get('segmentation', [])    # 多个polygon
                 area = anno_coco.get('area', None)                   # coco中，area是组面积，isat中是单个polygon面积
-                iscrowd = anno_coco.get('iscrowd', None)
+                iscrowd = anno_coco.get('iscrowd', False)
                 image_id = anno_coco.get('image_id', None)
                 bbox = anno_coco.get('bbox', [])
                 category_id = anno_coco.get('category_id', None)
                 id = anno_coco.get('id', None)
 
-                if iscrowd == 0:
+                if iscrowd is False:
                     # polygon
                     for segmentation in segmentations:
                         xs = segmentation[::2]
@@ -73,7 +86,7 @@ def read_from_coco(self, annotation_file):
 
                         objs.append(obj)
 
-                elif iscrowd == 1 and self.keep_crowd:
+                elif iscrowd is True and self.keep_crowd:
                     if isinstance(segmentations, dict) and 'counts' in segmentations:
                         # RLE
                         rles = mscoco_mask.frPyObjects(segmentations, height, width)
@@ -121,8 +134,14 @@ def read_from_coco(self, annotation_file):
             self.annos[self.remove_file_suffix(file_name)] = anno
         return True
 
-    def save_to_coco(self, annotation_file, cates: tuple=()):
+    def save_to_coco(self, annotation_file: str, cates: tuple=()) -> bool:
+        r"""
+        Save annotations to a COCO json file.
 
+        Arguments:
+            annotation_file (str): coco json annotations file.
+            cates (tuple): coco cates.
+        """
         coco_anno = {}
         # info
         coco_anno['info'] = {}