Bug fixes; move time_from_name to be topic specific

angel_system/data/common/kwcoco_utils.py

@@ -26,7 +26,6 @@
 from angel_system.data.common.load_data import (
     activities_from_dive_csv,
     objs_as_dataframe,
-    time_from_name,
     sanitize_str,
 )
 from angel_system.data.common.load_data import Re_order
@@ -49,7 +48,7 @@ def load_kwcoco(dset):
     return dset
 
 
-def add_activity_gt_to_kwcoco(task, dset):
+def add_activity_gt_to_kwcoco(topic, task, dset):
     """Takes an existing kwcoco file and fills in the "activity_gt"
     field on each image based on the activity annotations.
 
@@ -58,31 +57,38 @@ def add_activity_gt_to_kwcoco(task, dset):
 
     :param dset: kwcoco object or a string pointing to a kwcoco file
     """
+    if topic == "medical":
+        from angel_system.data.medical.load_bbn_data import time_from_name
+    elif topic == "cooking": 
+        from angel_system.data.cooking.load_kitware_data import time_from_name
+
     # Load kwcoco file
     dset = load_kwcoco(dset)
 
-    data_dir = f"/data/PTG/{task}/"
+    data_dir = f"/data/PTG/{topic}/"
     activity_gt_dir = f"{data_dir}/activity_anns"
 
+    # Load activity config
+    with open(
+        f"config/activity_labels/{topic}/{task}.yaml", "r"
+    ) as stream:
+        activity_config = yaml.safe_load(stream)
+    activity_labels = activity_config["labels"]
+    label_version = activity_config["version"]
+
+    activity_gt_dir = f"{activity_gt_dir}/{task}_labels/"
+
+    # Add ground truth to kwcoco
     for video_id in dset.index.videos.keys():
         video = dset.index.videos[video_id]
         video_name = video["name"]
         print(video_name)
 
         if "_extracted" in video_name:
             video_name = video_name.split("_extracted")[0]
-        video_skill = "m2"  # video["recipe"]
-
-        with open(
-            f"../config/activity_labels/{task}/task_{video_skill}.yaml", "r"
-        ) as stream:
-            recipe_activity_config = yaml.safe_load(stream)
-        recipe_activity_labels = recipe_activity_config["labels"]
 
-        recipe_activity_gt_dir = f"{activity_gt_dir}/{video_skill}_labels/"
-
-        activity_gt_fn = f"{recipe_activity_gt_dir}/{video_name}_activity_labels_v2.csv"
-        gt = activities_from_dive_csv(activity_gt_fn)
+        activity_gt_fn = f"{activity_gt_dir}/{video_name}_activity_labels_v{label_version}.csv"
+        gt = activities_from_dive_csv(topic, activity_gt_fn)
         gt = objs_as_dataframe(gt)
 
         image_ids = dset.index.vidid_to_gids[video_id]
@@ -92,7 +98,10 @@ def add_activity_gt_to_kwcoco(task, dset):
             im = dset.imgs[gid]
             frame_idx, time = time_from_name(im["file_name"])
 
-            matching_gt = gt.loc[(gt["start"] <= time) & (gt["end"] >= time)]
+            if time:
+                matching_gt = gt.loc[(gt["start"] <= time) & (gt["end"] >= time)]
+            else:
+                matching_gt = gt.loc[(gt["start_frame"] <= frame_idx) & (gt["end_frame"] >= frame_idx)]
 
             if matching_gt.empty:
                 label = "background"
@@ -106,7 +115,7 @@ def add_activity_gt_to_kwcoco(task, dset):
                 try:
                     activity = [
                         x
-                        for x in recipe_activity_labels
+                        for x in activity_labels
                         if int(x["id"]) == int(float(label))
                     ]
                 except:
@@ -131,6 +140,7 @@ def add_activity_gt_to_kwcoco(task, dset):
 
     # dset.fpath = dset.fpath.split(".")[0] + "_fixed.mscoco.json"
     dset.dump(dset.fpath, newlines=True)
+    return dset
 
 
 def visualize_kwcoco_by_label(dset=None, save_dir=""):

angel_system/data/common/load_data.py

@@ -43,34 +43,6 @@ def Re_order(image_list, image_number):
         new_list.append(image_list[idx])
     return new_list
 
