diff --git a/examples/EPM_to_benchmark.py b/examples/EPM_to_benchmark.py index 51a090a..941f193 100644 --- a/examples/EPM_to_benchmark.py +++ b/examples/EPM_to_benchmark.py @@ -13,6 +13,8 @@ import shutil from pathlib import Path +import sleap_io as sio + from poseinterface.io import annotations_to_coco # %% @@ -75,11 +77,8 @@ print(f"Video already exists at: {target_video_path}") # %% -# Convert DLC annotations to COCO format -# -------------------------------------- -# Here we use the :func:`annotations_to_coco` function from `poseinterface.io` -# which wraps around `sleap_io` functionality to perform the conversion. -# +# Define source annotations path +# ------------------------------ # The first attempt failed because the paths in the DLC annotations # csv file were given as # ``labeled-data,,.`` @@ -99,9 +98,43 @@ # Save COCO annotations inside the Frames directory target_annotations_path = target_frames_dir / f"{video_id}_framelabels.json" +# %% +# Build new filenames for COCO annotations +# ---------------------------------------- +# We need to determine the new filenames before calling annotations_to_coco. +# This ensures the JSON file references the renamed frame files, +# not the original source filenames. + +# Load annotations to get the order of labeled frames +labels = sio.load_file(source_annotations_path) + +# Build list of new filenames matching the order of labeled frames +coco_image_filenames = [] +for lf in labels.labeled_frames: + # Extract the original image path from the labeled frame + # For DLC, video.filename is a list of all image paths; + # use frame_idx to get the correct one for this labeled frame + filename = lf.video.filename + if isinstance(filename, list): + filename = filename[lf.frame_idx] + original_path = Path(filename) + # Extract frame number from filename, e.g. "img0042.png" -> "0042" + frame_number = original_path.stem.replace("img", "") + # Create new filename following the naming convention + new_filename = f"{video_id}_frame-{frame_number}.png" + coco_image_filenames.append(new_filename) + +# %% +# Convert DLC annotations to COCO format +# -------------------------------------- +# Here we use the :func:`annotations_to_coco` function from `poseinterface.io` +# which wraps around `sleap_io` functionality to perform the conversion. +# We pass the new filenames so the JSON references the renamed frames. + annotations_to_coco( input_path=source_annotations_path, output_json_path=target_annotations_path, + coco_image_filenames=coco_image_filenames, coco_visibility_encoding="ternary", ) print(f"Saved COCO annotations to: {target_annotations_path}")