wangtong/ISAT_with_sam
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

新版本2.0

Browse Source
This commit is contained in:
yatengLG 2023-07-10 10:33:42 +08:00
parent e38ffcfb63
commit bc8488bf8f
20 changed files with 2150 additions and 5789 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
annotation.py
View File

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

7
configs.py
View File

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

2
default.yaml
View File

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

2907
example/images/000000000113.json
View File

File diff suppressed because it is too large Load Diff

1135
example/images/000000002592.json
View File

File diff suppressed because it is too large Load Diff

3
example/images/isat.yaml
View File

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

1
icons.qrc
View File

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

3429
icons_rc.py
View File

File diff suppressed because it is too large Load Diff

4
isat.yaml
View File

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

40
segment_any/segment_any.py
View File

@ -1,7 +1,7 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
# @Author : LG # @Author : LG
from segment_anything import sam_model_registry, SamPredictor from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
import torch import torch
import numpy as np import numpy as np
@ -20,29 +20,30 @@ class SegAny:
self.device = 'cuda' if torch.cuda.is_available() else 'cpu' self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
sam = sam_model_registry[self.model_type](checkpoint=checkpoint) sam = sam_model_registry[self.model_type](checkpoint=checkpoint)
sam.to(device=self.device) sam.to(device=self.device)
self.predictor = SamPredictor(sam) self.predictor = SamAutomaticMaskGenerator(sam)
self.predictor_with_point_prompt = SamPredictor(sam)
self.image = None self.image = None
def set_image(self, image): def set_image(self, image):
self.image = image self.image = image
self.predictor.set_image(image) self.predictor_with_point_prompt.set_image(image)
def reset_image(self): def reset_image(self):
self.predictor.reset_image() self.predictor_with_point_prompt.reset_image()
self.image = None self.image = None
torch.cuda.empty_cache() 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_point = np.array(input_point)
input_label = np.array(input_label) 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_coords=input_point,
point_labels=input_label, point_labels=input_label,
multimask_output=True, multimask_output=True,
) )
mask_input = logits[np.argmax(scores), :, :] # Choose the model's best mask 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_coords=input_point,
point_labels=input_label, point_labels=input_label,
mask_input=mask_input[None, :, :], mask_input=mask_input[None, :, :],
@ -50,3 +51,28 @@ class SegAny:
) )
torch.cuda.empty_cache() torch.cuda.empty_cache()
return masks 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()

74
ui/MainWindow.py
View File

