Merge pull request #46 from CBroz1/dev

Add NWB export function

Author: tdincer

element_deeplabcut/export/README.md

@@ -0,0 +1,38 @@
+# Exporting data to NWB
+
+## Description
+
+The `export/nwb.py` module calls [DLC2NWB](https://github.com/DeepLabCut/DLC2NWB/) to
+save output generated by Element DeepLabCut as NWB files. 
+The main function, `dlc_session_to_nwb`, contains a flag to control calling a parallel 
+function in [Element Session](https://github.com/datajoint/element-session/blob/main/element_session/export/nwb.py).
+
+As DLC2NWB does not currently offer a separate function for generating `PoseEstimation`
+objects (see [ndx-pose](https://github.com/rly/ndx-pose)), the current solution is to
+allow DLC2NWB to write to disk, and optionally rewrite this file using metadata provided
+by the export function in Element Session.
+
+## Usage
+
+Before using, please install [DLC2NWB](https://github.com/DeepLabCut/DLC2NWB/)
+
+```bash
+pip install dlc2nwb
+```
+
+Then, call the export function using keys from the `PoseEstimation` table.
+
+```python
+from element_deeplabcut import model
+from element_session import session
+from element_deeplabcut.export import dlc_session_to_nwb
+
+session_key = (session.Session & CONDITION)
+pose_key = (model.PoseEstimation & session_key).fetch1('KEY')
+dlc_session_to_nwb(pose_key, use_element_session=True, session_kwargs=SESSION_KWARGS)
+```
+
+
+Here, `CONDITION` should uniquely identify a session and `SESSION_KWARGS` can be any of
+the items described in the docstring of `element_session.export.nwb.session_to_nwb`
+as a dictionary.

element_deeplabcut/export/__init__.py

@@ -0,0 +1 @@
+from .nwb import dlc_session_to_nwb

element_deeplabcut/export/nwb.py

@@ -0,0 +1,72 @@
+"""
+Portions of code adapted from DeepLabCut/DLC2NWB
+MIT License Copyright (c) 2022 Alexander Mathis
+DataJoint export methods for DeepLabCut 2.x
+"""
+import logging
+import warnings
+from pathlib import Path
+from collections import abc
+from pynwb import NWBHDF5IO
+from hdmf.build.warnings import DtypeConversionWarning
+from .. import model
+
+try:  # Not all users will want NWB export, so dependency not in requirements.
+    from dlc2nwb.utils import convert_h5_to_nwb, write_subject_to_nwb
+except ImportError:
+    raise ImportError(
+        "The package `dlc2nwb` is missing. Please run `pip install dlc2nwb`."
+    )
+
+logger = logging.getLogger("datajoint")
+
+
+def dlc_session_to_nwb(keys, use_element_session=True, session_kwargs=None):
+    """Using keys from PoseEstimation table, save DLC's h5 output to NWB.
+
+    Calls DLC2NWB to export NWB file using current h5 on disk. If use_element_session,
+    calls NWB export function from Elements for lab, animal and session, passing
+    session_kwargs. Saves output based on naming convention in DLC2NWB. If output path
+    already exists, returns output path without making changes to the file.
+    NOTE: does not support multianimal exports
+
+    Parameters
+    ----------
+    keys: One or more keys from model.PoseEstimation
+    use_element_session: Optional. If True, call NWB export from Element Session
+    session_kwargs: Optional. Additional keyword arguments for Element Session export
+
+    Returns output path of saved file
+    """
+    if not isinstance(keys, abc.Sequence):  # Ensure list for following loop
+        keys = [keys]
+
+    for key in keys:
+        write_file = True
+        subject_id = key["subject"]
+        output_dir = model.PoseEstimationTask.infer_output_dir(key)
+        config_file = str(output_dir / "dj_dlc_config.yaml")
+        video_name = Path((model.VideoRecording.File & key).fetch1("file_path")).stem
+        h5file = next(output_dir.glob(f"{video_name}*h5"))
+        output_path = h5file.replace(".h5", f"_{subject_id}.nwb")  # DLC2NWB convention
+
+        if Path(output_path).exists():
+            logger.warning(f"Skipping {subject_id}. NWB already exists.")
+            write_file = False
+
+        # Use standard DLC2NWB export
+        if write_file and not use_element_session:
+            output_path = convert_h5_to_nwb(config_file, h5file, subject_id)
+
+        # Pass Element Session export items in export
+        if write_file and use_element_session:
+            from element_session.export.nwb import session_to_nwb
+
+            session_nwb = session_to_nwb(key, **session_kwargs)  # call session export
+            dlc_nwb = write_subject_to_nwb(session_nwb, h5file, subject_id, config_file)
+            # warnings filter from DLC2NWB
+            with warnings.catch_warnings(), NWBHDF5IO(output_path, mode="w") as io:
+                warnings.filterwarnings("ignore", category=DtypeConversionWarning)
+                io.write(dlc_nwb)
+
+    return output_path

element_deeplabcut/readers/dlc_reader.py

@@ -1,9 +1,16 @@
 import re
+import logging
 import numpy as np
 import pandas as pd
 from pathlib import Path
 import pickle
 import ruamel.yaml as yaml
+from element_interface.utils import find_root_directory
+from .. import model
+from ..model import get_dlc_root_data_dir
+from datajoint.errors import DataJointError
+
+logger = logging.getLogger("datajoint")
 
 
 class PoseEstimation:
@@ -215,7 +222,11 @@ def do_pose_estimation(
     use_shelve=False,
     modelprefix="",  # need from paramset
 ):
-    """Launch DLC's analyze_videos within element-deeplabcut
+    """Launch DLC's analyze_videos within element-deeplabcut.
+
+    Also saves a copy of the current config in the output dir, with ensuring analyzed
+    videos in the video_set. NOTE: Config-specificed cropping not supported when adding
+    to config in this manner.
 
     Parameters
     ----------
@@ -232,6 +243,28 @@ def do_pose_estimation(
     dlc_project_path = Path(project_path)
     dlc_config["project_path"] = dlc_project_path.as_posix()
 
+    # ---- Add current video to config ---
+    for video_filepath in video_filepaths:
+        if video_filepath not in dlc_config["video_sets"]:
+            root_dir = find_root_directory(get_dlc_root_data_dir(), video_filepath)
+            relative_path = Path(video_filepath).relative_to(root_dir)
+            recording_id = (
+                model.VideoRecording.File & f'file_path="{relative_path}"'
+            ).fetch1("recording_id")
+            try:
+                px_width, px_height = (
+                    model.RecordingInfo & f'recording_id="{recording_id}"'
+                ).fetch1("px_width", "px_height")
+            except DataJointError:
+                logger.warn(
+                    f"Could not find RecordingInfo for {video_filepath.stem}"
+                    + "\n\tUsing zeros for crop value in config."
+                )
+                px_height, px_width = 0, 0
+            dlc_config["video_sets"].update(
+                {str(video_filepath): {"crop": f"0, {px_width}, 0, {px_height}"}}
+            )
+
     # ---- Write config files ----
     # To output dir: Important for loading/parsing output in datajoint
     _ = save_yaml(output_dir, dlc_config)