Read videos from HIM project files

data/box.py

@@ -0,0 +1,93 @@
+import json
+import os.path
+import re
+import zipfile
+from typing import NamedTuple, Optional, Iterable, List
+
+from numpy import ndarray
+
+from arkite_booster.cv2util import box_mask
+from arkite_booster.util import FloatPoint
+from arkite_booster.util import bounding_box_centre
+
+ArkiteBox = NamedTuple(typename='ArkiteBox',
+                       fields=[('name', str),
+                               ('points', List[FloatPoint]),
+                               ('depth', int),
+                               ('distance', int),
+                               ('centre', FloatPoint),
+                               ('mask', ndarray)])
+
+HEF_BOX_LIST_FILE = 'data/Model.Entities.Box.json'
+
+CALIBRATION_BOX_REGEX = re.compile('(Depth|IR) Cal ([0-9]+|- Center).*')
+GRAB_BOX_PREFIX = 'grab_'
+
+
+def load_boxes(model_entities_box_path: str,
+               project_ids: Optional[Iterable[int]] = None,
+               box_ids: Optional[Iterable[int]] = None,
+               keep_grab_boxes: bool = True,
+               keep_container_boxes: bool = True,
+               skip_calibration_boxes: bool = True) -> Iterable[ArkiteBox]:
+    model_entities_box_path = os.path.abspath(model_entities_box_path)
+
+    if model_entities_box_path.endswith('.json'):
+        with open(model_entities_box_path, 'r') as mebf:
+            all_boxes = json.load(mebf)
+    elif model_entities_box_path.endswith('.hef'):
+        with zipfile.ZipFile(model_entities_box_path, mode='r') as hef:
+            with hef.open(HEF_BOX_LIST_FILE) as mebf:
+                all_boxes = json.loads(mebf.read().decode('utf-8'))
+    else:
+        raise ValueError('Unknown file format: ' + model_entities_box_path)
+
+    boxes = []
+    for box_json in all_boxes:
+        if project_ids is not None and box_json['ProjectId'] not in project_ids:
+            continue
+
+        if box_ids is not None and box_json['Id'] not in box_ids:
+            continue
+
+        box_name = box_json['Name']
+        if skip_calibration_boxes and CALIBRATION_BOX_REGEX.match(box_name) is not None:
+            continue
+
+        if not keep_grab_boxes and box_name.startswith(GRAB_BOX_PREFIX):
+            continue
+
+        if not keep_container_boxes and not box_name.startswith(GRAB_BOX_PREFIX):
+            continue
+
+        points = [(coordinates["X"], coordinates["Y"]) for coordinates in box_json["DataPoints"]]
+        mask = box_mask(points)
+        centre = bounding_box_centre(points)
+
+        box = ArkiteBox(name=box_name,
+                        depth=box_json['Depth'],
+                        distance=box_json['Distance'],
+                        points=points,
+                        mask=mask,
+                        centre=centre)
+        boxes.append(box)
+
+    return boxes
+
+
+def boxes_for_use_case(box_metadata_root_dir: str,
+                       use_case: str,
+                       keep_grab_boxes: bool = True,
+                       keep_container_boxes: bool = True,
+                       skip_calibration_boxes: bool = True) -> Iterable[ArkiteBox]:
+    """Ad-hoc method to match the detection libraty with additional `.hef` files provided by Stijn."""
+    box_metadata_root_dir = os.path.abspath(box_metadata_root_dir)
+
+    use_case_name = use_case.split('/')[-1]
+    use_case_hef = os.path.join(box_metadata_root_dir, use_case_name + '.hef')
+    use_case_boxes = load_boxes(use_case_hef,
+                                keep_grab_boxes=keep_grab_boxes,
+                                keep_container_boxes=keep_container_boxes,
+                                skip_calibration_boxes=skip_calibration_boxes)
+
+    return use_case_boxes