Merge pull request #380 from hdefazio/dev/bbn_revival

Reviving BBN code

angel_system/data/common/cli/add_activity_gt_to_kwcoco.py

@@ -11,10 +11,16 @@ def main():
         default="obj_annotations.mscoco.json",
         help="kwcoco filename",
     )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="medical",
+        help="Task name",
+    )
 
     args = parser.parse_args()
 
-    add_activity_gt_to_kwcoco(args.dset)
+    add_activity_gt_to_kwcoco(args.task, args.dset)
 
 
 if __name__ == "__main__":

angel_system/data/common/hl_hands_utils.py

@@ -5,6 +5,8 @@
 
 import ubelt as ub
 
+from angel_system.data.common.kwcoco_utils import load_kwcoco
+
 
 def load_hl_hand_bboxes(extracted_dir):
     fn = extracted_dir + "/_hand_pose_2d_data.json"
@@ -127,3 +129,108 @@ def find_closest_hands(object_pair, detected_classes, preds):
 
     hand_label = close_hands if len(close_hands) > 0 else None
     return hand_label
+
+
+def add_hl_hands_to_kwcoco(dset, remove_existing=True, using_contact=True):
+    """Add bounding boxes for the hands based on the Hololen's joint positions
+
+    This requires the projected hololens joint information generated by
+    running ``scripts/hand_pose_converter.py``
+
+    :param dset: kwcoco object or a string pointing to a kwcoco file
+    :param using_contact: If True, adds contact state and contact conf
+        for obj-obj and obj-hand states for each detection
+    """
+    # Load kwcoco file
+    dset = load_kwcoco(dset)
+
+    gid_to_aids = dset.index.gid_to_aids
+
+    for video_id in ub.ProgIter(
+        dset.index.videos.keys(),
+        desc=f"Adding hololens hands for videos in {dset.fpath}",
+    ):
+        image_ids = dset.index.vidid_to_gids[video_id]
+
+        all_hand_pose_2d_image_space = None
+        remove_aids = []
+
+        for gid in sorted(image_ids):
+            im = dset.imgs[gid]
+            frame_idx, time = time_from_name(im["file_name"])
+
+            aids = gid_to_aids[gid]
+            anns = ub.dict_subset(dset.anns, aids)
+
+            # Mark hand detections to be removed
+            if remove_existing:
+                for aid, ann in anns.items():
+                    cat = dset.cats[ann["category_id"]]["name"]
+                    if "hand" in cat:
+                        remove_aids.append(aid)
+
+            # Find HL hands, should only run once per video
+            if not all_hand_pose_2d_image_space:
+                # <video_folder>/_extracted/images/<file_name>
+                extr_video_folder = im["file_name"].split("/")[:-2]
+                extr_video_folder = ("/").join(extr_video_folder)
+                all_hand_pose_2d_image_space = load_hl_hand_bboxes(extr_video_folder)
+
+            all_hands = (
+                all_hand_pose_2d_image_space[time]
+                if time in all_hand_pose_2d_image_space.keys()
+                else []
+            )
+
+            # Add HL hand bounding boxes if we have them
+            if all_hands != []:
+                # print("Adding hand bboxes from the hololens joints")
+                for joints in all_hands:
+                    keys = list(joints["joints"].keys())
+                    hand_label = joints["hand"]
+
+                    all_x_values = [joints["joints"][k]["projected"][0] for k in keys]
+                    all_y_values = [joints["joints"][k]["projected"][1] for k in keys]
+
+                    hand_bbox = [
+                        min(all_x_values),
+                        min(all_y_values),
+                        max(all_x_values),
+                        max(all_y_values),
+                    ]  # tlbr
+                    xywh = (
+                        kwimage.Boxes([hand_bbox], "tlbr")
+                        .toformat("xywh")
+                        .data[0]
+                        .tolist()
+                    )
+
+                    cat = dset.index.name_to_cat[hand_label]
+                    ann = {
+                        "area": xywh[2] * xywh[3],
+                        "image_id": im["id"],
+                        "category_id": cat["id"],
+                        "segmentation": [],
+                        "bbox": xywh,
+                        "confidence": 1,
+                    }
+
+                    if using_contact:
+                        ann["obj-obj_contact_state"] = False
+                        ann["obj-obj_contact_conf"] = 0
+
+                        ann["obj-hand_contact_state"] = False
+                        ann["obj-hand_contact_conf"] = 0
+
+                    dset.add_annotation(**ann)
+
+        # print(f"Removing annotations {remove_aids} in video {video_id}")
+        dset.remove_annotations(remove_aids)
+
+    fpath = dset.fpath.split(".mscoco")[0]
+    if remove_existing:
+        dset.fpath = f"{fpath}_hl_hands_only.mscoco.json"
+    else:
+        dset.fpath = f"{fpath}_plus_hl_hands.mscoco.json"
+    dset.dump(dset.fpath, newlines=True)
+    print(f"Saved predictions to {dset.fpath}")