添加基于标注框的**自动分割**功能

ISAT/segment_any/edge_sam/modeling/prompt_encoder.py

@@ -166,6 +166,7 @@ def forward(
             sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
         if boxes is not None:
             box_embeddings = self._embed_boxes(boxes)
+            sparse_embeddings = sparse_embeddings.to(box_embeddings.dtype)
             sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
 
         if masks is not None:

ISAT/segment_any/mobile_sam/modeling/prompt_encoder.py

@@ -166,6 +166,7 @@ def forward(
             sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
         if boxes is not None:
             box_embeddings = self._embed_boxes(boxes)
+            sparse_embeddings = sparse_embeddings.to(box_embeddings.dtype)
             sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
 
         if masks is not None:

ISAT/segment_any/segment_any.py

@@ -114,3 +114,11 @@ def predict_with_point_prompt(self, input_point, input_label):
         )
         torch.cuda.empty_cache()
         return masks
+
+    def predict_with_box_prompt(self, box):
+        masks, scores, logits = self.predictor_with_point_prompt.predict(
+            box = box,
+            multimask_output=False,
+        )
+        torch.cuda.empty_cache()
+        return masks

ISAT/segment_any/segment_anything/modeling/prompt_encoder.py

@@ -166,6 +166,7 @@ def forward(
             sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
         if boxes is not None:
             box_embeddings = self._embed_boxes(boxes)
+            sparse_embeddings = sparse_embeddings.to(box_embeddings.dtype)
             sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
 
         if masks is not None:

ISAT/segment_any/segment_anything_hq/modeling/prompt_encoder.py

@@ -166,6 +166,7 @@ def forward(
             sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
         if boxes is not None:
             box_embeddings = self._embed_boxes(boxes)
+            sparse_embeddings = sparse_embeddings.to(box_embeddings.dtype)
             sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
 
         if masks is not None:

ISAT/software.yaml

@@ -22,11 +22,11 @@ label:
 - color: '#204a87'
   name: hat
 software:
-  contour_mode: max_only
+  contour_mode: external
   language: zh
   mask_alpha: 0.6
   show_edge: true
-  show_prompt: false
-  use_bfloat16: false
-  use_polydp: false
+  show_prompt: true
+  use_bfloat16: true
+  use_polydp: true
   vertex_size: 2

ISAT/ui/MainWindow.py

@@ -2,7 +2,7 @@
 
 # Form implementation generated from reading ui file '/mnt/disk2/PycharmProjects/ISAT_with_segment_anything/ISAT/ui/MainWindow.ui'
 #
-# Created by: PyQt5 UI code generator 5.15.7
+# Created by: PyQt5 UI code generator 5.15.9
 #
 # WARNING: Any manual changes made to this file will be lost when pyuic5 is
 # run again.  Do not edit this file unless you know what you are doing.
@@ -305,6 +305,9 @@ def setupUi(self, MainWindow):
         icon25.addPixmap(QtGui.QPixmap(":/icon/icons/转换文件夹1_folder-conversion-one.svg"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
         self.actionConverter.setIcon(icon25)
         self.actionConverter.setObjectName("actionConverter")
+        self.actionAuto_segment = QtWidgets.QAction(MainWindow)
+        self.actionAuto_segment.setIcon(icon14)
+        self.actionAuto_segment.setObjectName("actionAuto_segment")
         self.menuFile.addAction(self.actionOpen_dir)
         self.menuFile.addAction(self.actionSave_dir)
         self.menuFile.addSeparator()
@@ -327,6 +330,7 @@ def setupUi(self, MainWindow):
         self.menuAbout.addAction(self.actionAbout)
         self.menuTools.addSeparator()
         self.menuTools.addAction(self.actionConverter)
+        self.menuTools.addAction(self.actionAuto_segment)
         self.menuEdit.addAction(self.actionSegment_anything)
         self.menuEdit.addAction(self.actionPolygon)
         self.menuEdit.addAction(self.actionBackspace)
@@ -481,3 +485,4 @@ def retranslateUi(self, MainWindow):
         self.actionModel_manage.setStatusTip(_translate("MainWindow", "Model manage."))
         self.actionModel_manage.setWhatsThis(_translate("MainWindow", "Model manage."))
         self.actionConverter.setText(_translate("MainWindow", "Converter"))
+        self.actionAuto_segment.setText(_translate("MainWindow", "Auto segment with bounding box"))

ISAT/ui/MainWindow.ui

@@ -152,6 +152,7 @@
     </property>
     <addaction name="separator"/>
     <addaction name="actionConverter"/>
+    <addaction name="actionAuto_segment"/>
    </widget>
    <widget class="QMenu" name="menuEdit">
     <property name="font">
@@ -821,6 +822,15 @@
     <string>Converter</string>
    </property>
   </action>
+  <action name="actionAuto_segment">
+   <property name="icon">
+    <iconset resource="../icons.qrc">
+     <normaloff>:/icon/icons/M_Favicon.ico</normaloff>:/icon/icons/M_Favicon.ico</iconset>
+   </property>
+   <property name="text">
+    <string>Auto segment with bounding box</string>
+   </property>
+  </action>
  </widget>
  <resources>
   <include location="../icons.qrc"/>

ISAT/ui/auto_segment.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+
+# Form implementation generated from reading ui file '/mnt/disk2/PycharmProjects/ISAT_with_segment_anything/ISAT/ui/auto_segment.ui'
+#
+# Created by: PyQt5 UI code generator 5.15.9
+#
+# WARNING: Any manual changes made to this file will be lost when pyuic5 is
+# run again.  Do not edit this file unless you know what you are doing.
+
+
+from PyQt5 import QtCore, QtGui, QtWidgets
+
+
+class Ui_Dialog(object):
+    def setupUi(self, Dialog):
+        Dialog.setObjectName("Dialog")
+        Dialog.resize(1000, 680)
+        font = QtGui.QFont()
+        font.setFamily("Times New Roman")
+        font.setPointSize(12)
+        Dialog.setFont(font)
+        self.verticalLayout = QtWidgets.QVBoxLayout(Dialog)
+        self.verticalLayout.setObjectName("verticalLayout")
+        self.widget = QtWidgets.QWidget(Dialog)
+        self.widget.setObjectName("widget")
+        self.horizontalLayout = QtWidgets.QHBoxLayout(self.widget)
+        self.horizontalLayout.setContentsMargins(0, 0, 0, 0)
+        self.horizontalLayout.setObjectName("horizontalLayout")
+        self.lineEdit_image_dir = QtWidgets.QLineEdit(self.widget)
+        self.lineEdit_image_dir.setObjectName("lineEdit_image_dir")
+        self.horizontalLayout.addWidget(self.lineEdit_image_dir)
+        self.pushButton_image_dir = QtWidgets.QPushButton(self.widget)
+        self.pushButton_image_dir.setObjectName("pushButton_image_dir")
+        self.horizontalLayout.addWidget(self.pushButton_image_dir)
+        self.verticalLayout.addWidget(self.widget)
+        self.widget_2 = QtWidgets.QWidget(Dialog)
+        self.widget_2.setObjectName("widget_2")
+        self.horizontalLayout_2 = QtWidgets.QHBoxLayout(self.widget_2)
+        self.horizontalLayout_2.setContentsMargins(0, 0, 0, 0)
+        self.horizontalLayout_2.setObjectName("horizontalLayout_2")
+        self.lineEdit_label_dir = QtWidgets.QLineEdit(self.widget_2)
+        self.lineEdit_label_dir.setObjectName("lineEdit_label_dir")
+        self.horizontalLayout_2.addWidget(self.lineEdit_label_dir)
+        self.pushButton_label_dir = QtWidgets.QPushButton(self.widget_2)
+        self.pushButton_label_dir.setObjectName("pushButton_label_dir")
+        self.horizontalLayout_2.addWidget(self.pushButton_label_dir)
+        self.verticalLayout.addWidget(self.widget_2)
+        self.widget_3 = QtWidgets.QWidget(Dialog)
+        self.widget_3.setObjectName("widget_3")
+        self.horizontalLayout_3 = QtWidgets.QHBoxLayout(self.widget_3)
+        self.horizontalLayout_3.setContentsMargins(0, 0, 0, 0)
+        self.horizontalLayout_3.setObjectName("horizontalLayout_3")
+        self.lineEdit_save_dir = QtWidgets.QLineEdit(self.widget_3)
+        self.lineEdit_save_dir.setObjectName("lineEdit_save_dir")
+        self.horizontalLayout_3.addWidget(self.lineEdit_save_dir)
+        self.pushButton_save_dir = QtWidgets.QPushButton(self.widget_3)
+        self.pushButton_save_dir.setObjectName("pushButton_save_dir")
+        self.horizontalLayout_3.addWidget(self.pushButton_save_dir)
+        self.verticalLayout.addWidget(self.widget_3)
+        self.label = QtWidgets.QLabel(Dialog)
+        self.label.setObjectName("label")
+        self.verticalLayout.addWidget(self.label)
+        self.textBrowser = QtWidgets.QTextBrowser(Dialog)
+        self.textBrowser.setMinimumSize(QtCore.QSize(0, 0))
+        font = QtGui.QFont()
+        font.setFamily("宋体")
+        self.textBrowser.setFont(font)
+        self.textBrowser.setObjectName("textBrowser")
+        self.verticalLayout.addWidget(self.textBrowser)
+        self.progressBar = QtWidgets.QProgressBar(Dialog)
+        self.progressBar.setProperty("value", 0)
+        self.progressBar.setObjectName("progressBar")
+        self.verticalLayout.addWidget(self.progressBar)
+        self.widget_4 = QtWidgets.QWidget(Dialog)
+        self.widget_4.setObjectName("widget_4")
+        self.horizontalLayout_4 = QtWidgets.QHBoxLayout(self.widget_4)
+        self.horizontalLayout_4.setContentsMargins(0, 0, 0, 0)
+        self.horizontalLayout_4.setObjectName("horizontalLayout_4")
+        spacerItem = QtWidgets.QSpacerItem(461, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum)
+        self.horizontalLayout_4.addItem(spacerItem)
+        self.pushButton_cancel = QtWidgets.QPushButton(self.widget_4)
+        self.pushButton_cancel.setObjectName("pushButton_cancel")
+        self.horizontalLayout_4.addWidget(self.pushButton_cancel)
+        self.pushButton_start = QtWidgets.QPushButton(self.widget_4)
+        self.pushButton_start.setObjectName("pushButton_start")
+        self.horizontalLayout_4.addWidget(self.pushButton_start)
+        self.verticalLayout.addWidget(self.widget_4)
+
+        self.retranslateUi(Dialog)
+        QtCore.QMetaObject.connectSlotsByName(Dialog)
+
+    def retranslateUi(self, Dialog):
+        _translate = QtCore.QCoreApplication.translate
+        Dialog.setWindowTitle(_translate("Dialog", "Auto segment"))
+        self.lineEdit_image_dir.setPlaceholderText(_translate("Dialog", "Image dir"))
+        self.pushButton_image_dir.setText(_translate("Dialog", "image root"))
+        self.lineEdit_label_dir.setPlaceholderText(_translate("Dialog", "VOC xmls dir (for object detection)"))
+        self.pushButton_label_dir.setText(_translate("Dialog", "xml root"))
+        self.lineEdit_save_dir.setPlaceholderText(_translate("Dialog", "ISAT jsons save dir"))
+        self.pushButton_save_dir.setText(_translate("Dialog", "save root"))
+        self.label.setText(_translate("Dialog", "Auto segment with bounding box.(\'contour_mode\' and \'use_polydp\' also effect the results.)"))
+        self.pushButton_cancel.setText(_translate("Dialog", "Cancel"))
+        self.pushButton_start.setText(_translate("Dialog", "Start"))