add bop_v1, bop_v2 tools

bop_toolkit_lib/dataset/plain_bop.py → bop_toolkit_lib/dataset/bop_v1.py

@@ -1,17 +1,49 @@
-from pathlib import Path
 import re
 import numpy as np
 import json
+import pathlib
 import bop_toolkit_lib.inout as inout
 
 
+def instance_id_from_mask_filename(fn):
+    return int(re.findall("\d+", fn)[-1])
+
+
+def load_masks(
+    scene_dir,
+    image_id,
+    mask_type='mask',
+    n_instances=None,
+    instance_ids=None
+):
+
+    if n_instances is not None and instance_ids is None:
+        instance_ids = range(n_instances)
+
+    if instance_ids is not None:
+        mask_paths = (
+            scene_dir / mask_type / f'{image_id:06d}_{instance_id}.png'
+            for instance_id in instance_ids)
+    else:
+        mask_paths = (
+            scene_dir / mask_type).glob(f'{image_id:06d}_*.png')
+        mask_paths = sorted(
+            mask_paths,
+            key=lambda p: instance_id_from_mask_filename(p.name)
+        )
+    masks = np.stack([
+        inout.load_im(p) for p in mask_paths
+    ], axis=-1)
+    return masks
+
+
 def read_scene_infos(
     scene_dir,
     read_image_ids=True,
     read_n_objects=True,
 ):
     # Outputs number of scenes, image ids for each scene.
-    scene_dir = Path(scene_dir)
+    scene_dir = pathlib.Path(scene_dir)
 
     infos = dict()
     infos['has_rgb'] = (scene_dir / 'rgb').exists()
@@ -41,7 +73,6 @@ def load_scene_data(
     load_scene_camera=True,
     load_scene_gt=True,
     load_scene_gt_info=True,
-
 ):
     scene_data = dict()
     if load_scene_camera:
