Skip to content

Commit

Permalink
add nwb table for lightning pose face parts results
Browse files Browse the repository at this point in the history
  • Loading branch information
arjunsridhar12345 committed Jul 12, 2024
1 parent 797cbf0 commit ed816a9
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 6 deletions.
42 changes: 42 additions & 0 deletions src/npc_sessions/sessions.py
Original file line number Diff line number Diff line change
Expand Up @@ -2716,6 +2716,48 @@ def _video_frame_times(
for path, timestamps in path_to_timestamps.items()
)

@npc_io.cached_property
def _LPFaceParts(self) -> tuple[pynwb.core.DynamicTable, ...]:
"""
Stores the lightning pose output as a dynamic table for each of the relevant cameras (side, face)
"""
LP_face_parts_dynamic_tables = []
if not self.is_video:
raise ValueError(f"{self.id} is not a session with video")

for video_path in self.video_paths:
camera_name = npc_mvr.get_camera_name(video_path.name)
if camera_name == "eye":
continue

nwb_camera_name = self.mvr_to_nwb_camera_name[camera_name]
timestamps = next(t for t in self._video_frame_times if nwb_camera_name in t.name).timestamps
assert len(timestamps) == npc_mvr.get_total_frames_in_video(video_path)
df = utils.get_LPFaceParts_predictions_dataframe(self.id, utils.LP_MAPPING[camera_name])
if len(timestamps) != len(df):
logger.warning(
f"{self.id} {camera_name} lightning pose face parts output has wrong shape {len(df)}, expected {len(timestamps)} frames."
"\nLightning pose face parts capsule was likely run with an additional data asset attached"
)
continue

df['timestamps'] = timestamps
name = f"Lightning_Pose_FaceParts_{nwb_camera_name}"
table_description = (
f"Lightning Pose tracking model fit to {len(utils.LP_VIDEO_FEATURES_MAPPING[utils.LP_MAPPING[camera_name]])} facial features for each frame of {nwb_camera_name} video. "
"Output for every frame is x,y coordinates in pixels along with the likelihood of the model for each feature in the frame. "
f"Features tracked are {utils.LP_VIDEO_FEATURES_MAPPING[utils.LP_MAPPING[camera_name]]} "
)
table = pynwb.core.DynamicTable.from_dataframe(
name=name,
table_description=table_description,
df=df
)
LP_face_parts_dynamic_tables.append(table)

return tuple(LP_face_parts_dynamic_tables)


@npc_io.cached_property
def _eye_tracking(self) -> pynwb.core.DynamicTable:
if not self.is_video:
Expand Down
42 changes: 36 additions & 6 deletions src/npc_sessions/utils/videos.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,11 @@
}

FACEMAP_CAMERA_NAMES: tuple[npc_mvr.CameraName, ...] = ("behavior", "face")

LP_CAMERA_NAMES = ("side", "face")
LP_MAPPING = {"behavior": "side", "face": "face"}
LP_VIDEO_FEATURES_MAPPING = {'side': ('eye_top_l', 'eye_bottom_l', 'whisker_pad_l_top', 'whisker_pad_l_side', 'ear_tip_l', 'ear_base_l', 'nostril_l', 'nose_tip', 'jaw', 'tongue_base_l', 'tongue_tip'),
'face': ('eye_top_l', 'eye_bottom_l', 'whisker_pad_l_top', 'whisker_pad_l_side', 'ear_tip_l', 'ear_base_l', 'nostril_l', 'nostril_r', 'nose_tip', 'tongue_base_l', 'tongue_tip', 'tongue_base_r')
}

def get_dlc_session_paf_graph(session: str, model_name: str) -> list:
"""
Expand Down Expand Up @@ -61,14 +65,31 @@ def h5_to_dataframe(h5_path: upath.UPath, key_name: str | None = None) -> pd.Dat


def get_LPFaceParts_predictions_dataframe(
session: str, camera: Literal["side", "face"]
session: str, camera: str
) -> pd.DataFrame:
"""
Gets the result dataframe with the lightning pose prediction for the facial features for the given camera
Gets the result dataframe with the lightning pose prediction for the facial features for the given camera.
Modifies the result dataframe for easier use
>>> df_predictions = get_LPFaceParts_predictions_dataframe('702136_2024-03-07', 'side')
>>> len(df_predictions.columns)
34
33
>>> df_predictions.columns
Index(['eye_top_l_x', 'eye_top_l_y', 'eye_top_l_likelihood', 'eye_bottom_l_x',
'eye_bottom_l_y', 'eye_bottom_l_likelihood', 'whisker_pad_l_top_x',
'whisker_pad_l_top_y', 'whisker_pad_l_top_likelihood',
'whisker_pad_l_side_x', 'whisker_pad_l_side_y',
'whisker_pad_l_side_likelihood', 'ear_tip_l_x', 'ear_tip_l_y',
'ear_tip_l_likelihood', 'ear_base_l_x', 'ear_base_l_y',
'ear_base_l_likelihood', 'nostril_l_x', 'nostril_l_y',
'nostril_l_likelihood', 'nose_tip_x', 'nose_tip_y',
'nose_tip_likelihood', 'jaw_x', 'jaw_y', 'jaw_likelihood',
'tongue_base_l_x', 'tongue_base_l_y', 'tongue_base_l_likelihood',
'tongue_tip_x', 'tongue_tip_y', 'tongue_tip_likelihood'],
dtype='object')
"""
if camera not in LP_CAMERA_NAMES:
raise ValueError(f'Undefined camera name {camera}. Must be one of either side or face')

session_LP_predictions_s3_paths = (
npc_lims.get_lpfaceparts_camera_predictions_s3_paths(session, camera)
)
Expand All @@ -84,8 +105,17 @@ def get_LPFaceParts_predictions_dataframe(

df_predictions = pd.read_csv(session_LP_prediction_csv_s3_path[0], low_memory=False)
data_array = df_predictions.to_numpy()
return pd.DataFrame(data_array[1:], columns=data_array[0])

df_predictions_rearranged = pd.DataFrame(data_array[1:], columns=data_array[0])
df_predictions_rearranged.drop(columns='bodyparts', inplace=True)

new_column_labels = []
for column in df_predictions_rearranged.columns.unique():
df_column = df_predictions_rearranged[column].iloc[0]
new_column_labels.append(f'{column}_{df_column.iloc[0]}')
new_column_labels.append(f'{column}_{df_column.iloc[1]}')
new_column_labels.append(f'{column}_{df_column.iloc[2]}')

return pd.DataFrame(data_array[1:, 1:], columns=new_column_labels).drop(0).reset_index(drop=True)

def get_ellipse_session_dataframe_from_h5(session: str) -> pd.DataFrame:
"""
Expand Down

0 comments on commit ed816a9

Please sign in to comment.