Merge pull request #386 from hdefazio/dev/m2_dataloader

Add M2 lab data to dataloader
PTG-Kitware · Feb 15, 2024 · 16d6baa · 16d6baa
2 parents 2944da9 + fc8d40f
commit 16d6baa
Show file tree

Hide file tree

Showing 2 changed files with 136 additions and 13 deletions.
diff --git a/angel_system/data/medical/data_paths.py b/angel_system/data/medical/data_paths.py
@@ -18,9 +18,7 @@
 m2_activity_gt_dir = f"{activity_gt_dir}/m2_labels/"
 m2_activity_config_fn = f"{activity_config_path}/recipe_m2.yaml"
 
-m2_ros_bags_dir = f"{ros_bags_dir}/m2/m2_extracted/"
 m2_bbn_data_dir = f"{bbn_data_dir}/Release_v0.5/v0.52/M2_Tourniquet/Data"
-
 m2_training_split = {
     "train_activity": [
         f"{m2_bbn_data_dir}/M2-{x}"
@@ -91,20 +89,131 @@
 m2_obj_dets_dir = f"{objects_dir}/m2"
 m2_obj_config = f"{object_config_path}/task_m2.yaml"
 
+# M2 Lab
+# ------
+m2_lab_bbn_data_dir = f"{bbn_data_dir}/M2_Lab_data/skills_by_frame"
+m2_lab_training_split = {
+    "train_activity": [],
+    "val": [],
+    "test": [
+        f"{m2_lab_bbn_data_dir}/tq_{x}"
+        for x in [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+        ]
+    ],
+}
+
 
-def grab_data(recipe, machine):
+# M2 Kitware
+# ----------
+m2_kitware_data_dir = f"{ros_bags_dir}/M2/M2_extracted/"
+m2_kitware_training_split = {
+    "train_activity": [],
+    "val": [],
+    "test": [
+        f"{m2_kitware_data_dir}/kitware_m2_video_{x}_extracted"
+        for x in [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31,
+            32,
+        ]
+    ],
+}
+
+
+def grab_data(skill, machine="gyges"):
     if machine == "gyges":
-        if recipe == "m2":
+        if skill == "m2":
             return (
                 ptg_root,
                 data_dir,
                 m2_activity_config_fn,
                 m2_activity_gt_dir,
-                m2_ros_bags_dir,
+                None,
                 m2_training_split,
                 m2_obj_dets_dir,
                 m2_obj_config,
             )
+        elif skill == "m2_lab":
+            return (
+                ptg_root,
+                data_dir,
+                None,
+                None,
+                m2_kitware_data_dir,
+                m2_lab_training_split,
+                None,
+                None,
+            )
+        elif skill == "m2_kitware":
+            return (
+                ptg_root,
+                data_dir,
+                None,
+                None,
+                None,
+                m2_kitware_training_split,
+                None,
+                None,
+            )
 
         else:
             raise NotImplementedError

diff --git a/angel_system/data/medical/load_bbn_data.py b/angel_system/data/medical/load_bbn_data.py
@@ -242,16 +242,31 @@ def activity_label_fixes(activity_label, target):
     return label, label_id
 
 
-def bbn_activity_txt_to_csv(root_dir):
+def bbn_activity_txt_to_csv(root_dir, output_dir):
     """
     Generate DIVE csv format activity annotations from BBN's text annotations
+
+    :param root_dir: Path to a folder containing video folders
+        Expected setup:
+            root_dir/
+                {VIDEO_NAME}/
+                    {action labels by frame}.txt
+                    images/
+                        {filename}.png
+                        ...
+
     """
-    task = "M2_Tourniquet"
-    print(f"{root_dir}/{task}/Data/*/*_action_labels_by_frame.txt")
+    print(f"{root_dir}/*/*_action_labels_by_frame.txt")
 
-    for action_txt_fn in glob.glob(
-        f"{root_dir}/{task}/Data/*/*.action_labels_by_frame.txt"
-    ):
+    action_fns = glob.glob(f"{root_dir}/*/*.action_labels_by_frame.txt")
+    if not action_fns:
+        # Lab videos
+        action_fns = glob.glob(f"{root_dir}/*/*_skills_frame.txt")
+    if not action_fns:
+        warnings.warn(f"No text annotations found in {root_dir}")
+        return
+
+    for action_txt_fn in action_fns:
         track_id = 0
         video_dir = os.path.dirname(action_txt_fn)
         video_name = os.path.basename(video_dir)
@@ -262,8 +277,7 @@ def bbn_activity_txt_to_csv(root_dir):
         lines = action_f.readlines()
 
         # Create output csv
-        task_dir = "m2_labels"
-        csv_fn = f"{activity_dir}/{task_dir}/{video_name}_activity_labels_v2.csv"
+        csv_fn = f"{output_dir}/{video_name}_activity_labels_v2.csv"
         csv_f = open(csv_fn, "w")
         csv_f.write(
             "# 1: Detection or Track-id,2: Video or Image Identifier,3: Unique Frame Identifier,4-7: Img-bbox(TL_x,TL_y,BR_x,BR_y),8: Detection or Length Confidence,9: Target Length (0 or -1 if invalid),10-11+: Repeated Species,Confidence Pairs or Attributes\n"