recommended

Author: Matt Sokoloff

labelbox/data/serialization/coco/annotation.py

@@ -1,10 +1,12 @@
-from pydantic import BaseModel
-from collections import defaultdict
 from typing import Tuple, List, Union
+from pathlib import Path
+from collections import defaultdict
+
+from pydantic import BaseModel
 import numpy as np
 
 
-def rle_decoding(rle_arr, w, h):
+def rle_decoding(rle_arr : List[int], w : int, h: int) -> np.ndarray:
     indices = []
     for idx, cnt in zip(rle_arr[0::2], rle_arr[1::2]):
         indices.extend(list(range(idx - 1,
@@ -51,5 +53,5 @@ class COCOObjectAnnotation(BaseModel):
 class PanopticAnnotation(BaseModel):
     # One to one relationship between image and panoptic annotation
     image_id: int
-    file_name: str
+    file_name: Path
     segments_info: List[SegmentInfo]

labelbox/data/serialization/coco/categories.py

@@ -1,6 +1,6 @@
+import sys
+
 from pydantic import BaseModel
-from typing import Optional
-import hashlib
 
 
 class Categories(BaseModel):
@@ -11,4 +11,4 @@ class Categories(BaseModel):
 
 
 def hash_category_name(name: str) -> int:
-    return int(hashlib.sha256(name.encode('utf-8')).hexdigest(), 16) % 10000
+    return hash(name) + sys.maxsize

labelbox/data/serialization/coco/converter.py

@@ -1,22 +1,23 @@
 from typing import Dict, Any
+from pathlib import Path
 import os
 
 from labelbox.data.annotation_types.collection import LabelCollection, LabelGenerator
 from labelbox.data.serialization.coco.instance_dataset import CocoInstanceDataset
 from labelbox.data.serialization.coco.panoptic_dataset import CocoPanopticDataset
 
 
-def create_path_if_not_exists(path: str, ignore_existing_data=False):
-    if not os.path.exists(path):
-        os.makedirs(path)
+def create_path_if_not_exists(path: Path, ignore_existing_data=False):
+    if not path.exists():
+        path.mkdir(parents=True, exist_ok=True)
     elif not ignore_existing_data and os.listdir(path):
         raise ValueError(
             f"Directory `{path}`` must be empty. Or set `ignore_existing_data=True`"
         )
 
 
-def validate_path(path, name):
-    if not os.path.exists(path):
+def validate_path(path: Path, name: str):
+    if not path.exists():
         raise ValueError(f"{name} `{path}` must exist")
 
 
@@ -28,10 +29,10 @@ class COCOConverter:
     Subclasses are currently ignored.
     To use subclasses, manually flatten them before using the converter.
     """
-
+    @staticmethod
     def serialize_instances(labels: LabelCollection,
-                            image_root: str,
-                            ignore_existing_data=False) -> Dict[str, Any]:
+                            image_root: Path,
+                            ignore_existing_data=False, max_workers = 8) -> Dict[str, Any]:
         """
         Convert a Labelbox LabelCollection into an mscoco dataset.
         This function will only convert masks, polygons, and rectangles.
@@ -43,18 +44,20 @@ def serialize_instances(labels: LabelCollection,
             image_root: Where to save images to
             ignore_existing_data: Whether or not to raise an exception if images already exist.
                 This exists only to support detectons panoptic fpn model which requires two mscoco payloads for the same images.
+            max_workers : Number of workers to process dataset with
         Returns:
             A dictionary containing labels in the coco object format.
         """
         create_path_if_not_exists(image_root, ignore_existing_data)
         return CocoInstanceDataset.from_common(labels=labels,
-                                               image_root=image_root).dict()
+                                               image_root=image_root, max_workers = max_workers).dict()
 
+    @staticmethod
     def serialize_panoptic(labels: LabelCollection,
-                           image_root: str,
-                           mask_root: str,
+                           image_root: Path,
+                           mask_root: Path,
                            all_stuff: bool = False,
-                           ignore_existing_data=False) -> Dict[str, Any]:
+                           ignore_existing_data=False, max_workers = 8) -> Dict[str, Any]:
         """
         Convert a Labelbox LabelCollection into an mscoco dataset.
         This function will only convert masks, polygons, and rectangles.
@@ -69,6 +72,7 @@ def serialize_panoptic(labels: LabelCollection,
                 To convert them to stuff class set `all_stuff=True`.
             ignore_existing_data: Whether or not to raise an exception if images already exist.
                 This exists only to support detectons panoptic fpn model which requires two mscoco payloads for the same images.
+            max_workers : Number of workers to process dataset with
         Returns:
             A dictionary containing labels in the coco panoptic format.
         """
@@ -77,11 +81,11 @@ def serialize_panoptic(labels: LabelCollection,
         return CocoPanopticDataset.from_common(labels=labels,
                                                image_root=image_root,
                                                mask_root=mask_root,
-                                               all_stuff=all_stuff).dict()
+                                               all_stuff=all_stuff , max_workers = max_workers).dict()
 
-    @classmethod
-    def deserialize_panoptic(cls, json_data: Dict[str, Any], image_root: str,
-                             mask_root: str) -> LabelGenerator:
+    @staticmethod
+    def deserialize_panoptic(json_data: Dict[str, Any], image_root: Path,
+                             mask_root: Path) -> LabelGenerator:
         """
         Convert coco panoptic data into the labelbox format (as a LabelGenerator).
 
@@ -98,9 +102,9 @@ def deserialize_panoptic(cls, json_data: Dict[str, Any], image_root: str,
         gen = objs.to_common(image_root, mask_root)
         return LabelGenerator(data=gen)
 
-    @classmethod
-    def deserialize_instances(cls, json_data: Dict[str, Any],
-                              image_root: str) -> LabelGenerator:
+    @staticmethod
+    def deserialize_instances(json_data: Dict[str, Any],
+                              image_root: Path) -> LabelGenerator:
         """
         Convert coco object data into the labelbox format (as a LabelGenerator).
 

labelbox/data/serialization/coco/image.py

@@ -1,51 +1,51 @@
-from labelbox.data.annotation_types import Label
-from pydantic import BaseModel
+from pathlib import Path
+
 from typing import Optional, Tuple
-import os
+from pydantic import BaseModel
 from PIL import Image
-import hashlib
 import imagesize
 
+from labelbox.data.annotation_types import Label
+
 
 class CocoImage(BaseModel):
     id: int
     width: int
     height: int
-    file_name: str
+    file_name: Path
     license: Optional[int] = None
     flickr_url: Optional[str] = None
     coco_url: Optional[str] = None
-    #date_captured: datetime
 
 
-def get_image_id(label: Label, idx: int):
+def get_image_id(label: Label, idx: int) -> int:
     if label.data.file_path is not None:
         file_name = label.data.file_path.replace(".jpg", "")
         if file_name.isdecimal():
-            return label.data.file_path.replace(".jpg", "")
+            return file_name
     return idx
 
 
-def get_image(label: Label, image_root, image_id):
-    path = os.path.join(image_root, f"{image_id}.jpg")
-    if not os.path.exists(path):
+def get_image(label: Label, image_root : Path, image_id: str) -> CocoImage:
+    path = Path(image_root, f"{image_id}.jpg")
+    if not path.exists():
         im = Image.fromarray(label.data.value)
         im.save(path)
         w, h = im.size
     else:
-        w, h = imagesize.get(path)
+        w, h = imagesize.get(str(path))
     return CocoImage(id=image_id,
                      width=w,
                      height=h,
-                     file_name=path.split(os.sep)[-1])
+                     file_name=Path(path.name))
 
 
 def id_to_rgb(id: int) -> Tuple[int, int, int]:
-    blue = id // (256 * 256)
-    rem = id - (blue * 256 * 256)
-    green = rem // 256
-    red = rem - (256 * green)
-    return red, green, blue
+    digits = []
+    for _ in range(3):
+        digits.append(id % 256)
+        id //= 256
+    return digits
 
 
 def rgb_to_id(red: int, green: int, blue: int) -> int:

labelbox/data/serialization/coco/instance_dataset.py

@@ -1,21 +1,22 @@
 # https://cocodataset.org/#format-data
 
-from labelbox.data.serialization.coco.categories import Categories, hash_category_name
-from labelbox.data.serialization.coco.annotation import COCOObjectAnnotation, RLE, get_annotation_lookup, rle_decoding
-from labelbox.data.serialization.coco.image import CocoImage, get_image, get_image_id
-from typing import Any, Dict, List
-from pydantic import BaseModel
-from ...annotation_types import ImageData, MaskData, Mask, ObjectAnnotation, Label, Polygon, Point, Rectangle
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from typing import Any, Dict, List, Tuple
+from pathlib import Path
+
 import numpy as np
-from PIL import Image
 from tqdm import tqdm
-import os
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from pydantic import BaseModel
+
+from ...annotation_types import ImageData, MaskData, Mask, ObjectAnnotation, Label, Polygon, Point, Rectangle
 from ...annotation_types.collection import LabelCollection
+from .categories import Categories, hash_category_name
+from .annotation import COCOObjectAnnotation, RLE, get_annotation_lookup, rle_decoding
+from .image import CocoImage, get_image, get_image_id
 
 
-def mask_to_coco_object_annotation(annotation: ObjectAnnotation, annot_idx,
-                                   image_id, category_id):
+def mask_to_coco_object_annotation(annotation: ObjectAnnotation, annot_idx : int,
+                                   image_id : int, category_id : int) -> COCOObjectAnnotation:
     # This is going to fill any holes into the multipolygon
     # If you need to support holes use the panoptic data format
     shapely = annotation.value.shapely.simplify(1).buffer(0)
@@ -37,8 +38,8 @@ def mask_to_coco_object_annotation(annotation: ObjectAnnotation, annot_idx,
         iscrowd=0)
 
 
-def vector_to_coco_object_annotation(annotation: ObjectAnnotation, annot_idx,
-                                     image_id: int, category_id):
+def vector_to_coco_object_annotation(annotation: ObjectAnnotation, annot_idx : int,
+                                     image_id: int, category_id: int) -> COCOObjectAnnotation:
     shapely = annotation.value.shapely
     xmin, ymin, xmax, ymax = shapely.bounds
     segmentation = []
@@ -61,15 +62,15 @@ def vector_to_coco_object_annotation(annotation: ObjectAnnotation, annot_idx,
                                 iscrowd=0)
 
 
-def rle_to_common(class_annotations, class_name):
+def rle_to_common(class_annotations : COCOObjectAnnotation, class_name : str) -> ObjectAnnotation:
     mask = rle_decoding(class_annotations.segmentation.counts,
                         *class_annotations.segmentation.size[::-1])
     return ObjectAnnotation(name=class_name,
                             value=Mask(mask=MaskData.from_2D_arr(mask),
                                        color=[1, 1, 1]))
 
 
-def segmentations_to_common(class_annotations, class_name):
+def segmentations_to_common(class_annotations : COCOObjectAnnotation, class_name: str) -> List[ObjectAnnotation]:
     # Technically it is polygons. But the key in coco is called segmentations..
     annotations = []
     for points in class_annotations.segmentation:
@@ -83,9 +84,9 @@ def segmentations_to_common(class_annotations, class_name):
 
 
 def process_label(label: Label,
-                  idx,
-                  image_root,
-                  max_annotations_per_image=10000):
+                  idx : int,
+                  image_root :str,
+                  max_annotations_per_image=10000) -> Tuple[np.ndarray, List[COCOObjectAnnotation], Dict[str, str]]:
     annot_idx = idx * max_annotations_per_image
     image_id = get_image_id(label, idx)
     image = get_image(label, image_root, image_id)
@@ -118,14 +119,14 @@ class CocoInstanceDataset(BaseModel):
     categories: List[Categories]
 
     @classmethod
-    def from_common(cls, labels: LabelCollection, image_root):
+    def from_common(cls, labels: LabelCollection, image_root : Path, max_workers = 8):
         all_coco_annotations = []
         categories = {}
         images = []
         futures = []
         coco_categories = {}
 
-        with ProcessPoolExecutor(max_workers=8) as exc:
+        with ProcessPoolExecutor(max_workers=max_workers) as exc:
             futures = [
                 exc.submit(process_label, label, idx, image_root)
                 for idx, label in enumerate(labels)
@@ -161,13 +162,13 @@ def to_common(self, image_root):
         annotation_lookup = get_annotation_lookup(self.annotations)
 
         for image in self.images:
-            im_path = os.path.join(image_root, image.file_name)
-            if not os.path.exists(im_path):
+            im_path = Path(image_root, image.file_name)
+            if not im_path.exists():
                 raise ValueError(
                     f"Cannot find file {im_path}. Make sure `image_root` is set properly"
                 )
 
-            data = ImageData(file_path=im_path)
+            data = ImageData(file_path=str(im_path))
             annotations = []
             for class_annotations in annotation_lookup[image.id]:
                 if isinstance(class_annotations.segmentation, RLE):