Skeleton templates

[talmo]

Problem background

• Defining the skeleton is an often confusing first step when starting a project.
• There are many best practices for skeleton design that are dependent on the specifics of the dataset and model type that will be used.
• Users defining their own skeletons contributes to the lack of standardization across labeling projects, making it harder to reuse and unify labeled data that have different skeleton definitions (see for example, this creative approach to dealing with this).

Feature proposal

• We have a number of skeletons that we've used for a variety of species and conditions, so let's make them readily available in the GUI so users can just pick from the list instead of having to define their own.
• Adding a basic visualization will help with UX (like we had in the original LEAP GUI).

Implementation details

PR 1: Base functionality (MVP)

1. Add the skeletons from the reference datasets in the SLEAP paper.

   • SLEAP Skeletons can be serialized to JSON files when you click Save from the skeleton pane in the GUI or via Skeleton.save_json().
   • Download one of the SLP files from each of the reference datasets and extract their skeleton. Name them by the dataset name (e.g., flies13, mice_hc, etc.).
   • Rename the skeleton to the dataset name within the JSON file (i.e., set the Skeleton.name attribute).
   • Add the JSON files to a new folder sleap/sleap/skeletons (analogous to sleap/sleap/training_profiles).

2. The Skeleton JSONs can be auto-discovered following a similar pattern to how we do it for training profiles:

   # Get local path to the training profiles (this may be different on different systems depending on how SLEAP is installed).
   profiles_folder = sleap.utils.get_package_file("sleap/training_profiles")
   
   # Search for JSON files within this folder and return them in a list.
   json_files = sleap.utils.find_files_by_suffix(profiles_folder, suffix=".json", depth=1)

   Alternatively, the paths can just be hardcoded since we won't be adding new skeletons all the time.

3. Add a GUI dropdown in the Skeleton panel for selecting and loading the built-in skeleton.

   • The code that defines the Skeleton panel elements is here:
     

     sleap/sleap/gui/app.py

     Lines 993 to 1067 in c4861e3

     	####### Skeleton #######
     	skeleton_layout = _make_dock(
     	"Skeleton", tab_with=videos_layout.parent().parent()
     	)
     	gb = QGroupBox("Nodes")
     	vb = QVBoxLayout()
     	self.skeletonNodesTable = GenericTableView(
     	state=self.state,
     	row_name="node",
     	model=SkeletonNodesTableModel(
     	items=self.state["skeleton"], context=self.commands
     	),
     	)
     	vb.addWidget(self.skeletonNodesTable)
     	hb = QHBoxLayout()
     	_add_button(hb, "New Node", self.commands.newNode)
     	_add_button(hb, "Delete Node", self.commands.deleteNode)
     	hbw = QWidget()
     	hbw.setLayout(hb)
     	vb.addWidget(hbw)
     	gb.setLayout(vb)
     	skeleton_layout.addWidget(gb)
     	def _update_edge_src():
     	self.skeletonEdgesDst.model().skeleton = self.state["skeleton"]
     	gb = QGroupBox("Edges")
     	vb = QVBoxLayout()
     	self.skeletonEdgesTable = GenericTableView(
     	state=self.state,
     	row_name="edge",
     	model=SkeletonEdgesTableModel(
     	items=self.state["skeleton"], context=self.commands
     	),
     	)
     	vb.addWidget(self.skeletonEdgesTable)
     	hb = QHBoxLayout()
     	self.skeletonEdgesSrc = QComboBox()
     	self.skeletonEdgesSrc.setEditable(False)
     	self.skeletonEdgesSrc.currentIndexChanged.connect(_update_edge_src)
     	self.skeletonEdgesSrc.setModel(SkeletonNodeModel(self.state["skeleton"]))
     	hb.addWidget(self.skeletonEdgesSrc)
     	hb.addWidget(QLabel("to"))
     	self.skeletonEdgesDst = QComboBox()
     	self.skeletonEdgesDst.setEditable(False)
     	hb.addWidget(self.skeletonEdgesDst)
     	self.skeletonEdgesDst.setModel(
     	SkeletonNodeModel(
     	self.state["skeleton"], lambda: self.skeletonEdgesSrc.currentText()
     	)
     	)
     	def new_edge():
     	src_node = self.skeletonEdgesSrc.currentText()
     	dst_node = self.skeletonEdgesDst.currentText()
     	self.commands.newEdge(src_node, dst_node)
     	_add_button(hb, "Add Edge", new_edge)
     	_add_button(hb, "Delete Edge", self.commands.deleteEdge)
     	hbw = QWidget()
     	hbw.setLayout(hb)
     	vb.addWidget(hbw)
     	gb.setLayout(vb)
     	skeleton_layout.addWidget(gb)
     	hb = QHBoxLayout()
     	_add_button(hb, "Load Skeleton", self.commands.openSkeleton)
     	_add_button(hb, "Save Skeleton", self.commands.saveSkeleton)
     	hbw = QWidget()
     	hbw.setLayout(hb)
     	skeleton_layout.addWidget(hbw)

     
     This is where a dropdown should be added.
   • Use a QComboBox to define the dropdown.
   • Populate the dropdown with step 2 above, but have it default to a blank item.
   • Modify the Load button functionality in the skeleton panel GUI to load from the list if the non-blank item is selected.
   • Optional: Play with the layout and possibly add text (QLabel) elements to make it clear to the user.

