class MainWindow(QMainWindow, Ui_MainWindow):
# 声明枚举量
# 当前缩放的方式
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = 0, 1, 2
CREATE_MODE, EDIT_MODE = 0, 1
DATABASE_PATH = osp.abspath(r'database')
####初始化####
def __init__(self, parent=None):
super(MainWindow, self).__init__(parent)
'''
ui文件内容包括
1.按键/菜单选项形式的Action的注册
2.菜单栏，工具栏中Action的分发
3.子窗口的注册
4.主窗口的整体布局
'''
self.setupUi(self)
'''注册和分发WidgetAction'''
self.raw_image = np.ndarray(0, dtype=int)
self.current_series = None
self.current_file = ''
self.current_file_wl = 0
self.current_file_ww = 0
# 注册zoom_action处理canvas以光标所在点为中心的缩放
self.zoom_widget = ZoomWidget()
self.zoom_action = QWidgetAction(self)
self.zoom_action.setObjectName('zoom_action')
self.zoom_action.setDefaultWidget(self.zoom_widget)
self.toolBar.insertAction(self.zoom_out_action, self.zoom_action)
self.zoom_mode = self.FIT_WINDOW
# 注册wlww_action处理canvas图像的窗位窗宽调整
self.wlww_widget = WlwwWidget()
self.wlww_action = QWidgetAction(self)
self.wlww_action.setObjectName('wlww_action')
self.wlww_action.setDefaultWidget(self.wlww_widget)
self.toolBar.insertAction(self.wlww_reset_action, self.wlww_action)
self.init_docks()
self.coupling_mainwindow_actions()
self.coupling_canvas()
self.coupling_series_list_widget()
self.coupling_annotations_list_widget()
self.coupling_label_edit_widget()
self.coupling_database_tree_widget()
# 设置窗口显示属性
self.setFocusPolicy(Qt.ClickFocus)
self.resize(1920, 1080)
self.showMaximized()
# self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)
def init_docks(self):
'''初始化的一部分，执行浮动子窗口和canvas的初始化，单列一个函数以提升可读性'''
self.database_widget = DatabaseWidget()
self.dataset_tree_dock.setWidget(self.database_widget)
self.menuView.addAction(self.dataset_tree_dock.toggleViewAction())
self.series_list_widget = SeriesListWidget()
self.series_list_dock.setWidget(self.series_list_widget)
self.menuView.addAction(self.series_list_dock.toggleViewAction())
self.annotations_list_widget = AnnotationsListWidget()
self.annotations_list_dock.setWidget(self.annotations_list_widget)
self.menuView.addAction(self.annotations_list_dock.toggleViewAction())
self.label_edit_widget = LabelEditWidget()
self.label_edit_dock.setWidget(self.label_edit_widget)
self.menuView.addAction(self.label_edit_dock.toggleViewAction())
self.canvas_area = QScrollArea()
# TODO: 原理？看了文档还是不懂，需要进行更多研究
self.canvas_area.setWidgetResizable(True)
self.scroll_bars = {
Qt.Vertical: self.canvas_area.verticalScrollBar(),
Qt.Horizontal: self.canvas_area.horizontalScrollBar()
}
self.canvas_widget = Canvas()
self.canvas_area.setWidget(self.canvas_widget)
self.setCentralWidget(self.canvas_area)
def coupling_mainwindow_actions(self):
self.save_file_action.triggered.connect(self.save_current_work)
self.wlww_reset_action.triggered.connect(self.wlww_reset_slot)
def coupling_database_tree_widget(self):
self.new_database_action.triggered.connect(self.new_database_slot)
self.open_database_action.triggered.connect(self.open_database_slot)
self.database_widget.series_selected_signal.connect(
self.change_series_slot)
# 之前初始化database_widget时,尚未绑定,也无从接收
self.database_widget.send_latest_modified_series()
def coupling_series_list_widget(self):
'''初始化的一部分，执行与当前文件序列列表耦合的指令，单列一个函数以提升可读性'''
'''从主窗口通过action对series_list_widget进行操作'''
self.action('open_dir').triggered.connect(self.open_dir_slot)
self.action('open_next_image').triggered.connect(
partial(self.series_list_widget.change_current_item_slot, 1))
self.action('open_prev_image').triggered.connect(
partial(self.series_list_widget.change_current_item_slot, -1))
'''响应series_list_widget信号'''
self.series_list_widget.currentItemChanged.connect(
self.change_current_file_slot)
def coupling_annotations_list_widget(self):
'''初始化的一部分，执行与标签列表耦合的指令，单列一个函数以提升可读性'''
self.annotations_list_widget.itemSelectionChanged.connect(
self.selected_annotations_changed_slot)
def coupling_label_edit_widget(self):
self.label_edit_widget.apply_label_signal.connect(
self.apply_label_slot)
def coupling_canvas(self):
'''初始化的一部分，执行与canvas耦合的指令，单列一个函数以提升可读性'''
'''从主窗口通过action对canvas进行操作'''
# 模式和创建类型切换
self.toggle_mode_action.triggered.connect(self.toggle_mode_slot)
self.create_polygon_action.triggered.connect(
partial(self.set_mode_slot, self.CREATE_MODE, 'polygon'))
self.create_rectangle_action.triggered.connect(
partial(self.set_mode_slot, self.CREATE_MODE, 'rectangle'))
self.create_circle_action.triggered.connect(
partial(self.set_mode_slot, self.CREATE_MODE, 'circle'))
self.create_polyline_action.triggered.connect(
partial(self.set_mode_slot, self.CREATE_MODE, 'polyline'))
self.action('create_line').triggered.connect(
partial(self.set_mode_slot, self.CREATE_MODE, 'line'))
self.action('create_point').triggered.connect(
partial(self.set_mode_slot, self.CREATE_MODE, 'point'))
# 缩放
self.action('zoom_in').triggered.connect(
partial(self.adjust_zoom_value, 1.1))
self.action('zoom_out').triggered.connect(
partial(self.adjust_zoom_value, 0.9))
self.action('fit_window').triggered.connect(self.fit_window_slot)
self.action('fit_window_width').triggered.connect(
self.fit_window_width_slot)
# 撤销
self.canvas_undo_action.triggered.connect(self.canvas_undo_slot)
# 复制选中的标记
self.copy_selected_annotations_action.triggered.connect(
self.canvas_widget.copy_selected_annotations)
# 删除选中的标记
self.delete_selected_annotataions_action.triggered.connect(
self.canvas_widget.delete_selected_annotations)
# 选中所有标记
self.select_all_annotations_action.triggered.connect(
self.canvas_widget.select_all_annotations)
# 增加点到邻近边
self.add_point_to_nearest_edge_action.triggered.connect(
self.canvas_widget.add_point_to_nearest_edge)
self.action('hide_selected_annotation').triggered.connect(
partial(self.canvas_widget.set_selected_annotations_visibility,
False))
self.action('toggle_all_annotatons_visibility').triggered.connect(
self.toggle_all_annotatons_visibility_slot)
# 分发actions到canvas菜单
actions = ['hide_selected_annotation', 'add_point_to_nearest_edge']
self.add_actions(self.canvas_widget.edit_menu, actions)
actions = [
'create_polygon', 'create_rectangle', 'create_circle',
'create_polyline', 'create_line', 'create_point'
]
self.add_actions(self.canvas_widget.create_menu, actions)
'''响应canvas信号'''
# 响应功能请求
self.zoom_widget.valueChanged.connect(self.zoom_action_slot)
self.wlww_widget.wlww_changed_signal.connect(self.wlww_action_slot)
self.canvas_widget.zoom_request.connect(self.zoom_requeset_slot)
self.canvas_widget.scroll_request.connect(self.scroll_request_slot)
self.canvas_widget.wlww_request.connect(self.wlww_request_slot)
# 响应状态变化信号
self.canvas_widget.annotations_changed_signal.connect(
self.annotations_list_widget.refresh)
self.canvas_widget.annotation_created_signal.connect(
self.label_new_annotation_slot)
self.canvas_widget.selected_annotations_changed_signal.connect(
self.selected_annotations_changed_slot)
# 响应功能可用性信号
self.canvas_widget.has_edge_tobe_added_signal.\
connect(lambda x: partial(
self.toggle_actions, ['add_point_to_nearest_edge'])(x))
self.canvas_widget.is_canvas_creating_signal.\
connect(lambda x: partial(
self.toggle_actions, ['add_point_to_nearest_edge'])(x))
####初始化完成####
'''下面的方法协助动作分发的结构化'''
def action(self, action_name: str) -> QAction:
return self.findChild(QAction, action_name + '_action')
def add_actions(self, widget, action_names: list):
'''
:param widget: 目标对象
:param action_names: action名称列表
'''
for action_name in action_names:
if action_name is None:
widget.addSeparator()
else:
widget.addAction(self.action(action_name))
def toggle_actions(self, action_names: List[str], value: bool) -> None:
for action_name in action_names:
self.action(action_name).setEnabled(value)
'''下面的方法与canvas进行交互'''
def change_current_file_slot(self, file_item: QListWidgetItem):
if file_item:
self.save_current_work()
self.current_file = file_item.text()
self.current_file_wl, self.current_file_ww, dicom_array = get_dicom_info(
self.current_file)
self.raw_image = dicom_array.copy()
if (self.wlww_widget.wl_spin.value()
== 0) and (self.wlww_widget.ww_spin.value() == 0):
self.wlww_widget.wl_spin.setValue(self.current_file_wl)
self.wlww_widget.ww_spin.setValue(self.current_file_ww)
else:
self.wlww_action_slot()
annotations_file = self.current_file.replace('.dcm', '.pkl')
# 读取标签文件
if osp.exists(annotations_file):
with open(annotations_file, 'rb') as annotations_pkl:
annotations = pickle.load(annotations_pkl)
self.canvas_widget.load_annotations(annotations)
# 保存当前文件
def scroll_request_slot(self, delta: int, orientation: int):
'''响应canvas的滚动请求'''
units = -delta * 0.1
bar = self.scroll_bars[orientation]
bar.setValue(bar.value() + bar.singleStep() * units)
def zoom_requeset_slot(self, delta: int, pos: QPoint):
'''响应canvas的缩放请求'''
canvas_width_old = self.canvas_widget.width()
coeff = 1.1 if delta > 0 else 0.9
# 按照倍数系数调整缩放值，触发缩放操作调整画布尺寸
self.adjust_zoom_value(coeff)
# 在画布上进行的缩放，以光标位置为中心，
# 调整尺寸之后，还要进行移动，保证光标在画布上的位置不变
canvas_width_new = self.canvas_widget.width()
canvas_scale_factor = canvas_width_new / canvas_width_old
x_shift = round(pos.x() * canvas_scale_factor) - pos.x()
y_shift = round(pos.y() * canvas_scale_factor) - pos.y()
self.scroll_bars[Qt.Horizontal].setValue(
self.scroll_bars[Qt.Horizontal].value() + x_shift)
self.scroll_bars[Qt.Vertical].setValue(
self.scroll_bars[Qt.Vertical].value() + y_shift)
def adjust_zoom_value(self, coeff: float):
self.set_zoom_value_mannually(self.zoom_widget.value() * coeff)
def fit_window_slot(self, value):
if value:
self.fit_window_width_action.setChecked(False)
self.zoom_mode = self.FIT_WINDOW
self.set_zoom_value_to_fit(self.FIT_WINDOW)
def fit_window_width_slot(self, value):
if value:
self.fit_window_action.setChecked(False)
self.zoom_mode = self.FIT_WIDTH
self.set_zoom_value_to_fit(self.FIT_WIDTH)
def set_zoom_value_to_fit(self, zoom_mode: int) -> None:
print('set_zoom_value_to_fit')
'''设置计算出的缩放比例，触发zoom_action'''
epsilon = 10.0
w1 = self.canvas_area.width() - epsilon
h1 = self.canvas_area.height() - epsilon
scroll_area_ratio = w1 / h1
w2 = self.canvas_widget.pixmap.width()
h2 = self.canvas_widget.pixmap.height()
pixmap_ratio = w2 / h2
if zoom_mode == self.FIT_WIDTH:
self.zoom_widget.setValue(w1 / w2 * 100)
elif zoom_mode == self.FIT_WINDOW:
if scroll_area_ratio > pixmap_ratio:
self.zoom_widget.setValue(h1 / h2 * 100)
else:
self.zoom_widget.setValue(w1 / w2 * 100)
def set_zoom_value_mannually(self, value: float) -> None:
'''设置缩放比例，触发zoom_action'''
self.fit_window_action.setChecked(False)
self.fit_window_width_action.setChecked(False)
self.zoom_widget.setValue(value)
self.zoom_mode = self.MANUAL_ZOOM
def zoom_action_slot(self):
'''缩放比例变化时触发，进行画布尺寸缩放'''
self.canvas_widget.scale = 0.01 * self.zoom_widget.value()
self.canvas_widget.adjustSize()
self.canvas_widget.update()
def wlww_reset_slot(self):
print('reset')
self.wlww_widget.wl_spin.setValue(self.current_file_wl)
self.wlww_widget.ww_spin.setValue(self.current_file_ww)
def wlww_request_slot(self, wl_delta, ww_delta):
'''响应canvas的窗位窗宽调整请求'''
self.wlww_widget.wl_spin.setValue(self.wlww_widget.wl_spin.value() +
wl_delta)
self.wlww_widget.ww_spin.setValue(self.wlww_widget.ww_spin.value() +
ww_delta)
def wlww_action_slot(self):
'''窗位窗宽数值变化时触发，按照新的窗位窗位窗宽生成图像，重绘画布'''
pixmap = dicom_array2pixmap(self.wlww_widget.wl_spin.value(),
self.wlww_widget.ww_spin.value(),
self.raw_image)
self.canvas_widget.change_pixmap(pixmap)
self.fit_window_slot(True)
def toggle_mode_slot(self) -> None:
'''实现模式切换'''
if not self.canvas_widget.create_mode:
self.canvas_widget.create_mode = True
else:
self.canvas_widget.create_mode = False
def set_mode_slot(self, mode: int, create_type='polygon') -> None:
'''实现所有进行模式和类型变更的action'''
if mode == self.EDIT_MODE:
self.canvas_widget.create_mode = False
elif mode == self.CREATE_MODE:
self.canvas_widget.create_mode = True
self.canvas_widget.create_type = create_type
def toggle_all_annotatons_visibility_slot(self) -> None:
'''显示/隐藏所有标记，具体逻辑是，若有标记被隐藏，显示所有标记，否则，隐藏所有标记'''
for annotation in self.canvas_widget.annotations:
if not annotation.is_visable:
self.canvas_widget.set_all_annotations_visibility(True)
return
self.canvas_widget.set_all_annotations_visibility(False)
return
def canvas_undo_slot(self):
'''撤销操作，根据具体状态进行不同操作'''
# 创建模式下
if self.canvas_widget.create_mode:
# （上一个）标记创建已经完成，撤销已完成标记的最后一个点
if not self.canvas_widget.current_annotation:
self.canvas_widget.undo_last_line()
# 有标记正在被创建，撤销当前标记的最后一个点
else:
self.canvas_widget.undo_last_point()
# 编辑模式下
else:
self.canvas_widget.restore_annotations()
def label_new_annotation_slot(self, new_annotation: Annotation):
new_annotation.label = LabelEditDialog.get_label()
print(new_annotation.label.segmentation)
def selected_annotations_changed_slot(self, selected_annotations=None):
'''
响应被选中标记变化信号，包括从canvas上选中和从annotation_list_widget上选中传出的信号
1.
2.更改标签编辑窗口内容为最后被选中的标记的标签内容
'''
if self.sender() == self.annotations_list_widget:
selected_annotations = []
for index in self.annotations_list_widget.selectedIndexes():
selected_annotations.append(
self.canvas_widget.annotations[index.row()])
self.canvas_widget.select_specific_annotations(
selected_annotations)
elif self.sender() == self.canvas_widget:
pass
if selected_annotations:
self.label_edit_widget.label = selected_annotations[-1].label
self.label_edit_widget.refresh()
else:
self.label_edit_widget.reset()
'''下面的方法与database_widget进行交互'''
def new_database_slot(self):
'''从指定目录新建dicom数据库'''
database_dir = QFileDialog.getExistingDirectory(self, '选择要扫描的目录')
if not database_dir:
return
database_name, ok = QInputDialog.getText(self, '输入数据库名称', '数据库名称：')
if ok:
self.database_widget.new_database(database_dir, database_name)
def open_database_slot(self):
'''从指定.xml文件打开dicom数据库'''
database_path, ok = QFileDialog.getOpenFileName(
self,
caption='选择要打开的数据库',
directory=self.DATABASE_PATH,
filter='数据库文件(*.xml)')
if ok:
self.database_widget.open_database(database_path)
def change_series_slot(self, series_item: QTreeWidgetItem,
files: List[str]) -> None:
'''
响应变更当前序列的请求,有以下来源
1.双击数据库窗口中的序列
2.从open_dir_action输入
'''
print('changing series')
self.auto_refresh_current_series_modified_time()
self.auto_refresh_current_series_state()
self.query_current_series_state()
self.input_files_slot(files)
self.current_series = series_item
self.database_widget.expand_recursively(self.current_series)
self.database_widget.setCurrentItem(self.current_series)
def query_current_series_state(self):
if self.current_series:
if self.current_series.checkState(0) != 2:
annotated = QMessageBox.question(
self, '', '标记上一个序列为已完成?', QMessageBox.Yes | QMessageBox.No,
QMessageBox.No)
if annotated == QMessageBox.Yes:
print(self.current_series)
self.current_series.setCheckState(0, 2)
def auto_refresh_current_series_state(self):
if self.current_series and self.annotations_list_widget.annotations:
self.current_series.setCheckState(0, 1)
def auto_refresh_current_series_modified_time(self):
if self.current_series:
self.current_series.setText(
4,
time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime(float(time.time()))))
'''下面的方法与series_list_widget进行交互'''
def open_dir_slot(self):
'''打开文件夹，将.dcm文件增加到数据库'''
files_dir = QFileDialog.getExistingDirectory(self, '选择要打开的目录')
if not files_dir:
return
files = get_dicom_files_path_from_dir(files_dir)
# 数据库被更改后,原先的current_series对应的节点会失效,需要在更改前提取uid,在更改后根据uid映射到新节点
current_series_uids = []
if self.current_series:
current_series_uids = self.database_widget.get_top_down_uid(
self.current_series)
self.database_widget.add_to_database(files)
if current_series_uids:
self.current_series = self.database_widget.get_item_from_top_down_uid(
current_series_uids)
self.database_widget.send_latest_imported_series()
def input_files_slot(self, files):
'''
响应输入新的文件更新当前文件序列的请求，有以下来源
1.从数据库窗口获取序列输入
2.从open_dir/open_file action输入
'''
files = sorted(files)
self.series_list_widget.refresh_files(files)
self.series_list_widget.change_current_item_slot(-9999)
if self.toggle_auto_wlww_action.isChecked():
self.wlww_reset_slot()
'''下面的方法与label_edit_widget进行交互'''
def apply_label_slot(self, label: LabelStruct):
if self.canvas_widget.selected_annotations:
self.canvas_widget.selected_annotations[-1].label = label
self.annotations_list_widget.refresh(
self.canvas_widget.annotations)
'''下面的方法实现全局功能'''
def save_current_work(self):
'''保存当前图像的所有标记为图像同名文件'''
if self.current_file and self.annotations_list_widget.annotations:
with open(self.current_file.replace('.dcm', '.pkl'),
'wb') as annotations_pkl:
pickle.dump(self.canvas_widget.annotations, annotations_pkl)
def closeEvent(self, *args, **kwargs):
'''退出前事件'''
super().closeEvent(*args, **kwargs)
self.auto_refresh_current_series_modified_time()
self.database_widget.save_item_states_and_modified_time()
class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
# Save as Pascal voc xml
self.defaultSaveDir = None
self.usingPascalVocFormat = True
if self.usingPascalVocFormat:
LabelFile.suffix = '.xml'
# For loading all image under a directory
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
# Whether we need to save or not.
self.dirty = False
# Enble auto saving if pressing next
self.autoSaving = True
self._noSelectionSlot = False
self._beginner = True
self.screencastViewer = "firefox"
self.screencast = "https://youtu.be/p0nR2YsCY_U"
self.loadPredefinedClasses()
# Main widgets and related state.
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.labelList = QListWidget()
self.itemsToShapes = {}
self.shapesToItems = {}
self.labelList.itemActivated.connect(self.labelSelectionChanged)
self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
self.labelList.itemDoubleClicked.connect(self.editLabel)
# Connect to itemChanged to detect checkbox changes.
self.labelList.itemChanged.connect(self.labelItemChanged)
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
listLayout.addWidget(self.labelList)
self.editButton = QToolButton()
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
self.labelListContainer = QWidget()
self.labelListContainer.setLayout(listLayout)
listLayout.addWidget(self.editButton) #, 0, Qt.AlignCenter)
listLayout.addWidget(self.labelList)
self.dock = QDockWidget(u'Box Labels', self)
self.dock.setObjectName(u'Labels')
self.dock.setWidget(self.labelListContainer)
# Tzutalin 20160906 : Add file list and dock to move faster
self.fileListWidget = QListWidget()
self.fileListWidget.itemDoubleClicked.connect(
self.fileitemDoubleClicked)
filelistLayout = QVBoxLayout()
filelistLayout.setContentsMargins(0, 0, 0, 0)
filelistLayout.addWidget(self.fileListWidget)
self.fileListContainer = QWidget()
self.fileListContainer.setLayout(filelistLayout)
self.filedock = QDockWidget(u'File List', self)
self.filedock.setObjectName(u'Files')
self.filedock.setWidget(self.fileListContainer)
self.zoomWidget = ZoomWidget()
self.colorDialog = ColorDialog(parent=self)
self.canvas = Canvas()
self.canvas.zoomRequest.connect(self.zoomRequest)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
# Tzutalin 20160906 : Add file list and dock to move faster
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.dockFeatures = QDockWidget.DockWidgetClosable\
| QDockWidget.DockWidgetFloatable
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
# Actions
action = partial(newAction, self)
quit = action('&Quit', self.close, 'Ctrl+Q', 'quit',
u'Quit application')
open = action('&Open', self.openFile, 'Ctrl+O', 'open',
u'Open image or label file')
opendir = action('&Open Dir', self.openDir, 'Ctrl+u', 'open',
u'Open Dir')
changeSavedir = action('&Change default saved Annotation dir',
self.changeSavedir, 'Ctrl+r', 'open',
u'Change default saved Annotation dir')
openAnnotation = action('&Open Annotation', self.openAnnotation,
'Ctrl+q', 'openAnnotation', u'Open Annotation')
openNextImg = action('&Next Image', self.openNextImg, 'd', 'next',
u'Open Next')
openPrevImg = action('&Prev Image', self.openPrevImg, 'a', 'prev',
u'Open Prev')
save = action('&Save',
self.saveFile,
'Ctrl+S',
'save',
u'Save labels to file',
enabled=False)
saveAs = action('&Save As',
self.saveFileAs,
'Ctrl+Shift+S',
'save-as',
u'Save labels to a different file',
enabled=False)
close = action('&Close', self.closeFile, 'Ctrl+W', 'close',
u'Close current file')
color1 = action('Box &Line Color', self.chooseColor1, 'Ctrl+L',
'color_line', u'Choose Box line color')
color2 = action('Box &Fill Color', self.chooseColor2, 'Ctrl+Shift+L',
'color', u'Choose Box fill color')
createMode = action('Create\nRectBox',
self.setCreateMode,
'Ctrl+N',
'new',
u'Start drawing Boxs',
enabled=False)
editMode = action('&Edit\nRectBox',
self.setEditMode,
'Ctrl+J',
'edit',
u'Move and edit Boxs',
enabled=False)
create = action('Create\nRectBox',
self.createShape,
'w',
'new',
u'Draw a new Box',
enabled=False)
delete = action('Delete\nRectBox',
self.deleteSelectedShape,
'Delete',
'delete',
u'Delete',
enabled=False)
copy = action('&Duplicate\nRectBox',
self.copySelectedShape,
'Ctrl+D',
'copy',
u'Create a duplicate of the selected Box',
enabled=False)
advancedMode = action('&Advanced Mode',
self.toggleAdvancedMode,
'Ctrl+Shift+A',
'expert',
u'Switch to advanced mode',
checkable=True)
hideAll = action('&Hide\nRectBox',
partial(self.togglePolygons, False),
'Ctrl+H',
'hide',
u'Hide all Boxs',
enabled=False)
showAll = action('&Show\nRectBox',
partial(self.togglePolygons, True),
'Ctrl+A',
'hide',
u'Show all Boxs',
enabled=False)
help = action('&Tutorial', self.tutorial, 'Ctrl+T', 'help',
u'Show demos')
zoom = QWidgetAction(self)
zoom.setDefaultWidget(self.zoomWidget)
self.zoomWidget.setWhatsThis(
u"Zoom in or out of the image. Also accessible with"\
" %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
fmtShortcut("Ctrl+Wheel")))
self.zoomWidget.setEnabled(False)
zoomIn = action('Zoom &In',
partial(self.addZoom, 10),
'Ctrl++',
'zoom-in',
u'Increase zoom level',
enabled=False)
zoomOut = action('&Zoom Out',
partial(self.addZoom, -10),
'Ctrl+-',
'zoom-out',
u'Decrease zoom level',
enabled=False)
zoomOrg = action('&Original size',
partial(self.setZoom, 100),
'Ctrl+=',
'zoom',
u'Zoom to original size',
enabled=False)
fitWindow = action('&Fit Window',
self.setFitWindow,
'Ctrl+F',
'fit-window',
u'Zoom follows window size',
checkable=True,
enabled=False)
fitWidth = action('Fit &Width',
self.setFitWidth,
'Ctrl+Shift+F',
'fit-width',
u'Zoom follows window width',
checkable=True,
enabled=False)
# Group zoom controls into a list for easier toggling.
zoomActions = (self.zoomWidget, zoomIn, zoomOut, zoomOrg, fitWindow,
fitWidth)
self.zoomMode = self.MANUAL_ZOOM
self.scalers = {
self.FIT_WINDOW:
self.scaleFitWindow,
self.FIT_WIDTH:
self.scaleFitWidth,
# Set to one to scale to 100% when loading files.
self.MANUAL_ZOOM:
lambda: 1,
}
edit = action('&Edit Label',
self.editLabel,
'Ctrl+E',
'edit',
u'Modify the label of the selected Box',
enabled=False)
self.editButton.setDefaultAction(edit)
shapeLineColor = action(
'Shape &Line Color',
self.chshapeLineColor,
icon='color_line',
tip=u'Change the line color for this specific shape',
enabled=False)
shapeFillColor = action(
'Shape &Fill Color',
self.chshapeFillColor,
icon='color',
tip=u'Change the fill color for this specific shape',
enabled=False)
labels = self.dock.toggleViewAction()
labels.setText('Show/Hide Label Panel')
labels.setShortcut('Ctrl+Shift+L')
# Lavel list context menu.
labelMenu = QMenu()
addActions(labelMenu, (edit, delete))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(
self.popLabelListMenu)
# Store actions for further handling.
self.actions = struct(
save=save,
saveAs=saveAs,
open=open,
close=close,
lineColor=color1,
fillColor=color2,
create=create,
delete=delete,
edit=edit,
copy=copy,
createMode=createMode,
editMode=editMode,
advancedMode=advancedMode,
shapeLineColor=shapeLineColor,
shapeFillColor=shapeFillColor,
zoom=zoom,
zoomIn=zoomIn,
zoomOut=zoomOut,
zoomOrg=zoomOrg,
fitWindow=fitWindow,
fitWidth=fitWidth,
zoomActions=zoomActions,
fileMenuActions=(open, opendir, save, saveAs, close, quit),
beginner=(),
advanced=(),
editMenu=(edit, copy, delete, None, color1, color2),
beginnerContext=(create, edit, copy, delete),
advancedContext=(createMode, editMode, edit, copy, delete,
shapeLineColor, shapeFillColor),
onLoadActive=(close, create, createMode, editMode),
onShapesPresent=(saveAs, hideAll, showAll))
self.menus = struct(file=self.menu('&File'),
edit=self.menu('&Edit'),
view=self.menu('&View'),
help=self.menu('&Help'),
recentFiles=QMenu('Open &Recent'),
labelList=labelMenu)
addActions(self.menus.file,
(open, opendir, changeSavedir, openAnnotation,
self.menus.recentFiles, save, saveAs, close, None, quit))
addActions(self.menus.help, (help, ))
addActions(self.menus.view,
(labels, advancedMode, None, hideAll, showAll, None, zoomIn,
zoomOut, zoomOrg, None, fitWindow, fitWidth))
self.menus.file.aboutToShow.connect(self.updateFileMenu)
# Custom context menu for the canvas widget:
addActions(self.canvas.menus[0], self.actions.beginnerContext)
addActions(self.canvas.menus[1],
(action('&Copy here', self.copyShape),
action('&Move here', self.moveShape)))
self.tools = self.toolbar('Tools')
self.actions.beginner = (open, opendir, openNextImg, openPrevImg, save,
None, create, copy, delete, None, zoomIn,
zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (open, save, None, createMode, editMode, None,
hideAll, showAll)
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
# Application state.
self.image = QImage()
self.filePath = u(defaultFilename)
self.recentFiles = []
self.maxRecent = 7
self.lineColor = None
self.fillColor = None
self.zoom_level = 100
self.fit_window = False
# XXX: Could be completely declarative.
# Restore application settings.
if have_qstring():
types = {
'filename': QString,
'recentFiles': QStringList,
'window/size': QSize,
'window/position': QPoint,
'window/geometry': QByteArray,
'line/color': QColor,
'fill/color': QColor,
'advanced': bool,
# Docks and toolbars:
'window/state': QByteArray,
'savedir': QString,
'lastOpenDir': QString,
}
else:
types = {
'filename': str,
'recentFiles': list,
'window/size': QSize,
'window/position': QPoint,
'window/geometry': QByteArray,
'line/color': QColor,
'fill/color': QColor,
'advanced': bool,
# Docks and toolbars:
'window/state': QByteArray,
'savedir': str,
'lastOpenDir': str,
}
self.settings = settings = Settings(types)
self.recentFiles = list(settings.get('recentFiles', []))
size = settings.get('window/size', QSize(600, 500))
position = settings.get('window/position', QPoint(0, 0))
self.resize(size)
self.move(position)
saveDir = u(settings.get('savedir', None))
self.lastOpenDir = u(settings.get('lastOpenDir', None))
if os.path.exists(saveDir):
self.defaultSaveDir = saveDir
self.statusBar().showMessage(
'%s started. Annotation will be saved to %s' %
(__appname__, self.defaultSaveDir))
self.statusBar().show()
# or simply:
#self.restoreGeometry(settings['window/geometry']
self.restoreState(settings.get('window/state', QByteArray()))
self.lineColor = QColor(settings.get('line/color', Shape.line_color))
self.fillColor = QColor(settings.get('fill/color', Shape.fill_color))
Shape.line_color = self.lineColor
Shape.fill_color = self.fillColor
def xbool(x):
if isinstance(x, QVariant):
return x.toBool()
return bool(x)
if xbool(settings.get('advanced', False)):
self.actions.advancedMode.setChecked(True)
self.toggleAdvancedMode()
# Populate the File menu dynamically.
self.updateFileMenu()
# Since loading the file may take some time, make sure it runs in the background.
self.queueEvent(partial(self.loadFile, self.filePath))
# Callbacks:
self.zoomWidget.valueChanged.connect(self.paintCanvas)
self.populateModeActions()
## Support Functions ##
def noShapes(self):
return not self.itemsToShapes
def toggleAdvancedMode(self, value=True):
self._beginner = not value
self.canvas.setEditing(True)
self.populateModeActions()
self.editButton.setVisible(not value)
if value:
self.actions.createMode.setEnabled(True)
self.actions.editMode.setEnabled(False)
self.dock.setFeatures(self.dock.features() | self.dockFeatures)
else:
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
def populateModeActions(self):
if self.beginner():
tool, menu = self.actions.beginner, self.actions.beginnerContext
else:
tool, menu = self.actions.advanced, self.actions.advancedContext
self.tools.clear()
addActions(self.tools, tool)
self.canvas.menus[0].clear()
addActions(self.canvas.menus[0], menu)
self.menus.edit.clear()
actions = (self.actions.create,) if self.beginner()\
else (self.actions.createMode, self.actions.editMode)
addActions(self.menus.edit, actions + self.actions.editMenu)
def setBeginner(self):
self.tools.clear()
addActions(self.tools, self.actions.beginner)
def setAdvanced(self):
self.tools.clear()
addActions(self.tools, self.actions.advanced)
def setDirty(self):
self.dirty = True
self.actions.save.setEnabled(True)
def setClean(self):
self.dirty = False
self.actions.save.setEnabled(False)
self.actions.create.setEnabled(True)
def toggleActions(self, value=True):
"""Enable/Disable widgets which depend on an opened image."""
for z in self.actions.zoomActions:
z.setEnabled(value)
for action in self.actions.onLoadActive:
action.setEnabled(value)
def queueEvent(self, function):
QTimer.singleShot(0, function)
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.itemsToShapes.clear()
self.shapesToItems.clear()
self.labelList.clear()
self.filePath = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
def currentItem(self):
items = self.labelList.selectedItems()
if items:
return items[0]
return None
def addRecentFile(self, filePath):
if filePath in self.recentFiles:
self.recentFiles.remove(filePath)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filePath)
def beginner(self):
return self._beginner
def advanced(self):
return not self.beginner()
## Callbacks ##
def tutorial(self):
subprocess.Popen([self.screencastViewer, self.screencast])
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
"""In the middle of drawing, toggling between modes should be disabled."""
self.actions.editMode.setEnabled(not drawing)
if not drawing and self.beginner():
# Cancel creation.
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.create.setEnabled(True)
def toggleDrawMode(self, edit=True):
self.canvas.setEditing(edit)
self.actions.createMode.setEnabled(edit)
self.actions.editMode.setEnabled(not edit)
def setCreateMode(self):
assert self.advanced()
self.toggleDrawMode(False)
def setEditMode(self):
assert self.advanced()
self.toggleDrawMode(True)
def updateFileMenu(self):
currFilePath = self.filePath
def exists(filename):
return os.path.exists(filename)
menu = self.menus.recentFiles
menu.clear()
files = [
f for f in self.recentFiles if f != currFilePath and exists(f)
]
for i, f in enumerate(files):
icon = newIcon('labels')
action = QAction(icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()),
self)
action.triggered.connect(partial(self.loadRecent, f))
menu.addAction(action)
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
def editLabel(self, item=None):
if not self.canvas.editing():
return
item = item if item else self.currentItem()
text = self.labelDialog.popUp(item.text())
if text is not None:
item.setText(text)
self.setDirty()
# Tzutalin 20160906 : Add file list and dock to move faster
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(u(item.text()))
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
# React to canvas signals.
def shapeSelectionChanged(self, selected=False):
if self._noSelectionSlot:
self._noSelectionSlot = False
else:
shape = self.canvas.selectedShape
if shape:
self.shapesToItems[shape].setSelected(True)
else:
self.labelList.clearSelection()
self.actions.delete.setEnabled(selected)
self.actions.copy.setEnabled(selected)
self.actions.edit.setEnabled(selected)
self.actions.shapeLineColor.setEnabled(selected)
self.actions.shapeFillColor.setEnabled(selected)
def addLabel(self, shape):
item = HashableQListWidgetItem(shape.label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.itemsToShapes[item] = shape
self.shapesToItems[shape] = item
self.labelList.addItem(item)
for action in self.actions.onShapesPresent:
action.setEnabled(True)
def remLabel(self, shape):
item = self.shapesToItems[shape]
self.labelList.takeItem(self.labelList.row(item))
del self.shapesToItems[shape]
del self.itemsToShapes[item]
def loadLabels(self, shapes):
s = []
for label, points, line_color, fill_color in shapes:
shape = Shape(label=label)
for x, y in points:
shape.addPoint(QPointF(x, y))
shape.close()
s.append(shape)
self.addLabel(shape)
if line_color:
shape.line_color = QColor(*line_color)
if fill_color:
shape.fill_color = QColor(*fill_color)
self.canvas.loadShapes(s)
def saveLabels(self, filePath):
lf = LabelFile()
def format_shape(s):
return dict(label=s.label,
line_color=s.line_color.getRgb()\
if s.line_color != self.lineColor else None,
fill_color=s.fill_color.getRgb()\
if s.fill_color != self.fillColor else None,
points=[(p.x(), p.y()) for p in s.points])
shapes = [format_shape(shape) for shape in self.canvas.shapes]
# Can add differrent annotation formats here
try:
unicodeFilePath = u(filePath)
if self.usingPascalVocFormat is True:
lf.savePascalVocFormat(unicodeFilePath, shapes,
unicodeFilePath, self.imageData,
self.lineColor.getRgb(),
self.fillColor.getRgb())
else:
lf.save(unicodeFilePath, shapes, unicodeFilePath,
self.imageData, self.lineColor.getRgb(),
self.fillColor.getRgb())
self.labelFile = lf
self.filePath = unicodeFilePath
return True
except LabelFileError as e:
self.errorMessage(u'Error saving label data', u'<b>%s</b>' % e)
return False
def copySelectedShape(self):
self.addLabel(self.canvas.copySelectedShape())
#fix copy and delete
self.shapeSelectionChanged(True)
def labelSelectionChanged(self):
item = self.currentItem()
if item and self.canvas.editing():
self._noSelectionSlot = True
self.canvas.selectShape(self.itemsToShapes[item])
def labelItemChanged(self, item):
shape = self.itemsToShapes[item]
label = item.text()
if label != shape.label:
shape.label = item.text()
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
## Callback functions:
def newShape(self):
"""Pop-up and give focus to the label editor.
position MUST be in global coordinates.
"""
if len(self.labelHist) > 0:
self.labelDialog = LabelDialog(parent=self,
listItem=self.labelHist)
# text = self.labelDialog.popUp()
text = QString(u'pockey') # label only pockey
if text is not None:
self.addLabel(self.canvas.setLastLabel(text))
if self.beginner(): # Switch to edit mode.
self.canvas.setEditing(True)
self.actions.create.setEnabled(True)
else:
self.actions.editMode.setEnabled(True)
self.setDirty()
if text not in self.labelHist:
self.labelHist.append(text)
else:
#self.canvas.undoLastLine()
self.canvas.resetAllLines()
def scrollRequest(self, delta, orientation):
units = -delta / (8 * 15)
bar = self.scrollBars[orientation]
bar.setValue(bar.value() + bar.singleStep() * units)
def setZoom(self, value):
self.actions.fitWidth.setChecked(False)
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.MANUAL_ZOOM
self.zoomWidget.setValue(value)
def addZoom(self, increment=10):
self.setZoom(self.zoomWidget.value() + increment)
def zoomRequest(self, delta):
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
def setFitWindow(self, value=True):
if value:
self.actions.fitWidth.setChecked(False)
self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
self.adjustScale()
def setFitWidth(self, value=True):
if value:
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
self.adjustScale()
def togglePolygons(self, value):
for item, shape in self.itemsToShapes.items():
item.setCheckState(Qt.Checked if value else Qt.Unchecked)
def loadFile(self, filePath=None):
"""Load the specified file, or the last opened file if None."""
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get('filename')
unicodeFilePath = u(filePath)
# Tzutalin 20160906 : Add file list and dock to move faster
# Highlight the file item
if unicodeFilePath and self.fileListWidget.count() > 0:
index = self.mImgList.index(unicodeFilePath)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
if unicodeFilePath and os.path.exists(unicodeFilePath):
if LabelFile.isLabelFile(unicodeFilePath):
try:
self.labelFile = LabelFile(unicodeFilePath)
except LabelFileError as e:
self.errorMessage(u'Error opening file',
(u"<p><b>%s</b></p>"
u"<p>Make sure <i>%s</i> is a valid label file.") \
% (e, unicodeFilePath))
self.status("Error reading %s" % unicodeFilePath)
return False
self.imageData = self.labelFile.imageData
self.lineColor = QColor(*self.labelFile.lineColor)
self.fillColor = QColor(*self.labelFile.fillColor)
else:
# Load image:
# read data first and store for saving into label file.
self.imageData = read(unicodeFilePath, None)
self.labelFile = None
image = QImage.fromData(self.imageData)
if image.isNull():
self.errorMessage(
u'Error opening file',
u"<p>Make sure <i>%s</i> is a valid image file." %
unicodeFilePath)
self.status("Error reading %s" % unicodeFilePath)
return False
self.status("Loaded %s" % os.path.basename(unicodeFilePath))
self.image = image
self.filePath = unicodeFilePath
self.canvas.loadPixmap(QPixmap.fromImage(image))
if self.labelFile:
self.loadLabels(self.labelFile.shapes)
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.addRecentFile(self.filePath)
self.toggleActions(True)
## Label xml file and show bound box according to its filename
if self.usingPascalVocFormat is True and \
self.defaultSaveDir is not None:
basename = os.path.basename(
os.path.splitext(self.filePath)[0]) + XML_EXT
xmlPath = os.path.join(self.defaultSaveDir, basename)
self.loadPascalXMLByFilename(xmlPath)
return True
return False
def resizeEvent(self, event):
if self.canvas and not self.image.isNull()\
and self.zoomMode != self.MANUAL_ZOOM:
self.adjustScale()
super(MainWindow, self).resizeEvent(event)
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.scale = 0.01 * self.zoomWidget.value()
self.canvas.adjustSize()
self.canvas.update()
def adjustScale(self, initial=False):
value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
self.zoomWidget.setValue(int(100 * value))
def scaleFitWindow(self):
"""Figure out the size of the pixmap in order to fit the main widget."""
e = 2.0 # So that no scrollbars are generated.
w1 = self.centralWidget().width() - e
h1 = self.centralWidget().height() - e
a1 = w1 / h1
# Calculate a new scale value based on the pixmap's aspect ratio.
w2 = self.canvas.pixmap.width() - 0.0
h2 = self.canvas.pixmap.height() - 0.0
a2 = w2 / h2
return w1 / w2 if a2 >= a1 else h1 / h2
def scaleFitWidth(self):
# The epsilon does not seem to work too well here.
w = self.centralWidget().width() - 2.0
return w / self.canvas.pixmap.width()
def closeEvent(self, event):
if not self.mayContinue():
event.ignore()
s = self.settings
# If it loads images from dir, don't load it at the begining
if self.dirname is None:
s['filename'] = self.filePath if self.filePath else ''
else:
s['filename'] = ''
s['window/size'] = self.size()
s['window/position'] = self.pos()
s['window/state'] = self.saveState()
s['line/color'] = self.lineColor
s['fill/color'] = self.fillColor
s['recentFiles'] = self.recentFiles
s['advanced'] = not self._beginner
if self.defaultSaveDir is not None and len(self.defaultSaveDir) > 1:
s['savedir'] = str(self.defaultSaveDir)
else:
s['savedir'] = ""
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
s['lastOpenDir'] = self.lastOpenDir
else:
s['lastOpenDir'] = ""
## User Dialogs ##
def loadRecent(self, filename):
if self.mayContinue():
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = ['.jpeg', '.jpg', '.png', '.bmp']
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relatviePath = os.path.join(root, file)
path = u(os.path.abspath(relatviePath))
images.append(path)
images.sort(key=lambda x: x.lower())
return images
def changeSavedir(self, _value=False):
if self.defaultSaveDir is not None:
path = str(self.defaultSaveDir)
else:
path = '.'
dirpath = str(
QFileDialog.getExistingDirectory(
self, '%s - Save to the directory' % __appname__, path,
QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
self.defaultSaveDir = dirpath
self.statusBar().showMessage(
'%s . Annotation will be saved to %s' %
('Change saved folder', self.defaultSaveDir))
self.statusBar().show()
def openAnnotation(self, _value=False):
if self.filePath is None:
return
path = os.path.dirname(u(self.filePath))\
if self.filePath else '.'
if self.usingPascalVocFormat:
formats = ['*.%s' % str(fmt).lower()\
for fmt in QImageReader.supportedImageFormats()]
filters = "Open Annotation XML file (%s)" % \
' '.join(formats + ['*.xml'])
filename = str(
QFileDialog.getOpenFileName(
self, '%s - Choose a xml file' % __appname__, path,
filters))
self.loadPascalXMLByFilename(filename)
def openDir(self, _value=False):
if not self.mayContinue():
return
path = os.path.dirname(self.filePath)\
if self.filePath else '.'
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
path = self.lastOpenDir
dirpath = u(
QFileDialog.getExistingDirectory(
self, '%s - Open Directory' % __appname__, path,
QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
self.lastOpenDir = dirpath
self.dirname = dirpath
self.filePath = None
self.fileListWidget.clear()
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def openPrevImg(self, _value=False):
if not self.mayContinue():
return
if len(self.mImgList) <= 0:
return
if self.filePath is None:
return
currIndex = self.mImgList.index(self.filePath)
if currIndex - 1 >= 0:
filename = self.mImgList[currIndex - 1]
if filename:
self.loadFile(filename)
def openNextImg(self, _value=False):
# Proceding next image without dialog if having any label
if self.autoSaving is True and self.defaultSaveDir is not None:
if self.dirty is True and self.hasLabels():
self.saveFile()
if not self.mayContinue():
return
if len(self.mImgList) <= 0:
return
filename = None
if self.filePath is None:
filename = self.mImgList[0]
else:
currIndex = self.mImgList.index(self.filePath)
if currIndex + 1 < len(self.mImgList):
filename = self.mImgList[currIndex + 1]
if filename:
self.loadFile(filename)
def openFile(self, _value=False):
if not self.mayContinue():
return
path = os.path.dirname(str(self.filePath))\
if self.filePath else '.'
formats = ['*.%s' % str(fmt).lower()\
for fmt in QImageReader.supportedImageFormats()]
filters = "Image & Label files (%s)" % \
' '.join(formats + ['*%s' % LabelFile.suffix])
filename = QFileDialog.getOpenFileName(
self, '%s - Choose Image or Label file' % __appname__, path,
filters)
if filename:
self.loadFile(filename)
def saveFile(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
if self.hasLabels():
if self.defaultSaveDir is not None and len(str(
self.defaultSaveDir)):
print('handle the image:' + self.filePath)
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(
imgFileName)[0] + LabelFile.suffix
savedPath = os.path.join(str(self.defaultSaveDir),
savedFileName)
self._saveFile(savedPath)
else:
self._saveFile(self.filePath if self.labelFile\
else self.saveFileDialog())
def saveFileAs(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
if self.hasLabels():
self._saveFile(self.saveFileDialog())
def saveFileDialog(self):
caption = '%s - Choose File' % __appname__
filters = 'File (*%s)' % LabelFile.suffix
openDialogPath = self.currentPath()
dlg = QFileDialog(self, caption, openDialogPath, filters)
dlg.setDefaultSuffix(LabelFile.suffix[1:])
dlg.setAcceptMode(QFileDialog.AcceptSave)
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
dlg.selectFile(filenameWithoutExtension)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
if dlg.exec_():
return dlg.selectedFiles()[0]
return ''
def _saveFile(self, annotationFilePath):
if annotationFilePath and self.saveLabels(annotationFilePath):
self.setClean()
self.statusBar().showMessage('Saved to %s' % annotationFilePath)
self.statusBar().show()
def closeFile(self, _value=False):
if not self.mayContinue():
return
self.resetState()
self.setClean()
self.toggleActions(False)
self.canvas.setEnabled(False)
self.actions.saveAs.setEnabled(False)
# Message Dialogs. #
def hasLabels(self):
if not self.itemsToShapes:
self.errorMessage(
u'No objects labeled',
u'You must label at least one object to save the file.')
return False
return True
def mayContinue(self):
return not (self.dirty and not self.discardChangesDialog())
def discardChangesDialog(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = u'You have unsaved changes, proceed anyway?'
return yes == QMessageBox.warning(self, u'Attention', msg, yes | no)
def errorMessage(self, title, message):
return QMessageBox.critical(self, title,
'<p><b>%s</b></p>%s' % (title, message))
def currentPath(self):
return os.path.dirname(self.filePath) if self.filePath else '.'
def chooseColor1(self):
color = self.colorDialog.getColor(self.lineColor,
u'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.lineColor = color
# Change the color for all shape lines:
Shape.line_color = self.lineColor
self.canvas.update()
self.setDirty()
def chooseColor2(self):
color = self.colorDialog.getColor(self.fillColor,
u'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.fillColor = color
Shape.fill_color = self.fillColor
self.canvas.update()
self.setDirty()
def deleteSelectedShape(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = u'You are about to permanently delete this Box, proceed anyway?'
if yes == QMessageBox.warning(self, u'Attention', msg, yes | no):
self.remLabel(self.canvas.deleteSelected())
self.setDirty()
if self.noShapes():
for action in self.actions.onShapesPresent:
action.setEnabled(False)
def chshapeLineColor(self):
color = self.colorDialog.getColor(self.lineColor,
u'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.canvas.selectedShape.line_color = color
self.canvas.update()
self.setDirty()
def chshapeFillColor(self):
color = self.colorDialog.getColor(self.fillColor,
u'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.canvas.selectedShape.fill_color = color
self.canvas.update()
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addLabel(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
def loadPredefinedClasses(self):
predefined_classes_path = os.path.join('data',
'predefined_classes.txt')
if os.path.exists(predefined_classes_path) is True:
with codecs.open(predefined_classes_path, 'r', 'utf8') as f:
for line in f:
line = line.strip()
if self.labelHist is None:
self.lablHist = [line]
else:
self.labelHist.append(line)
def loadPascalXMLByFilename(self, xmlPath):
if self.filePath is None:
return
if os.path.isfile(xmlPath) is False:
return
tVocParseReader = PascalVocReader(xmlPath)
shapes = tVocParseReader.getShapes()
self.loadLabels(shapes)
class MainWindow(QMainWindow):
def __init__(self, defaultFilename=None, defaultSaveDir=None):
super(MainWindow, self).__init__()
self.resize(800, 600)
self.setWindowTitle(__appname__)
self.filePath = defaultFilename
self.lastOpenDir = None
self.dirname = None
self.mImgList = []
self.defaultSaveDir = defaultSaveDir
self.filename = None
#最近打开文件
self.recentFiles = []
self.maxRecent = 7
#QAction关联到QToolBar并以QToolButton显示出来
action = partial(newAction, self)
open = action('&Open', self.openFile,
'Ctrl+O', 'open', u'Open image or label file')
opendir = action('&Open Dir', self.openDirDialog,
'Ctrl+u', 'open', u'Open Dir')
changeSavedir = action('&Change Save Dir', self.changeSavedirDialog,
'Ctrl+r', 'open', u'Change default saved Annotation dir')
openNextImg = action('&Next Image', self.openNextImg,
'd', 'next', u'Open Next')
openPrevImg = action('&Prev Image', self.openPrevImg,
'a', 'prev', u'Open Prev')
verify = action('&Verify Image', self.verifyImg,
'space', 'verify', u'Verify Image')
tools = {open : 'Open', opendir : 'Open Dir', changeSavedir : 'Change Save Dir',
openNextImg : 'Next Image', openPrevImg : 'Prev Image', verify : 'Verify Image'}
tools = tools.items()
for act, title in tools:
toolbar = ToolBar(title)
toolbar.setObjectName(u'%s ToolBar' % title)
toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
toolbar.addAction(act)
self.addToolBar(Qt.LeftToolBarArea, toolbar)
#中间显示图片区域
self.canvas = Canvas(parent=self)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollArea = scroll
self.setCentralWidget(scroll)
#右侧文件列表
self.fileListWidget = QListWidget()
self.fileListWidget.itemDoubleClicked.connect(self.fileitemDoubleClicked)
filelistLayout = QVBoxLayout()
filelistLayout.setContentsMargins(0, 0, 0, 0)
filelistLayout.addWidget(self.fileListWidget)
fileListContainer = QWidget()
fileListContainer.setLayout(filelistLayout)
self.filedock = QDockWidget(u'File List', self)
self.filedock.setObjectName(u'Files')
self.filedock.setWidget(fileListContainer)
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
#状态栏
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
self.labelCoordinates = QLabel('')
self.statusBar().addPermanentWidget(self.labelCoordinates)
def resetState(self):
self.filePath = None
self.imageData = None
self.canvas.resetState()
self.labelCoordinates.clear()
def errorMessage(self, title, message):
return QMessageBox.critical(self, title, '<p><b>%s</b></p>%s' % (title, message))
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.adjustSize()
self.canvas.update()
def addRecentFile(self, filePath):
if filePath in self.recentFiles:
self.recentFiles.remove(filePath)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filePath)
def loadFile(self, filePath=None):
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
return
print(filePath)
self.filename = os.path.split(filePath)[-1]
if self.fileListWidget.count() > 0:
index = self.mImgList.index(filePath)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
if os.path.exists(filePath):
self.imageData = read(filePath, None)
self.canvas.verified = False
image = QImage.fromData(self.imageData)
if image.isNull():
self.errorMessage(u'Error opening file', u"<p>Make sure <i>%s</i> is a valid image file." % filePath)
self.status("Error reading %s" % filePath)
return False
self.status("Loaded %s" % os.path.basename(filePath))
self.image = image
self.filePath = filePath
self.canvas.loadPixmap(QPixmap.fromImage(image))
self.canvas.setEnabled(True)
self.paintCanvas()
self.addRecentFile(self.filePath)
self.setWindowTitle(__appname__ + ' ' + filePath)
self.canvas.setFocus(True)
def openFile(self):
path = os.path.dirname(self.filePath) if self.filePath else '.'
formats = ['*.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
filters = "Image & Label files (%s)" % ' '.join(formats)
filename = QFileDialog.getOpenFileName(self, '%s - Choose Image or Label file' % __appname__, path, filters)
if filename:
if isinstance(filename, (tuple, list)):
filename = filename[0]
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relativePath = os.path.join(root, file)
path = os.path.abspath(relativePath)
images.append(path)
images.sort(key=lambda x: x.lower())
return images
def openNextImg(self):
if len(self.mImgList) <= 0:
return
filename = None
if self.filePath is None:
filename = self.mImgList[0]
else:
currIndex = self.mImgList.index(self.filePath)
if currIndex + 1 < len(self.mImgList):
filename = self.mImgList[currIndex + 1]
if filename:
self.loadFile(filename)
def importDirImages(self, dirpath):
self.lastOpenDir = dirpath
self.dirname = dirpath
self.filePath = None
self.fileListWidget.clear()
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def openDirDialog(self):
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
defaultOpenDirPath = self.lastOpenDir
else:
defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
targetDirPath = QFileDialog.getExistingDirectory(self, '%s - Open Directory' % __appname__,
defaultOpenDirPath, QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
self.importDirImages(targetDirPath)
def changeSavedirDialog(self):
if self.defaultSaveDir is not None:
path = self.defaultSaveDir
else:
path = '.'
dirpath = QFileDialog.getExistingDirectory(self, '%s - Save text to the directory' % __appname__, path,
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
if dirpath is not None and len(dirpath) > 1:
self.defaultSaveDir = dirpath
self.statusBar().showMessage('%s . Annotation will be saved to %s' % ('Change saved folder',
self.defaultSaveDir))
self.statusBar().show()
def openPrevImg(self):
if len(self.mImgList) <= 0:
return
if self.filePath is None:
return
currIndex = self.mImgList.index(self.filePath)
if currIndex - 1 >= 0:
filename = self.mImgList[currIndex - 1]
if filename:
self.loadFile(filename)
def verifyImg(self):
if self.defaultSaveDir is None:
self.changeSavedirDialog()
fname = self.filename.split('.')[0] + '.txt'
fname = os.path.join(self.defaultSaveDir, fname)
print(fname)
with open(fname, 'w') as f:
f.write(self.filename)
for point in self.canvas.points:
f.write(' ' + str(point.x()) + ' ' + str(point.y()))
self.openNextImg()
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(item.text())
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
class SWQT(QMainWindow):
def __init__(self, parent=None):
super(SWQT, self).__init__(parent)
self.initUI()
self.src = None
self.right_src = None
self.cfg = ConfigParser()
self.cfg.read('SWQT.ini')
self.canvas.setDrawColor(QColor(self.cfg.getint('file', 'drawColorRed'), self.cfg.getint('file', 'drawColorGreen'), self.cfg.getint('file', 'drawColorBlue')))
self.leftCameraId = self.cfg.getint('binocular_cameras', 'leftCameraId')
self.rightCameraId = self.cfg.getint('binocular_cameras', 'rightCameraId')
self.m_camera = []
self.bContinuouslyCaptured = False
if self.leftCameraId > -1:
self.left_cap = cv2.VideoCapture(self.leftCameraId)
if self.rightCameraId > -1:
self.right_cap = cv2.VideoCapture(self.rightCameraId)
def initUI(self):
self.setWindowTitle('SWQT')
self.setMinimumSize(700, 500)
self.showMaximized()
self.status = self.statusBar()
self.status.showMessage('Starting successfully', 5000)
self.setWindowIcon(QIcon('imgs/SWQT.png'))
self.dock = QDockWidget(self)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
dockScroll = QScrollArea(self.dock)
self.volumeMeasure = VolumeMeasure()
self.dockWidget = SWQTDocker(self.volumeMeasure)
dockScroll.setWidget(self.dockWidget)
dockScroll.setWidgetResizable(False)
self.dock.setWidget(dockScroll)
scroll = QScrollArea(self)
self.canvas = Canvas(self.dockWidget)
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(False)
self.setCentralWidget(scroll)
scroll.setAlignment(Qt.AlignCenter)
file = self.menuBar().addMenu('&File')
process = self.menuBar().addMenu('&Process')
settings = self.menuBar().addMenu('&Settings')
binocularCameras = self.menuBar().addMenu('&Binocular Cameras')
toolbarFile = self.addToolBar('File')
self.addToolBar(Qt.LeftToolBarArea, toolbarFile)
toolbarProcess = self.addToolBar('Process')
self.addToolBar(Qt.LeftToolBarArea, toolbarProcess)
toolbarSettings = self.addToolBar('Settings')
self.addToolBar(Qt.LeftToolBarArea, toolbarSettings)
toolbarBinocularCameras = self.addToolBar('Binocular Cameras')
self.addToolBar(Qt.LeftToolBarArea, toolbarBinocularCameras)
openAPicAction = create_action('imgs/openAPic.png', '&Open a picture', self, 'Ctrl+O', 'open a picture', file, toolbarFile)
openAPicAction.triggered.connect(self.openAPic)
quitAction = create_action('imgs/quit.jpg', '&Quit', self, 'Ctrl+Q', 'quit', file, toolbarFile)
quitAction.triggered.connect(self.close)
drawAction = create_action('imgs/draw.png', '&Draw', self, 'Ctrl+D', 'draw', file, toolbarFile)
drawAction.triggered.connect(self.draw)
openRightPicAction = create_action('imgs/openRightPic.jpg', '&Open right picture', self, 'Ctrl+2', 'open right picutre', file, toolbarFile)
openRightPicAction.triggered.connect(self.openRightPic)
refreshImageAction = create_action('imgs/refreshImage.jpg', '&Refresh image', self, 'Ctrl+R', 'refresh image', process, toolbarProcess)
refreshImageAction.triggered.connect(self.refreshImage)
scaleAction = create_action('imgs/scale.png', '&Scale image', self, 'Ctrl+S', 'scale image', process, toolbarProcess)
scaleAction.triggered.connect(self.scale)
someProcessAction = create_action('imgs/someProcess.jpg', '&Show some process', self, 'Ctrl+P', 'show some process', process, toolbarProcess)
someProcessAction.triggered.connect(self.someProcess)
settingsAction = create_action('imgs/settings.jpg', '&Settings', self, 'Ctrl+E', 'settings', settings, toolbarSettings)
settingsAction.triggered.connect(self.settings)
selectDrawColorAction = create_action('imgs/selectDrawColor.png', '&Select draw color', self, 'Ctrl+L', 'select draw color', file, toolbarFile)
selectDrawColorAction.triggered.connect(self.selectDrawColor)
clearAllShapesAction = create_action('imgs/clearAllShapes.png', '&Clear all shapes', self, 'Ctrl+1', 'clear all shapes', file, toolbarFile)
clearAllShapesAction.triggered.connect(self.clearAllShapes)
openCamerasAction = create_action('imgs/openBinocularCameras.jpg', '&Open binocular cameras', self, 'Ctrl+B', 'open binocular cameras', binocularCameras, toolbarBinocularCameras)
openCamerasAction.triggered.connect(self.openCameras)
measureAction = create_action('imgs/measure.png', '&Measure', self, 'Ctrl+M', 'measure', binocularCameras, toolbarBinocularCameras)
measureAction.triggered.connect(self.measure)
showCamerasAction = create_action('imgs/showCameras.jpg', '&Show binocular cameras', self, 'Ctrl+3', 'show binocular cameras', binocularCameras, toolbarBinocularCameras)
showCamerasAction.triggered.connect(self.showCameras)
def openAPic(self):
self.setBScale(False)
self.setBDraw(False)
pic_file, _ = QFileDialog.getOpenFileName(self, 'open a picture', self.cfg.get('file', 'lastOpenDir') if os.path.exists(self.cfg.get('file', 'lastOpenDir')) else '.', 'All files(*.*);;Image files(*.bmp *.jpg *.pbm *.png *.ppm *.xbm *.xpm)')
if os.path.isfile(pic_file):
self.cfg.set('file', 'lastOpenDir', os.path.dirname(pic_file))
with open('SWQT.ini', 'w') as configfile:
self.cfg.write(configfile)
self.resetState()
self.src = cv2.imread(pic_file)
if self.src is None:
messageBox_info('the path should include Englise symbol only or we could not read a .gif file', self)
else:
self.loadImage(npndarray2qimage(self.src))
def setPicScale(self, picScale):
pass
def resetState(self):
self.src = None
self.canvas.resetState()
def paintCanvas(self):
self.canvas.adjustSize()
self.canvas.update()
def refreshImage(self):
self.setBScale(False)
self.setBDraw(False)
if self.src is not None:
self.canvas.set_picScale(self.dockWidget.picScaleSlider.value()/10)
self.canvas.load_image(npndarray2qimage(self.src))
def loadImage(self, image):
self.status.showMessage('loading picture', 5000)
self.canvas.load_image(image)
self.paintCanvas()
def scale(self):
self.setBScale(True)
self.setBDraw(False)
def someProcess(self):
self.setBScale(False)
self.setBDraw(False)
if self.src != None:
toProcess = self.someProcessDetailed(self.src)
self.canvas.set_picScale(self.dockWidget.picScaleSlider.value()/10)
self.canvas.load_image(npndarray2qimage(toProcess))
def settings(self):
self.setBScale(False)
self.setBDraw(False)
settingsDialog = SettingsDialog(self.volumeMeasure, self)
settingsDialog.exec_()
def draw(self):
self.setBScale(False)
self.setBDraw(True)
def setBScale(self, bScale):
if bScale != self.dockWidget.bScaleCheckBox.isChecked():
self.canvas.bScale = bScale
if bScale == True:
self.dockWidget.bScaleCheckBox.setChecked(True)
else:
self.dockWidget.bScaleCheckBox.setChecked(False)
def setBDraw(self, bDraw):
if bDraw != self.dockWidget.bDrawCheckBox.isChecked():
self.canvas.bDraw = bDraw
if bDraw == True:
self.dockWidget.bDrawCheckBox.setChecked(True)
else:
self.dockWidget.bDrawCheckBox.setChecked(False)
def selectDrawColor(self):
self.setBScale(False)
colorDialog = QColorDialog(self)
colorDialog.setWindowTitle('Select Draw Color: ')
colorDialog.setCurrentColor(Qt.red)
if colorDialog.exec_() == QColorDialog.Accepted:
drawColor = colorDialog.selectedColor()
self.canvas.setDrawColor(drawColor)
red = drawColor.red()
green = drawColor.green()
blue = drawColor.blue()
self.cfg.set('file', 'drawColorRed', str(red))
self.cfg.set('file', 'drawColorGreen', str(green))
self.cfg.set('file', 'drawColorBlue', str(blue))
self.cfg.write(open('SWQT.ini', 'w'))
def clearAllShapes(self):
self.setBScale(False)
self.canvas.clearAllShapes()
self.paintCanvas()
def openCameras(self):
self.setBScale(False)
self.setBDraw(False)
for info in QCameraInfo.availableCameras():
self.m_camera.append(info)
if len(self.m_camera) < 2:
messageBox_info('please get your binocular cameras', self)
else:
self.camerasSelectDialog = QDialog(self)
self.camerasSelectDialog.setMinimumWidth(950)
self.camerasSelectDialog.setWindowTitle('Select your disired cameras')
camerasSelectDialogLayout = QGridLayout()
leftLabel = QLabel('left')
rightLabel = QLabel('right')
leftCameraButtonGroup = QButtonGroup(self.camerasSelectDialog)
rightCameraButtonGroup = QButtonGroup(self.camerasSelectDialog)
camerasSelectDialogLayout.addWidget(leftLabel, 1, 0)
camerasSelectDialogLayout.addWidget(rightLabel, 2, 0)
for i in range(len(self.m_camera)):
capLabel = QLabel()
capLabel.setFixedSize(400, 300)
cap = cv2.VideoCapture(i)
_, then_cap = cap.read()
cap.release()
cv2.destroyAllWindows()
capLabel.setPixmap(QPixmap.fromImage(npndarray2qimage(then_cap)))
camerasSelectDialogLayout.addWidget(capLabel, 0, i+1)
indexStr = 'Index: %d'%i
optionalLeftCameraRadioButton = QRadioButton(indexStr)
leftCameraButtonGroup.addButton(optionalLeftCameraRadioButton, i)
optionalRightCameraRadioButton = QRadioButton(indexStr)
rightCameraButtonGroup.addButton(optionalRightCameraRadioButton, i)
if i == 0:
optionalLeftCameraRadioButton.setChecked(True)
optionalRightCameraRadioButton.setChecked(True)
camerasSelectDialogLayout.addWidget(optionalLeftCameraRadioButton, 1, i+1)
camerasSelectDialogLayout.addWidget(optionalRightCameraRadioButton, 2, i+1)
button = QDialogButtonBox(self.camerasSelectDialog)
button.addButton('OK', QDialogButtonBox.YesRole)
button.addButton('Cancel', QDialogButtonBox.NoRole)
button.accepted.connect(self.camerasSelectDialog.accept)
button.rejected.connect(self.camerasSelectDialog.reject)
camerasSelectDialogLayout.addWidget(button, 3, 0, 1, -1)
self.camerasSelectDialog.setLayout(camerasSelectDialogLayout)
res = self.camerasSelectDialog.exec_()
if res == QDialog.Accepted:
if self.leftCameraId > -1:
self.left_cap.release()
if self.rightCameraId > -1:
self.right_cap.release()
cv2.destroyAllWindows()
self.leftCameraId = leftCameraButtonGroup.checkedId()
self.rightCameraId = rightCameraButtonGroup.checkedId()
self.cfg.set('binocular_cameras', 'leftCameraId', str(self.leftCameraId))
self.cfg.set('binocular_cameras', 'rightCameraId', str(self.rightCameraId))
self.cfg.write(open('SWQT.ini', 'w'))
self.left_cap = cv2.VideoCapture(self.leftCameraId)
self.right_cap = cv2.VideoCapture(self.rightCameraId)
self.m_camera.clear()
def openRightPic(self):
self.setBScale(False)
self.setBDraw(False)
picFile, _ = QFileDialog.getOpenFileName(self, 'open a right picture', self.cfg.get('file', 'lastOpenRightDir') if os.path.exists(self.cfg.get('file', 'lastOpenRightDir')) else '.', 'All files(*.*);;Image files(*.bmp *.jpg *.pbm *.png *.ppm *.xbm *.xpm)')
if os.path.isfile(picFile):
self.cfg.set('file', 'lastOpenRightDir', os.path.dirname(picFile))
with open('SWQT.ini', 'w') as configfile:
self.cfg.write(configfile)
self.right_src = cv2.imread(picFile)
if self.right_src is None:
messageBox_info('the path should include Englise symbol only or we could not read a .gif file', self)
else:
cv2.imshow('right pic', self.right_src)
def measure(self):
if self.dockWidget.picRadioButton.isChecked():
if self.right_src == None or self.src == None or self.src.shape != self.right_src.shape:
messageBox_info('please get two proper pictures', self)
else:
left = self.someProcessDetailed(self.src)
right = self.someProcessDetailed(self.right_src)
self.canvas.load_image(npndarray2qimage(left))
self.volumeMeasure.detect(left, right, self.src, self.right_src)
elif self.leftCameraId <= -1 or self.rightCameraId < 0:
messageBox_info('Please get your binocular cameras', self)
else:
if self.dockWidget.captureSelectedPicRadioButton.isChecked():
_, left_frame = self.left_cap.read()
_, right_frame = self.right_cap.read()
if left_frame.shape == right_frame.shape:
self.src = left_frame
self.right_src = right_frame
left = self.someProcessDetailed(left_frame)
right = self.someProcessDetailed(right_frame)
self.canvas.load_image(npndarray2qimage(left))
cv2.imshow('right pic', right_frame)
if not self.volumeMeasure.bAuto:
cv2.waitKey()
ret = self.volumeMeasure.detect(left, right, left_frame, right_frame, self.canvas.shapes)
if ret == 1:
messageBox_info('please binarify the image(eg. gray, canny or thresh)', self)
else:
messageBox_info('you possibly get wrong binocular cameras', self)
self.openCameras()
else:
self.bContinuouslyCaptured = not self.bContinuouslyCaptured
bChecked = False
while(self.bContinuouslyCaptured):
_, left_frame = self.left_cap.read()
_, right_frame = self.right_cap.read()
if not bChecked:
bChecked = True
if left_frame.shape != right_frame.shape:
messageBox_info('you possibly get wrong binocular cameras', self)
self.openCameras()
self.src = left_frame
self.right_src = right_frame
left = self.someProcessDetailed(left_frame)
right = self.someProcessDetailed(right_frame)
self.canvas.load_image(npndarray2qimage(left))
cv2.imshow('right pic', right_frame)
if not self.volumeMeasure.bAuto:
cv2.waitKey()
ret = self.volumeMeasure.detect(left, right, left_frame, right_frame, self.canvas.shapes)
ret = self.volumeMeasure.detect(left, right, left_frame, right_frame)
if ret == 1:
messageBox_info('please binarify the image(eg.gray, canny or thresh)', self)
if not self.bContinuouslyCaptured or ret != 0:
cv2.destroyWindow('right pic')
break
def someProcessDetailed(self, cv_image):
if cv_image.shape[2] == 4:
cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGRA2BGR)
if len(cv_image.shape) > 2 and self.dockWidget.grayCheckBox.isChecked():
cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
if self.dockWidget.gaussianBlurCheckBox.isChecked():
gaussianBlurKSize = self.cfg.getint('process', 'gaussianBlurKSize')
cv_image = cv2.GaussianBlur(cv_image, (gaussianBlurKSize, gaussianBlurKSize), 1, 1)
if self.dockWidget.CLAHECheckBox.isChecked():
clahe = cv2.createCLAHE(clipLimit=self.cfg.getfloat('process', 'CLAHECoeff'))
if len(cv_image.shape) == 2:
cv_image = clahe.apply(cv_image)
else:
clahe_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2Lab)
splited = cv2.split(clahe_image)
dst = clahe.apply(splited[0])
np.copyto(splited[0], dst)
cv_image = cv2.merge(splited)
cv_image = cv2.cvtColor(cv_image, cv2.COLOR_Lab2BGR)
if self.dockWidget.sharpenCheckBox.isChecked():
laplacian = np.ndarray([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
cv_image = cv2.filter2D(cv_image, -1, laplacian)
if self.dockWidget.laplacianCheckBox.isChecked():
laplacian = np.ndarray([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
cv_image = cv2.filter2D(cv_image, -1, laplacian)
if self.dockWidget.OTSUCheckBox.isChecked():
if len(cv_image.shape) == 2:
_, cv_image = cv2.threshold(cv_image, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
else:
splited = cv2.split(cv_image)
for i in range(3):
_, splited[i] = cv2.threshold(splited[i], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
cv_image = cv2.merge(splited)
if self.dockWidget.RCFContourDetectCheckBox.isChecked():
pass
if self.dockWidget.cannyCheckBox.isChecked():
cannyLowThresh = self.cfg.getint('process', 'cannyLowThresh')
cannyHighThresh = self.cfg.getint('process', 'cannyHighThresh')
if len(cv_image.shape) == 2:
cv_image = cv2.Canny(cv_image, cannyLowThresh, cannyHighThresh)
else:
splited = cv2.split(cv_image)
for i in range(3):
splited[i] = cv2.Canny(splited[i], cannyLowThresh, cannyHighThresh)
cv_image = cv2.merge(splited)
return cv_image
def showCameras(self):
if self.leftCameraId < 0 or self.rightCameraId < 0:
messageBox_info('please get your binocular cameras', self)
self.openCameras()
else:
while True:
_, left_frame = self.left_cap.read()
_, right_frame = self.right_cap.read()
cv2.imshow('left_frame', left_frame)
cv2.imshow('right_frame', right_frame)
c = cv2.waitKey(25)
if c == ord('q') or c == 27:
break
cv2.destroyAllWindows()
def closeEvent(self, *args, **kwargs):
if self.leftCameraId > -1:
self.left_cap.release()
if self.rightCameraId > -1:
self.right_cap.release()
cv2.destroyAllWindows()
if self.dockWidget.grayCheckBox.isChecked():
self.cfg.set('process', 'bGray', '1')
else:
self.cfg.set('process', 'bGray', '0')
if self.dockWidget.gaussianBlurCheckBox.isChecked():
self.cfg.set('process', 'bGaussianBlur', '1')
else:
self.cfg.set('process', 'bGaussianBlur', '0')
if self.dockWidget.CLAHECheckBox.isChecked():
self.cfg.set('process', 'bCLAHE', '1')
else:
self.cfg.set('process', 'bCLAHE', '0')
if self.dockWidget.sharpenCheckBox.isChecked():
self.cfg.set('process', 'bSharpen', '1')
else:
self.cfg.set('process', 'bSharpen', '0')
if self.dockWidget.laplacianCheckBox.isChecked():
self.cfg.set('process', 'bLaplacian', '1')
else:
self.cfg.set('process', 'bLaplacian', '0')
if self.dockWidget.OTSUCheckBox.isChecked():
self.cfg.set('process', 'bOTSU', '1')
else:
self.cfg.set('process', 'bOTSU', '0')
if self.dockWidget.RCFContourDetectCheckBox.isChecked():
self.cfg.set('process', 'bRCFContourDetect', '1')
else:
self.cfg.set('process', 'bRCFContourDetect', '0')
if self.dockWidget.cannyCheckBox.isChecked():
self.cfg.set('process', 'bCanny', '1')
else:
self.cfg.set('process', 'bCanny', '0')
if self.dockWidget.rectRadioButton.isChecked():
self.cfg.set('draw', 'drawWay', 'Rect')
elif self.dockWidget.circleRadioButton.isChecked():
self.cfg.set('draw', 'drawWay', 'Circle')
elif self.dockWidget.polyRadioButton.isChecked():
self.cfg.set('draw', 'drawWay', 'Poly')
else:
self.cfg.set('draw', 'drawWay', 'Eraser')
if self.dockWidget.picRadioButton.isChecked():
self.cfg.set('binocular_cameras', 'captureWay', 'Pic')
elif self.dockWidget.captureSelectedPicRadioButton.isChecked():
self.cfg.set('binocular_cameras', 'captureWay', 'CaptureSelectedPic')
else:
self.cfg.set('binocular_cameras', 'captureWay', 'Capture')
if self.dockWidget.cameraToGroundHeightCheckBox.isChecked():
self.cfg.set('binocular_cameras', 'bFromTop', '1')
else:
self.cfg.set('binocular_cameras', 'bFromTop', '0')
self.cfg.set('binocular_cameras', 'cameraToGroundHeight', self.dockWidget.cameraToGroundHeightLE.text())
if self.dockWidget.bVisualizedCheckBox.isChecked():
self.cfg.set('binocular_cameras', 'bVisualized', '1')
else:
self.cfg.set('binocular_cameras', 'bVisualized', '0')
if self.dockWidget.bCollectedCheckBox.isChecked():
self.cfg.set('binocular_cameras', 'bCollected', '1')
else:
self.cfg.set('binocular_cameras', 'bCollected', '0')
if self.dockWidget.bTimeRecordedCheckBox.isChecked():
self.cfg.set('binocular_cameras', 'bTimeRecorded', '1')
else:
self.cfg.set('binocular_cameras', 'bTimeRecorded', '0')
if self.dockWidget.bmmCheckBox.isChecked():
self.cfg.set('binocular_cameras', 'bmm', '1')
else:
self.cfg.set('binocular_cameras', 'bmm', '0')
if self.dockWidget.bAutoCheckBox.isChecked():
self.cfg.set('binocular_cameras', 'bAuto', '1')
else:
self.cfg.set('binocular_cameras', 'bAuto', '0')
self.cfg.set('binocular_cameras', 'processThing', self.dockWidget.processThingComboBox.currentText())
self.cfg.set('binocular_cameras', 'disparityStyle', self.dockWidget.disparityStyleComboBox.currentText())
self.cfg.set('binocular_cameras', 'constructionStyle', self.dockWidget.constructionStyleComboBox.currentText())
self.cfg.set('binocular_cameras', 'sizeCalculation', self.dockWidget.sizeCalculationComboBox.currentText())
self.cfg.write(open('SWQT.ini', 'w'))
class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None):
super().__init__()
self.dirty = True
self.mImgList = []
self.dirname = None
self._beginner = True
self.image_out_np = None
self.default_save_dir = None
# Application state
self.filePath = None
self.mattingFile = None
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
matResultShow = ResizedQWidget()
matResultShow.resize(150, 150)
self.pic = QLabel(matResultShow)
self.pic.resize(150, 150)
self.pic.setGeometry(50, 20, 150, 150)
# self.pic.resize(matResultShow.width(), matResultShow.height())
# self.pic.setScaledContents(True)
matResultShow.setLayout(listLayout)
self.resultdock = QDockWidget('Result Image', self)
# self.resultdock.adjustSize()
self.resultdock.setObjectName('result')
self.resultdock.setWidget(matResultShow)
self.resultdock.resize(150, 150)
self.fileListWidget = QListWidget()
self.fileListWidget.itemDoubleClicked.connect(
self.fileitemDoubleClicked)
fileListLayout = QVBoxLayout()
fileListLayout.setContentsMargins(0, 0, 0, 0)
fileListLayout.addWidget(self.fileListWidget)
fileListContainer = QWidget()
fileListContainer.setLayout(fileListLayout)
self.filedock = QDockWidget('File List', self)
self.filedock.setObjectName('Files')
self.filedock.setWidget(fileListContainer)
self.zoomWidget = ZoomWidget()
self.canvas = Canvas(parent=self)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.scrollArea = scroll
self.canvas.scrollRequest.connect(self.scrollRequest)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.resultdock)
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
self.resultdock.setFeatures(
self.resultdock.features() ^ self.dockFeatures)
# Actions
action = partial(newAction, self)
open_file = action('&Open', self.openFile, 'Ctrl+O', 'Open image')
open_dir = action('&Open Dir', self.openDir,
'Ctrl+D', 'Open image dir')
change_save_dir = action('&Change Save Dir', self.changeSavedirDialog)
# open_next_img = action('&Next Image', self.openNextImg,
# 'Ctrl+N', 'Open next image')
# open_pre_img = action('&Previous Image', self.openPreImg,
# 'Ctrl+M', 'Open previous image')
save = action('&Save', self.saveFile, 'Crl+S', 'Save output image')
create = action('Create\nRectBox', self.createShape,
'w', 'Draw a new Box')
matting = action('&Create\nMatting', self.grabcutMatting,
'e', 'GrabcutMatting')
self.scalers = {
self.FIT_WINDOW: self.scaleFitWindow,
self.FIT_WIDTH: self.scaleFitWidth,
# Set to one to scale to 100% when loading files.
self.MANUAL_ZOOM: lambda: 1,
}
# store actions for further handling
self.actions = struct(save=save, open_file=open_file,
open_dir=open_dir, change_save_dir=change_save_dir,
# open_next_img=open_next_img, open_pre_img=open_pre_img,
create=create, matting=matting)
# Auto saving: enable auto saving if pressing next
# self.autoSaving = QAction('Auto Saving', self)
# self.autoSaving.setCheckable(True)
# self.autoSaving.setChecked()
# set toolbar
self.tools = self.toolbar('Tools')
self.actions.all = (save, open_file, open_dir,
change_save_dir, create,
# open_pre_img, open_next_img,
matting)
addActions(self.tools, self.actions.all)
# set status
self.statusBar().showMessage('{} started.'.format(__appname__))
def okToContinue(self):
if self.dirty:
reply = QMessageBox.question(self, "Attention",
"you have unsaved changes, proceed anyway?",
QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel)
if reply == QMessageBox.Cancel:
return False
elif reply == QMessageBox.Yes:
return self.fileSave
return True
def resetState(self):
self.canvas.resetState()
def errorMessage(self, title, message):
return QMessageBox.critical(self, title,
'<p><b>%s</b></p>%s' % (title, message))
def beginner(self):
return self._beginner
def advanced(self):
return not self.beginner()
def openFile(self, _value=False):
path = os.path.dirname(self.filePath) if self.filePath else '.'
formats = ['*.%s' % fmt.data().decode("ascii").lower()
for fmt in QImageReader.supportedImageFormats()]
filters = "Image (%s)" % ' '.join(formats)
filename = QFileDialog.getOpenFileName(
self, '%s - Choose Image or Label file' % __appname__, path, filters)
if filename:
if isinstance(filename, (tuple, list)):
filename = filename[0]
self.loadFile(filename)
def openDir(self, dirpath=None):
defaultOpenDirPath = dirpth if dirpath else '.'
targetDirPath = QFileDialog.getExistingDirectory(self,
'%s - Open Directory' % __appname__, defaultOpenDirPath,
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
self.importDirImages(targetDirPath)
def importDirImages(self, dirpath):
self.fileListWidget.clear()
self.mImgList = self.scanAllImages(dirpath)
# self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def scanAllImages(self, folderPath):
extensions = ['.%s' % fmt.data().decode("ascii").lower()
for fmt in QImageReader.supportedImageFormats()]
imageList = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relativePath = os.path.join(root, file)
path = os.path.abspath(relativePath)
imageList.append(path)
imageList.sort(key=lambda x: x.lower())
return imageList
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(item.text())
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
def loadFile(self, filePath=None):
self.resetState()
self.canvas.setEnabled(False)
# highlight the file item
if filePath and self.fileListWidget.count() > 0:
index = self.mImgList.index(filePath)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
if filePath and os.path.exists(filePath):
# load image
self.imageData = read(filePath, None)
image = QImage.fromData(self.imageData)
if image.isNull():
self.errorMessage(u'Error opening file',
u'<p>Make sure <i>%s</i> is a valid image file.' % filePath)
self.status('Error reading %s' % filePath)
return False
self.status('Loaded %s' % os.path.basename(filePath))
self.image = image
self.filePath = filePath
self.canvas.loadPixmap(QPixmap.fromImage(image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
# self.toggleActions(True)
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def adjustScale(self, initial=False):
value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
self.zoomWidget.setValue(int(100 * value))
def scaleFitWindow(self):
"""Figure out the size of the pixmap in order to fit the main widget."""
e = 2.0 # So that no scrollbars are generated.
w1 = self.centralWidget().width() - e
h1 = self.centralWidget().height() - e
a1 = w1 / h1
# Calculate a new scale value based on the pixmap's aspect ratio.
w2 = self.canvas.pixmap.width() - 0.0
h2 = self.canvas.pixmap.height() - 0.0
a2 = w2 / h2
return w1 / w2 if a2 >= a1 else h1 / h2
def scaleFitWidth(self):
# The epsilon does not seem to work too well here.
w = self.centralWidget().width() - 2.0
return w / self.canvas.pixmap.width()
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.scale = 0.01 * self.zoomWidget.value()
self.canvas.adjustSize()
self.canvas.update()
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawMode(self, edit=True):
self.canvas.setEditing(edit)
self.actions.createMode.setEnabled(edit)
self.actions.editMode.setEnabled(not edit)
def grabcutMatting(self):
if self.mattingFile is None:
self.mattingFile = Grab_cut()
def format_shape(s):
return dict(line_color=s.line_color.getRgb(),
fill_color=s.fill_color.getRgb(),
points=[(p.x(), p.y()) for p in s.points])
shape = format_shape(self.canvas.shapes[-1])
self.image_out_np = self.mattingFile.image_matting(self.filePath,
shape, iteration=10)
self.showResultImg(self.image_out_np)
self.actions.save.setEnabled(True)
def showResultImg(self, image_np):
# resize to pic
factor = min(self.pic.width() /
image_np.shape[1], self.pic.height() / image_np.shape[0])
image_np = cv2.resize(image_np, None, fx=factor,
fy=factor, interpolation=cv2.INTER_AREA)
# image_np = cv2.resize((self.pic.height(), self.pic.width()))
image = QImage(image_np, image_np.shape[1],
image_np.shape[0], QImage.Format_ARGB32)
matImg = QPixmap(image)
self.pic.setPixmap(matImg)
def saveFile(self):
self._saveFile(self.saveFileDialog())
def _saveFile(self, saved_path):
print(saved_path)
if saved_path:
Grab_cut.resultSave(saved_path, self.image_out_np)
self.setClean()
self.statusBar().showMessage('Saved to %s' % saved_path)
self.statusBar().show()
def saveFileDialog(self):
caption = '%s - Choose File' % __appname__
filters = 'File (*%s)' % 'png'
if self.default_save_dir is not None and len(self.default_save_dir):
openDialogPath = self.default_save_dir
else:
openDialogPath = self.currentPath()
print(openDialogPath)
dlg = QFileDialog(self, caption, openDialogPath, filters)
dlg.setDefaultSuffix('png')
dlg.setAcceptMode(QFileDialog.AcceptSave)
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
dlg.selectFile(filenameWithoutExtension)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
if dlg.exec_():
return dlg.selectedFiles()[0]
return ''
def currentPath(self):
return os.path.dirname(self.filePath) if self.filePath else '.'
def changeSavedirDialog(self, _value=False):
if self.default_save_dir is not None:
path = self.default_save_dir
else:
path = '.'
dirpath = QFileDialog.getExistingDirectory(self,
'%s - Save annotations to the directory' % __appname__, path, QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks)
if dirpath is not None and len(dirpath) > 1:
self.default_save_dir = dirpath
self.statusBar().showMessage('%s . Annotation will be saved to %s' %
('Change saved folder', self.default_save_dir))
self.statusBar().show()
def setClean(self):
self.dirty = False
self.actions.save.setEnabled(False)
self.actions.create.setEnabled(True)
def openNextImg():
pass
def openPreImg():
pass
def scrollRequest(self, delta, orientation):
units = - delta / (8 * 15)
bar = self.scrollBars[orientation]
bar.setValue(bar.value() + bar.singleStep() * units)
class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = 0, 1, 2
def __init__(self, filename=None, output=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
# Whether we need to save or not.
self.dirty = False
self._noSelectionSlot = False
self._beginner = True
self.screencastViewer = "firefox"
self.screencast = "screencast.ogv"
# Main widgets and related state.
self.labelDialog = LabelDialog(parent=self)
self.labelList = QListWidget()
self.itemsToShapes = []
self.labelList.itemActivated.connect(self.labelSelectionChanged)
self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
self.labelList.itemDoubleClicked.connect(self.editLabel)
# Connect to itemChanged to detect checkbox changes.
self.labelList.itemChanged.connect(self.labelItemChanged)
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
listLayout.addWidget(self.labelList)
self.editButton = QToolButton()
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
self.labelListContainer = QWidget()
self.labelListContainer.setLayout(listLayout)
listLayout.addWidget(self.editButton)#, 0, Qt.AlignCenter)
listLayout.addWidget(self.labelList)
#define table view
self.table_view = QTableView(self)
self.table_view.setSelectionBehavior(QAbstractItemView.SelectRows)
self.table_view.show()
self.table_view.doubleClicked.connect(self.openCurrentData)
listLayout.addWidget(self.table_view)
self.dock = QDockWidget('Polygon Labels', self)
self.dock.setObjectName('Labels')
self.dock.setWidget(self.labelListContainer)
#self.dock.setWidget(self.infoLayout)
self.zoomWidget = ZoomWidget()
self.colorDialog = ColorDialog(parent=self)
self.canvas = Canvas()
self.canvas.zoomRequest.connect(self.zoomRequest)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
self.dockFeatures = QDockWidget.DockWidgetClosable\
| QDockWidget.DockWidgetFloatable
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
# Actions
action = partial(newAction, self)
quit = action('&Quit', self.close,
'Ctrl+Q', 'quit', 'Quit application')
open = action('&Open', self.openFile,
'Ctrl+O', 'open', 'Open image or label file')
save = action('&Save', self.saveFile,
'Ctrl+S', 'save', 'Save labels to file', enabled=False)
saveAs = action('&Save As', self.saveFileAs,
'Ctrl+Shift+S', 'save-as', 'Save labels to a different file',
enabled=False)
close = action('&Close', self.closeFile,
'Ctrl+W', 'close', 'Close current file')
color1 = action('Polygon &Line Color', self.chooseColor1,
'Ctrl+L', 'color_line', 'Choose polygon line color')
color2 = action('Polygon &Fill Color', self.chooseColor2,
'Ctrl+Shift+L', 'color', 'Choose polygon fill color')
createMode = action('Create\nPolygo&ns', self.setCreateMode,
'Ctrl+N', 'new', 'Start drawing polygons', enabled=False)
editMode = action('&Edit\nPolygons', self.setEditMode,
'Ctrl+J', 'edit', 'Move and edit polygons', enabled=False)
create = action('Create\nPolygo&n', self.createShape,
'Ctrl+N', 'new', 'Draw a new polygon', enabled=False)
delete = action('Delete\nPolygon', self.deleteSelectedShape,
'Delete', 'delete', 'Delete', enabled=False)
copy = action('&Duplicate\nPolygon', self.copySelectedShape,
'Ctrl+D', 'copy', 'Create a duplicate of the selected polygon',
enabled=False)
advancedMode = action('&Advanced Mode', self.toggleAdvancedMode,
'Ctrl+Shift+A', 'expert', 'Switch to advanced mode',
checkable=True)
hideAll = action('&Hide\nPolygons', partial(self.togglePolygons, False),
'Ctrl+H', 'hide', 'Hide all polygons',
enabled=False)
showAll = action('&Show\nPolygons', partial(self.togglePolygons, True),
'Ctrl+A', 'hide', 'Show all polygons',
enabled=False)
help = action('&Tutorial', self.tutorial, 'Ctrl+T', 'help',
'Show screencast of introductory tutorial')
zoom = QWidgetAction(self)
zoom.setDefaultWidget(self.zoomWidget)
self.zoomWidget.setWhatsThis(
"Zoom in or out of the image. Also accessible with"\
" %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
fmtShortcut("Ctrl+Wheel")))
self.zoomWidget.setEnabled(False)
zoomIn = action('Zoom &In', partial(self.addZoom, 10),
'Ctrl++', 'zoom-in', 'Increase zoom level', enabled=False)
zoomOut = action('&Zoom Out', partial(self.addZoom, -10),
'Ctrl+-', 'zoom-out', 'Decrease zoom level', enabled=False)
zoomOrg = action('&Original size', partial(self.setZoom, 100),
'Ctrl+=', 'zoom', 'Zoom to original size', enabled=False)
fitWindow = action('&Fit Window', self.setFitWindow,
'Ctrl+F', 'fit-window', 'Zoom follows window size',
checkable=True, enabled=False)
fitWidth = action('Fit &Width', self.setFitWidth,
'Ctrl+Shift+F', 'fit-width', 'Zoom follows window width',
checkable=True, enabled=False)
# Group zoom controls into a list for easier toggling.
zoomActions = (self.zoomWidget, zoomIn, zoomOut, zoomOrg, fitWindow, fitWidth)
self.zoomMode = self.MANUAL_ZOOM
self.scalers = {
self.FIT_WINDOW: self.scaleFitWindow,
self.FIT_WIDTH: self.scaleFitWidth,
# Set to one to scale to 100% when loading files.
self.MANUAL_ZOOM: lambda: 1,
}
edit = action('&Edit Label', self.editLabel,
'Ctrl+E', 'edit', 'Modify the label of the selected polygon',
enabled=False)
self.editButton.setDefaultAction(edit)
shapeLineColor = action('Shape &Line Color', self.chshapeLineColor,
icon='color_line', tip='Change the line color for this specific shape',
enabled=False)
shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor,
icon='color', tip='Change the fill color for this specific shape',
enabled=False)
labels = self.dock.toggleViewAction()
labels.setText('Show/Hide Label Panel')
labels.setShortcut('Ctrl+Shift+L')
# Lavel list context menu.
labelMenu = QMenu()
addActions(labelMenu, (edit, delete))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(self.popLabelListMenu)
# Store actions for further handling.
self.actions = struct(save=save, saveAs=saveAs, open=open, close=close,
lineColor=color1, fillColor=color2,
create=create, delete=delete, edit=edit, copy=copy,
createMode=createMode, editMode=editMode, advancedMode=advancedMode,
shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
fitWindow=fitWindow, fitWidth=fitWidth,
zoomActions=zoomActions,
fileMenuActions=(open,save,saveAs,close,quit),
beginner=(), advanced=(),
editMenu=(edit, copy, delete, None, color1, color2),
beginnerContext=(create, edit, copy, delete),
advancedContext=(createMode, editMode, edit, copy,
delete, shapeLineColor, shapeFillColor),
onLoadActive=(close, create, createMode, editMode),
onShapesPresent=(saveAs, hideAll, showAll))
self.menus = struct(
file=self.menu('&File'),
edit=self.menu('&Edit'),
view=self.menu('&View'),
help=self.menu('&Help'),
recentFiles=QMenu('Open &Recent'),
labelList=labelMenu)
addActions(self.menus.file,
(open, self.menus.recentFiles, save, saveAs, close, None, quit))
addActions(self.menus.help, (help,))
addActions(self.menus.view, (
labels, advancedMode, None,
hideAll, showAll, None,
zoomIn, zoomOut, zoomOrg, None,
fitWindow, fitWidth))
self.menus.file.aboutToShow.connect(self.updateFileMenu)
# Custom context menu for the canvas widget:
addActions(self.canvas.menus[0], self.actions.beginnerContext)
addActions(self.canvas.menus[1], (
action('&Copy here', self.copyShape),
action('&Move here', self.moveShape)))
self.tools = self.toolbar('Tools')
self.actions.beginner = (
open, save, None, create, copy, delete, None,
zoomIn, zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (
open, save, None,
createMode, editMode, None,
hideAll, showAll)
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
# Application state.
self.image = QImage()
self.filename = filename
self.labeling_once = output is not None
self.output = output
self.recentFiles = []
self.maxRecent = 7
self.lineColor = None
self.fillColor = None
self.zoom_level = 100
self.fit_window = False
# XXX: Could be completely declarative.
# Restore application settings.
self.settings = {}
self.recentFiles = self.settings.get('recentFiles', [])
size = self.settings.get('window/size', QSize(1000, 800))
position = self.settings.get('window/position', QPoint(0, 0))
self.resize(size)
self.move(position)
# or simply:
#self.restoreGeometry(settings['window/geometry']
self.restoreState(self.settings.get('window/state', QByteArray()))
self.lineColor = QColor(self.settings.get('line/color', Shape.line_color))
self.fillColor = QColor(self.settings.get('fill/color', Shape.fill_color))
Shape.line_color = self.lineColor
Shape.fill_color = self.fillColor
if self.settings.get('advanced', QVariant()):
self.actions.advancedMode.setChecked(True)
self.toggleAdvancedMode()
# Populate the File menu dynamically.
self.updateFileMenu()
# Since loading the file may take some time, make sure it runs in the background.
self.queueEvent(partial(self.loadFile, self.filename))
# Callbacks:
self.zoomWidget.valueChanged.connect(self.paintCanvas)
self.populateModeActions()
#self.firstStart = True
#if self.firstStart:
# QWhatsThis.enterWhatsThisMode()
## Support Functions ##
def openCurrentData(self,e):
row_idx, key = self.getCurrentData()
row_data = self.mymodel.getRow(row_idx)
try:
if row_data[1].lower()=="yes":
self.filename=self.labelFileDict[key]
else:
self.filename=self.imageFileDict[key]
self.loadFile(self.filename)
except KeyError as e:
print(e, " is not a dictionary key")
def getCurrentData(self):
index = self.table_view.selectionModel().currentIndex()
if index.column():
return index.row(), self.mymodel.getRow(index.row())[0]
else:
return index.row(), index.data()
def updateTable(self, table_data):
header = ["filename", "is labeled"]
self.mymodel=MyTableModel(table_data, header, self)
self.table_view.setModel(self.mymodel)
self.table_view.resizeColumnsToContents()
self.table_view.update()
def noShapes(self):
return not self.itemsToShapes
def toggleAdvancedMode(self, value=True):
self._beginner = not value
self.canvas.setEditing(True)
self.populateModeActions()
self.editButton.setVisible(not value)
if value:
self.actions.createMode.setEnabled(True)
self.actions.editMode.setEnabled(False)
self.dock.setFeatures(self.dock.features() | self.dockFeatures)
else:
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
def populateModeActions(self):
if self.beginner():
tool, menu = self.actions.beginner, self.actions.beginnerContext
else:
tool, menu = self.actions.advanced, self.actions.advancedContext
self.tools.clear()
addActions(self.tools, tool)
self.canvas.menus[0].clear()
addActions(self.canvas.menus[0], menu)
self.menus.edit.clear()
actions = (self.actions.create,) if self.beginner()\
else (self.actions.createMode, self.actions.editMode)
addActions(self.menus.edit, actions + self.actions.editMenu)
def setBeginner(self):
self.tools.clear()
addActions(self.tools, self.actions.beginner)
def setAdvanced(self):
self.tools.clear()
addActions(self.tools, self.actions.advanced)
def setDirty(self):
self.dirty = True
self.actions.save.setEnabled(True)
def setClean(self):
self.dirty = False
self.actions.save.setEnabled(False)
self.actions.create.setEnabled(True)
def toggleActions(self, value=True):
"""Enable/Disable widgets which depend on an opened image."""
for z in self.actions.zoomActions:
z.setEnabled(value)
for action in self.actions.onLoadActive:
action.setEnabled(value)
def queueEvent(self, function):
QTimer.singleShot(0, function)
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.itemsToShapes = []
self.labelList.clear()
self.filename = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
def currentItem(self):
items = self.labelList.selectedItems()
if items:
return items[0]
return None
def addRecentFile(self, filename):
if filename in self.recentFiles:
self.recentFiles.remove(filename)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filename)
def beginner(self):
return self._beginner
def advanced(self):
return not self.beginner()
## Callbacks ##
def tutorial(self):
subprocess.Popen([self.screencastViewer, self.screencast])
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
"""In the middle of drawing, toggling between modes should be disabled."""
self.actions.editMode.setEnabled(not drawing)
if not drawing and self.beginner():
# Cancel creation.
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.create.setEnabled(True)
def toggleDrawMode(self, edit=True):
self.canvas.setEditing(edit)
self.actions.createMode.setEnabled(edit)
self.actions.editMode.setEnabled(not edit)
def setCreateMode(self):
assert self.advanced()
self.toggleDrawMode(False)
def setEditMode(self):
assert self.advanced()
self.toggleDrawMode(True)
def updateFileMenu(self):
current = self.filename
def exists(filename):
return os.path.exists(str(filename))
menu = self.menus.recentFiles
menu.clear()
files = [f for f in self.recentFiles if f != current and exists(f)]
for i, f in enumerate(files):
icon = newIcon('labels')
action = QAction(
icon, '&%d %s' % (i+1, QFileInfo(f).fileName()), self)
action.triggered.connect(partial(self.loadRecent, f))
menu.addAction(action)
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
def editLabel(self, item=None):
if not self.canvas.editing():
return
item = item if item else self.currentItem()
text = self.labelDialog.popUp(item.text())
if text is not None:
item.setText(text)
self.setDirty()
# React to canvas signals.
def shapeSelectionChanged(self, selected=False):
if self._noSelectionSlot:
self._noSelectionSlot = False
else:
shape = self.canvas.selectedShape
if shape:
for item, shape_ in self.itemsToShapes:
if shape_ == shape:
break
item.setSelected(True)
else:
self.labelList.clearSelection()
self.actions.delete.setEnabled(selected)
self.actions.copy.setEnabled(selected)
self.actions.edit.setEnabled(selected)
self.actions.shapeLineColor.setEnabled(selected)
self.actions.shapeFillColor.setEnabled(selected)
def addLabel(self, shape):
item = QListWidgetItem(shape.label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.itemsToShapes.append((item, shape))
self.labelList.addItem(item)
for action in self.actions.onShapesPresent:
action.setEnabled(True)
def remLabel(self, shape):
for index, (item, shape_) in enumerate(self.itemsToShapes):
if shape_ == shape:
break
self.itemsToShapes.pop(index)
self.labelList.takeItem(self.labelList.row(item))
def loadLabels(self, shapes):
s = []
for label, points, line_color, fill_color in shapes:
shape = Shape(label=label)
for x, y in points:
shape.addPoint(QPointF(x, y))
shape.close()
s.append(shape)
self.addLabel(shape)
if line_color:
shape.line_color = QColor(*line_color)
if fill_color:
shape.fill_color = QColor(*fill_color)
self.canvas.loadShapes(s)
def saveLabels(self, filename):
lf = LabelFile()
def format_shape(s):
return dict(label=str(s.label),
line_color=s.line_color.getRgb()\
if s.line_color != self.lineColor else None,
fill_color=s.fill_color.getRgb()\
if s.fill_color != self.fillColor else None,
points=[(p.x(), p.y()) for p in s.points])
shapes = [format_shape(shape) for shape in self.canvas.shapes]
try:
lf.save(filename, shapes, str(self.filename), self.imageData,
self.lineColor.getRgb(), self.fillColor.getRgb())
self.labelFile = lf
self.filename = filename
return True
except LabelFileError as e:
self.errorMessage('Error saving label data',
'<b>%s</b>' % e)
return False
def copySelectedShape(self):
self.addLabel(self.canvas.copySelectedShape())
#fix copy and delete
self.shapeSelectionChanged(True)
def labelSelectionChanged(self):
item = self.currentItem()
if item and self.canvas.editing():
self._noSelectionSlot = True
for item_, shape in self.itemsToShapes:
if item_ == item:
break
self.canvas.selectShape(shape)
def labelItemChanged(self, item):
for item_, shape in self.itemsToShapes:
if item_ == item:
break
label = str(item.text())
if label != shape.label:
shape.label = str(item.text())
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
## Callback functions:
def newShape(self):
"""Pop-up and give focus to the label editor.
position MUST be in global coordinates.
"""
text = self.labelDialog.popUp()
if text is not None:
self.addLabel(self.canvas.setLastLabel(text))
if self.beginner(): # Switch to edit mode.
self.canvas.setEditing(True)
self.actions.create.setEnabled(True)
else:
self.actions.editMode.setEnabled(True)
self.setDirty()
else:
self.canvas.undoLastLine()
def scrollRequest(self, delta, orientation):
units = - delta * 0.1 # natural scroll
bar = self.scrollBars[orientation]
bar.setValue(bar.value() + bar.singleStep() * units)
def setZoom(self, value):
self.actions.fitWidth.setChecked(False)
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.MANUAL_ZOOM
self.zoomWidget.setValue(value)
def addZoom(self, increment=10):
self.setZoom(self.zoomWidget.value() + increment)
def zoomRequest(self, delta):
units = delta * 0.1
self.addZoom(units)
def setFitWindow(self, value=True):
if value:
self.actions.fitWidth.setChecked(False)
self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
self.adjustScale()
def setFitWidth(self, value=True):
if value:
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
self.adjustScale()
def togglePolygons(self, value):
for item, shape in self.itemsToShapes:
item.setCheckState(Qt.Checked if value else Qt.Unchecked)
def loadFile(self, filename=None):
"""Load the specified file, or the last opened file if None."""
self.resetState()
self.canvas.setEnabled(False)
if filename is None:
filename = self.settings.get('filename', '')
filename = str(filename)
if QFile.exists(filename):
if LabelFile.isLabelFile(filename):
try:
self.labelFile = LabelFile(filename)
# FIXME: PyQt4 installed via Anaconda fails to load JPEG
# and JSON encoded images.
# https://github.com/ContinuumIO/anaconda-issues/issues/131
if QImage.fromData(self.labelFile.imageData).isNull():
raise LabelFileError(
'Failed loading image data from label file.\n'
'Maybe this is a known issue of PyQt4 built on'
' Anaconda, and may be fixed by installing PyQt5.')
except LabelFileError as e:
self.errorMessage('Error opening file',
("<p><b>%s</b></p>"
"<p>Make sure <i>%s</i> is a valid label file.")\
% (e, filename))
self.status("Error reading %s" % filename)
return False
self.imageData = self.labelFile.imageData
self.lineColor = QColor(*self.labelFile.lineColor)
self.fillColor = QColor(*self.labelFile.fillColor)
else:
# Load image:
# read data first and store for saving into label file.
self.imageData = read(filename, None)
self.labelFile = None
image = QImage.fromData(self.imageData)
if image.isNull():
formats = ['*.{}'.format(fmt.data().decode())
for fmt in QImageReader.supportedImageFormats()]
self.errorMessage(
'Error opening file',
'<p>Make sure <i>{0}</i> is a valid image file.<br/>'
'Supported image formats: {1}</p>'
.format(filename, ','.join(formats)))
self.status("Error reading %s" % filename)
return False
self.status("Loaded %s" % os.path.basename(str(filename)))
self.image = image
self.filename = filename
self.canvas.loadPixmap(QPixmap.fromImage(image))
if self.labelFile:
self.loadLabels(self.labelFile.shapes)
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.addRecentFile(self.filename)
self.toggleActions(True)
return True
return False
def resizeEvent(self, event):
if self.canvas and not self.image.isNull()\
and self.zoomMode != self.MANUAL_ZOOM:
self.adjustScale()
super(MainWindow, self).resizeEvent(event)
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.scale = 0.01 * self.zoomWidget.value()
self.canvas.adjustSize()
self.canvas.update()
def adjustScale(self, initial=False):
value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
self.zoomWidget.setValue(int(100 * value))
def scaleFitWindow(self):
"""Figure out the size of the pixmap in order to fit the main widget."""
e = 2.0 # So that no scrollbars are generated.
w1 = self.centralWidget().width() - e
h1 = self.centralWidget().height() - e
a1 = w1/ h1
# Calculate a new scale value based on the pixmap's aspect ratio.
w2 = self.canvas.pixmap.width() - 0.0
h2 = self.canvas.pixmap.height() - 0.0
a2 = w2 / h2
return w1 / w2 if a2 >= a1 else h1 / h2
def scaleFitWidth(self):
# The epsilon does not seem to work too well here.
w = self.centralWidget().width() - 2.0
return w / self.canvas.pixmap.width()
def closeEvent(self, event):
if not self.mayContinue():
event.ignore()
s = self.settings
s['filename'] = self.filename if self.filename else ''
s['window/size'] = self.size()
s['window/position'] = self.pos()
s['window/state'] = self.saveState()
s['line/color'] = self.lineColor
s['fill/color'] = self.fillColor
s['recentFiles'] = self.recentFiles
s['advanced'] = not self._beginner
# ask the use for where to save the labels
#s['window/geometry'] = self.saveGeometry()
## User Dialogs ##
def loadRecent(self, filename):
if self.mayContinue():
self.loadFile(filename)
def openFile(self, _value=False):
if not self.mayContinue():
return
path = os.path.dirname(str(self.filename))\
if self.filename else '.'
formats = ['*.{}'.format(fmt.data().decode())
for fmt in QImageReader.supportedImageFormats()]
filters = "Image & Label files (%s)" % \
' '.join(formats + ['*%s' % LabelFile.suffix])
filename = QFileDialog.getOpenFileName(self,
'%s - Choose Image or Label file' % __appname__, path, filters)
if PYQT5:
filename, _ = filename
filename = str(filename)
self.filename=filename
self.makeImageAndLabelFileDict(os.path.split(filename)[0])
if hasattr(self, 'imageFileDict'):
self.updateTable(self.makeTableViewData())
if filename:
self.loadFile(filename)
def makeImageAndLabelFileDict(self, path):
try:
self.imageFileDict={}
self.labelFileDict={}
for fname in os.listdir(path):
if fname.endswith(".png"):
self.imageFileDict[fname.split('.')[0]]=os.path.join(path,fname)
elif fname.endswith(".json"):
self.labelFileDict[fname.split('.')[0]]=os.path.join(path,fname)
except FileNotFoundError:
msg=QMessageBox()
msg.setIcon(QMessageBox.Critical)
msg.setText("cannot find given path")
def makeTableViewData(self):
table_data=[]
if len(self.imageFileDict.keys())!=0:
for key in self.imageFileDict.keys():
if key in self.labelFileDict.keys():
table_data.append([key,'Yes'])
else:
table_data.append([key, 'no'])
else:
for key in self.labelFileDict.keys():
table_data.append([key,'Yes'])
return table_data
def saveFile(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
if self.hasLabels():
if self.labelFile:
self._saveFile(self.filename)
elif self.output:
self._saveFile(self.output)
else:
self._saveFile(self.saveFileDialog())
image = self.canvas.getLabelImage()
if image:
path, fname = os.path.split(self.filename)
if not os.path.exists(os.path.join(path,"label")):
os.mkdir(os.path.join(path,"label"))
fname=os.path.join("label","label_"+fname.replace("json","png"))
image.save(os.path.join(path,fname),"png")
src_filename=os.path.join("label",os.path.split(self.filename)[1])
self.image.save(os.path.join(path,src_filename.replace("json","png")),"png")
#table view update
row_idx, key = self.getCurrentData()
self.labelFileDict[key]=self.filename.replace("png", "json")
self.mymodel.setLabeled(row_idx)
self.table_view.update()
def saveFileAs(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
if self.hasLabels():
self._saveFile(self.saveFileDialog())
def saveFileDialog(self):
caption = '%s - Choose File' % __appname__
filters = 'Label files (*%s)' % LabelFile.suffix
dlg = QFileDialog(self, caption, self.currentPath(), filters)
dlg.setDefaultSuffix(LabelFile.suffix[1:])
dlg.setAcceptMode(QFileDialog.AcceptSave)
dlg.setOption(QFileDialog.DontConfirmOverwrite, False)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
basename = os.path.splitext(self.filename)[0]
default_labelfile_name = os.path.join(self.currentPath(),
basename + LabelFile.suffix)
filename = dlg.getSaveFileName(
self, 'Choose File', default_labelfile_name,
'Label files (*%s)' % LabelFile.suffix)
if PYQT5:
filename, _ = filename
filename = str(filename)
return filename
def _saveFile(self, filename):
if filename and self.saveLabels(filename):
self.addRecentFile(filename)
self.setClean()
if self.labeling_once:
self.close()
def closeFile(self, _value=False):
if not self.mayContinue():
return
self.resetState()
self.setClean()
self.toggleActions(False)
self.canvas.setEnabled(False)
self.actions.saveAs.setEnabled(False)
# Message Dialogs. #
def hasLabels(self):
if not self.itemsToShapes:
self.errorMessage('No objects labeled',
'You must label at least one object to save the file.')
return False
return True
def mayContinue(self):
return not (self.dirty and not self.discardChangesDialog())
def discardChangesDialog(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = 'You have unsaved changes, proceed anyway?'
return yes == QMessageBox.warning(self, 'Attention', msg, yes|no)
def errorMessage(self, title, message):
return QMessageBox.critical(self, title,
'<p><b>%s</b></p>%s' % (title, message))
def currentPath(self):
return os.path.dirname(str(self.filename)) if self.filename else '.'
def chooseColor1(self):
color = self.colorDialog.getColor(self.lineColor, 'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.lineColor = color
# Change the color for all shape lines:
Shape.line_color = self.lineColor
self.canvas.update()
self.setDirty()
def chooseColor2(self):
color = self.colorDialog.getColor(self.fillColor, 'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.fillColor = color
Shape.fill_color = self.fillColor
self.canvas.update()
self.setDirty()
def deleteSelectedShape(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = 'You are about to permanently delete this polygon, proceed anyway?'
if yes == QMessageBox.warning(self, 'Attention', msg, yes|no):
self.remLabel(self.canvas.deleteSelected())
self.setDirty()
if self.noShapes():
for action in self.actions.onShapesPresent:
action.setEnabled(False)
def chshapeLineColor(self):
color = self.colorDialog.getColor(self.lineColor, 'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.canvas.selectedShape.line_color = color
self.canvas.update()
self.setDirty()
def chshapeFillColor(self):
color = self.colorDialog.getColor(self.fillColor, 'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.canvas.selectedShape.fill_color = color
self.canvas.update()
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addLabel(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
class MainWindow(QMainWindow):
def __init__(self, defaultFilename=None):
super(MainWindow, self).__init__()
'''with fopen('log.txt', 'r') with f:
self.targetFile = f.readline().strip('\n')
targetSaveDir = f.readline().strip('\n')'''
self.resize(800, 600)
self.setWindowTitle(__appname__)
cf = configparser.ConfigParser()
cf.read('./TERSI-labelmkr.ini')
self.lastOpenDir = cf.get('env', 'lastOpenDir')
#with open('config/lastOpenDir.txt', 'r') as f:
#self.lastOpenDir = f.readline().strip('\n')
self.dirname = None
print('you have input some parameters: ' + str(defaultFilename))
self.filePath = defaultFilename
self.mImgList = []
'''if(targetSaveDir == 'None'):
self.defaultSaveDir = None
else:'''
self.defaultSaveDir = cf.get('env', 'defaultSaveDir')
self.cropSaveDir = cf.get('env', 'cropSaveDir')
'''with open('config/defaultSaveDir.txt', 'r') as f:
self.defaultSaveDir = f.readline().strip('\n')'''
self.filename = None
#最近打开文件
self.recentFiles = []
self.maxRecent = 7
self.picw = 0
self.pich = 0
#QAction关联到QToolBar并以QToolButton显示出来
action = partial(newAction, self)
open_ = action('&Open', self.openFile, 'Ctrl+O', 'open',
u'Open image or label file')
opendir = action('&Open Dir', self.openDirDialog, 'Ctrl+u', 'opendir',
u'Open Dir')
changeSavedir = action('&Change Save Dir', self.changeSavedirDialog,
'Ctrl+r', 'changesavedir',
u'Change default saved Annotation dir')
openNextImg = action('&Next Image', self.openNextImg, 'd', 'next',
u'Open Next')
openPrevImg = action('&Prev Image', self.openPrevImg, 'a', 'prev',
u'Open Prev')
verify = action('&Verify Image', self.verifyImg, 'space', 'verify',
u'Verify Image')
test = action('&Test Label', self.test, 't', 'test', u'Test Label')
test_sample = action('&Test Sample', self.testSample, 's',
'testsample', u'Test Sample')
delete_img = action('&Delete Img', self.delImg, 'del', 'delimg',
u'Delete Img')
crop_and_save = action('&Crop And Save', self.cropAndSave, 'c',
'cropandsave', u'Crop And Save')
clear = action('&Clear', self.clear, 'e', 'clear', u'Clear')
tools = {
open_: 'Open',
opendir: 'Open Dir',
changeSavedir: 'Change Save Dir',
openNextImg: 'Next Image',
openPrevImg: 'Prev Image',
verify: 'Verify Image',
test: 'Test',
test_sample: 'Test Sample',
delete_img: 'Delete Img',
crop_and_save: 'Crop And Save',
clear: 'Clear'
}
tools = tools.items()
toolbar = ToolBar('tools')
toolbar.setObjectName(u'%s ToolBar' % 'tools')
toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
for act, title in tools:
'''
toolbar = ToolBar(title)
toolbar.setObjectName(u'%s ToolBar' % title)
toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
toolbar.addAction(act)
self.addToolBar(Qt.LeftToolBarArea, toolbar)
'''
if act:
toolbar.addAction(act)
self.addToolBar(Qt.LeftToolBarArea, toolbar)
#中间显示图片区域
self.canvas = Canvas(parent=self)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollArea = scroll
self.setCentralWidget(scroll)
#右侧文件列表
self.fileListWidget = QListWidget()
self.fileListWidget.itemDoubleClicked.connect(
self.fileitemDoubleClicked)
filelistLayout = QVBoxLayout()
filelistLayout.setContentsMargins(0, 0, 0, 0)
filelistLayout.addWidget(self.fileListWidget)
fileListContainer = QWidget()
fileListContainer.setLayout(filelistLayout)
self.filedock = QDockWidget(u'File List', self)
self.filedock.setObjectName(u'Files')
self.filedock.setWidget(fileListContainer)
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
#状态栏
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
self.labelCoordinates = QLabel('')
self.statusBar().addPermanentWidget(self.labelCoordinates)
print(self.filePath)
if self.filePath:
if len(self.filePath) == 1:
if os.path.isdir(self.filePath[0]):
self.importDirImages(self.filePath[0])
elif os.path.isfile(self.filePath[0]):
self.mImgList = self.filePath
print('opening the file ' + self.filePath[0])
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
self.loadFile(self.filePath[0])
elif len(self.filePath) > 1:
self.mImgList = self.filePath
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
self.loadFile(self.mImgList[0])
else:
if os.path.exists(self.lastOpenDir):
messageBox = QMessageBox()
messageBox.setWindowTitle('询问框标题')
messageBox.setText('是否载入上次打开的图片目录？')
buttonY = messageBox.addButton('确定', QMessageBox.YesRole)
buttonN = messageBox.addButton('取消', QMessageBox.NoRole)
messageBox.exec_()
if messageBox.clickedButton() == buttonY:
self.importDirImages(self.lastOpenDir)
if os.path.exists(self.defaultSaveDir):
print('your labels\' default save dir is ', self.defaultSaveDir)
def resetState(self):
self.filePath = None
self.imageData = None
self.canvas.resetState()
self.labelCoordinates.clear()
def errorMessage(self, title, message):
return QMessageBox.critical(self, title,
'<p><b>%s</b></p>%s' % (title, message))
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.adjustSize()
self.canvas.update()
def addRecentFile(self, filePath):
if filePath in self.recentFiles:
self.recentFiles.remove(filePath)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filePath)
def loadFile(self, filePath=None):
self.resetState()
self.canvas.setEnabled(False)
if filePath is None or len(filePath) == 0:
pass
else:
print('loading ' + filePath)
self.filename = os.path.split(filePath)[-1]
if self.fileListWidget.count() > 0:
index = self.mImgList.index(filePath)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
if os.path.exists(filePath):
self.imageData = read(filePath, None)
self.canvas.verified = False
self.filePath = filePath
image = QImage.fromData(self.imageData)
if image.isNull():
self.errorMessage(
u'Error opening file',
u"<p>Make sure <i>%s</i> is a valid image file." %
filePath)
self.status("Error reading %s" % filePath)
else:
self.status("Loaded %s" % os.path.basename(filePath))
self.image = image
self.im = open_img(filePath)
self.pich, self.picw = self.im.shape[:2]
self.canvas.loadPixmap(QPixmap.fromImage(image))
self.canvas.setEnabled(True)
self.paintCanvas()
#self.addRecentFile(self.filePath)
self.setWindowTitle(__appname__ + ' ' + filePath)
self.canvas.setFocus(True)
def openFile(self):
path = os.path.dirname(self.filePath) if self.filePath else '.'
formats = [
'*.%s' % fmt.data().decode("ascii").lower()
for fmt in QImageReader.supportedImageFormats()
]
filters = "Image & Label files (%s)" % ' '.join(formats)
filename, _ = QFileDialog.getOpenFileName(
self, '%s - Choose Image or Label file' % __appname__, path,
filters)
if filename:
if isinstance(filename, (tuple, list)):
self.mImgList = filename
filename = filename[0]
else:
self.mImgList = [filename]
print('opening the file ' + filename)
self.fileListWidget.clear()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = [
'.%s' % fmt.data().decode("ascii").lower()
for fmt in QImageReader.supportedImageFormats()
]
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relativePath = os.path.join(root, file)
path = os.path.abspath(relativePath)
images.append(path)
images.sort(key=lambda x: x.lower())
return images
def openNextImg(self):
if len(self.mImgList) <= 0:
return
filename = None
if self.filePath is None:
filename = self.mImgList[0]
else:
currIndex = self.mImgList.index(self.filePath)
if currIndex + 1 < len(self.mImgList):
filename = self.mImgList[currIndex + 1]
if filename:
self.loadFile(filename)
def importDirImages(self, dirpath):
self.lastOpenDir = dirpath
cf = configparser.ConfigParser()
cf.read('./TERSI-labelmkr.ini')
cf.set('env', 'lastOpenDir', self.lastOpenDir)
with open('./TERSI-labelmkr.ini', 'w+') as f:
cf.write(f)
self.dirname = dirpath
self.filePath = None
self.fileListWidget.clear()
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def openDirDialog(self):
if self.lastOpenDir != 'None' and os.path.exists(self.lastOpenDir):
defaultOpenDirPath = self.lastOpenDir
else:
defaultOpenDirPath = os.path.dirname(
self.filePath) if self.filePath else '.'
#print(defaultOpenDirPath)
targetDirPath = QFileDialog.getExistingDirectory(
self, '%s - Open Directory' % __appname__, defaultOpenDirPath,
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
self.importDirImages(targetDirPath)
def changeSavedirDialog(self):
if self.defaultSaveDir != 'None' and os.path.exists(
self.defaultSaveDir):
path = self.defaultSaveDir
else:
path = '.'
dirpath = QFileDialog.getExistingDirectory(
self, '%s - Save xml to the directory' % __appname__, path,
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
if dirpath is not None and len(dirpath) > 1:
self.defaultSaveDir = dirpath
print('you have chosen save dir: ' + self.defaultSaveDir)
cf = configparser.ConfigParser()
cf.read('./TERSI-labelmkr.ini')
cf.set('env', 'defaultSaveDir', self.defaultSaveDir)
with open('./TERSI-labelmkr.ini', 'w+') as f:
cf.write(f)
self.statusBar().showMessage(
'%s . Annotation will be saved to %s' %
('Change saved folder', self.defaultSaveDir))
self.statusBar().show()
def cropSavedirDialog(self):
if self.cropSaveDir != 'None' and os.path.exists(self.cropSaveDir):
path = self.cropSaveDir
else:
path = '.'
dirpath = QFileDialog.getExistingDirectory(
self, '%s - Save the section to the directory' % __appname__, path,
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
if dirpath is not None and len(dirpath) > 1:
self.cropSaveDir = dirpath
print('you have chosen cropped pic save dir: ' + self.cropSaveDir)
cf = configparser.ConfigParser()
cf.read('./TERSI-labelmkr.ini')
cf.set('env', 'cropSaveDir', self.cropSaveDir)
with open('./TERSI-labelmkr.ini', 'w+') as f:
cf.write(f)
self.statusBar().showMessage(
'%s .jpg file will be saved to %s' %
('Change cropped pic saved folder', self.cropSaveDir))
self.statusBar().show()
def openPrevImg(self):
if len(self.mImgList) <= 0:
return
if self.filePath is None:
return
currIndex = self.mImgList.index(self.filePath)
if currIndex - 1 >= 0:
filename = self.mImgList[currIndex - 1]
if filename:
self.loadFile(filename)
def verifyImg(self):
if self.defaultSaveDir == 'None':
self.changeSavedirDialog()
fname = self.filename.split('.')[0] + '.xml'
fname = os.path.join(self.defaultSaveDir, fname)
print('verifying label info of ' + fname)
'''
with open(fname, 'w') as f:
f.write(self.filename)
for point in self.canvas.points:
f.write(' ' + str(point.x()) + ' ' + str(point.y()))
'''
root = ET.Element('TERSI-label')
subElement(root, 'picpath', os.path.abspath(self.filePath))
subElement(root, 'xmlpath', fname)
size = subElement(root, 'size')
subElement(size, 'picw', str(self.picw))
subElement(size, 'pich', str(self.pich))
for i in range(len(self.canvas.rectangles) // 2):
leftTop = self.canvas.rectangles[2 * i]
rightBottom = self.canvas.rectangles[2 * i + 1]
xmin = min(leftTop.x(), rightBottom.x())
xmax = max(leftTop.x(), rightBottom.x())
ymin = min(leftTop.y(), rightBottom.y())
ymax = max(leftTop.y(), rightBottom.y())
rectangle = subElement(root, 'rectangle')
subElement(rectangle, 'xmin',
str(self.regularize_point_pos(x=xmin)))
subElement(rectangle, 'xmax',
str(self.regularize_point_pos(x=xmax)))
subElement(rectangle, 'ymin',
str(self.regularize_point_pos(y=ymin)))
subElement(rectangle, 'ymax',
str(self.regularize_point_pos(y=ymax)))
for p in self.canvas.points:
point = subElement(root, 'point')
subElement(point, 'x', str(self.regularize_point_pos(x=p.x())))
subElement(point, 'y', str(self.regularize_point_pos(y=p.y())))
tree = ET.ElementTree(root)
tree.write(fname, encoding='utf-8', xml_declaration=True)
self.openNextImg()
def test(self):
if self.filename:
if os.path.exists(self.defaultSaveDir):
labelfile = os.path.join(self.defaultSaveDir,
self.filename.split('.')[0] + '.xml')
if os.path.exists(labelfile):
#self.canvas.test(labelfile)
tree = ET.parse(labelfile)
root = tree.getroot()
#print(root.find('picpath'))
#print(self.filePath)
if root.find('picpath').text == self.filePath:
print('testing ' + self.filePath)
for rect in root.findall('rectangle'):
x1 = int(float(rect.find('xmin').text))
y1 = int(float(rect.find('ymin').text))
x2 = int(float(rect.find('xmax').text))
y2 = int(float(rect.find('ymax').text))
#print(x1, y1, x2, y2)
cv2.rectangle(self.im, (x1, y1), (x2, y2),
(0, 255, 0), 3)
count = 1
for p in root.findall('point'):
x = int(float(p.find('x').text))
y = int(float(p.find('y').text))
#print(x, y)
cv2.circle(self.im, (x, y), 2, (0, 0, 255), -1)
cv2.putText(self.im, str(count), (x + 2, y + 2),
cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 255, 0), 2)
count += 1
cv2.imshow('Test Image %s' % self.filePath, self.im)
cv2.waitKey(0)
cv2.destroyAllWindows()
else:
'''
box = QMessageBox.information(self,
"消息框标题",
"请核查该图片标注文件存储目录是否正确或者请保存一个标注文件",
QMessageBox.Yes)'''
messageBox = QMessageBox()
messageBox.setWindowTitle('消息框标题')
messageBox.setText('请核查该图片标注文件存储目录是否正确或者请保存一个标注文件')
messageBox.addButton(QPushButton('确定'),
QMessageBox.YesRole)
messageBox.exec_()
else:
self.changeSavedirDialog()
def testSample(self):
if os.path.exists(self.defaultSaveDir):
for file in os.listdir(self.defaultSaveDir):
if file.endswith('.xml'):
tree = ET.parse(os.path.join(self.defaultSaveDir, file))
root = tree.getroot()
picpath = root.find('picpath').text
if os.path.exists(picpath):
img = open_img(picpath)
print('testing ' + picpath)
for rect in root.findall('rectangle'):
x1 = int(float(rect.find('xmin').text))
y1 = int(float(rect.find('ymin').text))
x2 = int(float(rect.find('xmax').text))
y2 = int(float(rect.find('ymax').text))
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0),
3)
count = 1
for p in root.findall('point'):
x = int(float(p.find('x').text))
y = int(float(p.find('y').text))
cv2.circle(img, (x, y), 2, (0, 0, 255), -1)
cv2.putText(img, str(count), (x + 2, y + 2),
cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 255, 0), 2)
count += 1
cv2.imshow('Test Sample - %s' % picpath, img)
keycode = cv2.waitKey(0)
cv2.destroyAllWindows()
if keycode == 27:
break
else:
self.changeSavedirDialog()
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(item.text())
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
def delImg(self):
print('deleting ' + self.filePath)
self.status('deleting ' + self.filePath)
os.remove(self.filePath)
to_remove = self.filePath
self.openNextImg()
self.mImgList.remove(to_remove)
self.fileListWidget.clear()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def regularize_point_pos(self, x=None, y=None):
if x is not None:
if x < 0:
x = 0
elif x >= self.picw:
x = self.picw - 1
return x
elif y is not None:
if y < 0:
y = 0
elif y >= self.pich:
y = self.pich - 1
return y
def cropAndSave(self):
if self.filename:
self.cropSavedirDialog()
if len(self.canvas.rectangles) == 2:
rect1 = self.canvas.rectangles[0]
rect2 = self.canvas.rectangles[1]
x1 = int(
self.regularize_point_pos(x=min(rect1.x(), rect2.x())))
y1 = int(
self.regularize_point_pos(y=min(rect1.y(), rect2.y())))
x2 = int(
self.regularize_point_pos(x=max(rect1.x(), rect2.x())))
y2 = int(
self.regularize_point_pos(y=max(rect1.y(), rect2.y())))
croppedim = self.im[y1:y2, x1:x2, :]
#cv2.imshow('img', croppedim)
#cv2.waitKey()
croppedimg = Image.fromarray(
cv2.cvtColor(croppedim, cv2.COLOR_BGR2RGB))
croppedimpath = os.path.join(
self.cropSaveDir,
self.filename.split('.')[-2] + '%d%d%d%d.jpg' %
(x1, y1, x2, y2))
#croppedimg.show()
#input()
#cv2.imwrite(croppedimpath, croppedim)
croppedimg.save(croppedimpath)
print('saved %s into %s' % (croppedimpath, self.cropSaveDir))
else:
messageBox = QMessageBox()
messageBox.setWindowTitle('消息框标题')
messageBox.setText('请标注一个矩形，如此才能截取')
messageBox.addButton(QPushButton('确定'), QMessageBox.YesRole)
messageBox.exec_()
def clear(self):
self.canvas.resetState()
class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, filename=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
# Save as Pascal voc xml
self.defaultSaveDir = None
self.usingPascalVocFormat = True
if self.usingPascalVocFormat:
LabelFile.suffix = '.xml'
# For loading all image under a directory
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
# Whether we need to save or not.
self.dirty = False
# Enble auto saving if pressing next
self.autoSaving = True
self._noSelectionSlot = False
self._beginner = True
self.screencastViewer = "firefox"
self.screencast = "https://youtu.be/p0nR2YsCY_U"
self.loadPredefinedClasses()
# Main widgets and related state.
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.labelList = QListWidget()
self.itemsToShapes = {}
self.shapesToItems = {}
self.labelList.itemActivated.connect(self.labelSelectionChanged)
self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
self.labelList.itemDoubleClicked.connect(self.editLabel)
# Connect to itemChanged to detect checkbox changes.
self.labelList.itemChanged.connect(self.labelItemChanged)
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
listLayout.addWidget(self.labelList)
self.editButton = QToolButton()
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
self.labelListContainer = QWidget()
self.labelListContainer.setLayout(listLayout)
listLayout.addWidget(self.editButton)#, 0, Qt.AlignCenter)
listLayout.addWidget(self.labelList)
self.dock = QDockWidget(u'Box Labels', self)
self.dock.setObjectName(u'Labels')
self.dock.setWidget(self.labelListContainer)
# Tzutalin 20160906 : Add file list and dock to move faster
self.fileListWidget = QListWidget()
self.fileListWidget.itemDoubleClicked.connect(self.fileitemDoubleClicked)
filelistLayout = QVBoxLayout()
filelistLayout.setContentsMargins(0, 0, 0, 0)
filelistLayout.addWidget(self.fileListWidget)
self.fileListContainer = QWidget()
self.fileListContainer.setLayout(filelistLayout)
self.filedock = QDockWidget(u'File List', self)
self.filedock.setObjectName(u'Files')
self.filedock.setWidget(self.fileListContainer)
self.zoomWidget = ZoomWidget()
self.colorDialog = ColorDialog(parent=self)
self.canvas = Canvas()
self.canvas.zoomRequest.connect(self.zoomRequest)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
# Tzutalin 20160906 : Add file list and dock to move faster
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.dockFeatures = QDockWidget.DockWidgetClosable\
| QDockWidget.DockWidgetFloatable
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
# Actions
action = partial(newAction, self)
quit = action('&Quit', self.close,
'Ctrl+Q', 'quit', u'Quit application')
open = action('&Open', self.openFile,
'Ctrl+O', 'open', u'Open image or label file')
opendir = action('&Open Dir', self.openDir,
'Ctrl+u', 'open', u'Open Dir')
changeSavedir = action('&Change default saved Annotation dir', self.changeSavedir,
'Ctrl+r', 'open', u'Change default saved Annotation dir')
openAnnotation = action('&Open Annotation', self.openAnnotation,
'Ctrl+q', 'openAnnotation', u'Open Annotation')
openNextImg = action('&Next Image', self.openNextImg,
'd', 'next', u'Open Next')
openPrevImg = action('&Prev Image', self.openPrevImg,
'a', 'prev', u'Open Prev')
save = action('&Save', self.saveFile,
'Ctrl+S', 'save', u'Save labels to file', enabled=False)
saveAs = action('&Save As', self.saveFileAs,
'Ctrl+Shift+S', 'save-as', u'Save labels to a different file',
enabled=False)
close = action('&Close', self.closeFile,
'Ctrl+W', 'close', u'Close current file')
color1 = action('Box &Line Color', self.chooseColor1,
'Ctrl+L', 'color_line', u'Choose Box line color')
color2 = action('Box &Fill Color', self.chooseColor2,
'Ctrl+Shift+L', 'color', u'Choose Box fill color')
createMode = action('Create\nRectBox', self.setCreateMode,
'Ctrl+N', 'new', u'Start drawing Boxs', enabled=False)
editMode = action('&Edit\nRectBox', self.setEditMode,
'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
create = action('Create\nRectBox', self.createShape,
'w', 'new', u'Draw a new Box', enabled=False)
delete = action('Delete\nRectBox', self.deleteSelectedShape,
'Delete', 'delete', u'Delete', enabled=False)
copy = action('&Duplicate\nRectBox', self.copySelectedShape,
'Ctrl+D', 'copy', u'Create a duplicate of the selected Box',
enabled=False)
advancedMode = action('&Advanced Mode', self.toggleAdvancedMode,
'Ctrl+Shift+A', 'expert', u'Switch to advanced mode',
checkable=True)
hideAll = action('&Hide\nRectBox', partial(self.togglePolygons, False),
'Ctrl+H', 'hide', u'Hide all Boxs',
enabled=False)
showAll = action('&Show\nRectBox', partial(self.togglePolygons, True),
'Ctrl+A', 'hide', u'Show all Boxs',
enabled=False)
help = action('&Tutorial', self.tutorial, 'Ctrl+T', 'help',
u'Show demos')
zoom = QWidgetAction(self)
zoom.setDefaultWidget(self.zoomWidget)
self.zoomWidget.setWhatsThis(
u"Zoom in or out of the image. Also accessible with"\
" %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
fmtShortcut("Ctrl+Wheel")))
self.zoomWidget.setEnabled(False)
zoomIn = action('Zoom &In', partial(self.addZoom, 10),
'Ctrl++', 'zoom-in', u'Increase zoom level', enabled=False)
zoomOut = action('&Zoom Out', partial(self.addZoom, -10),
'Ctrl+-', 'zoom-out', u'Decrease zoom level', enabled=False)
zoomOrg = action('&Original size', partial(self.setZoom, 100),
'Ctrl+=', 'zoom', u'Zoom to original size', enabled=False)
fitWindow = action('&Fit Window', self.setFitWindow,
'Ctrl+F', 'fit-window', u'Zoom follows window size',
checkable=True, enabled=False)
fitWidth = action('Fit &Width', self.setFitWidth,
'Ctrl+Shift+F', 'fit-width', u'Zoom follows window width',
checkable=True, enabled=False)
# Group zoom controls into a list for easier toggling.
zoomActions = (self.zoomWidget, zoomIn, zoomOut, zoomOrg, fitWindow, fitWidth)
self.zoomMode = self.MANUAL_ZOOM
self.scalers = {
self.FIT_WINDOW: self.scaleFitWindow,
self.FIT_WIDTH: self.scaleFitWidth,
# Set to one to scale to 100% when loading files.
self.MANUAL_ZOOM: lambda: 1,
}
edit = action('&Edit Label', self.editLabel,
'Ctrl+E', 'edit', u'Modify the label of the selected Box',
enabled=False)
self.editButton.setDefaultAction(edit)
shapeLineColor = action('Shape &Line Color', self.chshapeLineColor,
icon='color_line', tip=u'Change the line color for this specific shape',
enabled=False)
shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor,
icon='color', tip=u'Change the fill color for this specific shape',
enabled=False)
labels = self.dock.toggleViewAction()
labels.setText('Show/Hide Label Panel')
labels.setShortcut('Ctrl+Shift+L')
# Lavel list context menu.
labelMenu = QMenu()
addActions(labelMenu, (edit, delete))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(self.popLabelListMenu)
# Store actions for further handling.
self.actions = struct(save=save, saveAs=saveAs, open=open, close=close,
lineColor=color1, fillColor=color2,
create=create, delete=delete, edit=edit, copy=copy,
createMode=createMode, editMode=editMode, advancedMode=advancedMode,
shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
fitWindow=fitWindow, fitWidth=fitWidth,
zoomActions=zoomActions,
fileMenuActions=(open,opendir,save,saveAs,close,quit),
beginner=(), advanced=(),
editMenu=(edit, copy, delete, None, color1, color2),
beginnerContext=(create, edit, copy, delete),
advancedContext=(createMode, editMode, edit, copy,
delete, shapeLineColor, shapeFillColor),
onLoadActive=(close, create, createMode, editMode),
onShapesPresent=(saveAs, hideAll, showAll))
self.menus = struct(
file=self.menu('&File'),
edit=self.menu('&Edit'),
view=self.menu('&View'),
help=self.menu('&Help'),
recentFiles=QMenu('Open &Recent'),
labelList=labelMenu)
addActions(self.menus.file,
(open, opendir,changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, None, quit))
addActions(self.menus.help, (help,))
addActions(self.menus.view, (
labels, advancedMode, None,
hideAll, showAll, None,
zoomIn, zoomOut, zoomOrg, None,
fitWindow, fitWidth))
self.menus.file.aboutToShow.connect(self.updateFileMenu)
# Custom context menu for the canvas widget:
addActions(self.canvas.menus[0], self.actions.beginnerContext)
addActions(self.canvas.menus[1], (
action('&Copy here', self.copyShape),
action('&Move here', self.moveShape)))
self.tools = self.toolbar('Tools')
self.actions.beginner = (
open, opendir, openNextImg, openPrevImg, save, None, create, copy, delete, None,
zoomIn, zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (
open, save, None,
createMode, editMode, None,
hideAll, showAll)
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
# Application state.
self.image = QImage()
self.filename = filename
self.recentFiles = []
self.maxRecent = 7
self.lineColor = None
self.fillColor = None
self.zoom_level = 100
self.fit_window = False
# XXX: Could be completely declarative.
# Restore application settings.
if have_qstring():
types = {
'filename': QString,
'recentFiles': QStringList,
'window/size': QSize,
'window/position': QPoint,
'window/geometry': QByteArray,
'line/color': QColor,
'fill/color': QColor,
'advanced': bool,
# Docks and toolbars:
'window/state': QByteArray,
'savedir': QString,
'lastOpenDir': QString,
}
else:
types = {
'filename': str,
'recentFiles': list,
'window/size': QSize,
'window/position': QPoint,
'window/geometry': QByteArray,
'line/color': QColor,
'fill/color': QColor,
'advanced': bool,
# Docks and toolbars:
'window/state': QByteArray,
'savedir': str,
'lastOpenDir': str,
}
self.settings = settings = Settings(types)
self.recentFiles = list(settings.get('recentFiles', []))
size = settings.get('window/size', QSize(600, 500))
position = settings.get('window/position', QPoint(0, 0))
self.resize(size)
self.move(position)
saveDir = settings.get('savedir', None)
self.lastOpenDir = settings.get('lastOpenDir', None)
if os.path.exists(str(saveDir)):
self.defaultSaveDir = str(saveDir)
self.statusBar().showMessage('%s started. Annotation will be saved to %s' %(__appname__, self.defaultSaveDir))
self.statusBar().show()
# or simply:
#self.restoreGeometry(settings['window/geometry']
self.restoreState(settings.get('window/state', QByteArray()))
self.lineColor = QColor(settings.get('line/color', Shape.line_color))
self.fillColor = QColor(settings.get('fill/color', Shape.fill_color))
Shape.line_color = self.lineColor
Shape.fill_color = self.fillColor
def xbool(x):
if isinstance(x, QVariant):
return x.toBool()
return bool(x)
if xbool(settings.get('advanced', False)):
self.actions.advancedMode.setChecked(True)
self.toggleAdvancedMode()
# Populate the File menu dynamically.
self.updateFileMenu()
# Since loading the file may take some time, make sure it runs in the background.
self.queueEvent(partial(self.loadFile, self.filename))
# Callbacks:
self.zoomWidget.valueChanged.connect(self.paintCanvas)
self.populateModeActions()
## Support Functions ##
def noShapes(self):
return not self.itemsToShapes
def toggleAdvancedMode(self, value=True):
self._beginner = not value
self.canvas.setEditing(True)
self.populateModeActions()
self.editButton.setVisible(not value)
if value:
self.actions.createMode.setEnabled(True)
self.actions.editMode.setEnabled(False)
self.dock.setFeatures(self.dock.features() | self.dockFeatures)
else:
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
def populateModeActions(self):
if self.beginner():
tool, menu = self.actions.beginner, self.actions.beginnerContext
else:
tool, menu = self.actions.advanced, self.actions.advancedContext
self.tools.clear()
addActions(self.tools, tool)
self.canvas.menus[0].clear()
addActions(self.canvas.menus[0], menu)
self.menus.edit.clear()
actions = (self.actions.create,) if self.beginner()\
else (self.actions.createMode, self.actions.editMode)
addActions(self.menus.edit, actions + self.actions.editMenu)
def setBeginner(self):
self.tools.clear()
addActions(self.tools, self.actions.beginner)
def setAdvanced(self):
self.tools.clear()
addActions(self.tools, self.actions.advanced)
def setDirty(self):
self.dirty = True
self.actions.save.setEnabled(True)
def setClean(self):
self.dirty = False
self.actions.save.setEnabled(False)
self.actions.create.setEnabled(True)
def toggleActions(self, value=True):
"""Enable/Disable widgets which depend on an opened image."""
for z in self.actions.zoomActions:
z.setEnabled(value)
for action in self.actions.onLoadActive:
action.setEnabled(value)
def queueEvent(self, function):
QTimer.singleShot(0, function)
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.itemsToShapes.clear()
self.shapesToItems.clear()
self.labelList.clear()
self.filename = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
def currentItem(self):
items = self.labelList.selectedItems()
if items:
return items[0]
return None
def addRecentFile(self, filename):
if filename in self.recentFiles:
self.recentFiles.remove(filename)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filename)
def beginner(self):
return self._beginner
def advanced(self):
return not self.beginner()
## Callbacks ##
def tutorial(self):
subprocess.Popen([self.screencastViewer, self.screencast])
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
"""In the middle of drawing, toggling between modes should be disabled."""
self.actions.editMode.setEnabled(not drawing)
if not drawing and self.beginner():
# Cancel creation.
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.create.setEnabled(True)
def toggleDrawMode(self, edit=True):
self.canvas.setEditing(edit)
self.actions.createMode.setEnabled(edit)
self.actions.editMode.setEnabled(not edit)
def setCreateMode(self):
assert self.advanced()
self.toggleDrawMode(False)
def setEditMode(self):
assert self.advanced()
self.toggleDrawMode(True)
def updateFileMenu(self):
current = self.filename
def exists(filename):
return os.path.exists(filename)
menu = self.menus.recentFiles
menu.clear()
files = [f for f in self.recentFiles if f != current and exists(f)]
for i, f in enumerate(files):
icon = newIcon('labels')
action = QAction(
icon, '&%d %s' % (i+1, QFileInfo(f).fileName()), self)
action.triggered.connect(partial(self.loadRecent, f))
menu.addAction(action)
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
def editLabel(self, item=None):
if not self.canvas.editing():
return
item = item if item else self.currentItem()
text = self.labelDialog.popUp(item.text())
if text is not None:
item.setText(text)
self.setDirty()
# Tzutalin 20160906 : Add file list and dock to move faster
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(item.text())
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
# React to canvas signals.
def shapeSelectionChanged(self, selected=False):
if self._noSelectionSlot:
self._noSelectionSlot = False
else:
shape = self.canvas.selectedShape
if shape:
self.shapesToItems[shape].setSelected(True)
else:
self.labelList.clearSelection()
self.actions.delete.setEnabled(selected)
self.actions.copy.setEnabled(selected)
self.actions.edit.setEnabled(selected)
self.actions.shapeLineColor.setEnabled(selected)
self.actions.shapeFillColor.setEnabled(selected)
def addLabel(self, shape):
item = HashableQListWidgetItem(shape.label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.itemsToShapes[item] = shape
self.shapesToItems[shape] = item
self.labelList.addItem(item)
for action in self.actions.onShapesPresent:
action.setEnabled(True)
def remLabel(self, shape):
item = self.shapesToItems[shape]
self.labelList.takeItem(self.labelList.row(item))
del self.shapesToItems[shape]
del self.itemsToShapes[item]
def loadLabels(self, shapes):
s = []
for label, points, line_color, fill_color in shapes:
shape = Shape(label=label)
for x, y in points:
shape.addPoint(QPointF(x, y))
shape.close()
s.append(shape)
self.addLabel(shape)
if line_color:
shape.line_color = QColor(*line_color)
if fill_color:
shape.fill_color = QColor(*fill_color)
self.canvas.loadShapes(s)
def saveLabels(self, filename):
lf = LabelFile()
def format_shape(s):
return dict(label=s.label,
line_color=s.line_color.getRgb()\
if s.line_color != self.lineColor else None,
fill_color=s.fill_color.getRgb()\
if s.fill_color != self.fillColor else None,
points=[(p.x(), p.y()) for p in s.points])
shapes = [format_shape(shape) for shape in self.canvas.shapes]
# Can add differrent annotation formats here
try:
if self.usingPascalVocFormat is True:
print('savePascalVocFormat save to:' + filename)
lf.savePascalVocFormat(filename, shapes, str(self.filename), self.imageData,
self.lineColor.getRgb(), self.fillColor.getRgb())
else:
lf.save(filename, shapes, str(self.filename), self.imageData,
self.lineColor.getRgb(), self.fillColor.getRgb())
self.labelFile = lf
self.filename = filename
return True
except LabelFileError as e:
self.errorMessage(u'Error saving label data',
u'<b>%s</b>' % e)
return False
def copySelectedShape(self):
self.addLabel(self.canvas.copySelectedShape())
#fix copy and delete
self.shapeSelectionChanged(True)
def labelSelectionChanged(self):
item = self.currentItem()
if item and self.canvas.editing():
self._noSelectionSlot = True
self.canvas.selectShape(self.itemsToShapes[item])
def labelItemChanged(self, item):
shape = self.itemsToShapes[item]
label = item.text()
if label != shape.label:
shape.label = item.text()
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
## Callback functions:
def newShape(self):
"""Pop-up and give focus to the label editor.
position MUST be in global coordinates.
"""
if len(self.labelHist) > 0:
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
text = self.labelDialog.popUp()
if text is not None:
self.addLabel(self.canvas.setLastLabel(text))
if self.beginner(): # Switch to edit mode.
self.canvas.setEditing(True)
self.actions.create.setEnabled(True)
else:
self.actions.editMode.setEnabled(True)
self.setDirty()
if text not in self.labelHist:
self.labelHist.append(text)
else:
#self.canvas.undoLastLine()
self.canvas.resetAllLines()
def scrollRequest(self, delta, orientation):
units = - delta / (8 * 15)
bar = self.scrollBars[orientation]
bar.setValue(bar.value() + bar.singleStep() * units)
def setZoom(self, value):
self.actions.fitWidth.setChecked(False)
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.MANUAL_ZOOM
self.zoomWidget.setValue(value)
def addZoom(self, increment=10):
self.setZoom(self.zoomWidget.value() + increment)
def zoomRequest(self, delta):
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
def setFitWindow(self, value=True):
if value:
self.actions.fitWidth.setChecked(False)
self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
self.adjustScale()
def setFitWidth(self, value=True):
if value:
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
self.adjustScale()
def togglePolygons(self, value):
for item, shape in self.itemsToShapes.items():
item.setCheckState(Qt.Checked if value else Qt.Unchecked)
def loadFile(self, filename=None):
"""Load the specified file, or the last opened file if None."""
self.resetState()
self.canvas.setEnabled(False)
if filename is None:
filename = self.settings.get('filename')
# Tzutalin 20160906 : Add file list and dock to move faster
# Highlight the file item
if filename and self.fileListWidget.count() > 0:
index = self.mImgList.index(filename)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
if filename and QFile.exists(filename):
if LabelFile.isLabelFile(filename):
try:
self.labelFile = LabelFile(filename)
except LabelFileError as e:
self.errorMessage(u'Error opening file',
(u"<p><b>%s</b></p>"
u"<p>Make sure <i>%s</i> is a valid label file.")\
% (e, filename))
self.status("Error reading %s" % filename)
return False
self.imageData = self.labelFile.imageData
self.lineColor = QColor(*self.labelFile.lineColor)
self.fillColor = QColor(*self.labelFile.fillColor)
else:
# Load image:
# read data first and store for saving into label file.
self.imageData = read(filename, None)
self.labelFile = None
image = QImage.fromData(self.imageData)
if image.isNull():
self.errorMessage(u'Error opening file',
u"<p>Make sure <i>%s</i> is a valid image file." % filename)
self.status("Error reading %s" % filename)
return False
self.status("Loaded %s" % os.path.basename(str(filename)))
self.image = image
self.filename = filename
self.canvas.loadPixmap(QPixmap.fromImage(image))
if self.labelFile:
self.loadLabels(self.labelFile.shapes)
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.addRecentFile(self.filename)
self.toggleActions(True)
## Label xml file and show bound box according to its filename
if self.usingPascalVocFormat is True and \
self.defaultSaveDir is not None:
basename = os.path.basename(os.path.splitext(self.filename)[0]) + '.xml'
xmlPath = os.path.join(self.defaultSaveDir, basename)
self.loadPascalXMLByFilename(xmlPath)
return True
return False
def resizeEvent(self, event):
if self.canvas and not self.image.isNull()\
and self.zoomMode != self.MANUAL_ZOOM:
self.adjustScale()
super(MainWindow, self).resizeEvent(event)
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.scale = 0.01 * self.zoomWidget.value()
self.canvas.adjustSize()
self.canvas.update()
def adjustScale(self, initial=False):
value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
self.zoomWidget.setValue(int(100 * value))
def scaleFitWindow(self):
"""Figure out the size of the pixmap in order to fit the main widget."""
e = 2.0 # So that no scrollbars are generated.
w1 = self.centralWidget().width() - e
h1 = self.centralWidget().height() - e
a1 = w1/ h1
# Calculate a new scale value based on the pixmap's aspect ratio.
w2 = self.canvas.pixmap.width() - 0.0
h2 = self.canvas.pixmap.height() - 0.0
a2 = w2 / h2
return w1 / w2 if a2 >= a1 else h1 / h2
def scaleFitWidth(self):
# The epsilon does not seem to work too well here.
w = self.centralWidget().width() - 2.0
return w / self.canvas.pixmap.width()
def closeEvent(self, event):
if not self.mayContinue():
event.ignore()
s = self.settings
# If it loads images from dir, don't load it at the begining
if self.dirname is None:
s['filename'] = self.filename if self.filename else ''
else:
s['filename'] = ''
s['window/size'] = self.size()
s['window/position'] = self.pos()
s['window/state'] = self.saveState()
s['line/color'] = self.lineColor
s['fill/color'] = self.fillColor
s['recentFiles'] = self.recentFiles
s['advanced'] = not self._beginner
if self.defaultSaveDir is not None and len(self.defaultSaveDir) > 1:
s['savedir'] = str(self.defaultSaveDir)
else:
s['savedir'] = ""
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
s['lastOpenDir'] = str(self.lastOpenDir)
else:
s['lastOpenDir'] = ""
#ask the use for where to save the labels
#s['window/geometry'] = self.saveGeometry()
## User Dialogs ##
def loadRecent(self, filename):
if self.mayContinue():
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = ['.jpeg','.jpg', '.png', '.bmp']
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relatviePath = os.path.join(root, file)
path = u(os.path.abspath(relatviePath))
images.append(path)
images.sort(key=lambda x: x.lower())
return images
def changeSavedir(self, _value=False):
if self.defaultSaveDir is not None:
path = str(self.defaultSaveDir)
else:
path = '.'
dirpath = str(QFileDialog.getExistingDirectory(self,
'%s - Save to the directory' % __appname__, path, QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
self.defaultSaveDir = dirpath
self.statusBar().showMessage('%s . Annotation will be saved to %s' %('Change saved folder', self.defaultSaveDir))
self.statusBar().show()
def openAnnotation(self, _value=False):
if self.filename is None:
return
path = os.path.dirname(str(self.filename))\
if self.filename else '.'
if self.usingPascalVocFormat:
formats = ['*.%s' % str(fmt).lower()\
for fmt in QImageReader.supportedImageFormats()]
filters = "Open Annotation XML file (%s)" % \
' '.join(formats + ['*.xml'])
filename = str(QFileDialog.getOpenFileName(self,
'%s - Choose a xml file' % __appname__, path, filters))
self.loadPascalXMLByFilename(filename)
def openDir(self, _value=False):
if not self.mayContinue():
return
path = os.path.dirname(self.filename)\
if self.filename else '.'
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
path = self.lastOpenDir
dirpath = str(QFileDialog.getExistingDirectory(self,
'%s - Open Directory' % __appname__, path, QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
self.lastOpenDir = dirpath
self.dirname = dirpath
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def openPrevImg(self, _value=False):
if not self.mayContinue():
return
if len(self.mImgList) <= 0:
return
if self.filename is None:
return
currIndex = self.mImgList.index(self.filename)
if currIndex -1 >= 0:
filename = self.mImgList[currIndex-1]
if filename:
self.loadFile(filename)
def openNextImg(self, _value=False):
# Proceding next image without dialog if having any label
if self.autoSaving is True and self.defaultSaveDir is not None:
if self.dirty is True and self.hasLabels():
self.saveFile()
if not self.mayContinue():
return
if len(self.mImgList) <= 0:
return
filename = None
if self.filename is None:
filename = self.mImgList[0]
else:
currIndex = self.mImgList.index(self.filename)
if currIndex + 1 < len(self.mImgList):
filename = self.mImgList[currIndex+1]
if filename:
self.loadFile(filename)
def openFile(self, _value=False):
if not self.mayContinue():
return
path = os.path.dirname(str(self.filename))\
if self.filename else '.'
formats = ['*.%s' % str(fmt).lower()\
for fmt in QImageReader.supportedImageFormats()]
filters = "Image & Label files (%s)" % \
' '.join(formats + ['*%s' % LabelFile.suffix])
filename = str(QFileDialog.getOpenFileName(self,
'%s - Choose Image or Label file' % __appname__, path, filters))
if filename:
self.loadFile(filename)
def saveFile(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
if self.hasLabels():
if self.defaultSaveDir is not None and len(str(self.defaultSaveDir)):
print('handle the image:' + self.filename)
imgFileName = os.path.basename(self.filename)
savedFileName = os.path.splitext(imgFileName)[0] + LabelFile.suffix
savedPath = os.path.join(str(self.defaultSaveDir), savedFileName)
self._saveFile(savedPath)
else:
self._saveFile(self.filename if self.labelFile\
else self.saveFileDialog())
def saveFileAs(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
if self.hasLabels():
self._saveFile(self.saveFileDialog())
def saveFileDialog(self):
caption = '%s - Choose File' % __appname__
filters = 'File (*%s)' % LabelFile.suffix
openDialogPath = self.currentPath()
dlg = QFileDialog(self, caption, openDialogPath, filters)
dlg.setDefaultSuffix(LabelFile.suffix[1:])
dlg.setAcceptMode(QFileDialog.AcceptSave)
filenameWithoutExtension = os.path.splitext(self.filename)[0]
dlg.selectFile(filenameWithoutExtension)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
if dlg.exec_():
return dlg.selectedFiles()[0]
return ''
def _saveFile(self, filename):
if filename and self.saveLabels(filename):
self.addRecentFile(filename)
self.setClean()
self.statusBar().showMessage('Saved to %s' % filename)
self.statusBar().show()
def closeFile(self, _value=False):
if not self.mayContinue():
return
self.resetState()
self.setClean()
self.toggleActions(False)
self.canvas.setEnabled(False)
self.actions.saveAs.setEnabled(False)
# Message Dialogs. #
def hasLabels(self):
if not self.itemsToShapes:
self.errorMessage(u'No objects labeled',
u'You must label at least one object to save the file.')
return False
return True
def mayContinue(self):
return not (self.dirty and not self.discardChangesDialog())
def discardChangesDialog(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = u'You have unsaved changes, proceed anyway?'
return yes == QMessageBox.warning(self, u'Attention', msg, yes|no)
def errorMessage(self, title, message):
return QMessageBox.critical(self, title,
'<p><b>%s</b></p>%s' % (title, message))
def currentPath(self):
return os.path.dirname(str(self.filename)) if self.filename else '.'
def chooseColor1(self):
color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.lineColor = color
# Change the color for all shape lines:
Shape.line_color = self.lineColor
self.canvas.update()
self.setDirty()
def chooseColor2(self):
color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.fillColor = color
Shape.fill_color = self.fillColor
self.canvas.update()
self.setDirty()
def deleteSelectedShape(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = u'You are about to permanently delete this Box, proceed anyway?'
if yes == QMessageBox.warning(self, u'Attention', msg, yes|no):
self.remLabel(self.canvas.deleteSelected())
self.setDirty()
if self.noShapes():
for action in self.actions.onShapesPresent:
action.setEnabled(False)
def chshapeLineColor(self):
color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.canvas.selectedShape.line_color = color
self.canvas.update()
self.setDirty()
def chshapeFillColor(self):
color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.canvas.selectedShape.fill_color = color
self.canvas.update()
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addLabel(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
def loadPredefinedClasses(self):
predefined_classes_path = os.path.join('data', 'predefined_classes.txt')
if os.path.exists(predefined_classes_path) is True:
with codecs.open(predefined_classes_path, 'r', 'utf8') as f:
for line in f:
line = line.strip()
if self.labelHist is None:
self.lablHist = [line]
else:
self.labelHist.append(line)
def loadPascalXMLByFilename(self, xmlPath):
if self.filename is None:
return
if os.path.isfile(xmlPath) is False:
return
tVocParseReader = PascalVocReader(xmlPath)
shapes = tVocParseReader.getShapes()
self.loadLabels(shapes)
class LabelDialog(QDialog):
MODE_FIT_WINDOW, MODE_ADJUST_SIZE = ['F', 'A']
LIST_SIZE = [0, 0, 2200, 1200]
def __init__(self,
text="在此输入标签",
parent=None,
listItem=None,
pic_path=None):
super(LabelDialog, self).__init__(parent)
self.setWindowTitle(__moudlename__) # 标题
self.setGeometry(self.LIST_SIZE[0], self.LIST_SIZE[1],
self.LIST_SIZE[2], self.LIST_SIZE[3])
self.edit = QLineEdit()
self.edit.setText(text)
self.edit.setValidator(self.labelValidator())
self.edit.editingFinished.connect(self.postProcess)
model = QStringListModel()
# model.setStringList(listItem)
completer = QCompleter()
completer.setModel(model)
self.edit.setCompleter(completer)
layout = QVBoxLayout()
self.pic_path = pic_path
# 展示图片
try:
self.mode = self.MODE_FIT_WINDOW
self.canvas = Canvas(parent=self)
self.scrollArea = QScrollArea()
self.scrollArea.setGeometry(
QtCore.QRect(self.LIST_SIZE[0], self.LIST_SIZE[1],
self.LIST_SIZE[2], self.LIST_SIZE[3]))
image = QImage(self.pic_path)
self.canvas.load_pixmap(QPixmap.fromImage(image))
# print('self.scale_fit_window():', self.scale_fit_window())
self.file_or_dir_fit_window()
self.scrollArea.setWidget(self.canvas)
self.scrollArea.setWidgetResizable(True)
layout.addWidget(self.scrollArea)
self.setLayout(layout)
except Exception as e:
print('Error:', e)
def scale_fit_window(self):
e = 2.0 # So that no scrollbars are generated.
w1 = self.scrollArea.width() - e
h1 = self.scrollArea.height() - e
a1 = w1 / h1 # 宽高比a1 例如：16:9
w2 = self.canvas.pixmap.width() - 0.0
h2 = self.canvas.pixmap.height() - 0.0
a2 = w2 / h2
return w1 / w2 if a2 >= a1 else h1 / h2
def file_or_dir_fit_window(self):
if self.mode == self.MODE_FIT_WINDOW:
self.canvas.scale = self.scale_fit_window() # 随之变动
self.canvas.adjustSize()
self.repaint()
elif self.mode == self.MODE_ADJUST_SIZE:
self.canvas.adjustSize()
self.update()
@staticmethod
def labelValidator():
return QRegExpValidator(QRegExp(r'^[^ \t].+'), None)
def validate(self):
if self.edit.text().strip():
self.accept()
def postProcess(self):
self.edit.setText(self.edit.text())
def popUp(self, text='', move=True):
self.edit.setText(text)
self.edit.setSelection(0, len(text))
self.edit.setFocus(Qt.PopupFocusReason)
# if move:
# self.move(QCursor.pos())
return self.edit.text() if self.exec_() else None
def listItemClick(self, tQListWidgetItem):
text = tQListWidgetItem.text().strip()
self.edit.setText(text)
def listItemDoubleClick(self, tQListWidgetItem):
self.listItemClick(tQListWidgetItem)
self.validate()
