From fb0b6f847b992b036c7ecf34051bbb57c83be80b Mon Sep 17 00:00:00 2001 From: Hannah DeFazio Date: Wed, 24 Apr 2024 16:39:42 -0400 Subject: [PATCH] Bug fixes; move time_from_name to be topic specific --- angel_system/data/common/kwcoco_utils.py | 42 ++++++++++++------- angel_system/data/common/load_data.py | 35 +++------------- .../data/cooking/load_kitware_data.py | 31 +++++++++++++- angel_system/data/medical/load_bbn_data.py | 27 +++++++++++- 4 files changed, 86 insertions(+), 49 deletions(-) diff --git a/angel_system/data/common/kwcoco_utils.py b/angel_system/data/common/kwcoco_utils.py index 47fac3564..ace12e544 100644 --- a/angel_system/data/common/kwcoco_utils.py +++ b/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,12 +57,28 @@ 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"] @@ -71,18 +86,9 @@ def add_activity_gt_to_kwcoco(task, dset): 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=""): diff --git a/angel_system/data/common/load_data.py b/angel_system/data/common/load_data.py index 122d0f31a..92c2c6c3c 100644 --- a/angel_system/data/common/load_data.py +++ b/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\d+)_(?P\d+(?:_|.)\d+).(?P\w+)" -) - - -def time_from_name(fname): - """ - Extract the float timestamp from the filename. - - :param fname: Filename of an image in the format - frame___. - - :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) diff --git a/angel_system/data/cooking/load_kitware_data.py b/angel_system/data/cooking/load_kitware_data.py index 5c6d5d99b..befc42777 100644 --- a/angel_system/data/cooking/load_kitware_data.py +++ b/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\d+)_(?P\d+(?:_|.)\d+).(?P\w+)" +) + + +def time_from_name(fname): + """ + Extract the float timestamp from the filename. + + :param fname: Filename of an image in the format + frame___. + + :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"]: diff --git a/angel_system/data/medical/load_bbn_data.py b/angel_system/data/medical/load_bbn_data.py index 103d7b119..32ae5dead 100644 --- a/angel_system/data/medical/load_bbn_data.py +++ b/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\w+)-(?P\d+)_(?P\d+).(?P\w+)" +) + + +def time_from_name(fname): + """ + Extract the float timestamp from the filename. + + :param fname: Filename of an image in the format + frame___. + + :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