@ -77,6 +77,17 @@ class Ui_MainWindow(object):
font.setPointSize(12) font.setPointSize(12)
self.menuEdit.setFont(font) self.menuEdit.setFont(font)
self.menuEdit.setObjectName("menuEdit") 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) MainWindow.setMenuBar(self.menubar)
self.statusbar = QtWidgets.QStatusBar(MainWindow) self.statusbar = QtWidgets.QStatusBar(MainWindow)
self.statusbar.setLayoutDirection(QtCore.Qt.LeftToRight) self.statusbar.setLayoutDirection(QtCore.Qt.LeftToRight)
@ -100,20 +111,28 @@ class Ui_MainWindow(object):
self.dockWidgetContents_2.setObjectName("dockWidgetContents_2") self.dockWidgetContents_2.setObjectName("dockWidgetContents_2")
self.info_dock.setWidget(self.dockWidgetContents_2) self.info_dock.setWidget(self.dockWidgetContents_2)
MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.info_dock) 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 = QtWidgets.QWidget()
self.dockWidgetContents_3.setObjectName("dockWidgetContents_3") 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 = QtWidgets.QDockWidget(MainWindow)
self.files_dock.setObjectName("files_dock") self.files_dock.setObjectName("files_dock")
self.dockWidgetContents = QtWidgets.QWidget() self.dockWidgetContents = QtWidgets.QWidget()
self.dockWidgetContents.setObjectName("dockWidgetContents") self.dockWidgetContents.setObjectName("dockWidgetContents")
self.files_dock.setWidget(self.dockWidgetContents) self.files_dock.setWidget(self.dockWidgetContents)
MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.files_dock) 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) self.actionOpen_dir = QtWidgets.QAction(MainWindow)
icon2 = QtGui.QIcon() icon2 = QtGui.QIcon()
icon2.addPixmap(QtGui.QPixmap(":/icon/icons/照片_pic.svg"), QtGui.QIcon.Normal, QtGui.QIcon.Off) 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) icon26.addPixmap(QtGui.QPixmap(":/icon/icons/labelme_32x32.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
self.actionTo_LabelMe.setIcon(icon26) self.actionTo_LabelMe.setIcon(icon26)
self.actionTo_LabelMe.setObjectName("actionTo_LabelMe") 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.actionOpen_dir)
self.menuFile.addAction(self.actionSave_dir) self.menuFile.addAction(self.actionSave_dir)
self.menuFile.addSeparator() self.menuFile.addSeparator()
@ -305,9 +345,14 @@ class Ui_MainWindow(object):
self.menuEdit.addAction(self.actionEdit) self.menuEdit.addAction(self.actionEdit)
self.menuEdit.addAction(self.actionDelete) self.menuEdit.addAction(self.actionDelete)
self.menuEdit.addAction(self.actionSave) 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.menuFile.menuAction())
self.menubar.addAction(self.menuEdit.menuAction()) self.menubar.addAction(self.menuEdit.menuAction())
self.menubar.addAction(self.menuView.menuAction()) self.menubar.addAction(self.menuView.menuAction())
self.menubar.addAction(self.menuMode.menuAction())
self.menubar.addAction(self.menuTools.menuAction()) self.menubar.addAction(self.menuTools.menuAction())
self.menubar.addAction(self.menuAbout.menuAction()) self.menubar.addAction(self.menuAbout.menuAction())
self.toolBar.addAction(self.actionPrev) self.toolBar.addAction(self.actionPrev)
@ -344,10 +389,13 @@ class Ui_MainWindow(object):
self.menuLaguage.setTitle(_translate("MainWindow", "Laguage")) self.menuLaguage.setTitle(_translate("MainWindow", "Laguage"))
self.menuTools.setTitle(_translate("MainWindow", "Tools")) self.menuTools.setTitle(_translate("MainWindow", "Tools"))
self.menuEdit.setTitle(_translate("MainWindow", "Edit")) 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.toolBar.setWindowTitle(_translate("MainWindow", "toolBar"))
self.info_dock.setWindowTitle(_translate("MainWindow", "Info")) 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.files_dock.setWindowTitle(_translate("MainWindow", "Files"))
self.categories_dock.setWindowTitle(_translate("MainWindow", "Categories"))
self.actionOpen_dir.setText(_translate("MainWindow", "Images dir")) self.actionOpen_dir.setText(_translate("MainWindow", "Images dir"))
self.actionOpen_dir.setStatusTip(_translate("MainWindow", "Open images dir.")) self.actionOpen_dir.setStatusTip(_translate("MainWindow", "Open images dir."))
self.actionZoom_in.setText(_translate("MainWindow", "Zoom in")) 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.setStatusTip(_translate("MainWindow", "Accurately annotate by drawing polygon. "))
self.actionPolygon.setShortcut(_translate("MainWindow", "C")) self.actionPolygon.setShortcut(_translate("MainWindow", "C"))
self.actionVisible.setText(_translate("MainWindow", "Visible")) 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.actionVisible.setShortcut(_translate("MainWindow", "V"))
self.actionToCOCO.setText(_translate("MainWindow", "To COCO")) self.actionToCOCO.setText(_translate("MainWindow", "To COCO"))
self.actionToCOCO.setToolTip(_translate("MainWindow", "Convert ISAT 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.setText(_translate("MainWindow", "To LabelMe"))
self.actionTo_LabelMe.setToolTip(_translate("MainWindow", "Convert ISAT 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.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 import icons_rc

106
ui/MainWindow.ui
View File

@ -179,9 +179,33 @@
<addaction name="actionDelete"/> <addaction name="actionDelete"/>
<addaction name="actionSave"/> <addaction name="actionSave"/>
</widget> </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="menuFile"/>
<addaction name="menuEdit"/> <addaction name="menuEdit"/>
<addaction name="menuView"/> <addaction name="menuView"/>
<addaction name="menuMode"/>
<addaction name="menuTools"/> <addaction name="menuTools"/>
<addaction name="menuAbout"/> <addaction name="menuAbout"/>
</widget> </widget>
@ -258,7 +282,7 @@
</attribute> </attribute>
<widget class="QWidget" name="dockWidgetContents_2"/> <widget class="QWidget" name="dockWidgetContents_2"/>
</widget> </widget>
<widget class="QDockWidget" name="labels_dock"> <widget class="QDockWidget" name="annos_dock">
<property name="minimumSize"> <property name="minimumSize">
<size> <size>
<width>85</width> <width>85</width>
@ -269,7 +293,7 @@
<set>QDockWidget::AllDockWidgetFeatures</set> <set>QDockWidget::AllDockWidgetFeatures</set>
</property> </property>
<property name="windowTitle"> <property name="windowTitle">
<string>Labels</string> <string>Annos</string>
</property> </property>
<attribute name="dockWidgetArea"> <attribute name="dockWidgetArea">
<number>2</number> <number>2</number>
@ -285,6 +309,17 @@
</attribute> </attribute>
<widget class="QWidget" name="dockWidgetContents"/> <widget class="QWidget" name="dockWidgetContents"/>
</widget> </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"> <action name="actionOpen_dir">
<property name="icon"> <property name="icon">
<iconset resource="../icons.qrc"> <iconset resource="../icons.qrc">
@ -693,14 +728,17 @@
</action> </action>
<action name="actionVisible"> <action name="actionVisible">
<property name="icon"> <property name="icon">
<iconset> <iconset resource="../icons.qrc">
<normaloff>:/icon/icons/眼睛_eyes.svg</normaloff>:/icon/icons/眼睛_eyes.svg</iconset> <normaloff>:/icon/icons/眼睛_eyes.svg</normaloff>:/icon/icons/眼睛_eyes.svg</iconset>
</property> </property>
<property name="text"> <property name="text">
<string>Visible</string> <string>Visible</string>
</property> </property>
<property name="toolTip">
<string>Visible.</string>
</property>
<property name="statusTip"> <property name="statusTip">
<string>Visible</string> <string>Visible.</string>
</property> </property>
<property name="shortcut"> <property name="shortcut">
<string>V</string> <string>V</string>
@ -751,6 +789,66 @@
<string>Convert ISAT jsons to LabelMe jsons.</string> <string>Convert ISAT jsons to LabelMe jsons.</string>
</property> </property>
</action> </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> </widget>
<resources> <resources>
<include location="../icons.qrc"/> <include location="../icons.qrc"/>

2
ui/label_dock.py → ui/anno_dock.py
View File

@ -1,6 +1,6 @@
# -*- coding: utf-8 -*- # -*- 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 # Created by: PyQt5 UI code generator 5.15.7
# #

0
ui/label_dock.ui → ui/anno_dock.ui
View File

6
widgets/labels_dock_widget.py → widgets/annos_dock_widget.py
View File

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

112
widgets/canvas.py
View File

@ -2,9 +2,8 @@
# @Author : LG # @Author : LG
from PyQt5 import QtWidgets, QtGui, QtCore from PyQt5 import QtWidgets, QtGui, QtCore
from enum import Enum
from widgets.polygon import Polygon from widgets.polygon import Polygon
from configs import STATUSMode, CLICKMode, DRAWMode from configs import STATUSMode, CLICKMode, DRAWMode, CONTOURMode
from PIL import Image from PIL import Image
import numpy as np import numpy as np
import cv2 import cv2
@ -20,9 +19,11 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.mode = STATUSMode.VIEW self.mode = STATUSMode.VIEW
self.click = CLICKMode.POSITIVE self.click = CLICKMode.POSITIVE
self.draw_mode = DRAWMode.SEGMENTANYTHING # 默认使用segment anything进行快速标注 self.draw_mode = DRAWMode.SEGMENTANYTHING # 默认使用segment anything进行快速标注
self.contour_mode = CONTOURMode.SAVE_EXTERNAL # 默认SAM只保留外轮廓
self.click_points = [] self.click_points = []
self.click_points_mode = [] self.click_points_mode = []
self.masks:np.ndarray = None self.masks:np.ndarray = None
self.mask_alpha = 0.5
self.top_layer = 1 self.top_layer = 1
self.guide_line_x:QtWidgets.QGraphicsLineItem = None 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.image_data = np.repeat(self.image_data, 3, axis=2) # 转换为三通道
self.mainwindow.segany.set_image(self.image_data) self.mainwindow.segany.set_image(self.image_data)
else: 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 = QtWidgets.QGraphicsPixmapItem()
self.image_item.setZValue(0) self.image_item.setZValue(0)
@ -76,7 +77,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.mainwindow.actionSave.setEnabled(False) self.mainwindow.actionSave.setEnabled(False)
self.mainwindow.set_labels_visible(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): def change_mode_to_view(self):
self.mode = STATUSMode.VIEW self.mode = STATUSMode.VIEW
@ -101,7 +102,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.mainwindow.actionSave.setEnabled(self.mainwindow.can_be_annotated) self.mainwindow.actionSave.setEnabled(self.mainwindow.can_be_annotated)
self.mainwindow.set_labels_visible(True) 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): def change_mode_to_edit(self):
self.mode = STATUSMode.EDIT self.mode = STATUSMode.EDIT
@ -128,6 +129,15 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
def change_click_to_negative(self): def change_click_to_negative(self):
self.click = CLICKMode.NEGATIVE 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): def start_segment_anything(self):
self.draw_mode = DRAWMode.SEGMENTANYTHING self.draw_mode = DRAWMode.SEGMENTANYTHING
self.start_draw() self.start_draw()
@ -154,6 +164,11 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.change_mode_to_view() 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: if self.draw_mode == DRAWMode.SEGMENTANYTHING:
# mask to polygon # mask to polygon
# -------------- # --------------
@ -163,17 +178,54 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
h, w = masks.shape[-2:] h, w = masks.shape[-2:]
masks = masks.reshape(h, w) 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] contour = contours[0]
for cont in contours: for cont in contours:
if len(cont) > len(contour): if len(cont) > len(contour):
contour = cont contour = cont
contours = [contour]
for point in contour: for index, contour in enumerate(contours):
x, y = point[0] if self.current_graph is None:
self.current_graph.addPoint(QtCore.QPointF(x, y)) 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: elif self.draw_mode == DRAWMode.POLYGON:
if len(self.current_graph.points) < 1: if len(self.current_graph.points) < 1:
@ -198,9 +250,28 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.current_graph.addPoint(last_point) self.current_graph.addPoint(last_point)
self.current_graph.addPoint(QtCore.QPointF(last_point.x(), first_point.y())) 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清空 # mask清空
self.click_points.clear() self.click_points.clear()
@ -234,7 +305,7 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
for p in self.mainwindow.polygons: for p in self.mainwindow.polygons:
if p.zValue() > deleted_layer: if p.zValue() > deleted_layer:
p.setZValue(p.zValue() - 1) p.setZValue(p.zValue() - 1)
self.mainwindow.labels_dock_widget.update_listwidget() self.mainwindow.annos_dock_widget.update_listwidget()
def edit_polygon(self): def edit_polygon(self):
selectd_items = self.selectedItems() selectd_items = self.selectedItems()
@ -368,18 +439,21 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
def update_mask(self): def update_mask(self):
if not self.mainwindow.use_segment_anything: if not self.mainwindow.use_segment_anything:
return return
if self.image_data is None:
return
if not (self.image_data.ndim == 3 and self.image_data.shape[-1] == 3): if not (self.image_data.ndim == 3 and self.image_data.shape[-1] == 3):
return return
if len(self.click_points) > 0 and len(self.click_points_mode) > 0: 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 self.masks = masks
color = np.array([0, 0, 255]) color = np.array([0, 0, 255])
h, w = masks.shape[-2:] h, w = masks.shape[-2:]
mask_image = masks.reshape(h, w, 1) * color.reshape(1, 1, -1) mask_image = masks.reshape(h, w, 1) * color.reshape(1, 1, -1)
mask_image = mask_image.astype("uint8") mask_image = mask_image.astype("uint8")
mask_image = cv2.cvtColor(mask_image, cv2.COLOR_BGR2RGB) 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, mask_image = QtGui.QImage(mask_image[:], mask_image.shape[1], mask_image.shape[0], mask_image.shape[1] * 3,
QtGui.QImage.Format_RGB888) QtGui.QImage.Format_RGB888)
mask_pixmap = QtGui.QPixmap(mask_image) mask_pixmap = QtGui.QPixmap(mask_image)

