diff --git a/annolid/gui/app.py b/annolid/gui/app.py index cf4c11f..ce58a32 100644 --- a/annolid/gui/app.py +++ b/annolid/gui/app.py @@ -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 @@ -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, @@ -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,)) @@ -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_() diff --git a/annolid/gui/widgets/extract_keypoints_dialog.py b/annolid/gui/widgets/extract_keypoints_dialog.py new file mode 100644 index 0000000..26178bf --- /dev/null +++ b/annolid/gui/widgets/extract_keypoints_dialog.py @@ -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()