-
-RE_FILENAME_TIME = re.compile(
-    r"frame_(?P<frame>\d+)_(?P<ts>\d+(?:_|.)\d+).(?P<ext>\w+)"
-)
-
-
-def time_from_name(fname):
-    """
-    Extract the float timestamp from the filename.
-
-    :param fname: Filename of an image in the format
-        frame_<frame number>_<seconds>_<nanoseconds>.<extension>
-
-    :return: timestamp (float) in seconds
-    """
-    fname = os.path.basename(fname)
-    match = RE_FILENAME_TIME.match(fname)
-    time = match.group("ts")
-    if "_" in time:
-        time = time.split("_")
-        time = float(time[0]) + (float(time[1]) * 1e-9)
-    elif "." in time:
-        time = float(time)
-
-    frame = match.group("frame")
-    return int(frame), time
-
-
 def load_from_file(
     gt_fn, detections_fn
 ) -> Tuple[List[str], pd.DataFrame, pd.DataFrame]:
@@ -138,7 +110,7 @@ def load_from_file(
     return labels, gt, detections
 
 
-def activities_from_dive_csv(filepath: str) -> List[Activity]:
+def activities_from_dive_csv(topic, filepath: str) -> List[Activity]:
     """
     Load from a DIVE output CSV file a sequence of ground truth activity
     annotations.
@@ -149,6 +121,11 @@ def activities_from_dive_csv(filepath: str) -> List[Activity]:
     :param filepath: Filesystem path to the CSV file.
     :return: List of loaded activity annotations.
     """
+    if topic == "medical":
+        from angel_system.data.medical.load_bbn_data import time_from_name
+    elif topic == "cooking": 
+        from angel_system.data.cooking.load_kitware_data import time_from_name
+
     print(f"Loading ground truth activities from: {filepath}")
     df = pd.read_csv(filepath)
 

angel_system/data/cooking/load_kitware_data.py

@@ -1,13 +1,41 @@
+import os
 import kwcoco
 import glob
 import warnings
-
+import re
 import ubelt as ub
 
 from angel_system.data.common.load_data import Re_order
 from angel_system.data.common.kwcoco_utils import load_kwcoco
 
 
+
+RE_FILENAME_TIME = re.compile(
+    r"frame_(?P<frame>\d+)_(?P<ts>\d+(?:_|.)\d+).(?P<ext>\w+)"
+)
+
+
+def time_from_name(fname):
+    """
+    Extract the float timestamp from the filename.
+
+    :param fname: Filename of an image in the format
+        frame_<frame number>_<seconds>_<nanoseconds>.<extension>
+
+    :return: timestamp (float) in seconds
+    """
+    fname = os.path.basename(fname)
+    match = RE_FILENAME_TIME.match(fname)
+    time = match.group("ts")
+    if "_" in time:
+        time = time.split("_")
+        time = float(time[0]) + (float(time[1]) * 1e-9)
+    elif "." in time:
+        time = float(time)
+
+    frame = match.group("frame")
+    return int(frame), time
+
 def object_label_fixes(obj_cat):
     # Fix some deprecated labels
     if obj_cat in ["timer", "timer (20)", "timer (30)", "timer (else)"]:
@@ -47,7 +75,6 @@ def object_label_fixes(obj_cat):
 
     return obj_cat
 
-
 def activity_label_fixes(activity_label):
     # Temp fix until we can update the groundtruth labels
     if activity_label in ["microwave-30-sec", "microwave-60-sec"]:

angel_system/data/medical/load_bbn_data.py

@@ -22,6 +22,29 @@
 os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3"
 
 
+
+RE_FILENAME_TIME = re.compile(
+    r"(?P<task>\w+)-(?P<vid>\d+)_(?P<frame>\d+).(?P<ext>\w+)"
+)
+
+
+def time_from_name(fname):
+    """
+    Extract the float timestamp from the filename.
+
+    :param fname: Filename of an image in the format
+        frame_<frame number>_<seconds>_<nanoseconds>.<extension>
+
+    :return: timestamp (float) in seconds
+    """
+    fname = os.path.basename(fname)
+    match = RE_FILENAME_TIME.match(fname)
+
+    frame = match.group("frame")
+    time = None
+    return int(frame), time
+
+
 def dive_to_activity_file(videos_dir):
     """DIVE CSV to BBN TXT frame-level annotation file format"""
     for dive_csv in glob.glob(f"{videos_dir}/*/*.csv"):
@@ -313,10 +336,10 @@ def bbn_activity_txt_to_csv(task, root_dir, output_dir, label_version):
             end_frame = int(data[1])
 
             start_frame_fn = os.path.basename(
-                glob.glob(f"{video_dir}/images/*_{start_frame}*.png")[0]
+                glob.glob(f"{video_dir}/images/*_{start_frame}.png")[0]
             )
             end_frame_fn = os.path.basename(
-                glob.glob(f"{video_dir}/images/*_{end_frame}*.png")[0]
+                glob.glob(f"{video_dir}/images/*_{end_frame}.png")[0]
             )
 
             # Determine activity