新版本2.0

annotation.py

@@ -84,7 +84,7 @@ class Annotation:
                         if not is_polygon:
                             continue
                         category = shape.get('label', 'unknow')
-                        group = shape.get('group_id', '')
+                        group = shape.get('group_id', 0)
                         if group is None: group = ''
                         segmentation = shape.get('points', [])
                         iscrowd = shape.get('iscrowd', 0)

configs.py

@@ -33,3 +33,8 @@ class MAPMode(Enum):
     LABEL = 0
     SEMANTIC = 1
     INSTANCE = 2
+
+class CONTOURMode(Enum):
+    SAVE_MAX_ONLY = 0       # 只保留最多顶点的mask（一般为最大面积）
+    SAVE_EXTERNAL = 1       # 只保留外轮廓
+    SAVE_ALL = 2            # 保留所有轮廓

default.yaml

@@ -1,3 +1,5 @@
 label:
 - color: '#000000'
   name: __background__
+- color: '#00ff00'
+  name: aaaa

example/images/isat.yaml

@@ -1,3 +1,4 @@
+contour_mode: external
 label:
 - color: '#000000'
   name: __background__
@@ -19,3 +20,5 @@ label:
   name: cake
 - color: '#5c3566'
   name: fence
+language: en
+mask_alpha: 0.6

icons.qrc

@@ -1,6 +1,7 @@
 <RCC>
   <qresource prefix="icon">
     <file>icons/semantic.png</file>
+    <file>icons/眼睛_eyes.svg</file>
     <file>icons/VOC_32x32.png</file>
     <file>icons/labelme_32x32.png</file>
     <file>icons/coco.ico</file>

isat.yaml

@@ -1,3 +1,4 @@
+contour_mode: external
 label:
 - color: '#000000'
   name: __background__
@@ -19,4 +20,5 @@ label:
   name: cake
 - color: '#5c3566'
   name: fence
-language: zh
+language: en
+mask_alpha: 0.5

segment_any/segment_any.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # @Author  : LG
 
-from segment_anything import sam_model_registry, SamPredictor
+from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
 import torch
 import numpy as np
 
@@ -20,29 +20,30 @@ class SegAny:
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         sam = sam_model_registry[self.model_type](checkpoint=checkpoint)
         sam.to(device=self.device)
-        self.predictor = SamPredictor(sam)
+        self.predictor = SamAutomaticMaskGenerator(sam)
+        self.predictor_with_point_prompt = SamPredictor(sam)
         self.image = None
 
     def set_image(self, image):
         self.image = image
-        self.predictor.set_image(image)
+        self.predictor_with_point_prompt.set_image(image)
 
     def reset_image(self):
-        self.predictor.reset_image()
+        self.predictor_with_point_prompt.reset_image()
         self.image = None
         torch.cuda.empty_cache()
 
-    def predict(self, input_point, input_label):
+    def predict_with_point_prompt(self, input_point, input_label):
         input_point = np.array(input_point)
         input_label = np.array(input_label)
 
-        masks, scores, logits = self.predictor.predict(
+        masks, scores, logits = self.predictor_with_point_prompt.predict(
             point_coords=input_point,
             point_labels=input_label,
             multimask_output=True,
         )
         mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask
-        masks, _, _ = self.predictor.predict(
+        masks, _, _ = self.predictor_with_point_prompt.predict(
             point_coords=input_point,
             point_labels=input_label,
             mask_input=mask_input[None, :, :],
@@ -50,3 +51,28 @@ class SegAny:
         )
         torch.cuda.empty_cache()
         return masks
+
+    def predict(self, image):
+        self.image = image
+        masks = self.predictor.generate(image)
+        torch.cuda.empty_cache()
+        return masks
+
+
+if __name__ == '__main__':
+    from PIL import Image
+    import time
+    import matplotlib.pyplot as plt
+    time1 = time.time()
+    seg = SegAny('sam_vit_h_4b8939.pth')
+    image = np.array(Image.open('../example/images/000000000113.jpg'))
+    time2 = time.time()
+    print(time2-time1)
+    # seg.set_image()
+    masks = seg.predict(image)
+    print(time.time() - time2)
+    print(masks)
+    for mask in masks:
+        mask = mask['segmentation']
+        plt.imshow(mask)
+        plt.show()

ui/MainWindow.py

@@ -77,6 +77,17 @@ class Ui_MainWindow(object):
         font.setPointSize(12)
         self.menuEdit.setFont(font)
         self.menuEdit.setObjectName("menuEdit")
+        self.menuMode = QtWidgets.QMenu(self.menubar)
+        font = QtGui.QFont()
+        font.setFamily("Times New Roman")
+        font.setPointSize(12)
+        font.setBold(False)
+        font.setItalic(False)
+        font.setWeight(50)
+        self.menuMode.setFont(font)
+        self.menuMode.setObjectName("menuMode")
+        self.menuContour_mode = QtWidgets.QMenu(self.menuMode)
+        self.menuContour_mode.setObjectName("menuContour_mode")
         MainWindow.setMenuBar(self.menubar)
         self.statusbar = QtWidgets.QStatusBar(MainWindow)
         self.statusbar.setLayoutDirection(QtCore.Qt.LeftToRight)
@@ -100,20 +111,28 @@ class Ui_MainWindow(object):
         self.dockWidgetContents_2.setObjectName("dockWidgetContents_2")
         self.info_dock.setWidget(self.dockWidgetContents_2)
         MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.info_dock)
-        self.labels_dock = QtWidgets.QDockWidget(MainWindow)
+        self.annos_dock = QtWidgets.QDockWidget(MainWindow)
-        self.labels_dock.setMinimumSize(QtCore.QSize(85, 43))
+        self.annos_dock.setMinimumSize(QtCore.QSize(85, 43))
-        self.labels_dock.setFeatures(QtWidgets.QDockWidget.AllDockWidgetFeatures)
+        self.annos_dock.setFeatures(QtWidgets.QDockWidget.AllDockWidgetFeatures)
-        self.labels_dock.setObjectName("labels_dock")
+        self.annos_dock.setObjectName("annos_dock")
         self.dockWidgetContents_3 = QtWidgets.QWidget()
         self.dockWidgetContents_3.setObjectName("dockWidgetContents_3")
-        self.labels_dock.setWidget(self.dockWidgetContents_3)
+        self.annos_dock.setWidget(self.dockWidgetContents_3)
-        MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.labels_dock)
+        MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.annos_dock)
         self.files_dock = QtWidgets.QDockWidget(MainWindow)
         self.files_dock.setObjectName("files_dock")
         self.dockWidgetContents = QtWidgets.QWidget()
         self.dockWidgetContents.setObjectName("dockWidgetContents")
         self.files_dock.setWidget(self.dockWidgetContents)
         MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.files_dock)
