From d9ac5d8b422bddb2a6f79500e4ca654301f60d22 Mon Sep 17 00:00:00 2001
From: Jeremia Manurung <39850621+jeremia49@users.noreply.github.com>
Date: Sun, 10 Mar 2024 06:23:42 +0700
Subject: [PATCH] Add Flag Shortcuts

---
 labelme/app.py                     | 46 ++++++++++++++++++++++++++++++
 labelme/config/default_config.yaml |  3 ++
 2 files changed, 49 insertions(+)

diff --git a/labelme/app.py b/labelme/app.py
index 9a83a3711..4d1c792cf 100644
--- a/labelme/app.py
+++ b/labelme/app.py
@@ -205,6 +205,20 @@ def __init__(
         # Actions
         action = functools.partial(utils.newAction, self)
         shortcuts = self._config["shortcuts"]
+        setFlag1 = action(
+            "Set Flag 1",
+            functools.partial(self.toggleTrueFlag, 0),
+            shortcuts["flag_1"],
+            "labels",
+            "Set Flag 1",
+        )
+        setFlag2 = action(
+            "Set Flag 2",
+            functools.partial(self.toggleTrueFlag, 1),
+            shortcuts["flag_2"],
+            "labels",
+            "Set Flag 2",
+        )
         quit = action(
             self.tr("&Quit"),
             self.close,
@@ -656,6 +670,10 @@ def __init__(
             openPrevImg=openPrevImg,
             fileMenuActions=(open_, opendir, save, saveAs, close, quit),
             tool=(),
+
+            setFlag1=setFlag1,
+            setFlag2=setFlag2,
+
             # XXX: need to add some actions here to activate the shortcut
             editMenu=(
                 edit,
@@ -690,6 +708,10 @@ def __init__(
                 undo,
                 undoLastPoint,
                 removePoint,
+
+                setFlag1,
+                setFlag2,
+
             ),
             onLoadActive=(
                 close,
@@ -703,6 +725,10 @@ def __init__(
                 createAiMaskMode,
                 editMode,
                 brightnessContrast,
+                
+                setFlag1,
+                setFlag2,
+                
             ),
             onShapesPresent=(saveAs, hideAll, showAll, toggleAll),
         )
@@ -825,6 +851,9 @@ def __init__(
             zoom,
             None,
             selectAiModel,
+            None,
+            setFlag1,
+            setFlag2,
         )
 
         self.statusBar().showMessage(str(self.tr("%s started.")) % __appname__)
@@ -1290,6 +1319,23 @@ def loadFlags(self, flags):
             item.setCheckState(Qt.Checked if flag else Qt.Unchecked)
             self.flag_widget.addItem(item)
 
+    def toggleTrueFlag(self, index):
+        flags = {}
+        for i in range(self.flag_widget.count()):
+            item = self.flag_widget.item(i)
+            key = item.text()
+            flag = item.checkState() == Qt.Checked
+            flags[key] = flag
+        self.flag_widget.clear()
+        i = 0
+        for key, flag in flags.items():
+            item = QtWidgets.QListWidgetItem(key)
+            item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
+            item.setCheckState(Qt.Checked if i==index else Qt.Unchecked)
+            i += 1
+            self.flag_widget.addItem(item)
+        self.setDirty()
+    
     def saveLabels(self, filename):
         lf = LabelFile()
 
diff --git a/labelme/config/default_config.yaml b/labelme/config/default_config.yaml
index 128dc6d6a..bbf9b5de6 100644
--- a/labelme/config/default_config.yaml
+++ b/labelme/config/default_config.yaml
@@ -122,3 +122,6 @@ shortcuts:
   show_all_polygons: null
   hide_all_polygons: null
   toggle_all_polygons: T
+  
+  flag_1: 1
+  flag_2: 2
\ No newline at end of file