2
widgets/category_choice_dialog.py
View File

@ -88,7 +88,7 @@ class CategoryChoiceDialog(QtWidgets.QDialog, Ui_Dialog):
for vertex in self.scene.current_graph.vertexs: for vertex in self.scene.current_graph.vertexs:
vertex.setZValue(len(self.mainwindow.polygons)) 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.current_graph = None
self.scene.change_mode_to_view() self.scene.change_mode_to_view()

2
widgets/category_edit_dialog.py
View File

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

103
widgets/mainwindow.py
View File

@ -6,7 +6,8 @@ from ui.MainWindow import Ui_MainWindow
from widgets.setting_dialog import SettingDialog from widgets.setting_dialog import SettingDialog
from widgets.category_choice_dialog import CategoryChoiceDialog from widgets.category_choice_dialog import CategoryChoiceDialog
from widgets.category_edit_dialog import CategoryEditDialog 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.files_dock_widget import FilesDockWidget
from widgets.info_dock_widget import InfoDockWidget from widgets.info_dock_widget import InfoDockWidget
from widgets.right_button_menu import RightButtonMenu from widgets.right_button_menu import RightButtonMenu
@ -34,6 +35,8 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
super(MainWindow, self).__init__() super(MainWindow, self).__init__()
self.setupUi(self) self.setupUi(self)
self.init_ui() self.init_ui()
self.init_segment_anything()
self.image_root: str = None self.image_root: str = None
self.label_root:str = None self.label_root:str = None
@ -41,6 +44,9 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.current_index = None self.current_index = None
self.current_file_index: int = 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.config_file = CONFIG_FILE if os.path.exists(CONFIG_FILE) else DEFAULT_CONFIG_FILE
self.saved = True self.saved = True
self.can_be_annotated = True self.can_be_annotated = True
@ -57,7 +63,6 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.init_connect() self.init_connect()
self.reset_action() self.reset_action()
self.init_segment_anything()
def init_segment_anything(self): def init_segment_anything(self):
if os.path.exists('./segment_any/sam_vit_h_4b8939.pth'): 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.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_widget = FilesDockWidget(mainwindow=self)
self.files_dock.setWidget(self.files_dock_widget) self.files_dock.setWidget(self.files_dock_widget)
@ -135,6 +143,18 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.labelData.setFixedWidth(150) self.labelData.setFixedWidth(150)
self.statusbar.addPermanentWidget(self.labelData) 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() self.trans = QtCore.QTranslator()
def translate(self, language='zh'): def translate(self, language='zh'):
@ -148,7 +168,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
_app.installTranslator(self.trans) _app.installTranslator(self.trans)
self.retranslateUi(self) self.retranslateUi(self)
self.info_dock_widget.retranslateUi(self.info_dock_widget) 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.files_dock_widget.retranslateUi(self.files_dock_widget)
self.category_choice_widget.retranslateUi(self.category_choice_widget) self.category_choice_widget.retranslateUi(self.category_choice_widget)
self.category_edit_widget.retranslateUi(self.category_edit_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) self.show_image(self.current_index)
language = self.cfg.get('language', 'en') language = self.cfg.get('language', 'en')
self.cfg['language'] = language
self.translate(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): def set_saved_state(self, is_saved:bool):
self.saved = is_saved self.saved = is_saved
if self.files_list is not None and self.current_index is not None: 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 return
try: try:
self.polygons.clear() self.polygons.clear()
self.labels_dock_widget.listWidget.clear() self.annos_dock_widget.listWidget.clear()
self.scene.cancel_draw() self.scene.cancel_draw()
file_path = os.path.join(self.image_root, self.files_list[index]) file_path = os.path.join(self.image_root, self.files_list[index])
image_data = Image.open(file_path) image_data = Image.open(file_path)
@ -292,6 +323,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
# load label # load label
if self.can_be_annotated: if self.can_be_annotated:
self.current_group = 1
_, name = os.path.split(file_path) _, name = os.path.split(file_path)
label_path = os.path.join(self.label_root, '.'.join(name.split('.')[:-1]) + '.json') label_path = os.path.join(self.label_root, '.'.join(name.split('.')[:-1]) + '.json')
self.current_label = Annotation(file_path, label_path) self.current_label = Annotation(file_path, label_path)
@ -299,6 +331,11 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.current_label.load_annotation() self.current_label.load_annotation()
for object in self.current_label.objects: 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() polygon = Polygon()
self.scene.addItem(polygon) self.scene.addItem(polygon)
polygon.load_object(object) polygon.load_object(object)
@ -309,7 +346,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
else: else:
self.setWindowTitle('{}'.format(file_path)) self.setWindowTitle('{}'.format(file_path))
self.labels_dock_widget.update_listwidget() self.annos_dock_widget.update_listwidget()
self.info_dock_widget.update_widget() self.info_dock_widget.update_widget()
self.files_dock_widget.set_select(index) self.files_dock_widget.set_select(index)
self.current_index = 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.change_color(QtGui.QColor(self.category_color_dict.get(polygon.category, '#000000')))
polygon.color.setAlpha(255) polygon.color.setAlpha(255)
polygon.setBrush(polygon.color) 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.actionSegment_anything.setEnabled(False)
self.actionPolygon.setEnabled(False) self.actionPolygon.setEnabled(False)
self.actionVisible.setEnabled(False) self.actionVisible.setEnabled(False)
@ -428,14 +465,16 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
for vertex in polygon.vertexs: for vertex in polygon.vertexs:
vertex.setVisible(False) vertex.setVisible(False)
if polygon.group != '': 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: else:
rgb = self.instance_cmap[0] rgb = self.instance_cmap[0]
polygon.change_color(QtGui.QColor(rgb[0], rgb[1], rgb[2], 255)) polygon.change_color(QtGui.QColor(rgb[0], rgb[1], rgb[2], 255))
polygon.color.setAlpha(255) polygon.color.setAlpha(255)
polygon.setBrush(polygon.color) 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.actionSegment_anything.setEnabled(False)
self.actionPolygon.setEnabled(False) self.actionPolygon.setEnabled(False)
self.actionVisible.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.change_color(QtGui.QColor(self.category_color_dict.get(polygon.category, '#000000')))
polygon.color.setAlpha(polygon.nohover_alpha) polygon.color.setAlpha(polygon.nohover_alpha)
polygon.setBrush(polygon.color) 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.actionSegment_anything.setEnabled(self.use_segment_anything)
self.actionPolygon.setEnabled(True) self.actionPolygon.setEnabled(True)
self.actionVisible.setEnabled(True) self.actionVisible.setEnabled(True)
@ -468,9 +507,33 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
def set_labels_visible(self, visible=None): def set_labels_visible(self, visible=None):
if visible is 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): def ISAT_to_VOC(self):
self.ISAT_to_VOC_dialog.reset_gui() 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.actionBit_map.triggered.connect(self.change_bit_map)
self.actionVisible.triggered.connect(functools.partial(self.set_labels_visible, None)) 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.actionToVOC.triggered.connect(self.ISAT_to_VOC)
self.actionToCOCO.triggered.connect(self.ISAT_to_COCO) self.actionToCOCO.triggered.connect(self.ISAT_to_COCO)
self.actionTo_LabelMe.triggered.connect(self.ISAT_to_LABELME) 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.actionChinese.triggered.connect(self.translate_to_chinese)
self.actionEnglish.triggered.connect(self.translate_to_english) 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): def reset_action(self):
self.actionPrev.setEnabled(False) self.actionPrev.setEnabled(False)

2
widgets/polygon.py
View File

@ -133,7 +133,7 @@ class Polygon(QtWidgets.QGraphicsPolygonItem):
else: else:
self.color.setAlpha(self.nohover_alpha) self.color.setAlpha(self.nohover_alpha)
self.setBrush(self.color) 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 if change == QtWidgets.QGraphicsItem.GraphicsItemChange.ItemPositionChange: # ItemPositionHasChanged
bias = value bias = value