+        self.categories_dock = QtWidgets.QDockWidget(MainWindow)
+        self.categories_dock.setObjectName("categories_dock")
+        self.dockWidgetContents_4 = QtWidgets.QWidget()
+        self.dockWidgetContents_4.setObjectName("dockWidgetContents_4")
+        self.verticalLayout_2 = QtWidgets.QVBoxLayout(self.dockWidgetContents_4)
+        self.verticalLayout_2.setObjectName("verticalLayout_2")
+        self.categories_dock.setWidget(self.dockWidgetContents_4)
+        MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(1), self.categories_dock)
         self.actionOpen_dir = QtWidgets.QAction(MainWindow)
         icon2 = QtGui.QIcon()
         icon2.addPixmap(QtGui.QPixmap(":/icon/icons/照片_pic.svg"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
@@ -268,6 +287,27 @@ class Ui_MainWindow(object):
         icon26.addPixmap(QtGui.QPixmap(":/icon/icons/labelme_32x32.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
         self.actionTo_LabelMe.setIcon(icon26)
         self.actionTo_LabelMe.setObjectName("actionTo_LabelMe")
+        self.actionContour_Max_only = QtWidgets.QAction(MainWindow)
+        self.actionContour_Max_only.setCheckable(True)
+        font = QtGui.QFont()
+        font.setFamily("Times New Roman")
+        font.setPointSize(12)
+        self.actionContour_Max_only.setFont(font)
+        self.actionContour_Max_only.setObjectName("actionContour_Max_only")
+        self.actionContour_External = QtWidgets.QAction(MainWindow)
+        self.actionContour_External.setCheckable(True)
+        font = QtGui.QFont()
+        font.setFamily("Times New Roman")
+        font.setPointSize(12)
+        self.actionContour_External.setFont(font)
+        self.actionContour_External.setObjectName("actionContour_External")
+        self.actionContour_All = QtWidgets.QAction(MainWindow)
+        self.actionContour_All.setCheckable(True)
+        font = QtGui.QFont()
+        font.setFamily("Times New Roman")
+        font.setPointSize(12)
+        self.actionContour_All.setFont(font)
+        self.actionContour_All.setObjectName("actionContour_All")
         self.menuFile.addAction(self.actionOpen_dir)
         self.menuFile.addAction(self.actionSave_dir)
         self.menuFile.addSeparator()
@@ -305,9 +345,14 @@ class Ui_MainWindow(object):
         self.menuEdit.addAction(self.actionEdit)
         self.menuEdit.addAction(self.actionDelete)
         self.menuEdit.addAction(self.actionSave)
+        self.menuContour_mode.addAction(self.actionContour_Max_only)
+        self.menuContour_mode.addAction(self.actionContour_External)
+        self.menuContour_mode.addAction(self.actionContour_All)
+        self.menuMode.addAction(self.menuContour_mode.menuAction())
         self.menubar.addAction(self.menuFile.menuAction())
         self.menubar.addAction(self.menuEdit.menuAction())
         self.menubar.addAction(self.menuView.menuAction())
+        self.menubar.addAction(self.menuMode.menuAction())
         self.menubar.addAction(self.menuTools.menuAction())
         self.menubar.addAction(self.menuAbout.menuAction())
         self.toolBar.addAction(self.actionPrev)
@@ -344,10 +389,13 @@ class Ui_MainWindow(object):
         self.menuLaguage.setTitle(_translate("MainWindow", "Laguage"))
         self.menuTools.setTitle(_translate("MainWindow", "Tools"))
         self.menuEdit.setTitle(_translate("MainWindow", "Edit"))
+        self.menuMode.setTitle(_translate("MainWindow", "Mode"))
+        self.menuContour_mode.setTitle(_translate("MainWindow", "Contour mode"))
         self.toolBar.setWindowTitle(_translate("MainWindow", "toolBar"))
         self.info_dock.setWindowTitle(_translate("MainWindow", "Info"))
-        self.labels_dock.setWindowTitle(_translate("MainWindow", "Labels"))
+        self.annos_dock.setWindowTitle(_translate("MainWindow", "Annos"))
         self.files_dock.setWindowTitle(_translate("MainWindow", "Files"))
+        self.categories_dock.setWindowTitle(_translate("MainWindow", "Categories"))
         self.actionOpen_dir.setText(_translate("MainWindow", "Images dir"))
         self.actionOpen_dir.setStatusTip(_translate("MainWindow", "Open images dir."))
         self.actionZoom_in.setText(_translate("MainWindow", "Zoom in"))
@@ -422,7 +470,8 @@ class Ui_MainWindow(object):
         self.actionPolygon.setStatusTip(_translate("MainWindow", "Accurately annotate by drawing polygon. "))
         self.actionPolygon.setShortcut(_translate("MainWindow", "C"))
         self.actionVisible.setText(_translate("MainWindow", "Visible"))
-        self.actionVisible.setStatusTip(_translate("MainWindow", "Visible"))
+        self.actionVisible.setToolTip(_translate("MainWindow", "Visible."))
+        self.actionVisible.setStatusTip(_translate("MainWindow", "Visible."))
         self.actionVisible.setShortcut(_translate("MainWindow", "V"))
         self.actionToCOCO.setText(_translate("MainWindow", "To COCO"))
         self.actionToCOCO.setToolTip(_translate("MainWindow", "Convert ISAT to COCO"))
@@ -433,4 +482,13 @@ class Ui_MainWindow(object):
         self.actionTo_LabelMe.setText(_translate("MainWindow", "To LabelMe"))
         self.actionTo_LabelMe.setToolTip(_translate("MainWindow", "Convert ISAT to LabelMe"))
         self.actionTo_LabelMe.setStatusTip(_translate("MainWindow", "Convert ISAT jsons to LabelMe jsons."))
+        self.actionContour_Max_only.setText(_translate("MainWindow", "Max only"))
+        self.actionContour_Max_only.setStatusTip(_translate("MainWindow", "Contour save max only."))
+        self.actionContour_Max_only.setWhatsThis(_translate("MainWindow", "Contour save max only."))
+        self.actionContour_External.setText(_translate("MainWindow", "External"))
+        self.actionContour_External.setStatusTip(_translate("MainWindow", "Contour save external only."))
+        self.actionContour_External.setWhatsThis(_translate("MainWindow", "Contour save external only."))
+        self.actionContour_All.setText(_translate("MainWindow", "All"))
+        self.actionContour_All.setStatusTip(_translate("MainWindow", "Contour save all."))
+        self.actionContour_All.setWhatsThis(_translate("MainWindow", "Contour save all."))
 import icons_rc

ui/MainWindow.ui

@@ -179,9 +179,33 @@
     <addaction name="actionDelete"/>
     <addaction name="actionSave"/>
    </widget>
+   <widget class="QMenu" name="menuMode">
+    <property name="font">
+     <font>
+      <family>Times New Roman</family>
+      <pointsize>12</pointsize>
+      <weight>50</weight>
+      <italic>false</italic>
+      <bold>false</bold>
+     </font>
+    </property>
+    <property name="title">
+     <string>Mode</string>
+    </property>
+    <widget class="QMenu" name="menuContour_mode">
+     <property name="title">
+      <string>Contour mode</string>
+     </property>
+     <addaction name="actionContour_Max_only"/>
+     <addaction name="actionContour_External"/>
+     <addaction name="actionContour_All"/>
+    </widget>
+    <addaction name="menuContour_mode"/>
+   </widget>
    <addaction name="menuFile"/>
    <addaction name="menuEdit"/>
    <addaction name="menuView"/>
+   <addaction name="menuMode"/>
    <addaction name="menuTools"/>
    <addaction name="menuAbout"/>
   </widget>
@@ -258,7 +282,7 @@
    </attribute>
    <widget class="QWidget" name="dockWidgetContents_2"/>
   </widget>
-  <widget class="QDockWidget" name="labels_dock">
+  <widget class="QDockWidget" name="annos_dock">
    <property name="minimumSize">
     <size>
      <width>85</width>
@@ -269,7 +293,7 @@
     <set>QDockWidget::AllDockWidgetFeatures</set>
    </property>
    <property name="windowTitle">
-    <string>Labels</string>
+    <string>Annos</string>
    </property>
    <attribute name="dockWidgetArea">
     <number>2</number>
@@ -285,6 +309,17 @@
    </attribute>
    <widget class="QWidget" name="dockWidgetContents"/>
   </widget>
+  <widget class="QDockWidget" name="categories_dock">
+   <property name="windowTitle">
+    <string>Categories</string>
+   </property>
+   <attribute name="dockWidgetArea">
+    <number>1</number>
+   </attribute>
+   <widget class="QWidget" name="dockWidgetContents_4">
+    <layout class="QVBoxLayout" name="verticalLayout_2"/>
+   </widget>
+  </widget>
   <action name="actionOpen_dir">
    <property name="icon">
     <iconset resource="../icons.qrc">
@@ -693,14 +728,17 @@
   </action>
   <action name="actionVisible">
    <property name="icon">
-    <iconset>
+    <iconset resource="../icons.qrc">
      <normaloff>:/icon/icons/眼睛_eyes.svg</normaloff>:/icon/icons/眼睛_eyes.svg</iconset>
    </property>
    <property name="text">
     <string>Visible</string>
    </property>
+   <property name="toolTip">
+    <string>Visible.</string>
+   </property>
    <property name="statusTip">
-    <string>Visible</string>
+    <string>Visible.</string>
    </property>
    <property name="shortcut">
     <string>V</string>
@@ -751,6 +789,66 @@
     <string>Convert ISAT jsons to LabelMe jsons.</string>
    </property>
   </action>
+  <action name="actionContour_Max_only">
+   <property name="checkable">
+    <bool>true</bool>
+   </property>
+   <property name="text">
+    <string>Max only</string>
+   </property>
+   <property name="statusTip">
+    <string>Contour save max only.</string>
+   </property>
+   <property name="whatsThis">
+    <string>Contour save max only.</string>
+   </property>
+   <property name="font">
+    <font>
+     <family>Times New Roman</family>
+     <pointsize>12</pointsize>
+    </font>
+   </property>
+  </action>
+  <action name="actionContour_External">
+   <property name="checkable">
+    <bool>true</bool>
+   </property>
+   <property name="text">
+    <string>External</string>
+   </property>
+   <property name="statusTip">
+    <string>Contour save external only.</string>
+   </property>
+   <property name="whatsThis">
+    <string>Contour save external only.</string>
+   </property>
+   <property name="font">
+    <font>
+     <family>Times New Roman</family>
+     <pointsize>12</pointsize>
+    </font>
+   </property>
+  </action>
+  <action name="actionContour_All">
+   <property name="checkable">
+    <bool>true</bool>
+   </property>
+   <property name="text">
+    <string>All</string>
+   </property>
+   <property name="statusTip">
+    <string>Contour save all.</string>
+   </property>
+   <property name="whatsThis">
+    <string>Contour save all.</string>
+   </property>
+   <property name="font">
+    <font>
+     <family>Times New Roman</family>
+     <pointsize>12</pointsize>
+    </font>
+   </property>
+  </action>
  </widget>
  <resources>
   <include location="../icons.qrc"/>

ui/anno_dock.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Form implementation generated from reading ui file '/home/super/PycharmProjects/ISAT_with_segment_anything/ui/label_dock.ui'
+# Form implementation generated from reading ui file '/home/super/PycharmProjects/ISAT_with_segment_anything/ui/anno_dock.ui'
 #
 # Created by: PyQt5 UI code generator 5.15.7
 #

widgets/annos_dock_widget.py

@@ -2,13 +2,13 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore, QtGui
-from ui.label_dock import Ui_Form
+from ui.anno_dock import Ui_Form
 import functools
 
 
-class LabelsDockWidget(QtWidgets.QWidget, Ui_Form):
+class AnnosDockWidget(QtWidgets.QWidget, Ui_Form):
     def __init__(self, mainwindow):
-        super(LabelsDockWidget, self).__init__()
+        super(AnnosDockWidget, self).__init__()
         self.setupUi(self)
         self.mainwindow = mainwindow
         self.polygon_item_dict = {}

widgets/canvas.py

@@ -2,9 +2,8 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtGui, QtCore
-from enum import Enum
 from widgets.polygon import Polygon
-from configs import STATUSMode, CLICKMode, DRAWMode
+from configs import STATUSMode, CLICKMode, DRAWMode, CONTOURMode
 from PIL import Image
 import numpy as np
 import cv2
@@ -20,9 +19,11 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
         self.mode = STATUSMode.VIEW
         self.click = CLICKMode.POSITIVE
         self.draw_mode = DRAWMode.SEGMENTANYTHING   # 默认使用segment anything进行快速标注
+        self.contour_mode = CONTOURMode.SAVE_EXTERNAL       # 默认SAM只保留外轮廓
         self.click_points = []
         self.click_points_mode = []
         self.masks:np.ndarray = None
+        self.mask_alpha = 0.5
         self.top_layer = 1
 
         self.guide_line_x:QtWidgets.QGraphicsLineItem = None
@@ -42,7 +43,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
                 self.image_data = np.repeat(self.image_data, 3, axis=2) # 转换为三通道
                 self.mainwindow.segany.set_image(self.image_data)
             else:
-                self.mainwindow.statusbar.showMessage("Segment anything don't support the image with {} ndim.".format(self.image_data.ndim))
+                self.mainwindow.statusbar.showMessage("Segment anything don't support the image with shape {} .".format(self.image_data.shape))
                 
         self.image_item = QtWidgets.QGraphicsPixmapItem()
         self.image_item.setZValue(0)
@@ -76,7 +77,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
         self.mainwindow.actionSave.setEnabled(False)
 
         self.mainwindow.set_labels_visible(False)
-        self.mainwindow.labels_dock_widget.setEnabled(False)
+        self.mainwindow.annos_dock_widget.setEnabled(False)
 
     def change_mode_to_view(self):
         self.mode = STATUSMode.VIEW
@@ -101,7 +102,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
         self.mainwindow.actionSave.setEnabled(self.mainwindow.can_be_annotated)
 
         self.mainwindow.set_labels_visible(True)
-        self.mainwindow.labels_dock_widget.setEnabled(True)
+        self.mainwindow.annos_dock_widget.setEnabled(True)
 
     def change_mode_to_edit(self):
         self.mode = STATUSMode.EDIT
@@ -128,6 +129,15 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
     def change_click_to_negative(self):
         self.click = CLICKMode.NEGATIVE
 
+    def change_contour_mode_to_save_all(self):
+        self.contour_mode = CONTOURMode.SAVE_ALL
+
+    def change_contour_mode_to_save_max_only(self):
+        self.contour_mode = CONTOURMode.SAVE_MAX_ONLY
+
+    def change_contour_mode_to_save_external(self):
+        self.contour_mode = CONTOURMode.SAVE_EXTERNAL
+
     def start_segment_anything(self):
         self.draw_mode = DRAWMode.SEGMENTANYTHING
         self.start_draw()
@@ -154,6 +164,11 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
 
         self.change_mode_to_view()
 
+        category = self.mainwindow.current_category
+        group = self.mainwindow.current_group
+        is_crowd = False
+        note = ''
+
         if self.draw_mode == DRAWMode.SEGMENTANYTHING:
             # mask to polygon
             # --------------
@@ -163,18 +178,55 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
                 h, w = masks.shape[-2:]
                 masks = masks.reshape(h, w)
 
-                contours, _ = cv2.findContours(masks, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
+                if self.contour_mode == CONTOURMode.SAVE_ALL:
+                    # 当保留所有轮廓时，检测所有轮廓，并建立二层等级关系
+                    contours, hierarchy = cv2.findContours(masks, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_KCOS)
+                else:
+                    # 当只保留外轮廓或单个mask时，只检测外轮廓
+                    contours, hierarchy = cv2.findContours(masks, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
 
-                # 这里取轮廓点数最多的（可能返回多个轮廓）
+                if self.contour_mode == CONTOURMode.SAVE_MAX_ONLY:
                     contour = contours[0]
                     for cont in contours:
                         if len(cont) > len(contour):
                             contour = cont
+                    contours = [contour]
 
+                for index, contour in enumerate(contours):
+                    if self.current_graph is None:
+                        self.current_graph = Polygon()
+                        self.addItem(self.current_graph)
+
+                    if len(contour) < 3:
+                        continue
                     for point in contour:
                         x, y = point[0]
                         self.current_graph.addPoint(QtCore.QPointF(x, y))
 
+                    if self.contour_mode == CONTOURMode.SAVE_ALL and hierarchy[0][index][3] != -1:
+                        # 保存所有轮廓，且当前轮廓为子轮廓，则自轮廓类别设置为背景
+                        category = '__background__'
+                        group = 0
+                    else:
+                        category = self.mainwindow.current_category
+                        group = self.mainwindow.current_group
+
+                    self.current_graph.set_drawed(category,
+                                                  group,
+                                                  is_crowd,
+                                                  note,
+                                                  QtGui.QColor(self.mainwindow.category_color_dict[category]),
+                                                  self.top_layer)
+
+                    # 添加新polygon
+                    self.mainwindow.polygons.append(self.current_graph)
+                    # 设置为最高图层
+                    self.current_graph.setZValue(len(self.mainwindow.polygons))
+                    for vertex in self.current_graph.vertexs:
+                        vertex.setZValue(len(self.mainwindow.polygons))
+                    self.current_graph = None
+                self.mainwindow.current_group += 1
+
         elif self.draw_mode == DRAWMode.POLYGON:
             if len(self.current_graph.points) < 1:
                 return
@@ -198,9 +250,28 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
                 self.current_graph.addPoint(last_point)
                 self.current_graph.addPoint(QtCore.QPointF(last_point.x(), first_point.y()))
 
+            # 设置polygon 属性
+            self.current_graph.set_drawed(category,
+                                          group,
+                                          is_crowd,
+                                          note,
+                                          QtGui.QColor(self.mainwindow.category_color_dict[category]),
+                                          self.top_layer)
+            self.mainwindow.current_group += 1
+            # 添加新polygon
+            self.mainwindow.polygons.append(self.current_graph)
+            # 设置为最高图层
+            self.current_graph.setZValue(len(self.mainwindow.polygons))
+            for vertex in self.current_graph.vertexs:
+                vertex.setZValue(len(self.mainwindow.polygons))
         # 选择类别
-        self.mainwindow.category_choice_widget.load_cfg()
+        # self.mainwindow.category_choice_widget.load_cfg()
-        self.mainwindow.category_choice_widget.show()
+        # self.mainwindow.category_choice_widget.show()
+
+        self.mainwindow.annos_dock_widget.update_listwidget()
+
+        self.current_graph = None
+        self.change_mode_to_view()
 
         # mask清空
         self.click_points.clear()
@@ -234,7 +305,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
             for p in self.mainwindow.polygons:
                 if p.zValue() > deleted_layer:
                     p.setZValue(p.zValue() - 1)
-            self.mainwindow.labels_dock_widget.update_listwidget()
+            self.mainwindow.annos_dock_widget.update_listwidget()
 
     def edit_polygon(self):
         selectd_items = self.selectedItems()
@@ -368,18 +439,21 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
     def update_mask(self):
         if not self.mainwindow.use_segment_anything:
             return
+        if self.image_data is None:
+            return
         if not (self.image_data.ndim == 3 and self.image_data.shape[-1] == 3):
             return
 
         if len(self.click_points) > 0 and len(self.click_points_mode) > 0:
-            masks = self.mainwindow.segany.predict(self.click_points, self.click_points_mode)
+            masks = self.mainwindow.segany.predict_with_point_prompt(self.click_points, self.click_points_mode)
             self.masks = masks
             color = np.array([0, 0, 255])
             h, w = masks.shape[-2:]
             mask_image = masks.reshape(h, w, 1) * color.reshape(1, 1, -1)
             mask_image = mask_image.astype("uint8")
             mask_image = cv2.cvtColor(mask_image, cv2.COLOR_BGR2RGB)
-            mask_image = cv2.addWeighted(self.image_data, 0.5, mask_image, 0.9, 0)
+            # 这里通过调整原始图像的权重self.mask_alpha，来调整mask的明显程度。
+            mask_image = cv2.addWeighted(self.image_data, self.mask_alpha, mask_image, 1, 0)
             mask_image = QtGui.QImage(mask_image[:], mask_image.shape[1], mask_image.shape[0], mask_image.shape[1] * 3,
                                       QtGui.QImage.Format_RGB888)
             mask_pixmap = QtGui.QPixmap(mask_image)

widgets/category_choice_dialog.py

@@ -88,7 +88,7 @@ class CategoryChoiceDialog(QtWidgets.QDialog, Ui_Dialog):
         for vertex in self.scene.current_graph.vertexs:
             vertex.setZValue(len(self.mainwindow.polygons))
 
-        self.mainwindow.labels_dock_widget.update_listwidget()
+        self.mainwindow.annos_dock_widget.update_listwidget()
 
         self.scene.current_graph = None
         self.scene.change_mode_to_view()

widgets/category_edit_dialog.py

@@ -99,7 +99,7 @@ class CategoryEditDialog(QtWidgets.QDialog, Ui_Dialog):
         # 设置polygon 属性
         self.polygon.set_drawed(category, group, is_crowd, note,
                                             QtGui.QColor(self.mainwindow.category_color_dict.get(category, '#000000')))
-        self.mainwindow.labels_dock_widget.update_listwidget()
+        self.mainwindow.annos_dock_widget.update_listwidget()
 
         self.polygon = None
         self.scene.change_mode_to_view()

widgets/mainwindow.py

@@ -6,7 +6,8 @@ from ui.MainWindow import Ui_MainWindow
 from widgets.setting_dialog import SettingDialog
 from widgets.category_choice_dialog import CategoryChoiceDialog
 from widgets.category_edit_dialog import CategoryEditDialog
-from widgets.labels_dock_widget import LabelsDockWidget
+from widgets.category_dock_widget import CategoriesDockWidget
+from widgets.annos_dock_widget import AnnosDockWidget
 from widgets.files_dock_widget import FilesDockWidget
 from widgets.info_dock_widget import InfoDockWidget
 from widgets.right_button_menu import RightButtonMenu
@@ -34,6 +35,8 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         super(MainWindow, self).__init__()
         self.setupUi(self)
         self.init_ui()
+        self.init_segment_anything()
+
         self.image_root: str = None
         self.label_root:str = None
 
@@ -41,6 +44,9 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         self.current_index = None
         self.current_file_index: int = None
 
+        self.current_label = '__background__'
+        self.current_group = 1
+
         self.config_file = CONFIG_FILE if os.path.exists(CONFIG_FILE) else DEFAULT_CONFIG_FILE
         self.saved = True
         self.can_be_annotated = True
@@ -57,7 +63,6 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
 
         self.init_connect()
         self.reset_action()
-        self.init_segment_anything()
 
     def init_segment_anything(self):
         if os.path.exists('./segment_any/sam_vit_h_4b8939.pth'):
@@ -90,8 +95,11 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         #
         self.setting_dialog = SettingDialog(parent=self, mainwindow=self)
 
-        self.labels_dock_widget = LabelsDockWidget(mainwindow=self)
+        self.categories_dock_widget = CategoriesDockWidget(mainwindow=self)
-        self.labels_dock.setWidget(self.labels_dock_widget)
+        self.categories_dock.setWidget(self.categories_dock_widget)
+
+        self.annos_dock_widget = AnnosDockWidget(mainwindow=self)
+        self.annos_dock.setWidget(self.annos_dock_widget)
 
         self.files_dock_widget = FilesDockWidget(mainwindow=self)
         self.files_dock.setWidget(self.files_dock_widget)
@@ -135,6 +143,18 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         self.labelData.setFixedWidth(150)
         self.statusbar.addPermanentWidget(self.labelData)
 
+        #
+
+        self.toolBar.addSeparator()
+        self.mask_aplha = QtWidgets.QSlider(QtCore.Qt.Orientation.Horizontal, self)
+        self.mask_aplha.setFixedWidth(50)
+        self.mask_aplha.setStatusTip('Mask alpha.')
+        self.mask_aplha.setToolTip('Mask alpha ')
+        self.mask_aplha.setMaximum(10)
+        self.mask_aplha.setMinimum(3)
+        self.mask_aplha.valueChanged.connect(self.change_mask_aplha)
+        self.toolBar.addWidget(self.mask_aplha)
+
         self.trans = QtCore.QTranslator()
 
     def translate(self, language='zh'):
@@ -148,7 +168,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         _app.installTranslator(self.trans)
         self.retranslateUi(self)
         self.info_dock_widget.retranslateUi(self.info_dock_widget)
-        self.labels_dock_widget.retranslateUi(self.labels_dock_widget)
+        self.annos_dock_widget.retranslateUi(self.annos_dock_widget)
         self.files_dock_widget.retranslateUi(self.files_dock_widget)
         self.category_choice_widget.retranslateUi(self.category_choice_widget)
         self.category_edit_widget.retranslateUi(self.category_edit_widget)
@@ -182,8 +202,19 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
             self.show_image(self.current_index)
 
         language = self.cfg.get('language', 'en')
+        self.cfg['language'] = language
         self.translate(language)
 
+        contour_mode = self.cfg.get('contour_mode', 'max_only')
+        self.cfg['contour_mode'] = contour_mode
+        self.change_contour_mode(contour_mode)
+
+        mask_alpha = self.cfg.get('mask_alpha', 0.5)
+        self.cfg['mask_alpha'] = mask_alpha
+        self.mask_aplha.setValue(mask_alpha*10)
+
+        self.categories_dock_widget.update_widget()
+
     def set_saved_state(self, is_saved:bool):
         self.saved = is_saved
         if self.files_list is not None and self.current_index is not None:
@@ -255,7 +286,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
             return
         try:
             self.polygons.clear()
-            self.labels_dock_widget.listWidget.clear()
+            self.annos_dock_widget.listWidget.clear()
             self.scene.cancel_draw()
             file_path = os.path.join(self.image_root, self.files_list[index])
             image_data = Image.open(file_path)
@@ -292,6 +323,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
 
             # load label
             if self.can_be_annotated:
+                self.current_group = 1
                 _, name = os.path.split(file_path)
                 label_path = os.path.join(self.label_root, '.'.join(name.split('.')[:-1]) + '.json')
                 self.current_label = Annotation(file_path, label_path)
@@ -299,6 +331,11 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
                 self.current_label.load_annotation()
 
                 for object in self.current_label.objects:
+                    try:
+                        group = int(object.group)
+                        self.current_group = group+1 if group >= self.current_group else self.current_group
+                    except Exception as e:
+                        pass
                     polygon = Polygon()
                     self.scene.addItem(polygon)
                     polygon.load_object(object)
@@ -309,7 +346,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
             else:
                 self.setWindowTitle('{}'.format(file_path))
 
-            self.labels_dock_widget.update_listwidget()
+            self.annos_dock_widget.update_listwidget()
             self.info_dock_widget.update_widget()
             self.files_dock_widget.set_select(index)
             self.current_index = index
@@ -411,8 +448,8 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
                 polygon.change_color(QtGui.QColor(self.category_color_dict.get(polygon.category, '#000000')))
                 polygon.color.setAlpha(255)
                 polygon.setBrush(polygon.color)
-            self.labels_dock_widget.listWidget.setEnabled(False)
+            self.annos_dock_widget.listWidget.setEnabled(False)
-            self.labels_dock_widget.checkBox_visible.setEnabled(False)
+            self.annos_dock_widget.checkBox_visible.setEnabled(False)
             self.actionSegment_anything.setEnabled(False)
             self.actionPolygon.setEnabled(False)
             self.actionVisible.setEnabled(False)
@@ -428,14 +465,16 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
                 for vertex in polygon.vertexs:
                     vertex.setVisible(False)
                 if polygon.group != '':
-                    rgb = self.instance_cmap[int(polygon.group)]
+                    index = int(polygon.group)
+                    index = index % self.instance_cmap.shape[0]
+                    rgb = self.instance_cmap[index]
                 else:
                     rgb = self.instance_cmap[0]
                 polygon.change_color(QtGui.QColor(rgb[0], rgb[1], rgb[2], 255))
                 polygon.color.setAlpha(255)
                 polygon.setBrush(polygon.color)
-            self.labels_dock_widget.listWidget.setEnabled(False)
+            self.annos_dock_widget.listWidget.setEnabled(False)
-            self.labels_dock_widget.checkBox_visible.setEnabled(False)
+            self.annos_dock_widget.checkBox_visible.setEnabled(False)
             self.actionSegment_anything.setEnabled(False)
             self.actionPolygon.setEnabled(False)
             self.actionVisible.setEnabled(False)
@@ -454,8 +493,8 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
                 polygon.change_color(QtGui.QColor(self.category_color_dict.get(polygon.category, '#000000')))
                 polygon.color.setAlpha(polygon.nohover_alpha)
                 polygon.setBrush(polygon.color)
-            self.labels_dock_widget.listWidget.setEnabled(True)
+            self.annos_dock_widget.listWidget.setEnabled(True)
-            self.labels_dock_widget.checkBox_visible.setEnabled(True)
+            self.annos_dock_widget.checkBox_visible.setEnabled(True)
             self.actionSegment_anything.setEnabled(self.use_segment_anything)
             self.actionPolygon.setEnabled(True)
             self.actionVisible.setEnabled(True)
@@ -468,9 +507,33 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
 
     def set_labels_visible(self, visible=None):
         if visible is None:
-            visible = not self.labels_dock_widget.checkBox_visible.isChecked()
+            visible = not self.annos_dock_widget.checkBox_visible.isChecked()
-        self.labels_dock_widget.checkBox_visible.setChecked(visible)
+        self.annos_dock_widget.checkBox_visible.setChecked(visible)
-        self.labels_dock_widget.set_all_polygon_visible(visible)
+        self.annos_dock_widget.set_all_polygon_visible(visible)
+
+    def change_contour_mode(self, contour_mode='max_only'):
+        if contour_mode == 'max_only':
+            self.scene.change_contour_mode_to_save_max_only()
+        elif contour_mode == 'external':
+            self.scene.change_contour_mode_to_save_external()
+            self.statusbar.showMessage('Save all external contours will bring some noise.', 3000)
+        elif contour_mode == 'all':
+            self.scene.change_contour_mode_to_save_all()
+            self.statusbar.showMessage('Category of inner contour will be set _background__.', 3000)
+        else:
+            self.scene.change_contour_mode_to_save_max_only()
+            self.statusbar.showMessage('The contour mode [{}] not support.'.format(contour_mode), 3000)
+
+        self.actionContour_Max_only.setChecked(contour_mode == 'max_only')
+        self.actionContour_External.setChecked(contour_mode == 'external')
+        self.actionContour_All.setChecked(contour_mode == 'all')
+        self.cfg['contour_mode'] = contour_mode
+
+    def change_mask_aplha(self):
+        value = self.mask_aplha.value() / 10
+        self.scene.mask_alpha = value
+        self.scene.update_mask()
+        self.cfg['mask_alpha'] = value
 
     def ISAT_to_VOC(self):
         self.ISAT_to_VOC_dialog.reset_gui()
@@ -529,6 +592,10 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         self.actionBit_map.triggered.connect(self.change_bit_map)
         self.actionVisible.triggered.connect(functools.partial(self.set_labels_visible, None))
 
+        self.actionContour_Max_only.triggered.connect(functools.partial(self.change_contour_mode, 'max_only'))
+        self.actionContour_External.triggered.connect(functools.partial(self.change_contour_mode, 'external'))
+        self.actionContour_All.triggered.connect(functools.partial(self.change_contour_mode, 'all'))
+
         self.actionToVOC.triggered.connect(self.ISAT_to_VOC)
         self.actionToCOCO.triggered.connect(self.ISAT_to_COCO)
         self.actionTo_LabelMe.triggered.connect(self.ISAT_to_LABELME)
@@ -540,7 +607,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         self.actionChinese.triggered.connect(self.translate_to_chinese)
         self.actionEnglish.triggered.connect(self.translate_to_english)
 
-        self.labels_dock_widget.listWidget.doubleClicked.connect(self.scene.edit_polygon)
+        self.annos_dock_widget.listWidget.doubleClicked.connect(self.scene.edit_polygon)
 
     def reset_action(self):
         self.actionPrev.setEnabled(False)

widgets/polygon.py

@@ -133,7 +133,7 @@ class Polygon(QtWidgets.QGraphicsPolygonItem):
             else:
                 self.color.setAlpha(self.nohover_alpha)
                 self.setBrush(self.color)
-            self.scene().mainwindow.labels_dock_widget.set_selected(self) # 更新label面板
+            self.scene().mainwindow.annos_dock_widget.set_selected(self) # 更新label面板
 
         if change == QtWidgets.QGraphicsItem.GraphicsItemChange.ItemPositionChange: # ItemPositionHasChanged
             bias = value