更改为库方式导入

.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+    // 使用 IntelliSense 了解相关属性。 
+    // 悬停以查看现有属性的描述。
+    // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: main.py",
+            "type": "python",
+            "request": "launch",
+            "program": "main.py",
+            "console": "integratedTerminal",
+            "justMyCode": false
+        }
+    ]
+}

isat.yaml

@@ -21,5 +21,5 @@ label:
 - color: '#5c3566'
   name: fence
 language: en
-mask_alpha: 0.4
+mask_alpha: 0.7
 vertex_size: 2

main.py

@@ -3,7 +3,9 @@
 # import os
 
 from PyQt5 import QtWidgets
-from widgets.mainwindow import MainWindow
+import os
+os.environ['SAM_ANN_BASE_DIR'] = os.path.dirname(__file__)
+from sam_ann.widgets.mainwindow import MainWindow
 import sys
 
 

sam_ann/configs.py

@@ -1,6 +1,7 @@
 import yaml
 from enum import Enum
-
+import os
+BASE_DIR = os.environ['SAM_ANN_BASE_DIR']
 
 DEFAULT_CONFIG_FILE = 'default.yaml'
 CONFIG_FILE = 'isat.yaml'

sam_ann/segment_any/segment_any.py

@@ -7,38 +7,38 @@
 import torch
 import numpy as np
 import timm
-
+import os
 
 class SegAny:
     def __init__(self, checkpoint):
-        if 'mobile_sam' in checkpoint:
+        if 'mobile_sam' in os.path.basename(checkpoint):
             # mobile sam
             from mobile_sam import sam_model_registry, SamPredictor
             print('- mobile sam!')
             self.model_type = "vit_t"
-        elif 'sam_hq_vit' in checkpoint:
+        elif 'sam_hq_vit' in os.path.basename(checkpoint):
             # sam hq
             from segment_anything_hq import sam_model_registry, SamPredictor
             print('- sam hq!')
-            if 'vit_b' in checkpoint:
+            if 'vit_b' in os.path.basename(checkpoint):
                 self.model_type = "vit_b"
-            elif 'vit_l' in checkpoint:
+            elif 'vit_l' in os.path.basename(checkpoint):
                 self.model_type = "vit_l"
-            elif 'vit_h' in checkpoint:
+            elif 'vit_h' in os.path.basename(checkpoint):
                 self.model_type = "vit_h"
-            elif 'vit_tiny' in checkpoint:
+            elif 'vit_tiny' in os.path.basename(checkpoint):
                 self.model_type = "vit_tiny"
             else:
                 raise ValueError('The checkpoint named {} is not supported.'.format(checkpoint))
-        elif 'sam_vit' in checkpoint:
+        elif 'sam_vit' in os.path.basename(checkpoint):
             # sam
             from segment_anything import sam_model_registry, SamPredictor
             print('- sam!')
-            if 'vit_b' in checkpoint:
+            if 'vit_b' in os.path.basename(checkpoint):
                 self.model_type = "vit_b"
-            elif 'vit_l' in checkpoint:
+            elif 'vit_l' in os.path.basename(checkpoint):
                 self.model_type = "vit_l"
-            elif 'vit_h' in checkpoint:
+            elif 'vit_h' in os.path.basename(checkpoint):
                 self.model_type = "vit_h"
             else:
                 raise ValueError('The checkpoint named {} is not supported.'.format(checkpoint))

sam_ann/ui/COCO_to_ISAT_dialog.py

@@ -117,4 +117,4 @@ class Ui_Dialog(object):
         self.label.setText(_translate("Dialog", "Convert COCO json to ISAT jsons.All layer attr is 1."))
         self.pushButton_cancel.setText(_translate("Dialog", "cancel"))
         self.pushButton_apply.setText(_translate("Dialog", "convert"))
-import icons_rc
+

sam_ann/ui/ISAT_to_COCO_dialog.py

@@ -110,4 +110,4 @@ class Ui_Dialog(object):
         self.label.setText(_translate("Dialog", "Convert ISAT jsons to COCO json.The layer attr will be lost."))
         self.pushButton_cancel.setText(_translate("Dialog", "cancel"))
         self.pushButton_apply.setText(_translate("Dialog", "convert"))
-import icons_rc
+

sam_ann/ui/ISAT_to_VOC_dialog.py

@@ -124,4 +124,4 @@ class Ui_Dialog(object):
         self.label.setText(_translate("Dialog", "Convert ISAT annotations to VOC png."))
         self.pushButton_cancel.setText(_translate("Dialog", "cancel"))
         self.pushButton_apply.setText(_translate("Dialog", "convert"))
-import icons_rc
+

sam_ann/ui/MainWindow.py

@@ -494,5 +494,4 @@ class Ui_MainWindow(object):
         self.actionContour_External.setWhatsThis(_translate("MainWindow", "External contour save only."))
         self.actionContour_All.setText(_translate("MainWindow", "All"))
         self.actionContour_All.setStatusTip(_translate("MainWindow", "All contour save."))
-        self.actionContour_All.setWhatsThis(_translate("MainWindow", "All contour save."))
-import icons_rc
+        self.actionContour_All.setWhatsThis(_translate("MainWindow", "All contour save."))

sam_ann/ui/about_dialog.py

@@ -74,4 +74,4 @@ class Ui_Dialog(object):
         self.label_2.setText(_translate("Dialog", "ISAT with Segment anything."))
         self.label_4.setText(_translate("Dialog", "ISAT Copyright (C) 2022 yatengLG.\n"
 "http://www.yatenglg.cn/isat"))
-import icons_rc
+

