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:
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)

5
configs.py
View File

@ -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 # 保留所有轮廓

2
default.yaml
View File

@ -1,3 +1,5 @@
label:
- color: '#000000'
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:
- color: '#000000'
name: __background__
@ -19,3 +20,5 @@ label:
name: cake
- color: '#5c3566'
name: fence
language: en
mask_alpha: 0.6

1
icons.qrc
View File

@ -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>

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:
- color: '#000000'
name: __background__
@ -19,4 +20,5 @@ label:
name: cake
- color: '#5c3566'
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 -*-
# @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()

74
ui/MainWindow.py
View File

@ -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.labels_dock.setMinimumSize(QtCore.QSize(85, 43))
self.labels_dock.setFeatures(QtWidgets.QDockWidget.AllDockWidgetFeatures)
self.labels_dock.setObjectName("labels_dock")
self.annos_dock = QtWidgets.QDockWidget(MainWindow)
self.annos_dock.setMinimumSize(QtCore.QSize(85, 43))
self.annos_dock.setFeatures(QtWidgets.QDockWidget.AllDockWidgetFeatures)
self.annos_dock.setObjectName("annos_dock")
self.dockWidgetContents_3 = QtWidgets.QWidget()
self.dockWidgetContents_3.setObjectName("dockWidgetContents_3")
self.labels_dock.setWidget(self.dockWidgetContents_3)
MainWindow.addDockWidget(QtCore.Qt.DockWidgetArea(2), self.labels_dock)
self.annos_dock.setWidget(self.dockWidgetContents_3)
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

106
ui/MainWindow.ui
View File

@ -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"/>

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

@ -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
#

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
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 = {}

98
widgets/canvas.py
View File

@ -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.show()
# self.mainwindow.category_choice_widget.load_cfg()
# 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)

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:
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()

2
widgets/category_edit_dialog.py
View File

@ -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()

103
widgets/mainwindow.py
View File

@ -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.labels_dock.setWidget(self.labels_dock_widget)
self.categories_dock_widget = CategoriesDockWidget(mainwindow=self)
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.labels_dock_widget.checkBox_visible.setEnabled(False)
self.annos_dock_widget.listWidget.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.labels_dock_widget.checkBox_visible.setEnabled(False)
self.annos_dock_widget.listWidget.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.labels_dock_widget.checkBox_visible.setEnabled(True)
self.annos_dock_widget.listWidget.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()
self.labels_dock_widget.checkBox_visible.setChecked(visible)
self.labels_dock_widget.set_all_polygon_visible(visible)
visible = not self.annos_dock_widget.checkBox_visible.isChecked()
self.annos_dock_widget.checkBox_visible.setChecked(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)

2
widgets/polygon.py
View File

@ -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