# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging import os import time import os.path as osp from functools import partial import json from distutils.util import strtobool import webbrowser from easydict import EasyDict as edict from qtpy import QtGui, QtCore, QtWidgets from qtpy.QtWidgets import QMainWindow, QMessageBox, QTableWidgetItem, QApplication from qtpy.QtGui import QImage, QPixmap from qtpy.QtCore import Qt, QByteArray, QVariant, QCoreApplication, QThread, Signal, QTimer import cv2 import numpy as np from PIL import Image import paddle import paddle.nn.functional as F from eiseg import pjpath, __APPNAME__, logger from widget import ShortcutWidget, PolygonAnnotation from controller import InteractiveController from ui import Ui_EISeg import util from util import COCO from util import check_cn, normcase import plugin.remotesensing as rs from plugin.medical import med from plugin.remotesensing import Raster from plugin.n2grid import RSGrids, Grids, checkOpenGrid from plugin.video import InferenceCore, overlay_davis class APP_EISeg(QMainWindow, Ui_EISeg): IDILE, ANNING, EDITING = 0, 1, 2 # IDILE：网络，权重，图像三者任一没有加载 # EDITING：多边形编辑，可以交互式，但是多边形内部不能点 # ANNING：交互式标注，只能交互式，不能编辑多边形，多边形不接hover # 宫格标注背景颜色 GRID_COLOR = { "idle": QtGui.QColor(255, 255, 255), "current": QtGui.QColor(192, 220, 243), "finised": QtGui.QColor(185, 185, 225), "overlying": QtGui.QColor(51, 52, 227), } def __init__(self, parent=None): super(APP_EISeg, self).__init__(parent) self.settings = QtCore.QSettings( osp.join(pjpath, "config/setting.txt"), QtCore.QSettings.IniFormat) currentLang = self.settings.value("language") layoutdir = Qt.RightToLeft if currentLang == "Arabic" else Qt.LeftToRight self.setLayoutDirection(layoutdir) # 初始化界面 self.setupUi(self) # app变量 self._anning = False # self.status替代 self.isDirty = False # 是否需要保存 self.image = None # 可能先加载图片后加载模型，只用于暂存图片 self.predictor_params = { "brs_mode": "NoBRS", "with_flip": False, "zoom_in_params": { "skip_clicks": -1, "target_size": (400, 400), "expansion_ratio": 1.4, }, "predictor_params": { "net_clicks_limit": None, "max_size": 800, "with_mask": True, }, } self.controller = InteractiveController( predictor_params=self.predictor_params, prob_thresh=self.segThresh, ) self.video = InferenceCore() self.video_images = None self.video_masks = None # self.controller.labelList = util.LabelList() # 标签列表 self.save_status = { "gray_scale": True, "pseudo_color": True, "json": False, "coco": True, "cutout": True, } # 是否保存这几个格式 self.outputDir = None # 标签保存路径 self.labelPaths = [] # 所有outputdir中的标签文件路径 self.imagePaths = [] # 文件夹下所有待标注图片路径 self.currIdx = 0 # 文件夹标注当前图片下标 self.origExt = False # 是否使用图片本身拓展名，防止重名覆盖 if self.save_status["coco"]: self.coco = COCO() else: self.coco = None self.colorMap = util.colorMap if self.settings.value("cutout_background"): self.cutoutBackground = [ int(c) for c in self.settings.value("cutout_background") ] if len(self.cutoutBackground) == 3: self.cutoutBackground += tuple([255]) else: self.cutoutBackground = [0, 0, 128, 255] if self.settings.value("cross_color"): self.crossColor = [ int(c) for c in self.settings.value("cross_color") ] else: self.crossColor = [0, 0, 0, 127] self.scene.setPenColor(self.crossColor) # widget self.dockWidgets = { "model": [self.ModelDock], "data": [self.DataDock], "label": [self.LabelDock], "seg": [self.SegSettingDock], "rs": [self.RSDock], "med": [self.MedDock], "grid": [self.GridDock], "video": [self.VideoDock], "vseg": [self.VSTDock], "3d": [self.TDDock], } # self.display_dockwidget = [True, True, True, True, False, False, False, False, False, False] self.dockStatus = self.settings.value( "dock_status", QVariant([]), type=list) # 所有widget是否展示 if len(self.dockStatus) != len(self.dockWidgets): self.dockStatus = [True] * 4 + [False] * (len(self.dockWidgets) - 4) self.settings.setValue("dock_status", self.dockStatus) else: self.dockStatus = [strtobool(s) for s in self.dockStatus] self.layoutStatus = self.settings.value("layout_status", QByteArray()) # 界面元素位置 self.recentModels = self.settings.value( "recent_models", QVariant([]), type=list) self.video_recentModels = self.settings.value( "video_recent_models", QVariant([]), type=list) self.recentFiles = self.settings.value( "recent_files", QVariant([]), type=list) self.config = util.parse_configs(osp.join(pjpath, "config/config.yaml")) # 支持的图像格式 rs_ext = [".tif", ".tiff"] img_ext = [] for fmt in QtGui.QImageReader.supportedImageFormats(): fmt = ".{}".format(fmt.data().decode()) if fmt not in rs_ext: img_ext.append(fmt) video_ext = [ ".wmv", ".asf", ".asx", ".rm", ".rmvb", ".mp4", ".3gp", ".mov", ".m4v", ".avi", ".dat", ".mkv", ".flv", ".vob", ] self.video_ext = video_ext self.formats = [ img_ext, # 自然图像 [".dcm"], # 医学影像 rs_ext, # 遥感影像 video_ext, # 视频 ] # 遥感 self.raster = None self.grid = None self.rsRGB = [1, 1, 1] # 遥感索引 # 医疗参数 self.midx = 0 # 医疗切片索引 # 大图限制 self.thumbnail_min = 2000 # 初始化action self.initActions() # 更新近期记录 self.loadLayout() # 放前面 self.toggleWidget("all", warn=False) self.updateModelMenu() self.updateVideoModelMenu() self.updateRecentFile() # self.VideoDock.hide() # 窗口 ## 快捷键 self.ShortcutWidget = ShortcutWidget(self.actions, pjpath) ## 画布 self.scene.clickRequest.connect(self.canvasClick) self.canvas.zoomRequest.connect(self.viewZoomed) self.canvas.mousePosChanged.connect(self.scene.onMouseChanged) self.annImage = QtWidgets.QGraphicsPixmapItem() self.scene.addItem(self.annImage) ## 按钮点击 self.btnSave.clicked.connect(self.exportLabel) # 保存 self.listFiles.itemDoubleClicked.connect( self.imageListClicked) # 标签列表点击 self.btnAddClass.clicked.connect(self.addLabel) self.btnParamsSelect.clicked.connect(self.changeParam) # 模型参数选择 self.btn3DParamsSelect.clicked.connect(self.changePropgationParam) self.cheWithMask.stateChanged.connect(self.chooseMode) # with_mask self.btnPropagate.clicked.connect(self.on_propgation) ## 滑动 self.sldOpacity.valueChanged.connect(self.maskOpacityChanged) self.sldClickRadius.valueChanged.connect(self.clickRadiusChanged) self.sldThresh.valueChanged.connect(self.threshChanged) # self.sldBrush.valueChanged.connect(self.brushChanged) self.sldWw.valueChanged.connect(self.swwChanged) self.sldWc.valueChanged.connect(self.swcChanged) self.textWw.returnPressed.connect(self.twwChanged) self.textWc.returnPressed.connect(self.twcChanged) ## 标签列表点击 self.labelListTable.cellDoubleClicked.connect(self.labelListDoubleClick) self.labelListTable.cellClicked.connect(self.labelListClicked) self.labelListTable.cellChanged.connect(self.labelListItemChanged) ## 功能区选择 # self.rsShow.currentIndexChanged.connect(self.rsShowModeChange) # 显示模型 for bandCombo in self.bandCombos: bandCombo.currentIndexChanged.connect(self.rsBandSet) # 设置波段 # self.btnInitGrid.clicked.connect(self.initGrid) # 打开宫格 self.btnFinishedGrid.clicked.connect(self.saveGridLabel) ## 视频相关 self.timer.timeout.connect(self.on_time) self.videoPlay.clicked.connect(self.on_play) self.sldTime.valueChanged.connect(self.sframeChanged) self.textTime.returnPressed.connect(self.tframeChanged) self.ratio = 20 self.speedComboBox.currentIndexChanged.connect(self.on_speed) self.preFrameButton.clicked.connect(self.turnPreFrame) self.nextFrameButton.clicked.connect(self.turnNextFrame) def initActions(self): tr = partial(QtCore.QCoreApplication.translate, "APP_EISeg") action = partial(util.newAction, self) start = dir() # 打开/加载/保存 open_image = action( tr("&打开图像"), self.openImage, "open_image", "OpenImage", tr("打开一张图像进行标注"), ) open_folder = action( tr("&打开文件夹"), self.openFolder, "open_folder", "OpenFolder", tr("打开一个文件夹下所有的图像进行标注"), ) change_output_dir = action( tr("&改变标签保存路径"), partial(self.changeOutputDir, None), "change_output_dir", "ChangeOutputDir", tr("改变标签保存的文件夹路径"), ) load_param = action( tr("&加载模型参数"), self.changeParam, "load_param", "Model", tr("加载一个模型参数"), ) save = action( tr("&保存"), self.exportLabel, "save", "Save", tr("保存图像标签"), ) save_as = action( tr("&另存为"), partial( self.exportLabel, saveAs=True), "save_as", "SaveAs", tr("在指定位置另存为标签"), ) auto_save = action( tr("&自动保存"), self.toggleAutoSave, "auto_save", "AutoSave", tr("翻页同时自动保存"), checkable=True, ) # auto_save.setChecked(self.config.get("auto_save", False)) # 标注 turn_prev = action( tr("&上一张"), partial(self.turnImg, -1), "turn_prev", "Prev", tr("翻到上一张图片"), ) turn_next = action( tr("&下一张"), partial(self.turnImg, 1), "turn_next", "Next", tr("翻到下一张图片"), ) finish_object = action( tr("&完成当前目标"), self.finishObject, "finish_object", "Ok", tr("完成当前目标的标注"), ) clear = action( tr("&清除所有标注"), self.clearAll, "clear", "Clear", tr("清除所有标注信息"), ) undo = action( tr("&撤销"), self.undoClick, "undo", "Undo", tr("撤销一次点击"), ) redo = action( tr("&重做"), self.redoClick, "redo", "Redo", tr("重做一次点击"), ) del_active_polygon = action( tr("&删除多边形"), self.delActivePolygon, "del_active_polygon", "DeletePolygon", tr("删除当前选中的多边形"), ) del_all_polygon = action( tr("&删除所有多边形"), self.delAllPolygon, "del_all_polygon", "DeleteAllPolygon", tr("删除所有的多边形"), ) largest_component = action( tr("&保留最大连通块"), self.toggleLargestCC, "largest_component", "SaveLargestCC", tr("保留最大的连通块"), checkable=True, ) origional_extension = action( tr("&标签和图像使用相同拓展名"), self.toggleOrigExt, "origional_extension", "Same", tr("标签和图像使用相同拓展名，用于图像中有文件名相同但拓展名不同的情况，防止标签覆盖"), checkable=True, ) save_pseudo = action( tr("&伪彩色保存"), partial(self.toggleSave, "pseudo_color"), "save_pseudo", "SavePseudoColor", tr("保存为伪彩色图像"), checkable=True, ) save_pseudo.setChecked(self.save_status["pseudo_color"]) save_grayscale = action( tr("&灰度保存"), partial(self.toggleSave, "gray_scale"), "save_grayscale", "SaveGrayScale", tr("保存为灰度图像，像素的灰度为对应类型的标签"), checkable=True, ) save_grayscale.setChecked(self.save_status["gray_scale"]) save_json = action( tr("&JSON保存"), partial(self.toggleSave, "json"), "save_json", "SaveJson", tr("保存为JSON格式"), checkable=True, ) save_json.setChecked(self.save_status["json"]) save_coco = action( tr("&COCO保存"), partial(self.toggleSave, "coco"), "save_coco", "SaveCOCO", tr("保存为COCO格式"), checkable=True, ) save_coco.setChecked(self.save_status["coco"]) # test func self.show_rs_poly = action( tr("&显示遥感多边形"), None, "show_rs_poly", "Show", tr("显示遥感大图的多边形结果"), checkable=True, ) self.show_rs_poly.setChecked(False) self.grid_message = action( tr("&启用宫格检测"), None, "grid_message", "Show", tr("针对每张图片启用宫格检测"), checkable=True, ) self.grid_message.setChecked(True) eiseg_med3D = action( tr("&EISeg-Med3D"), self.enterEISegMed3D, "enterEISegMed3D", "EISegMed3D", tr("3D医疗交互式分割插件"), ) save_cutout = action( tr("&抠图保存"), partial(self.toggleSave, "cutout"), "save_cutout", "SaveCutout", tr("只保留前景，背景设置为背景色"), checkable=True, ) save_cutout.setChecked(self.save_status["cutout"]) set_cutout_background = action( tr("&设置抠图背景色"), self.setCutoutBackground, "set_cutout_background", self.cutoutBackground, tr("抠图后背景像素的颜色"), ) close = action( tr("&关闭"), partial(self.saveImage, True), "close", "Close", tr("关闭当前图像"), ) quit = action( tr("&退出"), self.close, "quit", "Quit", tr("退出软件"), ) export_label_list = action( tr("&导出标签列表"), partial(self.exportLabelList, None), "export_label_list", "ExportLabel", tr("将标签列表导出成标签配置文件"), ) import_label_list = action( tr("&载入标签列表"), partial(self.importLabelList, None), "import_label_list", "ImportLabel", tr("从标签配置文件载入标签列表"), ) clear_label_list = action( tr("&清空标签列表"), self.clearLabelList, "clear_label_list", "ClearLabel", tr("清空所有的标签"), ) clear_recent = action( tr("&清除近期文件记录"), self.clearRecentFile, "clear_recent", "ClearRecent", tr("清除近期标注文件记录"), ) model_widget = action( tr("&模型选择"), partial(self.toggleWidget, 0), "model_widget", "Net", tr("隐藏/展示模型选择面板"), checkable=True, ) data_widget = action( tr("&数据列表"), partial(self.toggleWidget, 1), "data_widget", "Data", tr("隐藏/展示数据列表面板"), checkable=True, ) label_widget = action( tr("&标签列表"), partial(self.toggleWidget, 2), "label_widget", "Label", tr("隐藏/展示标签列表面板"), checkable=True, ) segmentation_widget = action( tr("&分割设置"), partial(self.toggleWidget, 3), "segmentation_widget", "Setting", tr("隐藏/展示分割设置面板"), checkable=True, ) rs_widget = action( tr("&遥感设置"), partial(self.toggleWidget, 4), "rs_widget", "RemoteSensing", tr("隐藏/展示遥感设置面板"), checkable=True, ) mi_widget = action( tr("&医疗设置"), partial(self.toggleWidget, 5), "mi_widget", "MedicalImaging", tr("隐藏/展示医疗设置面板"), checkable=True, ) grid_ann_widget = action( tr("&N2宫格标注"), partial(self.toggleWidget, 6), "grid_ann_widget", "N2", tr("隐藏/展示N^2宫格细粒度标注面板"), checkable=True, ) video_play_widget = action( tr("&视频播放"), partial(self.toggleWidget, 7), "video_play_widget", "Video", tr("隐藏/展示视频播放面板"), checkable=True, ) video_anno_widget = action( tr("&视频标注"), partial(self.toggleWidget, 8), "video_anno_widget", "VideoAnno", tr("隐藏/展示视频标注面板"), checkable=True, ) td_widget = action( tr("&3D显示"), partial(self.toggleWidget, 9), "td_widget", "3D", tr("隐藏/展示3D显示面板"), checkable=True, ) quick_start = action( tr("&快速入门"), self.quickStart, "quick_start", "Use", tr("主要功能使用介绍"), ) report_bug = action( tr("&反馈问题"), self.reportBug, "report_bug", "ReportBug", tr("通过Github Issue反馈使用过程中遇到的问题。我们会尽快进行修复"), ) edit_shortcuts = action( tr("&编辑快捷键"), self.editShortcut, "edit_shortcuts", "Shortcut", tr("编辑软件快捷键"), ) toggle_logging = action( tr("&调试日志"), self.toggleLogging, "toggle_logging", "Log", tr("用于观察软件执行过程和进行debug。我们不会自动收集任何日志，可能会希望您在反馈问题时间打开此功能，帮助我们定位问题。"), checkable=True, ) toggle_logging.setChecked(bool(self.settings.value("log", False))) use_qt_widget = action( tr("&使用QT文件窗口"), self.useQtWidget, "use_qt_widget", "Qt", tr("如果使用文件选择窗口时遇到问题可以选择使用Qt窗口"), checkable=True, ) # print( # "use_qt_widget", # self.settings.value("use_qt_widget", type=bool), # ) use_qt_widget.setChecked( self.settings.value( "use_qt_widget", False, type=bool)) self.actions = util.struct() for name in dir(): if name not in start: self.actions.append(eval(name)) def newWidget(text, icon, showAction): widget = QtWidgets.QMenu(text) widget.setIcon(util.newIcon(icon)) widget.aboutToShow.connect(showAction) return widget recent_files = newWidget(self.tr("近期文件"), "Data", self.updateRecentFile) recent_params = newWidget( self.tr("近期模型及参数"), "Net", self.updateModelMenu) video_recent_params = newWidget( self.tr("近期视频传播模型及参数"), "Net", self.updateVideoModelMenu) languages = newWidget(self.tr("语言"), "Language", self.updateLanguage) self.menus = util.struct( recent_files=recent_files, recent_params=recent_params, video_recent_params=video_recent_params, languages=languages, fileMenu=( open_image, open_folder, change_output_dir, load_param, clear_recent, recent_files, recent_params, video_recent_params, None, save, save_as, auto_save, None, turn_next, turn_prev, close, None, quit, ), labelMenu=( export_label_list, import_label_list, clear_label_list, ), functionMenu=( largest_component, del_active_polygon, del_all_polygon, None, origional_extension, save_pseudo, save_grayscale, save_cutout, set_cutout_background, None, save_json, save_coco, None, # test self.show_rs_poly, None, self.grid_message, ), showMenu=( model_widget, data_widget, label_widget, segmentation_widget, rs_widget, mi_widget, grid_ann_widget, video_play_widget, video_anno_widget, td_widget, ), helpMenu=( languages, use_qt_widget, quick_start, report_bug, edit_shortcuts, toggle_logging, ), expandMenu=(eiseg_med3D, ), toolBar=( finish_object, clear, undo, redo, turn_prev, turn_next, None, save_pseudo, save_grayscale, save_cutout, save_json, save_coco, origional_extension, None, largest_component, ), ) def menu(title, actions=None): menu = self.menuBar().addMenu(title) if actions: util.addActions(menu, actions) return menu menu(tr("文件"), self.menus.fileMenu) menu(tr("标注"), self.menus.labelMenu) menu(tr("功能"), self.menus.functionMenu) menu(tr("显示"), self.menus.showMenu) menu(tr("帮助"), self.menus.helpMenu) menu(tr("更多"), self.menus.expandMenu) util.addActions(self.toolBar, self.menus.toolBar) def __setColor(self, action, setting_name): c = action color = QtWidgets.QColorDialog.getColor( QtGui.QColor(*c), self, options=QtWidgets.QColorDialog.ShowAlphaChannel, ) action = color.getRgb() self.settings.setValue(setting_name, [int(c) for c in action]) return action def on_speed(self, sender): text = self.speedComboBox.currentText() self.ratio = int(20 * float(text[4:-1])) if self.timer.isActive(): self.timer.stop() self.timer.start(1000 // self.ratio) def setCutoutBackground(self): self.cutoutBackground = self.__setColor(self.cutoutBackground, "cutout_background") self.actions.set_cutout_background.setIcon( util.newIcon(self.cutoutBackground)) def editShortcut(self): self.ShortcutWidget.center() self.ShortcutWidget.show() # 多语言 def updateLanguage(self): self.menus.languages.clear() langs = os.listdir(osp.join(pjpath, "util/translate")) langs = [n.split(".")[0] for n in langs if n.endswith("qm")] langs.append("中文") for lang in langs: if lang == self.currLanguage: continue entry = util.newAction( self, lang, partial(self.changeLanguage, lang), None, lang if lang != "Arabic" else "Egypt", ) self.menus.languages.addAction(entry) def changeLanguage(self, lang): self.settings.setValue("language", lang) self.warn(self.tr("切换语言"), self.tr("切换语言需要重启软件才能生效")) # 近期图像 def updateRecentFile(self): menu = self.menus.recent_files menu.clear() recentFiles = self.settings.value( "recent_files", QVariant([]), type=list) files = [f for f in recentFiles if osp.exists(f)] for i, f in enumerate(files): icon = util.newIcon("File") action = QtWidgets.QAction(icon, "&【%d】 %s" % (i + 1, QtCore.QFileInfo(f).fileName()), self) action.triggered.connect(partial(self.openRecentImage, f)) menu.addAction(action) if len(files) == 0: menu.addAction(self.tr("无近期文件")) self.settings.setValue("recent_files", files) def addRecentFile(self, path): if not osp.exists(path): return paths = self.settings.value("recent_files", QVariant([]), type=list) if path not in paths: paths.append(path) if len(paths) > 15: del paths[0] self.settings.setValue("recent_files", paths) self.updateRecentFile() def clearRecentFile(self): self.settings.remove("recent_files") self.statusbar.showMessage(self.tr("已清除最近打开文件"), 10000) # 模型加载 def updateModelMenu(self): menu = self.menus.recent_params menu.clear() self.recentModels = [ m for m in self.recentModels if osp.exists(m["param_path"]) ] for idx, m in enumerate(self.recentModels): icon = util.newIcon("Model") action = QtWidgets.QAction( icon, f"{osp.basename(m['param_path'])}", self, ) action.triggered.connect( partial(self.setModelParam, m["param_path"])) menu.addAction(action) if len(self.recentModels) == 0: menu.addAction(self.tr("无近期模型记录")) self.settings.setValue("recent_params", self.recentModels) def updateVideoModelMenu(self): menu = self.menus.video_recent_params menu.clear() self.video_recentModels = [ m for m in self.video_recentModels if osp.exists(m["param_path"]) ] for idx, m in enumerate(self.video_recentModels): icon = util.newIcon("Model") action = QtWidgets.QAction( icon, f"{osp.basename(m['param_path'])}", self, ) action.triggered.connect( partial(self.setVideoModelParam, m["param_path"])) menu.addAction(action) if len(self.video_recentModels) == 0: menu.addAction(self.tr("无近期视频传播模型记录")) self.settings.setValue("video_recent_params", self.video_recentModels) def setModelParam(self, paramPath): res = self.changeParam(paramPath) if res: return True return False def setVideoModelParam(self, paramPath): res = self.changePropgationParam(paramPath) if res: return True return False def changeParam(self, param_path: str=None): if not param_path: filters = self.tr("Paddle静态模型权重文件(*.pdiparams)") start_path = ("." if len(self.recentModels) == 0 else osp.dirname(self.recentModels[-1]["param_path"])) if self.settings.value("use_qt_widget", False, type=bool): options = QtWidgets.QFileDialog.DontUseNativeDialog else: options = QtWidgets.QFileDialog.ReadOnly param_path, _ = QtWidgets.QFileDialog.getOpenFileName( self, self.tr("选择传播模型参数") + " - " + __APPNAME__, start_path, filters, options=options, ) # QtWidgets.QFileDialog.DontUseNativeDialog if not param_path: return False # 中文路径打不开 if check_cn(param_path): self.warn(self.tr("参数路径存在中文"), self.tr("请修改参数路径为非中文路径！")) return False success, res = self.controller.setModel(param_path) if success: model_dict = {"param_path": param_path} if model_dict not in self.recentModels: self.recentModels.insert(0, model_dict) if len(self.recentModels) > 10: del self.recentModels[-1] else: # 如果存在移动位置，确保加载最近模型的正确 self.recentModels.remove(model_dict) self.recentModels.insert(0, model_dict) self.settings.setValue("recent_models", self.recentModels) self.statusbar.showMessage( osp.basename(param_path) + self.tr(" 模型加载成功"), 10000) return True else: self.warnException(res) return False def changePropgationParam(self, param_path: str=None): if not param_path: filters = self.tr("Paddle静态模型权重文件(*.pdiparams)") start_path = ( ".") if len(self.video_recentModels) == 0 else osp.dirname( self.video_recentModels[-1]["param_path"]) if self.settings.value("use_qt_widget", False, type=bool): options = QtWidgets.QFileDialog.DontUseNativeDialog else: options = QtWidgets.QFileDialog.ReadOnly param_path, _ = QtWidgets.QFileDialog.getOpenFileName( self, self.tr("选择模型参数") + " - " + __APPNAME__, start_path, filters, options=options, ) # QtWidgets.QFileDialog.DontUseNativeDialog if not param_path: return False # 中文路径打不开 if check_cn(param_path): self.warn(self.tr("参数路径存在中文"), self.tr("请修改参数路径为非中文路径！")) return False success, res = self.video.set_model(param_path) if success: model_dict = {"param_path": param_path} if model_dict not in self.video_recentModels: self.video_recentModels.insert(0, model_dict) if len(self.recentModels) > 10: del self.recentModels[-1] else: # 如果存在移动位置，确保加载最近模型的正确 self.video_recentModels.remove(model_dict) self.video_recentModels.insert(0, model_dict) self.settings.setValue("video_recent_models", self.video_recentModels) self.statusbar.showMessage( osp.basename(param_path) + self.tr("视频传播模型加载成功"), 10000) return True else: self.warnException(res) return False def chooseMode(self): self.predictor_params["predictor_params"][ "with_mask"] = self.cheWithMask.isChecked() self.controller.reset_predictor(predictor_params=self.predictor_params) if self.cheWithMask.isChecked(): self.statusbar.showMessage(self.tr("掩膜已启用"), 10000) else: self.statusbar.showMessage(self.tr("掩膜已关闭"), 10000) def loadRecentModelParam(self): if len(self.recentModels) == 0: self.statusbar.showMessage(self.tr("没有最近使用模型信息，请加载模型"), 10000) return m = self.recentModels[0] param_path = m["param_path"] self.setModelParam(param_path) def loadVideoRecentModelParam(self): if len(self.video_recentModels) == 0: self.statusbar.showMessage(self.tr("没有最近使用的视频传播模型信息，请加载模型"), 10000) return m = self.video_recentModels[0] param_path = m["param_path"] self.setVideoModelParam(param_path) # 标签列表 def importLabelList(self, filePath=None): if filePath is None: if self.settings.value("use_qt_widget", False, type=bool): options = QtWidgets.QFileDialog.DontUseNativeDialog else: options = QtWidgets.QFileDialog.ReadOnly filters = self.tr("标签配置文件") + " (*.txt)" filePath, _ = QtWidgets.QFileDialog.getOpenFileName( self, self.tr("选择标签配置文件路径") + " - " + __APPNAME__, ".", filters, options=options, ) filePath = normcase(filePath) if not osp.exists(filePath): return self.controller.importLabel(filePath) logger.info(f"Loaded label list: {self.controller.labelList.labelList}") self.refreshLabelList() def exportLabelList(self, savePath: str=None): if len(self.controller.labelList) == 0: self.warn(self.tr("没有需要保存的标签"), self.tr("请先添加标签之后再进行保存！")) return if savePath is None: filters = self.tr("标签配置文件") + "(*.txt)" dlg = QtWidgets.QFileDialog( self, self.tr("保存标签配置文件"), ".", filters, ) dlg.setOption(QtWidgets.QFileDialog.DontConfirmOverwrite, False) if self.settings.value("use_qt_widget", False, type=bool): options = QtWidgets.QFileDialog.DontUseNativeDialog else: options = QtWidgets.QFileDialog.DontUseCustomDirectoryIcons dlg.setDefaultSuffix("txt") dlg.setAcceptMode(QtWidgets.QFileDialog.AcceptSave) savePath, _ = dlg.getSaveFileName( self, self.tr("选择保存标签配置文件路径") + " - " + __APPNAME__, ".", filters, options=options, ) self.controller.exportLabel(savePath) def addLabel(self, idx=None, txt="", c=None): c = self.colorMap.get_color() table = self.labelListTable idx = table.rowCount() table.insertRow(table.rowCount()) self.controller.addLabel(idx + 1, txt, c) numberItem = QTableWidgetItem(str(idx + 1)) numberItem.setFlags(QtCore.Qt.ItemIsEnabled) table.setItem(idx, 0, numberItem) table.setItem(idx, 1, QTableWidgetItem()) colorItem = QTableWidgetItem() colorItem.setBackground(QtGui.QColor(c[0], c[1], c[2])) colorItem.setFlags(QtCore.Qt.ItemIsEnabled) table.setItem(idx, 2, colorItem) delItem = QTableWidgetItem() delItem.setIcon(util.newIcon("Clear")) delItem.setTextAlignment(Qt.AlignCenter) delItem.setFlags(QtCore.Qt.ItemIsEnabled) table.setItem(idx, 3, delItem) self.adjustTableSize() self.labelListClicked(self.labelListTable.rowCount() - 1, 0) def adjustTableSize(self): self.labelListTable.horizontalHeader().setDefaultSectionSize(25) self.labelListTable.horizontalHeader().setSectionResizeMode( 0, QtWidgets.QHeaderView.Fixed) self.labelListTable.horizontalHeader().setSectionResizeMode( 3, QtWidgets.QHeaderView.Fixed) self.labelListTable.horizontalHeader().setSectionResizeMode( 2, QtWidgets.QHeaderView.Fixed) self.labelListTable.setColumnWidth(2, 50) def clearLabelList(self): if len(self.controller.labelList) == 0: return True res = self.warn( self.tr("清空标签列表?"), self.tr("请确认是否要清空标签列表"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Cancel: return False self.controller.labelList.clear() if self.controller: self.controller.label_list = [] self.controller.curr_label_number = 0 self.labelListTable.clear() self.labelListTable.setRowCount(0) return True def refreshLabelList(self): table = self.labelListTable table.clearContents() table.setRowCount(len(self.controller.labelList)) table.setColumnCount(4) for idx, lab in enumerate(self.controller.labelList): numberItem = QTableWidgetItem(str(lab.idx)) numberItem.setFlags(QtCore.Qt.ItemIsEnabled) table.setItem(idx, 0, numberItem) table.setItem(idx, 1, QTableWidgetItem(lab.name)) c = lab.color colorItem = QTableWidgetItem() colorItem.setBackground(QtGui.QColor(c[0], c[1], c[2])) colorItem.setFlags(QtCore.Qt.ItemIsEnabled) table.setItem(idx, 2, colorItem) delItem = QTableWidgetItem() delItem.setIcon(util.newIcon("Clear")) delItem.setTextAlignment(Qt.AlignCenter) delItem.setFlags(QtCore.Qt.ItemIsEnabled) table.setItem(idx, 3, delItem) self.adjustTableSize() cols = [0, 1, 3] for idx in cols: table.resizeColumnToContents(idx) self.adjustTableSize() def labelListDoubleClick(self, row, col): if col != 2: return table = self.labelListTable color = QtWidgets.QColorDialog.getColor() if color.getRgb() == (0, 0, 0, 255): return table.item(row, col).setBackground(color) self.controller.labelList[row].color = color.getRgb()[:3] if self.controller: self.controller.label_list = self.controller.labelList for p in self.scene.polygon_items: idlab = self.controller.labelList.getLabelById(p.labelIndex) if idlab is not None: color = idlab.color p.setColor(color, color) self.labelListClicked(row, 0) @property def currLabelIdx(self): return self.controller.curr_label_number - 1 def labelListClicked(self, row, col): table = self.labelListTable if col == 3: labelIdx = int(table.item(row, 0).text()) if self.status == self.EDITING: if self.checkLabel(labelIdx): self.controller.labelList.remove(labelIdx) table.removeRow(row) else: self.warn( self.tr("无法删除"), self.tr("当前多边形中存在此标签"), QMessageBox.Yes) elif self.status == self.ANNING: self.warn( self.tr("无法删除"), self.tr("交互式标注模式无法删除标签"), QMessageBox.Yes) if col == 0 or col == 1: for cl in range(2): for idx in range(len(self.controller.labelList)): table.item(idx, cl).setBackground(QtGui.QColor(255, 255, 255)) table.item(row, cl).setBackground(QtGui.QColor(48, 140, 198)) table.item(row, 0).setSelected(True) if self.controller: self.controller.setCurrLabelIdx(int(table.item(row, 0).text())) self.controller.label_list = self.controller.labelList def labelListItemChanged(self, row, col): self.colorMap.usedColors = self.controller.labelList.colors try: if col == 1: name = self.labelListTable.item(row, col).text() self.controller.labelList[row].name = name except: pass # 多边形标注 def createPoly(self, curr_polygon, color): if curr_polygon is None: return for points in curr_polygon: if len(points) < 3: continue poly = PolygonAnnotation( self.controller.labelList[self.currLabelIdx].idx, self.controller.image.shape, self.delPolygon, self.setDirty, color, color, self.opacity, ) poly.labelIndex = self.controller.labelList[self.currLabelIdx].idx self.scene.addItem(poly) self.scene.polygon_items.append(poly) for p in points: poly.addPointLast(QtCore.QPointF(p[0], p[1])) self.setDirty(True) def delActivePolygon(self): for idx, polygon in enumerate(self.scene.polygon_items): if polygon.hasFocus(): res = self.warn( self.tr("确认删除？"), self.tr("确认要删除当前选中多边形标注？"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Yes: self.delPolygon(polygon) def delPolygon(self, polygon): polygon.remove() if self.save_status["coco"]: if polygon.coco_id: self.coco.delAnnotation( polygon.coco_id, self.coco.imgNameToId[osp.basename(self.imagePath)], ) self.setDirty(True) def delAllPolygon(self): for p in self.scene.polygon_items[::-1]: # 删除所有多边形 self.delPolygon(p) def delActivePoint(self): for polygon in self.scene.polygon_items: polygon.removeFocusPoint() # 图片/标签 io def getMask(self): if not self.controller or self.controller.image is None: return s = self.controller.imgShape pesudo = np.zeros([s[0], s[1]]) # 覆盖顺序，从上往下 # TODO: 是标签数值大的会覆盖小的吗? # A: 是列表中上面的覆盖下面的，由于标签可以移动，不一定是大小按顺序覆盖 # RE: 我们做医学的时候覆盖比较多，感觉一般是数值大的标签覆盖数值小的标签。按照上面覆盖下面的话可能跟常见的情况正好是反过来的，感觉可能从下往上覆盖会比较好 len_lab = self.labelListTable.rowCount() for i in range(len_lab - 1, -1, -1): idx = int(self.labelListTable.item(len_lab - i - 1, 0).text()) for poly in self.scene.polygon_items: if poly.labelIndex == idx: pts = np.int32([np.array(poly.scnenePoints)]) cv2.fillPoly(pesudo, pts=pts, color=idx) return pesudo def openRecentImage(self, file_path): self.queueEvent(partial(self.loadImage, file_path)) self.listFiles.addItems([file_path.replace("\\", "/")]) self.currIdx = self.listFiles.count() - 1 self.listFiles.setCurrentRow(self.currIdx) # 移动位置 self.imagePaths.append(file_path) def openImage(self, filePath: str=None): # 在triggered.connect中使用不管默认filePath为什么返回值都为False if not isinstance(filePath, str) or filePath is False: prompts = ["图片", "医学影像", "遥感影像", "视频"] filters = "" for fmts, p in zip(self.formats, prompts): filters += f"{p} ({' '.join(['*' + f for f in fmts])}) ;; " filters = filters[:-3] recentPath = self.settings.value("recent_files", []) if len(recentPath) == 0: recentPath = "." else: recentPath = osp.dirname(recentPath[0]) if self.settings.value("use_qt_widget", False, type=bool): options = QtWidgets.QFileDialog.DontUseNativeDialog else: options = QtWidgets.QFileDialog.ReadOnly filePath, _ = QtWidgets.QFileDialog.getOpenFileName( self, self.tr("选择待标注图片") + " - " + __APPNAME__, recentPath, filters, options=options, ) if len(filePath) == 0: # 用户没选就直接关闭窗口 return if osp.splitext(filePath)[-1] in self.video_ext: if not paddle.device.is_compiled_with_cuda( ): # TODO: 可以使用GPU却返回False self.warn( self.tr("请在gpu电脑上进行视频标注"), self.tr("准备进行视频标注，由于视频标注需要一定计算，请尽量确保在gpu的电脑上进行操作!")) filePath = normcase(filePath) if not self.loadImage(filePath): return False # 3. 添加记录 self.listFiles.addItems([filePath]) self.currIdx = self.listFiles.count() - 1 self.listFiles.setCurrentRow(self.currIdx) # 移动位置 self.imagePaths.append(filePath) return True def openFolder(self, inputDir: str=None): # 1. 如果没传文件夹，弹框让用户选 if not isinstance(inputDir, str): recentPath = self.settings.value("recent_files", []) if len(recentPath) == 0: recentPath = "." else: recentPath = osp.dirname(recentPath[-1]) options = (QtWidgets.QFileDialog.ShowDirsOnly | QtWidgets.QFileDialog.DontResolveSymlinks) if self.settings.value("use_qt_widget", False, type=bool): options = options | QtWidgets.QFileDialog.DontUseNativeDialog inputDir = QtWidgets.QFileDialog.getExistingDirectory( self, self.tr("选择待标注图片文件夹") + " - " + __APPNAME__, recentPath, options, ) if not osp.exists(inputDir): return # 2. 关闭当前图片，清空文件列表 self.saveImage(close=True) self.imagePaths = [] self.listFiles.clear() # 3. 扫描文件夹下所有图片 # 3.1 获取所有文件名 imagePaths = os.listdir(inputDir) exts = tuple(f for fmts in self.formats for f in fmts) imagePaths = [n for n in imagePaths if n.lower().endswith(exts)] imagePaths = [n for n in imagePaths if not n[0] == "."] imagePaths.sort() if len(imagePaths) == 0: return # 3.2 设置默认输出路径 if self.outputDir is None: # 没设置为文件夹下的 label 文件夹 self.outputDir = osp.join(inputDir, "label") if not osp.exists(self.outputDir): os.makedirs(self.outputDir) # 3.3 有重名图片，标签保留原来拓展名 names = [] for name in imagePaths: name = osp.splitext(name)[0] if name not in names: names.append(name) else: self.toggleOrigExt(True) break imagePaths = [osp.join(inputDir, n) for n in imagePaths] for p in imagePaths: p = normcase(p) self.imagePaths.append(p) self.listFiles.addItem(p) # 3.4 加载已有的标注 if self.outputDir is not None and osp.exists(self.outputDir): self.changeOutputDir(self.outputDir) if len(self.imagePaths) != 0: self.currIdx = 0 self.turnImg(0) self.inputDir = inputDir def loadImage(self, path): if self.controller.model is None: self.warn("未检测到模型", "请先加载模型参数") return # 1. 拒绝None和不存在的路径，关闭当前图像 if not path: return path = normcase(path) if not osp.exists(path): return self.imagePath = path self.saveImage(True) # 关闭当前图像 self.eximgsInit() # TODO: 将grid的部分整合到saveImage里 # 2. 判断图像类型，打开 # TODO: 加用户指定类型的功能 image = None # 直接if会报错，因为打开遥感图像后多波段不存在，现在把遥感图像的单独抽出来了 # 自然图像 if path.lower().endswith(tuple(self.formats[0])): image = cv2.imdecode(np.fromfile(path, dtype=np.uint8), 1) image = image[:, :, ::-1] # BGR转RGB if self.grid_message.isChecked(): if checkOpenGrid(image, self.thumbnail_min): if self.loadGrid(image, False): image, _ = self.grid.getGrid(0, 0) # 自然图像不进行缩小 else: if self.dockWidgets["grid"][0].isVisible() is True: self.grid = Grids(image) self.initGrid() image, _ = self.grid.getGrid(0, 0) # 医学影像 if path.lower().endswith(tuple(self.formats[1])): if not self.dockStatus[5]: res = self.warn( self.tr("未启用医疗组件"), self.tr("加载医疗影像需启用医疗组件，是否立即启用？"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Cancel: return False self.toggleWidget(5) if not self.dockStatus[5]: return False image = med.dcm_reader(path) # TODO: 添加多层支持 if image.shape[-1] != 1: self.warn("医学影像打开错误", "暂不支持打开多层医学影像") return False maxValue = np.max(image) # 根据数据模态自适应窗宽窗位 minValue = np.min(image) if minValue == 0: ww = maxValue wc = int(maxValue / 2) else: ww = maxValue + int(abs(minValue)) wc = int((minValue + maxValue) / 2) self.sldWw.setValue(int(ww)) self.textWw.setText(str(ww)) self.sldWc.setValue(int(wc)) self.textWc.setText(str(wc)) self.controller.rawImage = self.image = image image = med.windowlize(image, self.ww, self.wc) # 遥感图像 if path.lower().endswith(tuple(self.formats[ 2])): # imghdr.what(path) == "tiff": if not self.dockStatus[4]: res = self.warn( self.tr("未打开遥感组件"), self.tr("打开遥感图像需启用遥感组件，是否立即启用？"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Cancel: return False self.toggleWidget(4) if not self.dockStatus[4]: return False self.raster = Raster(path) gi = self.raster.showGeoInfo() self.edtGeoinfo.setText( self.tr("● 波段数：") + gi[0] + "\n" + self.tr("● 数据类型：") + gi[1] + "\n" + self.tr("● 行数：") + gi[2] + "\n" + self.tr("● 列数：") + gi[ 3] + "\n" + "● EPSG：" + gi[4]) if max(self.rsRGB) > self.raster.geoinfo.count: self.rsRGB = [1, 1, 1] self.raster.setBand(self.rsRGB) if self.grid_message.isChecked(): if self.raster.checkOpenGrid(self.thumbnail_min): if self.loadGrid(self.raster): image, _ = self.raster.getGrid(0, 0) else: image, _ = self.raster.getArray() else: image, _ = self.raster.getArray() else: if self.dockWidgets["grid"][0].isVisible() is True: self.grid = RSGrids(self.raster) self.raster.open_grid = True self.initGrid() image, _ = self.raster.getGrid(0, 0) else: image, _ = self.raster.getArray() self.updateBandList() # self.updateSlideSld(True) else: self.edtGeoinfo.setText(self.tr("无")) # 视频 if path.lower().endswith(tuple(self.formats[3])): # mp4 if not self.dockStatus[7]: res = self.warn( self.tr("未启用视频组件"), self.tr("加载视频需启用视频组件，是否立即启用？"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Cancel: return False self.toggleWidget(7) self.toggleWidget(8) if not self.dockStatus[7]: return False if not self.dockStatus[8]: return False # self.video_masks = None self.video_images, self.fps = self.video.set_video(path) self.video_masks = np.zeros( (self.video.num_frames, self.video.height, self.video.width), dtype=np.uint8) self.sldTime.setMaximum(self.video.num_frames - 1) image = self.video_images[self.video.cursur] self.sldTime.setProperty("value", 0) # 清空3d显示 if self.TDDock.isVisible(): self.vtkWidget.init() # TODO: 处理 # 如果没找到图片的reader if image is None: self.warn("打开图像失败", f"未找到{path}文件对应的读取程序") return self.image = image self.controller.setImage(image) self.updateImage(True) # 2. 加载标签 self.loadLabel(path) self.addRecentFile(path) return True def loadLabel(self, imgPath): if imgPath == "": return None if self.video_images is not None: videoName = osp.splitext(osp.basename(imgPath))[0] maskPath = None for path in self.labelPaths: if osp.basename(path) == videoName: maskPath = osp.join(path, 'mask') if not maskPath: return for cursur in range(self.video.num_frames): h, w = self.video_masks[cursur].shape frame_mask = np.zeros([h, w]) pseudo = cv2.imread( osp.join(maskPath, '{:05d}.png'.format(cursur))) for lab in self.controller.labelList: frame_mask[(pseudo == lab.color[::-1])[:, :, 0]] = lab.idx self.video_masks[cursur] = frame_mask return # 1. 读取json格式标签 if self.save_status["json"]: def getName(path): return osp.splitext(osp.basename(path))[0] imgName = getName(imgPath) labelPath = None for path in self.labelPaths: if not path.endswith(".json"): continue if self.origExt: if getName(path) == osp.basename(imgPath): labelPath = path break else: if getName(path) == imgName: labelPath = path break if not labelPath: return labels = json.loads(open(labelPath, "r").read()) for label in labels: color = label["color"] labelIdx = label["labelIdx"] points = label["points"] poly = PolygonAnnotation( labelIdx, self.controller.image.shape, self.delPolygon, self.setDirty, color, color, self.opacity, ) self.scene.addItem(poly) self.scene.polygon_items.append(poly) for p in points: poly.addPointLast(QtCore.QPointF(p[0], p[1])) # 2. 读取coco格式标签 if self.save_status["coco"]: imgId = self.coco.imgNameToId.get(osp.basename(imgPath), None) if imgId is None: return anns = self.coco.imgToAnns[imgId] for ann in anns: xys = ann["segmentation"][0] points = [] for idx in range(0, len(xys), 2): points.append([xys[idx], xys[idx + 1]]) labelIdx = ann["category_id"] idlab = self.controller.labelList.getLabelById(labelIdx) if idlab is not None: color = idlab.color poly = PolygonAnnotation( ann["category_id"], self.controller.image.shape, self.delPolygon, self.setDirty, color, color, self.opacity, ann["id"], ) self.scene.addItem(poly) self.scene.polygon_items.append(poly) for p in points: poly.addPointLast(QtCore.QPointF(p[0], p[1])) def turnImg(self, delta, list_click=False): if (self.grid is None or self.grid.curr_idx is None) or list_click: # 1. 检查是否有图可翻，保存标签 self.currIdx += delta if self.currIdx >= len(self.imagePaths) or self.currIdx < 0: self.currIdx -= delta if delta == 1: self.statusbar.showMessage(self.tr(f"没有后一张图片")) else: self.statusbar.showMessage(self.tr(f"没有前一张图片")) self.saveImage(False) return else: self.saveImage(True) # 2. 打开新图 self.loadImage(self.imagePaths[self.currIdx]) self.listFiles.setCurrentRow(self.currIdx) else: self.turnGrid(delta) self.setDirty(False) def imageListClicked(self): if not self.controller: self.warn(self.tr("模型未加载"), self.tr("尚未加载模型，请先加载模型！")) self.changeParam() if not self.controller: return if self.controller.is_incomplete_mask: self.exportLabel() toRow = self.listFiles.currentRow() delta = toRow - self.currIdx self.turnImg(delta, True) def finishObject(self): if not self.controller or self.image is None: return current_mask, curr_polygon = self.controller.finishObject( building=self.boundaryRegular.isChecked()) if curr_polygon is not None: self.updateImage() if current_mask is not None: # current_mask = current_mask.astype(np.uint8) * 255 # polygon = util.get_polygon(current_mask) color = self.controller.labelList[self.currLabelIdx].color self.createPoly(curr_polygon, color) # 状态改变 if self.status == self.EDITING: self.status = self.ANNING for p in self.scene.polygon_items: p.setAnning(isAnning=True) else: self.status = self.EDITING for p in self.scene.polygon_items: p.setAnning(isAnning=False) current_mask = self.getMask() if self.video_images is not None: if current_mask.max() != 0: self.video_masks[self.video.cursur] = current_mask def completeLastMask(self): # 返回最后一个标签是否完成，false就是还有带点的 if not self.controller or self.controller.image is None: return True if not self.controller.is_incomplete_mask: return True res = self.warn( self.tr("完成最后一个目标？"), self.tr("是否完成最后一个目标的标注，不完成不会进行保存。"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Yes: self.finishObject() self.exportLabel() self.setDirty(False) return True return False def saveImage(self, close=False): if self.controller and self.controller.image is not None: # 1. 完成正在交互式标注的标签 self.completeLastMask() # 2. 进行保存 if self.isDirty: if self.actions.auto_save.isChecked(): self.exportLabel() else: res = self.warn( self.tr("保存标签？"), self.tr("标签尚未保存，是否保存标签"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Yes: self.exportLabel() self.setDirty(False) if close: # 3. 清空多边形标注，删掉图片 for p in self.scene.polygon_items[::-1]: p.remove() self.scene.polygon_items = [] self.controller.resetLastObject() self.updateImage() self.controller.image = None if close: self.annImage.setPixmap(QPixmap()) if self.video_images is not None and self.video_masks is not None: self.reset_video() def reset_video(self): self.video_images = None self.video_masks = None self.timer.stop() self.textTime.setText(str(0)) self.videoPlay.setText(self.tr("播放")) self.videoPlay.setIcon( QtGui.QIcon(osp.join(pjpath, "resource/Play.png"))) self.ratio = 20 self.speedComboBox.setCurrentIndex(2) self.video.reset() def exportLabel(self, saveAs=False, savePath=None, lab_input=None): # 1. 需要处于标注状态 if not self.controller or self.controller.image is None: return # 2. 完成正在交互式标注的标签 self.completeLastMask() # 3. 确定保存路径 # 3.1 如果参数指定了保存路径直接存到savePath if not savePath: if not saveAs and self.outputDir is not None: # 3.2 指定了标签文件夹，而且不是另存为：根据标签文件夹和文件名出保存路径 name, ext = osp.splitext(osp.basename(self.imagePath)) if not self.origExt: ext = ".png" savePath = osp.join( self.outputDir, name + ext, ) if self.video_images is not None and self.video_masks is not None: savePath = osp.join(self.outputDir, name) os.makedirs(savePath, exist_ok=True) else: # 3.3 没有指定标签存到哪，或者是另存为：弹框让用户选 savePath = self.chooseSavePath() if savePath is None or not osp.exists(osp.dirname(savePath)): return if savePath not in self.labelPaths: self.labelPaths.append(savePath) # 视频帧保存&视频保存 if self.video_masks is not None: if osp.exists(savePath): res = self.warn( self.tr("文件夹已经存在"), self.tr("该文件夹下不为空，您确定继续保存在此路径下吗？"), QMessageBox.Yes | QMessageBox.Cancel, ) if res == QMessageBox.Cancel: return os.makedirs(savePath, exist_ok=True) if osp.isdir(savePath): mask_dir = osp.join(savePath, 'mask') overlay_dir = osp.join(savePath, 'overlay') os.makedirs(mask_dir, exist_ok=True) os.makedirs(overlay_dir, exist_ok=True) progress = QtWidgets.QProgressDialog(self) progress.setWindowTitle("请稍等") progress.setLabelText("正在保存...") progress.setCancelButtonText("取消") progress.setMinimumDuration(5) progress.setWindowModality(Qt.WindowModal) progress.setRange(0, self.video.num_frames) videoname = savePath + "_overlay.mp4" fourcc = cv2.VideoWriter_fourcc(*'mp4v') h, w = self.video_masks[0].shape videoWrite = cv2.VideoWriter(videoname, fourcc, self.fps, (w, h)) for i in range(0, self.video.num_frames): # Save mask mask = self.video_masks[i].astype('uint8') pseudo = np.zeros([h, w, 3]) # mask = self.controller.result_mask # print(pseudo.shape, mask.shape) for lab in self.controller.labelList: pseudo[mask == lab.idx, :] = lab.color[::-1] cv2.imwrite( os.path.join(mask_dir, '{:05d}.png'.format(i)), pseudo) # Save overlay overlay = overlay_davis(self.video_images[i], self.video_masks[i], self.opacity, self.controller.palette) videoWrite.write(overlay[:, :, ::-1]) # write video overlay = Image.fromarray(overlay) overlay.save( os.path.join(overlay_dir, '{:05d}.png'.format(i))) progress.setValue(i) if progress.wasCanceled(): # QMessageBox.warning(self, "提示", "保存失败") break progress.setValue(self.video.num_frames) videoWrite.release() self.setDirty(False) self.statusbar.showMessage( self.tr("视频帧成功保存至") + " " + savePath, 5000) return if lab_input is None: mask_output = self.getMask() s = self.controller.imgShape else: mask_output = lab_input s = lab_input.shape # BUG: 如果用了多边形标注从多边形生成mask # 4.1 保存灰度图 if self.save_status["gray_scale"]: if self.raster is not None: # FIXME: when big map saved, self.raster is None, # so adjust polygon can't saved in tif's mask. pathHead, _ = osp.splitext(savePath) # if self.rsSave.isChecked(): tifPath = pathHead + "_mask.tif" self.raster.saveMask(mask_output, tifPath) if self.shpSave.isChecked(): shpPath = pathHead + ".shp" print(rs.save_shp(shpPath, tifPath)) else: ext = osp.splitext(savePath)[1] cv2.imencode(ext, mask_output)[1].tofile(savePath) # self.labelPaths.append(savePath) # 4.2 保存伪彩色 if self.save_status["pseudo_color"]: pseudoPath, ext = osp.splitext(savePath) pseudoPath = pseudoPath + "_pseudo" + ext pseudo = np.zeros([s[0], s[1], 3], dtype="uint8") # mask = self.controller.result_mask mask = mask_output # print(pseudo.shape, mask.shape) for lab in self.controller.labelList: pseudo[mask == lab.idx, :] = lab.color[::-1] cv2.imencode(ext, pseudo)[1].tofile(pseudoPath) # 4.3 保存前景抠图 if self.save_status["cutout"]: mattingPath, ext = osp.splitext(savePath) mattingPath = mattingPath + "_cutout" + ext img = np.ones([s[0], s[1], 4], dtype="uint8") * 255 cim = cv2.resize(self.controller.image.copy(), (s[1], s[0])) img[:, :, :3] = cim img[mask_output == 0] = self.cutoutBackground img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA) cv2.imencode(ext, img)[1].tofile(mattingPath) # 4.4 保存json if self.save_status["json"]: polygons = self.scene.polygon_items labels = [] for polygon in polygons: l = self.controller.labelList[polygon.labelIndex - 1] label = { "name": l.name, "labelIdx": l.idx, "color": l.color, "points": [], } for p in polygon.scnenePoints: label["points"].append(p) labels.append(label) if self.origExt: jsonPath = savePath + ".json" else: jsonPath = osp.splitext(savePath)[0] + ".json" open(jsonPath, "w", encoding="utf-8").write(json.dumps(labels)) self.labelPaths.append(jsonPath) # 4.5 保存coco if self.save_status["coco"]: if not self.coco.hasImage(osp.basename(self.imagePath)): imgId = self.coco.addImage( osp.basename(self.imagePath), s[1], s[0]) else: imgId = self.coco.imgNameToId[osp.basename(self.imagePath)] for polygon in self.scene.polygon_items: points = [] for p in polygon.scnenePoints: for val in p: points.append(val) if not polygon.coco_id: annId = self.coco.addAnnotation(imgId, polygon.labelIndex, points) polygon.coco_id = annId else: self.coco.updateAnnotation(polygon.coco_id, imgId, points) for lab in self.controller.labelList: if self.coco.hasCat(lab.idx): self.coco.updateCategory(lab.idx, lab.name, lab.color) else: self.coco.addCategory(lab.idx, lab.name, lab.color) saveDir = (self.outputDir if self.outputDir is not None else osp.dirname(savePath)) cocoPath = osp.join(saveDir, "annotations.json") open( cocoPath, "w", encoding="utf-8").write(json.dumps(self.coco.dataset)) self.setDirty(False) self.statusbar.showMessage(self.tr("标签成功保存至") + " " + savePath, 5000) def chooseSavePath(self): formats = [ "*.{}".format(fmt.data().decode()) for fmt in QtGui.QImageReader.supportedImageFormats() ] filters = "Label file (%s)" % " ".join(formats) dlg = QtWidgets.QFileDialog( self, self.tr("保存标签文件路径"), osp.dirname(self.imagePath), filters, ) dlg.setDefaultSuffix("png") dlg.setAcceptMode(QtWidgets.QFileDialog.AcceptSave) dlg.setOption(QtWidgets.QFileDialog.DontConfirmOverwrite, False) dlg.setOption(QtWidgets.QFileDialog.DontUseNativeDialog, False) if self.video_masks is not None: savePath = dlg.getExistingDirectory( self, self.tr("选择标签文件保存路径"), osp.splitext(osp.basename(self.imagePath))[0], ) name, ext = osp.splitext(osp.basename(self.imagePath)) savePath = osp.join(savePath, name) else: savePath, _ = dlg.getSaveFileName( self, self.tr("选择标签文件保存路径"), osp.splitext(osp.basename(self.imagePath))[0] + ".png", ) return savePath def eximgsInit(self): self.gridTable.setRowCount(0) self.gridTable.clearContents() # 清零 self.raster = None self.grid = None def setDirty(self, isDirty): self.isDirty = isDirty def changeOutputDir(self, outputDir=None): # 1. 弹框选择标签路径 if outputDir is None: options = (QtWidgets.QFileDialog.ShowDirsOnly | QtWidgets.QFileDialog.DontResolveSymlinks) if self.settings.value("use_qt_widget", False, type=bool): options = options | QtWidgets.QFileDialog.DontUseNativeDialog outputDir = QtWidgets.QFileDialog.getExistingDirectory( self, self.tr("选择标签保存路径") + " - " + __APPNAME__, self.settings.value("output_dir", "."), options, ) if not osp.exists(outputDir): return False self.settings.setValue("output_dir", outputDir) self.outputDir = outputDir # 2. 加载标签 # 2.1 如果保存coco格式，加载coco标签 if self.save_status["coco"]: defaultPath = osp.join(self.outputDir, "annotations.json") if osp.exists(defaultPath): self.initCoco(defaultPath) # 2.2 如果保存json格式，获取所有json文件名 if self.save_status["json"]: labelPaths = os.listdir(outputDir) labelPaths = [n for n in labelPaths if n.endswith(".json")] labelPaths = [osp.join(outputDir, n) for n in labelPaths] self.labelPaths = labelPaths # 加载对应的标签列表 lab_auto_save = osp.join(self.outputDir, "autosave_label.txt") if osp.exists(lab_auto_save) == False: lab_auto_save = osp.join(self.outputDir, "label/autosave_label.txt") if osp.exists(lab_auto_save): try: self.importLabelList(lab_auto_save) except: pass return True def maskOpacityChanged(self): self.sldOpacity.textLab.setText(str(self.opacity)) if not self.controller or self.controller.image is None: return for polygon in self.scene.polygon_items: polygon.setOpacity(self.opacity) self.updateImage() if self.video_images is not None and self.video_masks is not None: self.show_current_frame() def clickRadiusChanged(self): self.sldClickRadius.textLab.setText(str(self.clickRadius)) if not self.controller or self.controller.image is None: return self.updateImage() if self.video_images is not None and self.video_masks is not None: self.show_current_frame() def threshChanged(self): self.sldThresh.textLab.setText(str(self.segThresh)) if not self.controller or self.controller.image is None: return self.controller.prob_thresh = self.segThresh self.updateImage() if self.video_images is not None and self.video_masks is not None: self.show_current_frame() # def slideChanged(self): # self.sldMISlide.textLab.setText(str(self.slideMi)) # if not self.controller or self.controller.image is None: # return # self.midx = int(self.slideMi) - 1 # self.miSlideSet() # self.updateImage() def undoClick(self): if self.image is None: return if not self.controller: return self.controller.undoClick() self.updateImage() if not self.controller.is_incomplete_mask: self.setDirty(False) def clearAll(self): if not self.controller or self.controller.image is None: return self.controller.resetLastObject() self.updateImage() self.setDirty(False) def redoClick(self): if self.image is None: return if not self.controller: return self.controller.redoClick() self.updateImage() def canvasClick(self, x, y, isLeft): c = self.controller if c.image is None: return if not c.inImage(x, y): return if not c.modelSet: self.warn(self.tr("未选择模型", self.tr("尚未选择模型，请先在右上角选择模型"))) return if self.status == self.IDILE: return currLabel = self.controller.curr_label_number if not currLabel or currLabel == 0: self.warn(self.tr("未选择当前标签"), self.tr("请先在标签列表中单击点选标签")) return self.controller.addClick(x, y, isLeft) self.updateImage() self.status = self.ANNING def updateImage(self, reset_canvas=False): if not self.controller: return image = self.controller.get_visualization( alpha_blend=self.opacity, click_radius=self.clickRadius, ) height, width, _ = image.shape bytesPerLine = 3 * width image = QImage(image.data, width, height, bytesPerLine, QImage.Format_RGB888) if reset_canvas: self.resetZoom(width, height) self.annImage.setPixmap(QPixmap(image)) def update_interact_viz(self): height, width, channel = self.viz.shape bytesPerLine = 3 * width qImg = QImage(self.viz.data, width, height, bytesPerLine, QImage.Format_RGB888) self.annImage.setPixmap(QPixmap(qImg)) def viewZoomed(self, scale): self.scene.scale = scale self.scene.updatePolygonSize() # 界面缩放重置 def resetZoom(self, width, height): # 每次加载图像前设定下当前的显示框，解决图像缩小后不在中心的问题 self.scene.setSceneRect(0, 0, width, height) # 缩放清除 self.canvas.scale(1 / self.canvas.zoom_all, 1 / self.canvas.zoom_all) # 重置缩放 self.canvas.zoom_all = 1 # 最佳缩放 s_eps = 0.98 scr_cont = [ (self.scrollArea.width() * s_eps) / width, (self.scrollArea.height() * s_eps) / height, ] if scr_cont[0] * height > self.scrollArea.height(): self.canvas.zoom_all = scr_cont[1] else: self.canvas.zoom_all = scr_cont[0] self.canvas.scale(self.canvas.zoom_all, self.canvas.zoom_all) self.scene.scale = self.canvas.zoom_all def keyReleaseEvent(self, event): # print(event.key(), Qt.Key_Control) # 释放ctrl的时候刷新图像，对应自适应点大小在缩放后刷新 if not self.controller or self.controller.image is None: return if event.key() == Qt.Key_Control: self.updateImage() def queueEvent(self, function): QtCore.QTimer.singleShot(0, function) def toggleOrigExt(self, dst=None): if dst: self.origExt = dst else: self.origExt = not self.origExt self.actions.origional_extension.setChecked(self.origExt) def toggleAutoSave(self, save): if save and not self.outputDir: self.changeOutputDir(None) if save and not self.outputDir: save = False self.actions.auto_save.setChecked(save) self.settings.setValue("auto_save", save) def toggleSave(self, type): self.save_status[type] = not self.save_status[type] if type == "coco" and self.save_status["coco"]: self.initCoco() if type == "coco": self.save_status["json"] = not self.save_status["coco"] self.actions.save_json.setChecked(self.save_status["json"]) if type == "json": self.save_status["coco"] = not self.save_status["json"] self.actions.save_coco.setChecked(self.save_status["coco"]) def initCoco(self, coco_path: str=None): if not coco_path: if not self.outputDir or not osp.exists(self.outputDir): coco_path = None else: coco_path = osp.join(self.outputDir, "annotations.json") else: if not osp.exists(coco_path): coco_path = None self.coco = COCO(coco_path) if self.clearLabelList(): self.controller.labelList = util.LabelList(self.coco.dataset[ "categories"]) self.refreshLabelList() def toggleWidget(self, index=None, warn=True): # TODO: 输入从数字改成名字 # 1. 改变 if isinstance(index, int): self.dockStatus[index] = not self.dockStatus[index] # 2. 判断widget是否可以开启 # 2.1 遥感 if self.dockStatus[4] and not (rs.check_gdal() and rs.check_rasterio()): if warn: self.warn( self.tr("无法导入GDAL或rasterio"), self.tr("使用遥感工具需要安装GDAL和rasterio！"), QMessageBox.Yes, ) self.statusbar.showMessage(self.tr("打开遥感工具失败，请安装GDAL和rasterio")) self.dockStatus[4] = False # 2.2 医疗 if self.dockStatus[5] and not med.has_sitk(): if warn: self.warn( self.tr("无法导入SimpleITK"), self.tr("使用医疗工具需要安装SimpleITK！"), QMessageBox.Yes, ) self.statusbar.showMessage(self.tr("打开医疗工具失败，请安装SimpleITK")) self.dockStatus[5] = False # 2.3 3D显示 if self.dockStatus[9] and not self.vtkWidget.convert_vtk(): if warn: self.warn( self.tr("无法导入VTK"), self.tr("使用3D显示工具需要安装VTK！"), QMessageBox.Yes, ) self.statusbar.showMessage(self.tr("打开3D显示工具失败，请安装VTK")) self.dockStatus[9] = False widgets = list(self.dockWidgets.values()) for idx, s in enumerate(self.dockStatus): self.menus.showMenu[idx].setChecked(s) if s: for w in widgets[idx]: w.show() else: for w in widgets[idx]: w.hide() self.settings.setValue("dock_status", self.dockStatus) # self.display_dockwidget[index] = bool(self.display_dockwidget[index] - 1) # self.toggleDockWidgets() self.saveLayout() # def toggleDockWidgets(self, is_init=False): # if is_init == True: # if self.dockStatus != []: # if len(self.dockStatus) != len(self.menus.showMenu): # self.settings.remove("dock_status") # else: # self.display_dockwidget = [strtobool(w) for w in self.dockStatus] # for i in range(len(self.menus.showMenu)): # self.menus.showMenu[i].setChecked(bool(self.display_dockwidget[i])) # else: # self.settings.setValue("dock_status", self.display_dockwidget) # for t, w in zip(self.display_dockwidget, self.dockWidgets.values()): # if t == True: # w.show() # else: # w.hide() def rsBandSet(self, idx): if self.raster is None: return for i in range(len(self.bandCombos)): self.rsRGB[i] = self.bandCombos[i].currentIndex() + 1 # 从1开始 self.raster.setBand(self.rsRGB) if self.grid is not None: if isinstance(self.grid.curr_idx, (list, tuple)): row, col = self.grid.curr_idx image, _ = self.raster.getGrid(row, col) else: image, _ = self.raster.getArray() else: image, _ = self.raster.getArray() self.image = image self.controller.image = image self.updateImage() # def miSlideSet(self): # image = rs.slice_img(self.controller.rawImage, self.midx) # self.test_show(image) # def changeWorkerShow(self, index): # self.display_dockwidget[index] = bool(self.display_dockwidget[index] - 1) # self.toggleDockWidgets() def updateBandList(self, clean=False): if clean: for i in range(len(self.bandCombos)): try: # 避免打开jpg后再打开tif报错 self.bandCombos[i].currentIndexChanged.disconnect() except TypeError: pass self.bandCombos[i].clear() self.bandCombos[i].addItems(["band_1"]) return bands = self.raster.geoinfo.count for i in range(len(self.bandCombos)): try: # 避免打开jpg后再打开tif报错 self.bandCombos[i].currentIndexChanged.disconnect() except TypeError: pass self.bandCombos[i].clear() self.bandCombos[i].addItems( [("band_" + str(j + 1)) for j in range(bands)]) try: self.bandCombos[i].setCurrentIndex(self.rsRGB[i] - 1) except IndexError: pass for bandCombo in self.bandCombos: bandCombo.currentIndexChanged.connect(self.rsBandSet) # 设置波段 # def updateSlideSld(self, clean=False): # if clean: # self.sldMISlide.setMaximum(1) # return # C = self.controller.rawImage.shape[-1] if len(self.controller.rawImage.shape) == 3 else 1 # self.sldMISlide.setMaximum(C) def toggleLargestCC(self, on): try: self.controller.filterLargestCC(on) except: pass # 宫格标注 def initGrid(self): self.delAllPolygon() grid_row_count, grid_col_count = self.grid.createGrids() self.gridTable.setRowCount(grid_row_count) self.gridTable.setColumnCount(grid_col_count) for r in range(grid_row_count): for c in range(grid_col_count): self.gridTable.setItem(r, c, QtWidgets.QTableWidgetItem()) self.gridTable.item(r, c).setBackground(self.GRID_COLOR["idle"]) self.gridTable.item(r, c).setFlags( Qt.ItemIsSelectable) # 无法高亮选择 # 初始显示第一个 self.grid.curr_idx = (0, 0) self.gridTable.item(0, 0).setBackground(self.GRID_COLOR["overlying"]) # 事件注册 self.gridTable.cellClicked.connect(self.changeGrid) # load polygon if self.outputDir is not None: name = osp.splitext(osp.basename(self.imagePath))[0] json_path = osp.join(self.outputDir, name + "_grid_saved.json") if osp.exists(json_path): self.grid.json_labels = json.loads(open(json_path, "r").read()) # load label for jlab in self.grid.json_labels: is_add = True for label in self.controller.labelList.labelList: if jlab["labelIdx"] == label.idx and jlab[ "name"] == label.name: is_add = False break if is_add is True: self.addLabel(jlab["labelIdx"], jlab["name"], jlab["color"]) self.changeGrid(0, 0) # load mask for jlab in self.grid.json_labels: pts = np.int32([np.array(jlab["points"])]) cv2.fillPoly( self.grid.mask_grids[jlab["row"]][jlab["col"]], pts=pts, color=jlab["labelIdx"]) def changeGrid(self, row, col): def find_in_json(r, c, json_labels): idxs = [] for idx, json_label in enumerate(json_labels): if json_label["row"] == r and json_label["col"] == c: idxs.append(idx) return idxs # 清除未保存的切换 self.finishObject() # TODO: 这块应该通过dirty判断? if self.grid.curr_idx is not None: self.saveGrid() # 切换时自动保存上一块 last_r, last_c = self.grid.curr_idx if self.grid.mask_grids[last_r][last_c] is None: self.gridTable.item( last_r, last_c).setBackground(self.GRID_COLOR["idle"]) else: self.gridTable.item( last_r, last_c).setBackground(self.GRID_COLOR["finised"]) self.delAllPolygon() image, _ = self.grid.getGrid(row, col) self.controller.setImage(image) self.grid.curr_idx = (row, col) idxs = find_in_json(row, col, self.grid.json_labels) if len(idxs) != 0: # 加载之前的标注 self.gridTable.item(row, col).setBackground(self.GRID_COLOR["overlying"]) for idx in idxs: label = self.grid.json_labels[idx] color = label["color"] labelIdx = label["labelIdx"] points = label["points"] poly = PolygonAnnotation( labelIdx, self.controller.image.shape, self.delPolygon, self.setDirty, color, color, self.opacity, ) self.scene.addItem(poly) self.scene.polygon_items.append(poly) for p in points: poly.addPointLast(QtCore.QPointF(p[0], p[1])) [self.grid.json_labels.remove(celement) for \ celement in [self.grid.json_labels[i] for i in idxs]] else: self.gridTable.item(row, col).setBackground(self.GRID_COLOR["current"]) # 刷新 self.updateImage(True) def saveGrid(self): row, col = self.grid.curr_idx if self.grid.curr_idx is None: return self.gridTable.item(row, col).setBackground(self.GRID_COLOR["overlying"]) # if len(np.unique(self.grid.mask_grids[row][col])) == 1: self.grid.mask_grids[row][col] = np.array(self.getMask()) # save grid label to load polygons = self.scene.polygon_items for polygon in polygons: l = self.controller.labelList[polygon.labelIndex - 1] label = { "row": row, "col": col, "name": l.name, "labelIdx": l.idx, "color": l.color, "points": [], } for p in polygon.scnenePoints: label["points"].append(p) self.grid.json_labels.append(label) # save every blocks or not if self.cheSaveEvery.isChecked(): _, fullflname = osp.split(self.listFiles.currentItem().text()) fname, _ = os.path.splitext(fullflname) if self.outputDir is None: if self.changeOutputDir() is False: self.cheSaveEvery.setChecked(False) return save_ima_path = osp.join( self.outputDir, (fname + "_data_" + str(row) + "_" + str(col) + ".tif")) save_lab_path = osp.join( self.outputDir, (fname + "_mask_" + str(row) + "_" + str(col) + ".tif")) im, tf = self.raster.getGrid(row, col) h, w = im.shape[:2] geoinfo = edict() geoinfo.xsize = w geoinfo.ysize = h geoinfo.dtype = self.raster.geoinfo.dtype geoinfo.crs = self.raster.geoinfo.crs geoinfo.geotf = tf self.raster.saveMask(self.grid.mask_grids[row][col], save_lab_path, geoinfo) # 保存mask self.raster.saveMask(im, save_ima_path, geoinfo, 3) # 保存图像 def turnGrid(self, delta): # 切换下一个宫格 r, c = self.grid.curr_idx if self.grid.curr_idx is not None else (0, -1) c += delta if c >= self.grid.grid_count[1]: c = 0 r += 1 if r >= self.grid.grid_count[0]: r = 0 if c < 0: c = self.grid.grid_count[1] - 1 r -= 1 if r < 0: r = self.grid.grid_count[0] - 1 self.changeGrid(r, c) def closeGrid(self): self.grid = None self.gridTable.setRowCount(0) self.gridTable.clearContents() def saveGridLabel(self): if self.grid is None: return if self.outputDir is not None: name, ext = osp.splitext(osp.basename(self.imagePath)) if not self.origExt: ext = ".png" save_path = osp.join(self.outputDir, name + ext) else: save_path = self.chooseSavePath() if save_path == "": return try: self.finishObject() self.saveGrid() # 先保存当前 except: pass self.delAllPolygon() # 清理 mask = self.grid.splicingList(save_path) json_path = save_path.replace(".png", "_grid_saved.json") open( json_path, "w", encoding="utf-8").write(json.dumps(self.grid.json_labels)) if self.grid.__class__.__name__ == "RSGrids": self.image, geo_tf = self.raster.getArray() if geo_tf is None: self.statusbar.showMessage(self.tr("图像过大，已显示缩略图")) else: self.image = self.grid.detimg self.controller.image = self.image self.controller._result_mask = mask self.exportLabel(savePath=save_path, lab_input=mask) # -- RS Show polygon demo -- if self.show_rs_poly.isChecked(): h, w = self.image.shape[:2] th_mask = cv2.resize( mask, dsize=(w, h), interpolation=cv2.INTER_NEAREST) indexs = np.unique(th_mask)[1:] for i in indexs: i_mask = np.zeros_like(th_mask, dtype="uint8") i_mask[th_mask == i] = 255 curr_polygon = util.get_polygon(i_mask) color = self.controller.labelList[i - 1].color self.createPoly(curr_polygon, color) for p in self.scene.polygon_items: p.setAnning(isAnning=False) # -- RS Show polygon demo -- # 刷新 grid_row_count = self.gridTable.rowCount() grid_col_count = self.gridTable.colorCount() for r in range(grid_row_count): for c in range(grid_col_count): try: self.gridTable.item( r, c).setBackground(self.GRID_COLOR["idle"]) except: pass self.raster = None self.closeGrid() self.updateBandList(True) self.controller.setImage(self.image) self.updateImage(True) self.setDirty(False) @property def opacity(self): return self.sldOpacity.value() / 100 @property def clickRadius(self): return self.sldClickRadius.value() @property def segThresh(self): return self.sldThresh.value() / 100 # @property # def slideMi(self): # return self.sldMISlide.value() def warnException(self, e): e = str(e) title = e.split("。")[0] self.warn(title, e) def warn(self, title, text, buttons=QMessageBox.Yes): msg = QMessageBox() # msg.setIcon(QMessageBox.Warning) msg.setWindowTitle(title) msg.setText(text) msg.setStandardButtons(buttons) return msg.exec_() @property def status(self): # TODO: 图片，模型 if not self.controller: return self.IDILE c = self.controller if c.model is None or c.image is None: return self.IDILE if self._anning: return self.ANNING return self.EDITING @status.setter def status(self, status): if status not in [self.ANNING, self.EDITING]: return if status == self.ANNING: self._anning = True else: self._anning = False def loadGrid(self, img, is_rs=True): res = self.warn(self.tr("图像过大"), self.tr("图像过大，将启用宫格功能！"), \ buttons=QMessageBox.Yes | QMessageBox.No) if res == QMessageBox.Yes: # 打开宫格功能 if self.dockWidgets["grid"][0].isVisible() is False: # TODO: 改成self.dockStatus self.menus.showMenu[-1].setChecked(True) # self.display_dockwidget[-1] = True self.dockWidgets["grid"][0].show() self.grid = RSGrids(img) if is_rs else Grids(img) self.initGrid() return True return False # 界面布局 def loadLayout(self): self.restoreState(self.layoutStatus) # TODO: 这里检查环境，判断是不是开医疗和遥感widget def saveLayout(self): # 保存界面 self.settings.setValue("layout_status", QByteArray(self.saveState())) self.settings.setValue( "save_status", [(k, self.save_status[k]) for k in self.save_status.keys()]) # # 如果设置了保存路径，把标签也保存下 # if self.outputDir is not None and len(self.controller.labelList) != 0: # self.exportLabelList(osp.join(self.outputDir, "autosave_label.txt")) def closeEvent(self, event): self.saveImage() self.saveLayout() QCoreApplication.quit() # sys.exit(0) def reportBug(self): webbrowser.open("https://github.com/PaddlePaddle/PaddleSeg/issues") def enterEISegMed3D(self): webbrowser.open( "https://github.com/PaddlePaddle/PaddleSeg/tree/develop/EISeg/med3d") def quickStart(self): # self.saveImage(True) # self.canvas.setStyleSheet(self.note_style) webbrowser.open( "https://github.com/PaddlePaddle/PaddleSeg/tree/release/2.6/EISeg") def toggleLogging(self, s): if s: logger.setLevel(logging.DEBUG) else: logger.setLevel(logging.CRITICAL) self.settings.setValue("log", s) def toBeImplemented(self): self.statusbar.showMessage(self.tr("功能尚在开发")) # 医疗 def wwChanged(self): if not self.controller or self.image is None: return try: # 那种jpg什么格式的医疗图像调整窗宽等会造成崩溃 self.textWw.selectAll() self.controller.image = med.windowlize(self.controller.rawImage, self.ww, self.wc) self.updateImage() except: pass def wcChanged(self): if not self.controller or self.image is None: return try: self.textWc.selectAll() self.controller.image = med.windowlize(self.controller.rawImage, self.ww, self.wc) self.updateImage() except: pass @property def ww(self): return int(self.textWw.text()) @property def wc(self): return int(self.textWc.text()) def twwChanged(self): if self.ww > self.sldWw.maximum(): self.textWw.setText(str(self.sldWw.maximum())) if self.ww < self.sldWw.minimum(): self.textWw.setText(str(self.sldWw.minimum())) self.sldWw.setProperty("value", self.ww) self.wwChanged() def swwChanged(self): self.textWw.setText(str(self.sldWw.value())) self.wwChanged() def twcChanged(self): if self.wc > self.sldWc.maximum(): self.textWc.setText(str(self.sldWc.maximum())) if self.wc < self.sldWc.minimum(): self.textWc.setText(str(self.sldWc.minimum())) self.sldWc.setProperty("value", self.wc) self.wcChanged() def swcChanged(self): self.textWc.setText(str(self.sldWc.value())) self.wcChanged() # 视频 def tframeChanged(self): if self.video_images is None: return if self.video.cursur > self.sldTime.maximum(): self.textTime.setText(str(self.sldTime.maximum())) if self.video.cursur < self.sldTime.minimum(): self.textTime.setText(str(self.sldTime.minimum())) self.sldTime.setProperty("value", int(self.textTime.text())) def sframeChanged(self): if self.video_images is None: return self.textTime.setText(str(self.sldTime.value())) self.video.cursur = int(self.textTime.text()) self.controller.setImage(self.video_images[self.video.cursur]) self.delAllPolygon() self.show_current_frame() # print('current_frame:',self.video.cursur) def turnPreFrame(self): if self.video_images is None: return self.video.cursur -= 1 if self.video.cursur < 0: self.video.cursur = self.video.num_frames - 1 self.sldTime.setProperty("value", self.video.cursur) def turnNextFrame(self): if self.video_images is None: return self.video.cursur += 1 if self.video.cursur > self.video.num_frames - 1: self.video.cursur = 0 self.sldTime.setProperty("value", self.video.cursur) def show_current_frame(self): self.viz = overlay_davis(self.video_images[self.video.cursur], self.video_masks[self.video.cursur], self.opacity, self.controller.palette) self.update_interact_viz() self.sldTime.setProperty("value", self.video.cursur) def brushChanged(self): self.textBrush.setText(str(self.sldBrush.value())) def on_time(self): self.video.cursur += 1 if self.video.cursur > self.video.num_frames - 1: self.video.cursur = 0 self.sldTime.setProperty("value", self.video.cursur) def on_play(self): if self.video_images is None: self.warn(self.tr("图片格式无法播放"), self.tr("请先加载视频")) return if self.timer.isActive(): self.timer.stop() self.videoPlay.setText(self.tr("播放")) self.videoPlay.setIcon( QtGui.QIcon(osp.join(pjpath, "resource/Play.png"))) else: # self.delAllPolygon() self.timer.start(1000 // self.ratio) self.videoPlay.setText(self.tr("暂停")) self.videoPlay.setIcon( QtGui.QIcon(osp.join(pjpath, "resource/Stop.png"))) def getVideoMask(self): if self.video_masks is not None: return self.video_masks[self.video.cursur] else: return None def on_propgation(self): self.finishObject() if self.video_images is None: self.warn(self.tr("未加载视频"), self.tr("请先在加载图像按钮中加载视频")) return if self.video.prop_net_segm is None: self.warn(self.tr("传播模型未加载"), self.tr("尚未加载视频传播模型，请先加载模型!")) return if self.video.fuse_net is None: self.warn(self.tr("融合模型未加载"), self.tr("尚未加载视频融合模型，请先加载模型!")) return current_mask = self.getMask() if current_mask is None: self.warn(self.tr("未提供传播参考帧"), self.tr("请先在标注传播参考帧再进行传播")) return if current_mask.max() == 0: current_mask = self.video_masks[self.video.cursur] # self.warn(self.tr("未新增标注"), self.tr("请先添加新标注再进行传播")) # return print('-------------start propgation----------------') self.statusbar.showMessage(self.tr("开始传播")) # set object self.video.set_objects(int(max(self.video.k, current_mask.max()))) self.video.set_images(self.video_images) one_hot_mask = F.one_hot( paddle.to_tensor(current_mask).astype('int32'), int(self.video.k + 1)) self.one_hot_mask = one_hot_mask.transpose([2, 0, 1]).unsqueeze(1) start = time.time() self.video_masks = self.video.interact( self.one_hot_mask, self.video.cursur, self.progress_total_cb, self.progress_step_cb) end = time.time() print("propagation time cost", end - start) self.statusbar.showMessage(self.tr("传播完成!"), 5000) # 传播进度条重置 self.proPropagete.setValue(0) self.proPropagete.setFormat('0%') self.delAllPolygon() self.show_current_frame() # 3d显示 color_map = [] for lab in self.controller.labelList: color_map.append(lab.color) if self.TDDock.isVisible(): self.vtkWidget.show_array( np.uint8(self.video_masks), (1., 1., 1.), color_map) def progress_step_cb(self): self.progress_num += 1 ratio = self.progress_num / self.progress_max self.proPropagete.setValue(int(ratio * 100)) self.proPropagete.setFormat('%2.1f%%' % (ratio * 100)) QApplication.processEvents() def progress_total_cb(self, total): self.progress_max = total self.progress_num = -1 self.progress_step_cb() def useQtWidget(self, s): self.settings.setValue("use_qt_widget", s) def checkLabel(self, labelIndex): for p in self.scene.polygon_items: if p.labelIndex == labelIndex: return False return True