DeepLabCut / DLC2NWB

Utilities to convert DeepLabCut (DLC), output to/from Neurodata Without Borders (NWB) format.
MIT License
14 stars 7 forks source link

Add optional parameters to `PoseEstimation` #24

Open vigji opened 3 months ago

vigji commented 3 months ago

As described in https://github.com/catalystneuro/neuroconv/issues/915, I would need to pass over a name value to PoseEstimation to support multiple DLC pose estimations in the same nwb file.

I am currently fixing it by piping an optional dictionary of additional kwargs from the args in this way:

 def _write_pes_to_nwbfile(
    nwbfile,
    animal,
    df_animal,
    scorer,
    video,  # Expects this to be a tuple; first index is string path, second is the image shape as "0, width, 0, height"
    paf_graph,
    timestamps,
    exclude_nans,
    **optional_kwargs,
):  
    ...

    pe = PoseEstimation(
        pose_estimation_series=pose_estimation_series,
        description="2D keypoint coordinates estimated using DeepLabCut.",
        original_videos=[video[0]],
        # TODO check if this is a mandatory arg in ndx-pose (can skip if video is not found_
        dimensions=[list(map(int, video[1].split(",")))[1::2]],
        scorer=scorer,
        source_software="DeepLabCut",
        source_software_version=deeplabcut_version,
        nodes=[pes.name for pes in pose_estimation_series],
        edges=paf_graph if paf_graph else None,
        **optional_kwargs,
    )

If you would be open to support this fix, I can submit a new PR or add this solutions as well to #23, as you prefer!