Skip to content

Commit

Permalink
camera update
Browse files Browse the repository at this point in the history
  • Loading branch information
keyaloding committed Aug 13, 2024
1 parent e7b9fd3 commit e932340
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion sleap_io/io/nwb.py
Original file line number Diff line number Diff line change
Expand Up @@ -559,6 +559,7 @@ def append_nwb_data(
default_metadata["source_software_version"] = sleap_version

labels_data_df = convert_predictions_to_dataframe(labels)
cameras = []

# For every video create a processing module
for video_index, video in enumerate(labels.videos):
Expand All @@ -567,6 +568,13 @@ def append_nwb_data(
nwb_processing_module = get_processing_module_for_video(
processing_module_name, nwbfile
)

camera = nwbfile.create_device(
name=f"camera {video_index}",
description=f"Camera used to record video {video_index}",
manufacturer="No specified manufacturer",
)
cameras.append(camera)

# Propagate video metadata
default_metadata["original_videos"] = [f"{video.filename}"]
Expand All @@ -589,6 +597,7 @@ def append_nwb_data(
track_name,
video,
default_metadata,
nwbfile,
)
nwb_processing_module.add(pose_estimation_container)

Expand Down Expand Up @@ -762,6 +771,7 @@ def build_pose_estimation_container_for_track(
track_name: str,
video: Video,
pose_estimation_metadata: dict,
nwbfile: NWBFile,
) -> PoseEstimation:
"""Create a PoseEstimation container for a track.
Expand All @@ -772,6 +782,7 @@ def build_pose_estimation_container_for_track(
track_name (str): The name of the track in labels.tracks
video (Video): The video to which data belongs to
pose_estimation_metadata (dict): Metadata for the pose estimation.
nwbfile (NWBFile): The nwbfile.
Returns:
PoseEstimation: A PoseEstimation multicontainer where the time series
Expand Down Expand Up @@ -821,6 +832,7 @@ def build_pose_estimation_container_for_track(
nodes=skeleton.node_names,
edges=np.array(skeleton.edge_inds).astype("uint64"),
source_software="SLEAP",
devices=[list(nwbfile.devices.values())[0]],
# dimensions=np.array([[video.backend.height, video.backend.width]]),
)

Expand All @@ -831,7 +843,7 @@ def build_pose_estimation_container_for_track(


def build_track_pose_estimation_list(
track_data_df: pd.DataFrame, timestamps: ArrayLike # type: ignore[return]
track_data_df: pd.DataFrame, timestamps: ArrayLike # type: ignore[return]
) -> List[PoseEstimationSeries]:
"""Build a list of PoseEstimationSeries from tracks.
Expand Down

0 comments on commit e932340

Please sign in to comment.