PR 2: Preview

Once the core functionality is implemented, add a couple of UI elements which will help with the UX. A plain dropdown with no information other than the filename is not very descriptive, so new users may not know which skeleton to select. One way to deal with this is to add two additional elements to the skeleton: a text description and a preview image.

1. Modify the Skeleton class to also contain two optional fields:

   class Skeleton:
       # ...
       preview_image: Optional[bytes] = None
       description: Optional[str] = None

2. Make sure that serialization/deserialization is handled properly in Skeleton.to_json() and Skeleton.from_json()

   • The JSON serializer might handle byte strings just fine, but we should also check for backwards compatibility with older skeletons that will not have this field.

3. Add images to the built in skeletons.

   • This can just be screenshots that you take from the GUI after loading the reference datasets. They should have node label names visible.
   • Encode the screenshots as byte strings so they can be stored in the JSON file in the preview_image field. This can be done as such:

   import base64
   from io import BytesIO
   from PIL import Image
   
   img = Image.open("test.png")
   img_stream = BytesIO()
   img.save(img_stream, format="PNG")  # encode the image as PNG (TODO: compress/resize)
   img_bytes = img_stream.getvalue()  # image in binary format
   img_b64 = base64.b64encode(img_bytes)  # bytes

   • This is basically a thumbnail and should be <<5 KB.
   • Make sure the bytestring is as small as possible, so explore different resolutions, lossy compression values or maybe just use JPEG instead of PNG to encode the image.
   • Do this once per skeleton.
   • Optionally add the serialization code above in sleap.utils for future use.

4. Implement image decoding from base64 strings so we can convert them back into images we can display in the GUI:

   import base64
   from io import BytesIO
   from PIL import Image
   import numpy as np
   
   img_b64 = skeleton.preview_image  # this should be stored in the file originally
   img = Image.open(BytesIO(base64.b64decode(img_b64)))  # decode the image into a PIL.Image
   # np.array(img)   # convert it into a (height, width, channels) uint8 array
   # img.toqimage()  # convert it into a QImage for display in Qt
   # img.toqpixmap()  # convert it into a QPixmap for display in Qt

   This is basically the inverse of step 3.

5. Add the preview and description elements to the Skeleton panel in the GUI.

   • These should be fetched from the JSON files by loading the skeleton when the user selects one from the dropdown.
   • The description can be displayed using a QLabel.
   • The image display is a little trickier. It probably is best rendered via a QPixmap. Search the SLEAP code for how we use it in other places for examples.

[roomrys]

This feature is now available in the (pre) release 1.3.0a0, to install, first uninstall and then:
conda (Windows/Linux/GPU):

conda create -y -n sleap -c sleap -c sleap/label/dev -c nvidia -c conda-forge sleap=1.3.0a0

pip (any OS except Apple Silicon):

pip install sleap==1.3.0a0

Warning: This is a pre-release! Expect bugs and strange behavior when testing.