Skip to content

Commit 077e42d

Browse files
Saksham20h-mayorquinjeylau
authored
create a separate method to write to existing nwbfile (#10)
* separate out methods to write to nwb per subject * add the option to provide a config_dict instead of .yml file path * scrapt using config dict * Update dlc2nwb/utils.py Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com> * Update dlc2nwb/utils.py Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Co-authored-by: Heberto Mayorquin <h.mayorquin@gmail.com> Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com>
1 parent 5c70444 commit 077e42d

1 file changed

Lines changed: 99 additions & 59 deletions

File tree

‎dlc2nwb/utils.py‎

Lines changed: 99 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -51,36 +51,11 @@ def _ensure_individuals_in_header(df, dummy_name):
5151
return df
5252

5353

54-
def convert_h5_to_nwb(config, h5file, individual_name="ind1"):
55-
"""
56-
Convert a DeepLabCut (DLC) video prediction, h5 data file to Neurodata Without Borders (NWB). Also
57-
takes project config, to store relevant metadata.
58-
59-
Parameters
60-
----------
61-
config : str
62-
Path to a project config.yaml file
63-
64-
h5file : str
65-
Path to a h5 data file
66-
67-
individual_name : str
68-
Name of the subject (whose pose is predicted) for single-animal DLC project.
69-
For multi-animal projects, the names from the DLC project will be used directly.
70-
71-
TODO: allow one to overwrite those names, with a mapping?
72-
73-
Returns
74-
-------
75-
list of str
76-
List of paths to the newly created NWB data files.
77-
By default NWB files are stored in the same folder as the h5file.
78-
79-
"""
54+
def _get_pes_args(config_file, h5file, individual_name):
8055
if "DLC" not in h5file or not h5file.endswith(".h5"):
8156
raise IOError("The file passed in is not a DeepLabCut h5 data file.")
8257

83-
cfg = auxiliaryfunctions.read_config(config)
58+
cfg = auxiliaryfunctions.read_config(config_file)
8459

8560
vidname, scorer = os.path.split(h5file)[-1].split("DLC")
8661
scorer = "DLC" + os.path.splitext(scorer)[0]
@@ -123,37 +98,105 @@ def convert_h5_to_nwb(config, h5file, individual_name="ind1"):
12398
) # setting timestamps to dummy TODO: extract timestamps in DLC?
12499
else:
125100
timestamps = get_movie_timestamps(video[0])
101+
return scorer, df, video, paf_graph, timestamps, cfg
102+
103+
104+
def _write_pes_to_nwbfile(nwbfile, animal, df_animal, scorer, video, paf_graph, timestamps):
105+
pose_estimation_series = []
106+
for kpt, xyp in df_animal.groupby(level="bodyparts", axis=1, sort=False):
107+
data = xyp.to_numpy()
108+
109+
pes = PoseEstimationSeries(
110+
name=f"{animal}_{kpt}",
111+
description=f"Keypoint {kpt} from individual {animal}.",
112+
data=data[:, :2],
113+
unit="pixels",
114+
reference_frame="(0,0) corresponds to the bottom left corner of the video.",
115+
timestamps=timestamps,
116+
confidence=data[:, 2],
117+
confidence_definition="Softmax output of the deep neural network.",
118+
)
119+
pose_estimation_series.append(pes)
120+
121+
pe = PoseEstimation(
122+
pose_estimation_series=pose_estimation_series,
123+
description="2D keypoint coordinates estimated using DeepLabCut.",
124+
original_videos=[video[0]],
125+
# TODO check if this is a mandatory arg in ndx-pose (can skip if video is not found_
126+
dimensions=[list(map(int, video[1].split(",")))[1::2]],
127+
scorer=scorer,
128+
source_software="DeepLabCut",
129+
source_software_version=__version__,
130+
nodes=[pes.name for pes in pose_estimation_series],
131+
edges=paf_graph,
132+
)
133+
if 'behavior' in nwbfile.processing:
134+
behavior_pm = nwbfile.processing["behavior"]
135+
else:
136+
behavior_pm = nwbfile.create_processing_module(
137+
name="behavior", description="processed behavioral data"
138+
)
139+
behavior_pm.add(pe)
140+
return nwbfile
141+
142+
143+
def write_subject_to_nwb(nwbfile, h5file, individual_name, config_file):
144+
"""
145+
Given, subject name, write h5file to an existing nwbfile.
146+
147+
Parameters
148+
----------
149+
nwbfile: pynwb.NWBFile
150+
nwbfile to write the subject specific pose estimation series.
151+
h5file : str
152+
Path to a h5 data file
153+
individual_name : str
154+
Name of the subject (whose pose is predicted) for single-animal DLC project.
155+
For multi-animal projects, the names from the DLC project will be used directly.
156+
config_file : str
157+
Path to a project config.yaml file
158+
config_dict : dict
159+
dict containing configuration options. Provide this as alternative to config.yml file.
160+
161+
Returns
162+
-------
163+
nwbfile: pynwb.NWBFile
164+
nwbfile with pes written in the behavior module
165+
"""
166+
scorer, df, video, paf_graph, timestamps, _ = _get_pes_args(config_file, h5file, individual_name)
167+
df_animal = df.groupby(level="individuals", axis=1).get_group(individual_name)
168+
return _write_pes_to_nwbfile(nwbfile, individual_name, df_animal, scorer, video, paf_graph, timestamps)
169+
170+
171+
def convert_h5_to_nwb(config, h5file, individual_name="ind1"):
172+
"""
173+
Convert a DeepLabCut (DLC) video prediction, h5 data file to Neurodata Without Borders (NWB). Also
174+
takes project config, to store relevant metadata.
175+
176+
Parameters
177+
----------
178+
config : str
179+
Path to a project config.yaml file
180+
181+
h5file : str
182+
Path to a h5 data file
183+
184+
individual_name : str
185+
Name of the subject (whose pose is predicted) for single-animal DLC project.
186+
For multi-animal projects, the names from the DLC project will be used directly.
187+
188+
TODO: allow one to overwrite those names, with a mapping?
189+
190+
Returns
191+
-------
192+
list of str
193+
List of paths to the newly created NWB data files.
194+
By default NWB files are stored in the same folder as the h5file.
126195
196+
"""
197+
scorer, df, video, paf_graph, timestamps, cfg = _get_pes_args(config, h5file, individual_name)
127198
output_paths = []
128199
for animal, df_ in df.groupby(level="individuals", axis=1):
129-
pose_estimation_series = []
130-
for kpt, xyp in df_.groupby(level="bodyparts", axis=1, sort=False):
131-
data = xyp.to_numpy()
132-
133-
pes = PoseEstimationSeries(
134-
name=f"{animal}_{kpt}",
135-
description=f"Keypoint {kpt} from individual {animal}.",
136-
data=data[:, :2],
137-
unit="pixels",
138-
reference_frame="(0,0) corresponds to the bottom left corner of the video.",
139-
timestamps=timestamps,
140-
confidence=data[:, 2],
141-
confidence_definition="Softmax output of the deep neural network.",
142-
)
143-
pose_estimation_series.append(pes)
144-
145-
pe = PoseEstimation(
146-
pose_estimation_series=pose_estimation_series,
147-
description="2D keypoint coordinates estimated using DeepLabCut.",
148-
original_videos=[video[0]],
149-
dimensions=[list(map(int, video[1].split(",")))[1::2]],
150-
scorer=scorer,
151-
source_software="DeepLabCut",
152-
source_software_version=__version__,
153-
nodes=[pes.name for pes in pose_estimation_series],
154-
edges=paf_graph,
155-
)
156-
157200
nwbfile = NWBFile(
158201
session_description=cfg["Task"],
159202
experimenter=cfg["scorer"],
@@ -162,10 +205,7 @@ def convert_h5_to_nwb(config, h5file, individual_name="ind1"):
162205
)
163206

164207
# TODO Store the test_pose_config as well?
165-
behavior_pm = nwbfile.create_processing_module(
166-
name="behavior", description="processed behavioral data"
167-
)
168-
behavior_pm.add(pe)
208+
nwbfile = _write_pes_to_nwbfile(nwbfile, animal, df_, scorer, video, paf_graph, timestamps)
169209
output_path = h5file.replace(".h5", f"_{animal}.nwb")
170210
with warnings.catch_warnings(), NWBHDF5IO(output_path, mode="w") as io:
171211
warnings.filterwarnings("ignore", category=DtypeConversionWarning)

0 commit comments

Comments
 (0)