sam_ann/ui/anno_dock.py

@@ -68,4 +68,4 @@ class Ui_Form(object):
         _translate = QtCore.QCoreApplication.translate
         Form.setWindowTitle(_translate("Form", "Form"))
         self.checkBox_visible.setText(_translate("Form", "Visible"))
-import icons_rc
+

sam_ann/ui/category_choice.py

@@ -125,4 +125,4 @@ class Ui_Dialog(object):
         self.checkBox_iscrowded.setText(_translate("Dialog", "is crowded"))
         self.pushButton_cancel.setText(_translate("Dialog", "cancel"))
         self.pushButton_apply.setText(_translate("Dialog", "apply"))
-import icons_rc
+

sam_ann/ui/setting_dialog.py

@@ -108,4 +108,3 @@ class Ui_Dialog(object):
         self.pushButton_export.setText(_translate("Dialog", "Export"))
         self.cancel_button.setText(_translate("Dialog", "Cancel"))
         self.apply_button.setText(_translate("Dialog", "Apply"))
-import icons_rc

sam_ann/ui/shortcut_dialog.py

@@ -289,4 +289,4 @@ class Ui_Dialog(object):
         self.label_8.setText(_translate("Dialog", "To top"))
         self.label_28.setText(_translate("Dialog", "Polygons Visible"))
         self.label_29.setText(_translate("Dialog", "V"))
-import icons_rc
+

sam_ann/widgets/COCO_to_ISAT_dialog.py

@@ -2,8 +2,8 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore
-from ui.COCO_to_ISAT_dialog import Ui_Dialog
-from tools.fromCOCO import FROMCOCO
+from sam_ann.ui.COCO_to_ISAT_dialog import Ui_Dialog
+from sam_ann.tools.fromCOCO import FROMCOCO
 
 
 class COCOtoISATDialog(QtWidgets.QDialog, Ui_Dialog):

sam_ann/widgets/ISAT_to_COCO_dialog.py

@@ -2,8 +2,8 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore, QtGui
-from ui.ISAT_to_COCO_dialog import Ui_Dialog
-from tools.toCOCO import TOCOCO
+from sam_ann.ui.ISAT_to_COCO_dialog import Ui_Dialog
+from sam_ann.tools.toCOCO import TOCOCO
 
 
 class ISATtoCOCODialog(QtWidgets.QDialog, Ui_Dialog):

sam_ann/widgets/ISAT_to_LABELME_dialog.py

@@ -2,8 +2,8 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore, QtGui
-from ui.ISAT_to_LABELME_dialog import Ui_Dialog
-from tools.toLABELME import TOLABELME
+from sam_ann.ui.ISAT_to_LABELME_dialog import Ui_Dialog
+from sam_ann.tools.toLABELME import TOLABELME
 import os
 
 class ISATtoLabelMeDialog(QtWidgets.QDialog, Ui_Dialog):

sam_ann/widgets/ISAT_to_VOC_dialog.py

@@ -2,9 +2,9 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore, QtGui
-from ui.ISAT_to_VOC_dialog import Ui_Dialog
-import tools.toVOC as toVOC
-from configs import load_config
+from sam_ann.ui.ISAT_to_VOC_dialog import Ui_Dialog
+import sam_ann.tools.toVOC as toVOC
+from sam_ann.configs import load_config
 import os
 
 

sam_ann/widgets/about_dialog.py

@@ -2,7 +2,7 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore, QtGui
-from ui.about_dialog import Ui_Dialog
+from sam_ann.ui.about_dialog import Ui_Dialog
 
 class AboutDialog(QtWidgets.QDialog, Ui_Dialog):
     def __init__(self, parent):

sam_ann/widgets/annos_dock_widget.py

@@ -2,7 +2,7 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtCore, QtGui
-from ui.anno_dock import Ui_Form
+from sam_ann.ui.anno_dock import Ui_Form
 import functools
 import re
 

sam_ann/widgets/canvas.py

@@ -2,8 +2,8 @@
 # @Author  : LG
 
 from PyQt5 import QtWidgets, QtGui, QtCore
-from widgets.polygon import Polygon, Vertex
-from configs import STATUSMode, CLICKMode, DRAWMode, CONTOURMode
+from sam_ann.widgets.polygon import Polygon, Vertex
+from sam_ann.configs import STATUSMode, CLICKMode, DRAWMode, CONTOURMode
 from PIL import Image
 import numpy as np
 import cv2
@@ -39,6 +39,11 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
             self.clear()
 
             self.image_data = np.array(Image.open(image_path))
+            if len(self.image_data.shape) == 2:
+                self.image_data = np.stack([self.image_data] * 3, axis=-1)
+            if self.image_data.shape[-1] >= 4:
+                self.image_data = self.image_data[:, :, :3]
+                
        
             self.image_item = QtWidgets.QGraphicsPixmapItem()
             self.image_item.setZValue(0)
@@ -138,7 +143,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
 
     def start_segment_anything(self):
         self.draw_mode = DRAWMode.SEGMENTANYTHING
-        
+
         self.start_draw()
 
     def start_draw_polygon(self):

sam_ann/widgets/category_choice_dialog.py

@@ -2,9 +2,9 @@
 # @Author  : LG
 
 
-from ui.category_choice import Ui_Dialog
+from sam_ann.ui.category_choice import Ui_Dialog
 from PyQt5 import QtWidgets, QtGui, QtCore
-from configs import load_config, CONFIG_FILE, DEFAULT_CONFIG_FILE
+from sam_ann.configs import load_config, CONFIG_FILE, DEFAULT_CONFIG_FILE
 import os