Skip to content

Commit

Permalink
Add feature to extract shape keypoints from JSON files
Browse files Browse the repository at this point in the history
This commit introduces a new feature to the file menu, enabling users to extract shape keypoints from JSON files. Users can now select a folder containing JSON files and input a space-separated list of instance names (e.g., 'mouse rat') to extract the shape keypoints for each instance. This functionality is accessible from the GUI, providing a convenient way for users to process their data.
  • Loading branch information
healthonrails committed Jun 11, 2024
1 parent 2dc5c17 commit a9dcbc0
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 0 deletions.
14 changes: 14 additions & 0 deletions annolid/gui/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
from annolid.gui.widgets.step_size_widget import StepSizeWidget
from annolid.gui.widgets.downsample_videos_dialog import VideoRescaleWidget
from annolid.gui.widgets.convert_sleap_dialog import ConvertSleapDialog
from annolid.gui.widgets.extract_keypoints_dialog import ExtractShapeKeyPointsDialog
from annolid.gui.widgets.convert_labelme2csv_dialog import LabelmeJsonToCsvDialog
from annolid.postprocessing.quality_control import pred_dict_to_labelme
from annolid.annotation.timestamps import convert_frame_number_to_time
Expand Down Expand Up @@ -357,6 +358,14 @@ def __init__(self,
self.tr("Save CSV")
)

extract_shape_keypoints = action(
self.tr("&Extract Shape Keypoints"),
self.extract_and_save_shape_keypoints,
None,
"Extract Shape Keypoints",
self.tr("Extract Shape Keypoints")
)

convert_sleap = action(
self.tr("&Load SLEAP h5"),
self.convert_sleap_h5_to_labelme,
Expand Down Expand Up @@ -532,6 +541,7 @@ def __init__(self,
utils.addActions(self.menus.file, (quality_control,))
utils.addActions(self.menus.file, (downsample_video,))
utils.addActions(self.menus.file, (convert_csv,))
utils.addActions(self.menus.file, (extract_shape_keypoints,))
utils.addActions(self.menus.file, (convert_sleap,))
utils.addActions(self.menus.file, (place_perference,))
utils.addActions(self.menus.file, (advance_params,))
Expand Down Expand Up @@ -648,6 +658,10 @@ def convert_sleap_h5_to_labelme(self):
convert_sleap_h5_widget = ConvertSleapDialog()
convert_sleap_h5_widget.exec_()

def extract_and_save_shape_keypoints(self):
extract_shape_keypoints_dialog = ExtractShapeKeyPointsDialog()
extract_shape_keypoints_dialog.exec_()

def place_preference_analyze(self):
place_preference_analyze_widget = TrackingAnalyzerDialog()
place_preference_analyze_widget.exec_()
Expand Down
46 changes: 46 additions & 0 deletions annolid/gui/widgets/extract_keypoints_dialog.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
from qtpy.QtWidgets import (QVBoxLayout, QPushButton,
QFileDialog, QLineEdit, QLabel, QDialog,
QMessageBox)

from annolid.postprocessing.skeletonization import main as extract_shape_keypoints


class ExtractShapeKeyPointsDialog(QDialog):
def __init__(self):
super().__init__()
self.setWindowTitle("Extract Keypoints")
self.initUI()

def initUI(self):
layout = QVBoxLayout()

self.selectFolderBtn = QPushButton("Select Folder")
self.selectFolderBtn.clicked.connect(self.selectFolder)
layout.addWidget(self.selectFolderBtn)

self.folderLabel = QLabel()
layout.addWidget(self.folderLabel)

self.instanceNamesEdit = QLineEdit()
self.instanceNamesEdit.setPlaceholderText(
"Enter instance names separated by space")
layout.addWidget(self.instanceNamesEdit)

self.runBtn = QPushButton("Run")
self.runBtn.clicked.connect(self.extract_keypoints)
layout.addWidget(self.runBtn)

self.setLayout(layout)

def selectFolder(self):
folder = QFileDialog.getExistingDirectory(self, "Select Folder")
if folder:
self.folderLabel.setText(folder)

def extract_keypoints(self):
input_folder = self.folderLabel.text()
instance_names = self.instanceNamesEdit.text().split()
extract_shape_keypoints(input_folder, instance_names)
# Display message to the user
QMessageBox.information(self, "Processing Complete", "Processing is complete.")
self.accept()

0 comments on commit a9dcbc0

Please sign in to comment.