Accessing sleap.labels for a conversion function to Lighting Pose

[JacopoRazzauti]

Hello everyone,

I am trying to create a conversion function that converts SLEAP labels into a format that is readable by lightning pose (which uses this data directory format: https://lightning-pose.readthedocs.io/en/latest/source/user_guide/directory_structure.html). This means I need to export the frames and store them in the labelled data folder and then generate a csv that indexes the frames and the locations for each instance and for each node (as you can see here: https://github.com/danbider/lightning-pose/blob/main/data/mirror-mouse-example/CollectedData.csv). For this purpose, I wrote a function that currently works on the exported csv files of a set of frames labelled in SLEAP. Ideally, I would like this function to work on the slp.labels but I have a hard time figuring out how the frames and the coordinates of the labels are stored in the sleap format to adapt my function to them. Do you have any suggestions on how to optimize it to receive a .slp file as an input and output a csv + a set of indexed frames?

Thank you!

Here is the function:

import cv2
import pandas as pd
import os

def convert_sleap_to_lightingpose(sleap_csv_path, video_directory, labeled_data_directory, output_csv_path):
    '''This function converts a SLEAP CSV file to a CSV file readable by LightingPose.'''
    
    sleap_df = pd.read_csv(sleap_csv_path) # Read the SLEAP CSV
    
    num_of_tracks = sleap_df['track'].max() # this determins the number of tracks in the sleap csv
    
    body_parts_headers = ['scorer'] # this prepare the column headers for body parts
    coords_header = ['coords'] # this prepare the column headers for coordinates
    
    for track_id in range(1, num_of_tracks + 1):
        body_parts_headers.extend([f'track{track_id}_thorax', f'track{track_id}_thorax'])
        coords_header.extend(['x', 'y']) # this adds the x and y coordinates to the column headers as per LP format
    
    grouped = sleap_df.groupby('frame_idx') # this groups the data by frame index
    
    video_path = os.path.join(video_directory, 'mfemale1_mirror.mp4') # this prepares the video path
    
    frame = cv2.VideoCapture(video_path) # this opens the video file
    
    if not os.path.exists(labeled_data_directory):
        os.makedirs(labeled_data_directory)
        
    rows = []
    image_idx = 1
    
    for frame_idx, frame_group in grouped: # this processes each frame
        frame.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) # this sets the video to the correct frame
        ret, frame = frame.read() # this reads the frame
        
        if not ret:
            continue
        
        image_path = f'img{image_idx}.png' # this saves the frame as an image
        full_image_path = os.path.join(labeled_data_directory, image_path) # this prepares the full image path
        cv2.imwrite(full_image_path, frame) # this writes the image to the labeled data directory
        
        row = [full_image_path] # this initializes the row data with the image path
        
        for track_id in range(1, num_of_tracks + 1): # this fills in the coordinates for all tracks
            track_data = frame_group[frame_group['track'] == track_id]
            if not track_data.empty:
                row.extend([track_data.iloc[0]['Thorax.x'], track_data.iloc[0]['Thorax.y']])
            else:
                row.extend([None, None])
                
        rows.append(row)
        image_idx += 1
    
    frame.release()
    
    with open(output_csv_path, 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(body_parts_headers)
        writer.writerow(coords_headers)
        writer.writerows(data_rows)

[talmo]

Hi @JacopoRazzauti,

For these sorts of conversion tasks, we recommend using sleap-io, which is a standalone package that can read SLP files and do conversions.

As an example, here's a notebook using sleap-io for converting to Ultralytics format for training.

You may also want to check out:

• neuroconv from NWB which supports both loaders.
• movement which has both loaders.

The two above may not yet support exporting training data though (we're working to add this to NWB FWIW).

Let us know if you have any questions!

Talmo