@@ -64,9 +95,10 @@ def load_image_data(
     load_depth=True,
     load_mask_visib=True,
     load_mask=False,
-    load_scene_gt=False,
-    load_scene_gt_info=False,
+    load_gt=False,
+    load_gt_info=False,
     rescale_depth=True,
+    instance_ids=None,
 ):
     """Loads all data for one image including images and annotations
 
@@ -83,7 +115,7 @@ def load_image_data(
 
     image_data = dict()
 
-    scene_dir = Path(scene_dir)
+    scene_dir = pathlib.Path(scene_dir)
     if isinstance(image_id, str):
         image_id = int(image_id)
 
@@ -109,40 +141,44 @@ def load_image_data(
         if rescale_depth:
             im_depth *= camera['depth_scale']
         image_data['im_depth'] = im_depth
-
-    if load_scene_gt:
-        scene_gt = inout.load_scene_gt(scene_dir / 'scene_gt.json')
-        scene_gt = scene_gt[image_id]
-        image_data['scene_gt'] = scene_gt
 
-    if load_scene_gt_info:
+    if load_gt:
+        scene_gt = inout.load_json(scene_dir / 'scene_gt.json')
+        gt = scene_gt[image_id]
+        if instance_ids is not None:
+            gt = [gt_n for n, gt_n in enumerate(gt) if n in instance_ids]
+        gt = [inout._gt_as_json(gt_n) for gt_n in gt]
+        image_data['gt'] = gt
+
+    if load_gt_info:
         scene_gt_info = inout.load_json(
             scene_dir / 'scene_gt_info.json', keys_to_int=True)
-        scene_gt_info = scene_gt_info[str(image_id)]
-        image_data['scene_gt_info'] = scene_gt_info
+        gt_info = scene_gt_info[image_id]
+        if instance_ids is not None:
+            gt_info = [
+                gt_info_n for n, gt_info_n in enumerate(gt_info)
+                if n in instance_ids]
+        gt_info = [inout._gt_as_json(gt_info) for gt_info_n in gt_info]
+        image_data['gt_info'] = inout._gt_as_json(gt_info)
 
     if load_mask_visib:
-        mask_visib_paths = (
-            scene_dir / 'mask_visib').glob(f'{image_id:06d}_*.png')
-        mask_visib_paths = sorted(
-            mask_visib_paths,
-            key=lambda p: int(re.findall("\d+", p.name)[-1])
+        mask_visib = load_masks(
+            scene_dir,
+            image_id,
+            mask_type='mask_visib',
+            n_instances=len(gt) if gt is not None else None,
+            instance_ids=instance_ids
         )
-        mask_visib = np.stack([
-            inout.load_im(p) for p in mask_visib_paths
-        ], axis=-1)
         image_data['mask_visib'] = mask_visib
 
     if load_mask:
-        mask_paths = (
-            scene_dir / 'mask').glob(f'{image_id:06d}_*.png')
-        mask_paths = sorted(
-            mask_paths,
-            key=lambda p: int(re.findall("\d+", p.name)[-1])
+        mask = load_masks(
+            scene_dir,
+            image_id,
+            mask_type='mask',
+            n_instances=len(gt) if gt is not None else None,
+            instance_ids=instance_ids
         )
-        mask = np.stack([
-            inout.load_im(p) for p in mask_paths
-        ], axis=-1)
         image_data['mask'] = mask
 
     return image_data

bop_toolkit_lib/dataset/bop_v2.py

@@ -0,0 +1,161 @@
+import pathlib
+import numpy as np
+import json
+from bop_toolkit_lib import inout
+from bop_toolkit_lib import pycoco_utils
+
+
+def _save_scene_dict(
+    scene_dict,
+    image_tpath,
+    json_converter,
+):
+    for image_id, image_dict in scene_dict.items():
+        image_dict = json_converter(image_dict)
+        path = image_tpath.format(image_id=image_id)
+        inout.save_json(path, image_dict)
+    return
+
+
+def save_scene_camera(
+    scene_camera,
+    image_camera_tpath,
+):
+    _save_scene_dict(
+        scene_camera,
+        image_camera_tpath,
+        inout._camera_as_json
+    )
+    return
+
+
+def save_scene_gt(
+    scene_gt,
+    image_gt_tpath,
+):
+    _save_scene_dict(
+        scene_gt,
+        image_gt_tpath,
+        lambda lst: [inout._gt_as_json(d) for d in lst]
+    )
+    return
+
+
+def save_masks(
+    masks,
+    masks_path,
+):
+    masks_rle = dict()
+    for instance_id, mask in enumerate(masks):
+        mask_rle = pycoco_utils.binary_mask_to_rle(mask)
+        masks_rle[instance_id] = mask_rle
+    inout.save_json(masks_path, masks_rle)
+    return
+
+
+def io_load_masks(
+    mask_file,
+    instance_ids=None
+):
+    masks_rle = json.load(mask_file)
+    if instance_ids is not None:
+        instance_ids = masks_rle.keys()
+        instance_ids = sorted(instance_ids)
+    masks = np.stack([
+        pycoco_utils.rle_to_binary_mask(mask_rle)[:, :, None]
+        for mask_rle in masks_rle], axis=-1)
+    return masks
+
+
+def io_load_gt(
+    gt_file,
+    instance_ids=None,
+):
+    gt = json.load(gt_file)
+    if instance_ids is not None:
+        gt = [gt_n for n, gt_n in enumerate(gt) if n in instance_ids]
+    gt = [inout._gt_as_json(gt_n) for gt_n in gt]
+    return gt
+
+
+def load_image_infos(
+    dataset_dir,
+    image_key,
+):
+    def _file_path(ext):
+        return dataset_dir / f'{image_key}.{ext}'
+
+    infos = dict()
+    dataset_dir = pathlib.Path(dataset_dir)
+    has_rgb = _file_path('rgb.png').exists()
+    has_rgb = has_rgb or _file_path('rgb.jpg').exists()
+    infos['has_rgb'] = has_rgb
+    infos['has_depth'] = _file_path('depth.png').exsits()
+    infos['has_gray'] = _file_path('gray.tiff').exists()
+    infos['has_mask'] = _file_path('mask.png').exists()
+    infos['has_mask_visib'] = _file_path('mask_visib.png').exists()
+    infos['has_gt'] = _file_path('gt.json').exists()
+    infos['has_gt_info'] = _file_path('gt_info.json').exists()
+    return infos
+
+
+def load_image_data(
+    dataset_dir,
+    image_key,
+    load_rgb=True,
+    load_gray=False,
+    load_depth=True,
+    load_mask_visib=True,
+    load_mask=False,
+    load_gt=False,
+    load_gt_info=False,
+    rescale_depth=True,
+    instance_ids=None,
+):
+
+    def _file_path(ext):
+        return dataset_dir / f'{image_key}.{ext}'
+
+    image_data = dict()
+
+    camera = inout.load_json(_file_path('camera.json'))
+    camera = inout._camera_as_numpy(camera)
+    image_data['camera'] = camera
+
+    if load_rgb:
+        rgb_path = _file_path('rgb.jpg')
+        if not rgb_path.exists():
+            rgb_path = _file_path('rgb.png')
+        image_data['im_rgb'] = inout.load_im(rgb_path).astype(np.uint8)
+
+    if load_gray:
+        gray_path = _file_path('gray.tiff') 
+        im_gray = inout.load_im(gray_path).astype(np.uint8)
+        image_data['im_gray'] = im_gray
+
+    if load_depth:
+        depth_path = _file_path('depth.png')
+        im_depth = inout.load_im(depth_path).astype(np.float32)
+        if rescale_depth:
+            im_depth *= camera['depth_scale']
+        image_data['im_depth'] = im_depth
+
+    if load_gt:
+        with open(_file_path('gt.json'), 'r') as f:
+            image_data['gt'] = io_load_gt(f, instance_ids=instance_ids)
+
+    if load_gt_info:
+        with open(_file_path('gt_info.json'), 'r') as f:
+            image_data['gt_info'] = io_load_gt(f, instance_ids=instance_ids)
+
+    if load_mask_visib:
+        with open(_file_path('mask_visib.json'), 'r') as f:
+            image_data['mask_visib'] = io_load_masks(
+                f, instance_ids=instance_ids)
+
+    if load_mask:
+        with open(_file_path('mask.json'), 'r') as f:
+            image_data['mask'] = io_load_masks(
+                f, instance_ids=instance_ids)
+
+    return image_data

bop_toolkit_lib/inout.py

@@ -137,6 +137,29 @@ def load_cam_params(path):
   return cam
 
 
+def _camera_as_numpy(camera):
+  if 'cam_K' in camera.keys():
+    camera['cam_K'] = \
+      np.array(camera['cam_K'], np.float64).reshape((3, 3))
+  if 'cam_R_w2c' in camera.keys():
+    camera['cam_R_w2c'] = \
+      np.array(camera['cam_R_w2c'], np.float64).reshape((3, 3))
+  if 'cam_t_w2c' in camera.keys():
+    camera['cam_t_w2c'] = \
+      np.array(camera['cam_t_w2c'], np.float64).reshape((3, 1))
+  return camera
+
+
+def _camera_as_json(camera):
+  if 'cam_K' in camera.keys():
+    camera['cam_K'] = camera['cam_K'].flatten().tolist()
+  if 'cam_R_w2c' in camera.keys():
+    camera['cam_R_w2c'] = camera['cam_R_w2c'].flatten().tolist()
+  if 'cam_t_w2c' in camera.keys():
+    camera['cam_t_w2c'] = camera['cam_t_w2c'].flatten().tolist()
+  return camera
+
+
 def load_scene_camera(path):
   """Loads content of a JSON file with information about the scene camera.
 
