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

update

Browse Source
This commit is contained in:
yatenglg 2023-04-27 21:59:14 +08:00
parent 2347f7d402
commit 5e76d380c6
16 changed files with 3068 additions and 1148 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
UpdateLog.md
View File

@ -8,3 +8,14 @@
- 添加了显示/隐藏按钮（快捷键V），用于显示或隐藏所有多边形
- 添加了GPU显存占用
#
- 标注时隐藏所有多边形
- 修改windows中文路径下，图片打开的bug
- bug修复
#
- 添加转换voc格式png图片的功能 finish
- 添加转换coco格式json的功能 ing
- 添加转换coco格式json转ISAT格式json的功能 ing

27
annotation.py
View File

@ -2,7 +2,8 @@
# @Author : LG
import os
import cv2
from PIL import Image
import numpy as np
from json import load, dump
from typing import List
@ -27,8 +28,16 @@ class Annotation:
self.img_name = img_name
self.label_path = label_path
self.note = ''
image = cv2.imread(image_path)
self.height, self.width, self.depth = image.shape
image = np.array(Image.open(image_path))
if image.ndim == 3:
self.height, self.width, self.depth = image.shape
elif image.ndim == 2:
self.height, self.width = image.shape
self.depth = 0
else:
self.height, self.width, self.depth = image.shape[:, :3]
print('Warning: Except image has 2 or 3 ndim, but get {}.'.format(image.ndim))
del image
self.objects:List[Object,...] = []
@ -43,9 +52,15 @@ class Annotation:
# ISAT格式json
objects = dataset.get('objects', [])
self.img_name = info.get('name', '')
self.width = info.get('width', 0)
self.height = info.get('height', 0)
self.depth = info.get('depth', 0)
width = info.get('width', None)
if width is not None:
self.width = width
height = info.get('height', None)
if height is not None:
self.height = height
depth = info.get('depth', None)
if depth is not None:
self.depth = depth
self.note = info.get('note', '')
for obj in objects:
category = obj.get('category', 'unknow')

2
icons.qrc
View File

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

3737
icons_rc.py
View File

File diff suppressed because it is too large Load Diff

2
isat.yaml
View File

@ -19,4 +19,4 @@ label:
name: cake
- color: '#5c3566'
name: fence
language: en
language: zh

2
segment_any/segment_any.py
View File

@ -24,7 +24,7 @@ class SegAny:
self.image = None
def set_image(self, image):
self.image = image
self.predictor.set_image(image)
def reset_image(self):

3
tools/label_convert.py → tools/toVOC.py
View File

@ -7,7 +7,8 @@ from PIL import Image, ImageDraw, ImageColor
import mahotas
import imgviz
class Converter:
class ToVOC:
def __init__(self, cfg, is_segmentation):
self.is_segmentation = is_segmentation

200
tools/tococo.py Normal file
View File