@@ -148,15 +171,7 @@ def load_scene_camera(path):
   scene_camera = load_json(path, keys_to_int=True)
 
   for im_id in scene_camera.keys():
-    if 'cam_K' in scene_camera[im_id].keys():
-      scene_camera[im_id]['cam_K'] = \
-        np.array(scene_camera[im_id]['cam_K'], np.float64).reshape((3, 3))
-    if 'cam_R_w2c' in scene_camera[im_id].keys():
-      scene_camera[im_id]['cam_R_w2c'] = \
-        np.array(scene_camera[im_id]['cam_R_w2c'], np.float64).reshape((3, 3))
-    if 'cam_t_w2c' in scene_camera[im_id].keys():
-      scene_camera[im_id]['cam_t_w2c'] = \
-        np.array(scene_camera[im_id]['cam_t_w2c'], np.float64).reshape((3, 1))
+    scene_camera[im_id] = _camera_as_numpy(scene_camera[im_id])
   return scene_camera
 
 
@@ -169,16 +184,29 @@ def save_scene_camera(path, scene_camera):
   :param scene_camera: Dictionary to save to the JSON file.
   """
   for im_id in sorted(scene_camera.keys()):
-    im_camera = scene_camera[im_id]
-    if 'cam_K' in im_camera.keys():
-      im_camera['cam_K'] = im_camera['cam_K'].flatten().tolist()
-    if 'cam_R_w2c' in im_camera.keys():
-      im_camera['cam_R_w2c'] = im_camera['cam_R_w2c'].flatten().tolist()
-    if 'cam_t_w2c' in im_camera.keys():
-      im_camera['cam_t_w2c'] = im_camera['cam_t_w2c'].flatten().tolist()
+    scene_camera[im_id] = _camera_as_json(scene_camera[im_id])
   save_json(path, scene_camera)
 
 
+def _gt_as_numpy(gt):
+  if 'cam_R_m2c' in gt.keys():
+    gt['cam_R_m2c'] = \
+      np.array(gt['cam_R_m2c'], np.float64).reshape((3, 3))
+  if 'cam_t_m2c' in gt.keys():
+    gt['cam_t_m2c'] = \
+      np.array(gt['cam_t_m2c'], np.float64).reshape((3, 1))
+  return gt
+
+def _gt_as_json(gt):
+  if 'cam_R_m2c' in gt.keys():
+    gt['cam_R_m2c'] = gt['cam_R_m2c'].flatten().tolist()
+  if 'cam_t_m2c' in gt.keys():
+    gt['cam_t_m2c'] = gt['cam_t_m2c'].flatten().tolist()
+  if 'obj_bb' in gt.keys():
+    gt['obj_bb'] = [int(x) for x in gt['obj_bb']]
+  return gt
+
+
 def load_scene_gt(path):
   """Loads content of a JSON file with ground-truth annotations.