@ -0,0 +1,200 @@
# -*- coding: utf-8 -*-
# @Author : LG
from json import load, dump
import os
from pycocotools import mask as coco_mask
import cv2
import imgviz
class COCOConverter:
def __init__(self, ):
pass
# labels = cfg.get('label', [])
# for index, label_dict in enumerate(labels):
# category = label_dict.get('name', 'unknow')
# color = label_dict.get('color', '#000000')
def convert_to_coco(self, isat_json_root:str, to_path:str):
coco_anno = {}
# info
coco_anno['info'] = {}
coco_anno['info']['description'] = 'coco from ISAT'
coco_anno['info']['version'] = None
coco_anno['info']['year'] = None
coco_anno['info']['contributor'] = None
coco_anno['info']['date_created'] = None
# licenses
coco_anno['licenses'] = []
coco_anno['licenses'][0] = {}
coco_anno['licenses'][0]['url'] = None
coco_anno['licenses'][0]['id'] = 0
coco_anno['licenses'][0]['name'] = None
# images and annotations
coco_anno['images'] = []
coco_anno['annotations'] = []
jsons = [f for f in os.listdir(isat_json_root) if f.endswith('.json')]
for json in jsons:
coco_image_info = {}
dataset = load(os.path.join(isat_json_root, json))
info = dataset.get('info', {})
description = info.get('description', '')
if not description.startswith('ISAT'):
# 不是ISAT格式json
continue
img_name = info.get('name', '')
width = info.get('width', None)
height = info.get('height', None)
depth = info.get('depth', None)
note = info.get('note', '')
objects = dataset.get('objects', [])
for obj in objects:
category = obj.get('category', 'unknow')
group = obj.get('group', 0)
if group is None: group = 0
segmentation = obj.get('segmentation', [])
iscrowd = obj.get('iscrowd', 0)
note = obj.get('note', '')
area = obj.get('area', 0)
layer = obj.get('layer', 2)
bbox = obj.get('bbox', [])
def convert_from_coco(self, coco_json:str, to_path:str, keep_crowd:bool=False):
assert coco_json.endswith('.json')
annos = {}
if os.path.exists(coco_json):
with open(coco_json, 'r') as f:
dataset = load(f)
images = {image.get('id', None):{
'file_name': image.get('file_name', ''),
'height': image.get('height', ''),
'width': image.get('width', ''),
} for image in dataset.get('images', [])}
annotations = dataset.get('annotations', [])
categories = {categorie.get('id', None): {'name': categorie.get('name', '')} for categorie in dataset.get('categories', [])}
for index, annotation in enumerate(annotations):
annotation_index = annotation.get('id')
annotation_image_id = annotation.get('image_id')
annotation_category_id = annotation.get('category_id')
file_name = images[annotation_image_id].get('file_name')
height = images[annotation_image_id].get('height')
width = images[annotation_image_id].get('width')
iscrowd = annotation["iscrowd"]
if file_name == '000000279278.jpg':
continue
if annotation_image_id not in annos:
annos[annotation_image_id] = {}
objects = annos[annotation_image_id].get('objects', [])
if iscrowd == 0:
# polygon
segmentations = annotation.get('segmentation')
for segmentation in segmentations:
xs = segmentation[::2]
ys = segmentation[1::2]
points = [[x, y] for x ,y in zip(xs, ys)]
obj = {
'category': categories.get(annotation_category_id).get('name'),
'group': annotation_index,
'area': None,
'segmentation': points,
'layer': 1,
'bbox': None,
'iscrowd': iscrowd,
'note': ''
}
objects.append(obj)
elif iscrowd == 1 and keep_crowd:
# RLE
segmentations = annotation.get('segmentation', {})
rles = coco_mask.frPyObjects(segmentations, height, width)
masks = coco_mask.decode(rles)
contours, _ = cv2.findContours(masks, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
for contour in contours:
points = []
for point in contour:
x, y = point[0]
points.append([float(x), float(y)])
obj = {
'category': categories.get(annotation_category_id).get('name'),
'group': annotation_index,
'area': None,
'segmentation': points,
'layer': 1,
'bbox': None,
'iscrowd': iscrowd,
'note': ''
}
objects.append(obj)
else:
pass
annos[annotation_image_id]['objects'] = objects
for image_id, values in annos.items():
image_path = images[image_id].get('file_name')
folder, name = os.path.split(image_path)
height = images[image_id].get('height')
width = images[image_id].get('width')
objects = values.get('objects', [])
isat_anno = {}
isat_anno['info'] = {}
isat_anno['info']['description'] = 'ISAT'
isat_anno['info']['folder'] = folder
isat_anno['info']['name'] = name
isat_anno['info']['width'] = width
isat_anno['info']['height'] = height
isat_anno['info']['depth'] = None
isat_anno['info']['note'] = ''
isat_anno['objects'] = []
# coco annotation的id 太大了，这里缩一下，每张图片重新开始计数
groups_dict = {}
for obj in objects:
group = obj.get('group', 0)
if group not in groups_dict:
groups_dict[group] = len(groups_dict)+1
for obj in objects:
object = {}
object['category'] = obj.get('category', '')
object['group'] = groups_dict.get(obj.get('group', 0))
object['segmentation'] = obj.get('segmentation', [])
object['area'] = obj.get('area', None)
object['layer'] = obj.get('layer', None)
object['bbox'] = obj.get('bbox', None)
object['iscrowd'] = obj.get('iscrowd', 0)
object['note'] = obj.get('note', '')
isat_anno['objects'].append(object)
json_name = '.'.join(name.split('.')[:-1]) + '.json'
save_json = os.path.join(to_path, json_name)
with open(save_json, 'w') as f:
try:
dump(isat_anno, f)
print('Converted coco to ISAT: {}'.format(json_name))
except Exception as e:
print('Convert coco to ISAT {} ,error: {}'.format(json_name, e))
### 类别文件
cmap = imgviz.label_colormap()
if __name__ == '__main__':
coco = COCOConverter()
coco.convert_from_coco(
'/mnt/disk/coco/annotations/instances_val2017.json',
'/mnt/disk/coco/isat_json',
True
)

8
ui/convert_dialog.py → ui/ISAT_to_VOC_dialog.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/convert_dialog.ui'
# Form implementation generated from reading ui file 'ISAT_to_VOC_dialog.ui'
#
# Created by: PyQt5 UI code generator 5.15.7
#
@ -15,7 +15,7 @@ class Ui_Dialog(object):
def setupUi(self, Dialog):
Dialog.setObjectName("Dialog")
Dialog.setWindowModality(QtCore.Qt.NonModal)
Dialog.resize(450, 175)
Dialog.resize(509, 175)
font = QtGui.QFont()
font.setFamily("Times New Roman")
font.setPointSize(12)
@ -103,7 +103,7 @@ class Ui_Dialog(object):
def retranslateUi(self, Dialog):
_translate = QtCore.QCoreApplication.translate
Dialog.setWindowTitle(_translate("Dialog", "ISAT to png converter"))
Dialog.setWindowTitle(_translate("Dialog", "ISAT to VOC png"))
self.lineEdit_save_root.setPlaceholderText(_translate("Dialog", "png save root"))
self.pushButton_save_root.setText(_translate("Dialog", "Save root"))
self.checkBox_is_instance.setText(_translate("Dialog", "Is Instance"))
@ -111,7 +111,7 @@ class Ui_Dialog(object):
self.pushButton_label_root.setText(_translate("Dialog", "Label root"))
self.label_3.setAccessibleName(_translate("Dialog", "aaa"))
self.label_3.setText(_translate("Dialog", "/"))
self.label.setText(_translate("Dialog", "Convert ISAT annotations to png."))
self.label.setText(_translate("Dialog", "Convert ISAT annotations to VOC png."))
self.pushButton_cache.setText(_translate("Dialog", "cache"))
self.pushButton_apply.setText(_translate("Dialog", "convert"))
import icons_rc

10
ui/convert_dialog.ui → ui/ISAT_to_VOC_dialog.ui
View File

@ -9,7 +9,7 @@
<rect>
<x>0</x>
<y>0</y>
<width>450</width>
<width>509</width>
<height>175</height>
</rect>
</property>
@ -20,7 +20,7 @@
</font>
</property>
<property name="windowTitle">
<string>ISAT to png converter</string>
<string>ISAT to VOC png</string>
</property>
<property name="sizeGripEnabled">
<bool>false</bool>
@ -173,7 +173,7 @@
<string notr="true">color: rgb(255, 0, 0);</string>
</property>
<property name="text">
<string>Convert ISAT annotations to png.</string>
<string>Convert ISAT annotations to VOC png.</string>
</property>
</widget>
</item>
@ -196,7 +196,7 @@
<string>cache</string>
</property>
<property name="icon">
<iconset resource="../icons.qrc">
<iconset>
<normaloff>:/icons/icons/关闭_close-one.svg</normaloff>:/icons/icons/关闭_close-one.svg</iconset>
</property>
</widget>
@ -207,7 +207,7 @@
<string>convert</string>
</property>
<property name="icon">
<iconset resource="../icons.qrc">
<iconset>
<normaloff>:/icons/icons/校验_check-one.svg</normaloff>:/icons/icons/校验_check-one.svg</iconset>
</property>
</widget>

25
ui/MainWindow.py
View File

@ -210,12 +210,12 @@ class Ui_MainWindow(object):
icon18.addPixmap(QtGui.QPixmap(":/icon/icons/去底部_to-bottom.svg"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
self.actionTo_bottom.setIcon(icon18)
self.actionTo_bottom.setObjectName("actionTo_bottom")
self.actionConverter = QtWidgets.QAction(MainWindow)
self.actionToVOC = QtWidgets.QAction(MainWindow)
icon19 = QtGui.QIcon()
icon19.addPixmap(QtGui.QPixmap(":/icon/icons/转换文件夹1_folder-conversion-one.svg"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
self.actionConverter.setIcon(icon19)
self.actionConverter.setWhatsThis("")
self.actionConverter.setObjectName("actionConverter")
icon19.addPixmap(QtGui.QPixmap(":/icon/icons/voc.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
self.actionToVOC.setIcon(icon19)
self.actionToVOC.setWhatsThis("")
self.actionToVOC.setObjectName("actionToVOC")
self.actionChinese = QtWidgets.QAction(MainWindow)
self.actionChinese.setCheckable(True)
font = QtGui.QFont()
@ -255,6 +255,11 @@ class Ui_MainWindow(object):
icon24.addPixmap(QtGui.QPixmap(":/icon/icons/眼睛_eyes.svg"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
self.actionVisible.setIcon(icon24)
self.actionVisible.setObjectName("actionVisible")
self.actionToCOCO = QtWidgets.QAction(MainWindow)
icon25 = QtGui.QIcon()
icon25.addPixmap(QtGui.QPixmap(":/icon/icons/coco.ico"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
self.actionToCOCO.setIcon(icon25)
self.actionToCOCO.setObjectName("actionToCOCO")
self.menuFile.addAction(self.actionOpen_dir)
self.menuFile.addAction(self.actionSave_dir)
self.menuFile.addSeparator()
@ -275,7 +280,8 @@ class Ui_MainWindow(object):
self.menuAbout.addAction(self.menuLaguage.menuAction())
self.menuAbout.addAction(self.actionShortcut)
self.menuAbout.addAction(self.actionAbout)
self.menuTools.addAction(self.actionConverter)
self.menuTools.addAction(self.actionToVOC)
self.menuTools.addAction(self.actionToCOCO)
self.menuEdit.addAction(self.actionSegment_anything)
self.menuEdit.addAction(self.actionPolygon)
self.menuEdit.addAction(self.actionBackspace)
@ -383,8 +389,9 @@ class Ui_MainWindow(object):
self.actionTo_bottom.setToolTip(_translate("MainWindow", "Move polygon to bottom layer"))
self.actionTo_bottom.setStatusTip(_translate("MainWindow", "Move polygon to bottom layer."))
self.actionTo_bottom.setShortcut(_translate("MainWindow", "B"))
self.actionConverter.setText(_translate("MainWindow", "Label converter"))
self.actionConverter.setStatusTip(_translate("MainWindow", "Convert annotations to png image."))
self.actionToVOC.setText(_translate("MainWindow", "ToVOC"))
self.actionToVOC.setToolTip(_translate("MainWindow", "Convert ISAT to VOC"))
self.actionToVOC.setStatusTip(_translate("MainWindow", "Convert annotations to png image."))
self.actionChinese.setText(_translate("MainWindow", "中文"))
self.actionEnglish.setText(_translate("MainWindow", "English"))
self.actionBackspace.setText(_translate("MainWindow", "Backspace"))
@ -406,4 +413,6 @@ class Ui_MainWindow(object):
self.actionVisible.setText(_translate("MainWindow", "Visible"))
self.actionVisible.setStatusTip(_translate("MainWindow", "Visible"))
self.actionVisible.setShortcut(_translate("MainWindow", "V"))
self.actionToCOCO.setText(_translate("MainWindow", "ToCOCO"))
self.actionToCOCO.setToolTip(_translate("MainWindow", "Convert ISAT to COCO"))
import icons_rc

24
ui/MainWindow.ui
View File

@ -150,7 +150,8 @@
<property name="title">
<string>Tools</string>
</property>
<addaction name="actionConverter"/>
<addaction name="actionToVOC"/>
<addaction name="actionToCOCO"/>
</widget>
<widget class="QMenu" name="menuEdit">
<property name="font">
@ -569,13 +570,16 @@
<string>B</string>
</property>
</action>
<action name="actionConverter">
<action name="actionToVOC">
<property name="icon">
<iconset resource="../icons.qrc">
<normaloff>:/icon/icons/转换文件夹1_folder-conversion-one.svg</normaloff>:/icon/icons/转换文件夹1_folder-conversion-one.svg</iconset>
<normaloff>:/icon/icons/voc.png</normaloff>:/icon/icons/voc.png</iconset>
</property>
<property name="text">
<string>Label converter</string>
<string>ToVOC</string>
</property>
<property name="toolTip">
<string>Convert ISAT to VOC</string>
</property>
<property name="statusTip">
<string>Convert annotations to png image.</string>
@ -699,6 +703,18 @@
<string>V</string>
</property>
</action>
<action name="actionToCOCO">
<property name="icon">
<iconset resource="../icons.qrc">
<normaloff>:/icon/icons/coco.ico</normaloff>:/icon/icons/coco.ico</iconset>
</property>
<property name="text">
<string>ToCOCO</string>
</property>
<property name="toolTip">
<string>Convert ISAT to COCO</string>
</property>
</action>
</widget>
<resources>
<include location="../icons.qrc"/>

86
widgets/ISAT_to_VOC_dialog.py Normal file
View File

@ -0,0 +1,86 @@
# -*- coding: utf-8 -*-
# @Author : LG
from PyQt5 import QtWidgets, QtCore, QtGui
from ui.ISAT_to_VOC_dialog import Ui_Dialog
from tools.toVOC import ToVOC
import os
class ISATtoVOCDialog(QtWidgets.QDialog, Ui_Dialog):
def __init__(self, parent, mainwindow):
super(ISATtoVOCDialog, self).__init__(parent)
self.setupUi(self)
self.mainwindow = mainwindow
self.label_root = None
self.save_root = None
self.pause = False
self.setWindowModality(QtCore.Qt.WindowModality.WindowModal)
self.init_connect()
def reset_gui(self):
self.widget_process.setVisible(False)
self.lineEdit_label_root.clear()
self.lineEdit_save_root.clear()
def _label_root(self):
dir = QtWidgets.QFileDialog.getExistingDirectory(self)
if dir:
self.label_root = dir
self.lineEdit_label_root.setText(dir)
else:
self.lineEdit_label_root.clear()
def _save_root(self):
dir = QtWidgets.QFileDialog.getExistingDirectory(self)
if dir:
self.save_root = dir
self.lineEdit_save_root.setText(dir)
else:
self.lineEdit_save_root.clear()
def cache(self):
self.pause = True
self.close()
def apply(self):
self.pause = False
if self.label_root is None or self.save_root is None:
return
# 语义分割，保存类别文件
if not self.checkBox_is_instance.isChecked():
with open(os.path.join(self.save_root, 'classesition.txt'), 'w') as f:
for index, label in enumerate(self.mainwindow.cfg.get('label', [])):
f.write('{} {}\n'.format(label.get('name'), index))
converter = ToVOC(self.mainwindow.cfg, self.checkBox_is_instance.isChecked())
jsons = [f for f in os.listdir(self.label_root) if f.endswith('.json')]
self.pushButton_label_root.setEnabled(False)
self.pushButton_save_root.setEnabled(False)
self.checkBox_is_instance.setEnabled(False)
self.widget_process.setVisible(True)
self.progressBar.setMaximum(len(jsons))
self.all_num.setText('{}'.format(len(jsons)))
for index, json in enumerate(jsons):
if self.pause:
break
label_path = os.path.join(self.label_root, json)
save_path = os.path.join(self.save_root, json[:-5]+'.png')
converter.convert(label_path, save_path)
self.progressBar.setValue(index+1)
self.current_num.setText('{}'.format(index+1))
self.pushButton_label_root.setEnabled(True)
self.pushButton_save_root.setEnabled(True)
self.checkBox_is_instance.setEnabled(True)
def init_connect(self):
self.pushButton_label_root.clicked.connect(self._label_root)
self.pushButton_save_root.clicked.connect(self._save_root)
self.pushButton_apply.clicked.connect(self.apply)
self.pushButton_cache.clicked.connect(self.cache)

24
widgets/canvas.py
View File

@ -30,18 +30,15 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
def load_image(self, image_path:str):
self.clear()
self.image_data = np.array(Image.open(image_path))
if self.mainwindow.use_segment_anything:
self.mainwindow.segany.reset_image()
self.image_data = np.array(Image.open(image_path))
if self.mainwindow.use_segment_anything and self.mainwindow.can_be_annotated:
if self.image_data.ndim == 3 and self.image_data.shape[-1] == 3:
self.mainwindow.segany.set_image(self.image_data)
elif self.image_data.ndim == 2 and image_path.endswith('.png'):
# 单通道图标签图
pass
else:
QtWidgets.QMessageBox.warning(self.mainwindow, 'Warning', 'Segment anything only support 3 channel rgb image.')
self.mainwindow.statusbar.showMessage('Segment anything only support 3 channel rgb image.')
self.image_item = QtWidgets.QGraphicsPixmapItem()
self.image_item.setZValue(0)
@ -74,6 +71,9 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.mainwindow.actionDelete.setEnabled(False)
self.mainwindow.actionSave.setEnabled(False)
self.mainwindow.set_labels_visible(False)
self.mainwindow.labels_dock_widget.setEnabled(False)
def change_mode_to_view(self):
self.mode = STATUSMode.VIEW
self.image_item.setCursor(QtGui.QCursor(QtCore.Qt.CursorShape.ArrowCursor))
@ -81,8 +81,11 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.mainwindow.actionPrev.setEnabled(True)
self.mainwindow.actionNext.setEnabled(True)
self.mainwindow.actionSegment_anything.setEnabled(self.mainwindow.use_segment_anything)
self.mainwindow.actionPolygon.setEnabled(True)
self.mainwindow.actionSegment_anything.setEnabled(self.mainwindow.use_segment_anything and self.mainwindow.can_be_annotated)
if self.mainwindow.use_segment_anything and self.mainwindow.segany.image is None:
self.mainwindow.actionSegment_anything.setEnabled(False)
self.mainwindow.actionPolygon.setEnabled(self.mainwindow.can_be_annotated)
self.mainwindow.actionBackspace.setEnabled(False)
self.mainwindow.actionFinish.setEnabled(False)
self.mainwindow.actionCancel.setEnabled(False)
@ -91,7 +94,10 @@ class AnnotationScene(QtWidgets.QGraphicsScene):
self.mainwindow.actionTo_bottom.setEnabled(False)
self.mainwindow.actionEdit.setEnabled(False)
self.mainwindow.actionDelete.setEnabled(False)
self.mainwindow.actionSave.setEnabled(True)
self.mainwindow.actionSave.setEnabled(self.mainwindow.can_be_annotated)
self.mainwindow.set_labels_visible(True)
self.mainwindow.labels_dock_widget.setEnabled(True)
def change_mode_to_edit(self):
self.mode = STATUSMode.EDIT

10
widgets/converter.py
View File

@ -2,14 +2,14 @@
# @Author : LG
from PyQt5 import QtWidgets, QtCore, QtGui
from ui.convert_dialog import Ui_Dialog
from tools.label_convert import Converter
from ui.ISAT_to_VOC_dialog import Ui_Dialog
from tools.toVOC import ToVOC
import os
class ConvertDialog(QtWidgets.QDialog, Ui_Dialog):
class TOVOCDialog(QtWidgets.QDialog, Ui_Dialog):
def __init__(self, parent, mainwindow):
super(ConvertDialog, self).__init__(parent)
super(TOVOCDialog, self).__init__(parent)
self.setupUi(self)
self.mainwindow = mainwindow
self.label_root = None
@ -55,7 +55,7 @@ class ConvertDialog(QtWidgets.QDialog, Ui_Dialog):
for index, label in enumerate(self.mainwindow.cfg.get('label', [])):
f.write('{} {}\n'.format(label.get('name'), index))
converter = Converter(self.mainwindow.cfg, self.checkBox_is_instance.isChecked())
converter = ToVOC(self.mainwindow.cfg, self.checkBox_is_instance.isChecked())
jsons = [f for f in os.listdir(self.label_root) if f.endswith('.json')]
self.pushButton_label_root.setEnabled(False)

45
widgets/mainwindow.py
View File

@ -12,7 +12,7 @@ from widgets.info_dock_widget import InfoDockWidget
from widgets.right_button_menu import RightButtonMenu
from widgets.shortcut_dialog import ShortcutDialog
from widgets.about_dialog import AboutDialog
from widgets.converter import ConvertDialog
from widgets.ISAT_to_VOC_dialog import ISATtoVOCDialog
from widgets.canvas import AnnotationScene, AnnotationView
from configs import STATUSMode, MAPMode, load_config, save_config, CONFIG_FILE, DEFAULT_CONFIG_FILE
from annotation import Object, Annotation
@ -40,6 +40,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.config_file = CONFIG_FILE if os.path.exists(CONFIG_FILE) else DEFAULT_CONFIG_FILE
self.saved = True
self.can_be_annotated = True
self.load_finished = False
self.polygons:list = []
@ -83,6 +84,9 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.labelGPUResource.setText('segment anything unused.')
def init_ui(self):
# 待完成
self.actionToCOCO.setEnabled(False)
#
self.setting_dialog = SettingDialog(parent=self, mainwindow=self)
self.labels_dock_widget = LabelsDockWidget(mainwindow=self)
@ -98,7 +102,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.category_choice_widget = CategoryChoiceDialog(self, mainwindow=self, scene=self.scene)
self.category_edit_widget = CategoryEditDialog(self, self, self.scene)
self.convert_dialog = ConvertDialog(self, mainwindow=self)
self.ISAT_to_VOC_dialog = ISATtoVOCDialog(self, mainwindow=self)
self.view = AnnotationView(parent=self)
self.view.setScene(self.scene)
@ -147,7 +151,7 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.setting_dialog.retranslateUi(self.setting_dialog)
self.about_dialog.retranslateUi(self.about_dialog)
self.shortcut_dialog.retranslateUi(self.shortcut_dialog)
self.convert_dialog.retranslateUi(self.convert_dialog)
self.ISAT_to_VOC_dialog.retranslateUi(self.ISAT_to_VOC_dialog)
def translate_to_chinese(self):
self.translate('zh')
@ -248,22 +252,37 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
image_data = Image.open(file_path)
self.png_palette = image_data.getpalette()
if self.png_palette is not None:
if self.png_palette is not None and file_path.endswith('.png'):
self.statusbar.showMessage('This is a label file.')
self.actionSegment_anything.setEnabled(False)
self.actionPolygon.setEnabled(False)
self.actionSave.setEnabled(False)
self.actionBit_map.setEnabled(False)
self.can_be_annotated = False
else:
self.can_be_annotated = True
if self.can_be_annotated:
self.actionSegment_anything.setEnabled(self.use_segment_anything)
self.actionPolygon.setEnabled(True)
self.actionSave.setEnabled(True)
self.actionBit_map.setEnabled(True)
self.actionBackspace.setEnabled(True)
self.actionFinish.setEnabled(True)
self.actionCancel.setEnabled(True)
self.actionVisible.setEnabled(True)
else:
self.actionSegment_anything.setEnabled(False)
self.actionPolygon.setEnabled(False)
self.actionSave.setEnabled(False)
self.actionBit_map.setEnabled(False)
self.actionBackspace.setEnabled(False)
self.actionFinish.setEnabled(False)
self.actionCancel.setEnabled(False)
self.actionVisible.setEnabled(False)
self.scene.load_image(file_path)
self.view.zoomfit()
# load label
if self.png_palette is None:
if self.can_be_annotated:
_, 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)
@ -444,9 +463,9 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
self.labels_dock_widget.checkBox_visible.setChecked(visible)
self.labels_dock_widget.set_all_polygon_visible(visible)
def label_converter(self):
self.convert_dialog.reset_gui()
self.convert_dialog.show()
def ISAT_to_VOC(self):
self.ISAT_to_VOC_dialog.reset_gui()
self.ISAT_to_VOC_dialog.show()
def help(self):
self.shortcut_dialog.show()
@ -489,7 +508,7 @@ 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.actionConverter.triggered.connect(self.label_converter)
self.actionToVOC.triggered.connect(self.ISAT_to_VOC)
self.actionShortcut.triggered.connect(self.help)
self.actionAbout.triggered.connect(self.about)