add new model

EISeg/eiseg/util/regularization/rs_regularization.py

+# 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.
+"""
+This code is based on https://github.com/niecongchong/RS-building-regularization
+Ths copyright of niecongchong/RS-building-regularization is as follows:
+Apache License [see LICENSE for details]
+"""
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+from .rdp_alg import rdp
+from .cal_point import cal_ang, cal_dist, cal_azimuth
+from .rotate_ang import Nrotation_angle_get_coor_coordinates, Srotation_angle_get_coor_coordinates
+from .cal_line import line, intersection, par_line_dist, point_in_line
+
+
+def boundary_regularization(contours, img_shape, epsilon=6):
+    h, w = img_shape[0:2]
+    # 轮廓定位
+    contours = np.squeeze(contours)
+    # 轮廓精简DP
+    contours = rdp(contours, epsilon=epsilon)
+    contours[:, 1] = h - contours[:, 1]
+    # 轮廓规则化
+    dists = []
+    azis = []
+    azis_index = []
+    # 获取每条边的长度和方位角
+    for i in range(contours.shape[0]):
+        cur_index = i
+        next_index = i + 1 if i < contours.shape[0] - 1 else 0
+        prev_index = i - 1
+        cur_point = contours[cur_index]
+        nest_point = contours[next_index]
+        prev_point = contours[prev_index]
+        dist = cal_dist(cur_point, nest_point)
+        azi = cal_azimuth(cur_point, nest_point)
+        dists.append(dist)
+        azis.append(azi)
+        azis_index.append([cur_index, next_index])
+    # 以最长的边的方向作为主方向
+    longest_edge_idex = np.argmax(dists)
+    main_direction = azis[longest_edge_idex]
+    # 方向纠正，绕中心点旋转到与主方向垂直或者平行
+    correct_points = []
+    para_vetr_idxs = []  # 0平行 1垂直
+    for i, (azi, (point_0_index,
+                  point_1_index)) in enumerate(zip(azis, azis_index)):
+        if i == longest_edge_idex:
+            correct_points.append(
+                [contours[point_0_index], contours[point_1_index]])
+            para_vetr_idxs.append(0)
+        else:
+            # 确定旋转角度
+            rotate_ang = main_direction - azi
+            if np.abs(rotate_ang) < 180 / 4:
+                rotate_ang = rotate_ang
+                para_vetr_idxs.append(0)
+            elif np.abs(rotate_ang) >= 90 - 180 / 4:
+                rotate_ang = rotate_ang + 90
+                para_vetr_idxs.append(1)
+            # 执行旋转任务
+            point_0 = contours[point_0_index]
+            point_1 = contours[point_1_index]
+            point_middle = (point_0 + point_1) / 2
+            if rotate_ang > 0:
+                rotate_point_0 = Srotation_angle_get_coor_coordinates(
+                    point_0, point_middle, np.abs(rotate_ang))
+                rotate_point_1 = Srotation_angle_get_coor_coordinates(
+                    point_1, point_middle, np.abs(rotate_ang))
+            elif rotate_ang < 0:
+                rotate_point_0 = Nrotation_angle_get_coor_coordinates(
+                    point_0, point_middle, np.abs(rotate_ang))
+                rotate_point_1 = Nrotation_angle_get_coor_coordinates(
+                    point_1, point_middle, np.abs(rotate_ang))
+            else:
+                rotate_point_0 = point_0
+                rotate_point_1 = point_1
+            correct_points.append([rotate_point_0, rotate_point_1])
+    correct_points = np.array(correct_points)
+    # 相邻边校正，垂直取交点，平行平移短边或者加线
+    final_points = []
+    final_points.append(correct_points[0][0])
+    for i in range(correct_points.shape[0] - 1):
+        cur_index = i
+        next_index = i + 1 if i < correct_points.shape[0] - 1 else 0
+        cur_edge_point_0 = correct_points[cur_index][0]
+        cur_edge_point_1 = correct_points[cur_index][1]
+        next_edge_point_0 = correct_points[next_index][0]
+        next_edge_point_1 = correct_points[next_index][1]
+        cur_para_vetr_idx = para_vetr_idxs[cur_index]
+        next_para_vetr_idx = para_vetr_idxs[next_index]
+        if cur_para_vetr_idx != next_para_vetr_idx:
+            # 垂直取交点
+            L1 = line(cur_edge_point_0, cur_edge_point_1)
+            L2 = line(next_edge_point_0, next_edge_point_1)
+            point_intersection = intersection(L1, L2)
+            final_points.append(point_intersection)
+        elif cur_para_vetr_idx == next_para_vetr_idx:
+            # 平行分两种，一种加短线，一种平移，取决于距离阈值
+            L1 = line(cur_edge_point_0, cur_edge_point_1)
+            L2 = line(next_edge_point_0, next_edge_point_1)
+            marg = par_line_dist(L1, L2)
+            if marg < 3:
+                # 平移
+                point_move = point_in_line(
+                    next_edge_point_0[0], next_edge_point_0[1],
+                    cur_edge_point_0[0], cur_edge_point_0[1],
+                    cur_edge_point_1[0], cur_edge_point_1[1])
+                final_points.append(point_move)
+                # 更新平移之后的下一条边
+                correct_points[next_index][0] = point_move
+                correct_points[next_index][1] = point_in_line(
+                    next_edge_point_1[0], next_edge_point_1[1],
+                    cur_edge_point_0[0], cur_edge_point_0[1],
+                    cur_edge_point_1[0], cur_edge_point_1[1])
+            else:
+                # 加线
+                add_mid_point = (cur_edge_point_1 + next_edge_point_0) / 2
+                add_point_1 = point_in_line(
+                    add_mid_point[0], add_mid_point[1], cur_edge_point_0[0],
+                    cur_edge_point_0[1], cur_edge_point_1[0],
+                    cur_edge_point_1[1])
+                add_point_2 = point_in_line(
+                    add_mid_point[0], add_mid_point[1], next_edge_point_0[0],
+                    next_edge_point_0[1], next_edge_point_1[0],
+                    next_edge_point_1[1])
+                final_points.append(add_point_1)
+                final_points.append(add_point_2)
+    final_points.append(final_points[0])
+    final_points = np.array(final_points)
+    final_points[:, 1] = h - final_points[:, 1]
+    final_points = final_points[np.newaxis, :].transpose((1, 0, 2))
+    return final_points
+
+
+# def rs_build_re(mask):
+#     # 中值滤波，去噪
+#     ori_img = cv2.medianBlur(mask, 5)
+#     ori_img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2GRAY)
+#     ret, ori_img = cv2.threshold(ori_img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
+#     # 连通域分析
+#     num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(ori_img, connectivity=8)
+#     # 遍历联通域
+#     for i in range(1, num_labels):
+#         img = np.zeros_like(labels)
+#         index = np.where(labels==i)
+#         img[index] = 255
+#         img = np.array(img, dtype=np.uint8)
+#         regularization_contour = boundary_regularization(img).astype(np.int32)
+#         cv2.polylines(img=mask, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)
+#         single_out = np.zeros_like(mask)
+#         cv2.polylines(img=single_out, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)
+#         cv2.imwrite('single_out_{}.jpg'.format(i), single_out)

EISeg/eiseg/util/serialization.py

+from functools import wraps
+from copy import deepcopy
+import inspect
+import paddle.nn as nn
+
+
+def serialize(init):
+    parameters = list(inspect.signature(init).parameters)
+
+    @wraps(init)
+    def new_init(self, *args, **kwargs):
+        params = deepcopy(kwargs)
+        for pname, value in zip(parameters[1:], args):
+            params[pname] = value
+
+        config = {"class": get_classname(self.__class__), "params": dict()}
+        specified_params = set(params.keys())
+
+        for pname, param in get_default_params(self.__class__).items():
+            if pname not in params:
+                params[pname] = param.default
+
+        for name, value in list(params.items()):
+            param_type = "builtin"
+            if inspect.isclass(value):
+                param_type = "class"
+                value = get_classname(value)
+
+            config["params"][name] = {
+                "type": param_type,
+                "value": value,
+                "specified": name in specified_params,
+            }
+
+        setattr(self, "_config", config)
+        init(self, *args, **kwargs)
+
+    return new_init
+
+
+def load_model(config, **kwargs):
+    model_class = get_class_from_str(config["class"])
+    model_default_params = get_default_params(model_class)
+
+    model_args = dict()
+    for pname, param in config["params"].items():
+        value = param["value"]
+        if param["type"] == "class":
+            value = get_class_from_str(value)
+
+        if pname not in model_default_params and not param["specified"]:
+            continue
+
+        assert pname in model_default_params
+        if not param["specified"] and model_default_params[
+                pname].default == value:
+            continue
+        model_args[pname] = value
+
+    model_args.update(kwargs)
+
+    return model_class(**model_args)
+
+
+def get_config_repr(config):
+    config_str = f'Model: {config["class"]}\n'
+    for pname, param in config["params"].items():
+        value = param["value"]
+        if param["type"] == "class":
+            value = value.split(".")[-1]
+        param_str = f"{pname:<22} = {str(value):<12}"
+        if not param["specified"]:
+            param_str += " (default)"
+        config_str += param_str + "\n"
+    return config_str
+
+
+def get_default_params(some_class):
+    params = dict()
+    for mclass in some_class.mro():
+        if mclass is nn.Layer or mclass is object:
+            continue
+
+        mclass_params = inspect.signature(mclass.__init__).parameters
+        for pname, param in mclass_params.items():
+            if param.default != param.empty and pname not in params:
+                params[pname] = param
+
+    return params
+
+
+def get_classname(cls):
+    module = cls.__module__
+    name = cls.__qualname__
+    if module is not None and module != "__builtin__":
+        name = module + "." + name
+    return name
+
+
+def get_class_from_str(class_str):
+    components = class_str.split(".")
+    mod = __import__(".".join(components[:-1]))
+    for comp in components[1:]:
+        mod = getattr(mod, comp)
+    return mod

EISeg/eiseg/util/vis.py

+from functools import lru_cache
+
+import cv2
+import numpy as np
+
+
+def visualize_instances(imask,
+                        bg_color=255,
+                        boundaries_color=None,
+                        boundaries_width=1,
+                        boundaries_alpha=0.8):
+    num_objects = imask.max() + 1
+    palette = get_palette(num_objects)
+    if bg_color is not None:
+        palette[0] = bg_color
+
+    result = palette[imask].astype(np.uint8)
+    if boundaries_color is not None:
+        boundaries_mask = get_boundaries(
+            imask, boundaries_width=boundaries_width)
+        tresult = result.astype(np.float32)
+        tresult[boundaries_mask] = boundaries_color
+        tresult = tresult * boundaries_alpha + (1 - boundaries_alpha) * result
+        result = tresult.astype(np.uint8)
+
+    return result
+
+
+@lru_cache(maxsize=16)
+def get_palette(num_cls):
+    return np.array([[0, 0, 0], [128, 0, 0], [0, 128, 0], [0, 0, 128],
+                     [128, 128, 0], [0, 128, 128], [128, 0, 128]])
+
+
+def visualize_mask(mask, num_cls):
+    palette = get_palette(num_cls)
+    mask[mask == -1] = 0
+
+    return palette[mask].astype(np.uint8)
+
+
+def visualize_proposals(proposals_info, point_color=(255, 0, 0),
+                        point_radius=1):
+    proposal_map, colors, candidates = proposals_info
+
+    proposal_map = draw_probmap(proposal_map)
+    for x, y in candidates:
+        proposal_map = cv2.circle(proposal_map, (y, x), point_radius,
+                                  point_color, -1)
+
+    return proposal_map
+
+
+def draw_probmap(x):
+    return cv2.applyColorMap((x * 255).astype(np.uint8), cv2.COLORMAP_HOT)
+
+
+def draw_points(image, points, color, radius=3):
+    image = image.copy()
+    for p in points:
+        image = cv2.circle(image, (int(p[1]), int(p[0])), radius, color, -1)
+
+    return image
+
+
+def draw_instance_map(x, palette=None):
+    num_colors = x.max() + 1
+    if palette is None:
+        palette = get_palette(num_colors)
+
+    return palette[x].astype(np.uint8)
+
+
+def blend_mask(image, mask, alpha=0.6):
+    if mask.min() == -1:
+        mask = mask.copy() + 1
+
+    imap = draw_instance_map(mask)
+    result = (image * (1 - alpha) + alpha * imap).astype(np.uint8)
+    return result
+
+
+def get_boundaries(instances_masks, boundaries_width=1):
+    boundaries = np.zeros(
+        (instances_masks.shape[0], instances_masks.shape[1]), dtype=np.bool)
+
+    for obj_id in np.unique(instances_masks.flatten()):
+        if obj_id == 0:
+            continue
+
+        obj_mask = instances_masks == obj_id
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        inner_mask = cv2.erode(
+            obj_mask.astype(np.uint8), kernel,
+            iterations=boundaries_width).astype(np.bool)
+
+        obj_boundary = np.logical_xor(obj_mask,
+                                      np.logical_and(inner_mask, obj_mask))
+        boundaries = np.logical_or(boundaries, obj_boundary)
+    return boundaries
+
+
+def draw_with_blend_and_clicks(
+        img,
+        mask=None,
+        alpha=0.6,
+        clicks_list=None,
+        pos_color=(0, 255, 0),
+        neg_color=(255, 0, 0),
+        radius=4,
+        palette=None, ):
+    result = img.copy()
+
+    if mask is not None:
+        if not palette:
+            palette = get_palette(np.max(mask) + 1)
+        palette = np.array(palette)
+        rgb_mask = palette[mask.astype(np.uint8)]
+
+        mask_region = (mask > 0).astype(np.uint8)
+        result = (result * (1 - mask_region[:, :, np.newaxis]) + (1 - alpha) *
+                  mask_region[:, :, np.newaxis] * result + alpha * rgb_mask)
+        result = result.astype(np.uint8)
+
+    if clicks_list is not None and len(clicks_list) > 0:
+        pos_points = [
+            click.coords for click in clicks_list if click.is_positive
+        ]
+        neg_points = [
+            click.coords for click in clicks_list if not click.is_positive
+        ]
+
+        result = draw_points(result, pos_points, pos_color, radius=radius)
+        result = draw_points(result, neg_points, neg_color, radius=radius)
+
+    return result

EISeg/eiseg/widget/__init__.py

+from .shortcut import ShortcutWidget
+from .loading import LoadingWidget
+from .line import LineItem
+from .grip import GripItem
+from .bbox import BBoxAnnotation
+from .polygon import PolygonAnnotation
+from .scene import AnnotationScene
+from .view import AnnotationView
+from .create import (create_text, create_button, create_slider, DockWidget,
+                     creat_dock)
+from .table import TableWidget

EISeg/eiseg/widget/bbox.py

+# 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.
+
+from qtpy import QtWidgets, QtGui, QtCore
+from qtpy.QtCore import Qt
+
+
+# Note, bbox annotation is more convenient than the default boundingBox generated by QGrpaphicItem
+class BBoxAnnotation(QtWidgets.QGraphicsPathItem):
+    def __init__(
+            self,
+            labelIndex,
+            polyline,
+            borderColor=[0, 0, 255],
+            cocoIndex=None,
+            parent=None, ):
+        super(BBoxAnnotation, self).__init__(parent)
+        self.polyline = polyline
+        self.corner_points = []
+        self.upper_right = QtCore.QPointF()
+        self.bottom_left = QtCore.QPointF()
+        self.w = -1.0
+        self.h = -1.0
+
+        self.parent = parent
+        self.is_added = False
+        if self.parent is not None:
+            self.is_added = True
+        self.labelIndex = labelIndex
+        self.coco_id = cocoIndex
+        self.bbox_hovering = True
+
+        # set rendering attributes
+        self.setZValue(10)
+
+        # b = borderColor
+        # self.borderColor = QtGui.QColor(b[0], b[1], b[2])
+        self.borderColor = QtGui.QColor(128, 128, 128)
+        self.borderColor.setAlphaF(0.8)
+        pen = QtGui.QPen(self.borderColor, 1.2)
+        pen.setStyle(Qt.DashDotLine)
+        self.setPen(pen)
+
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable, False)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable, False)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemSendsGeometryChanges, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsFocusable, False)
+        self.setAcceptHoverEvents(False)
+        # self.setCursor(QtGui.QCursor(QtCore.Qt.PointingHandCursor))
+
+    @property
+    def scnenePoints(self):
+        # return 4 corner points
+        raise Exception("Not Implemented Yet!")
+
+    def setAnning(self, isAnning=True):
+        raise Exception("Not Implemented Yet!")
+
+    def remove(self):
+        raise Exception("Not Implemented Yet!")
+
+    # ===== generate geometry info
+
+    def create_corners(self):
+        bbox_rect_geo = self.polyline.boundingRect()
+        self.bottom_left = bbox_rect_geo.bottomLeft()
+        self.upper_right = bbox_rect_geo.topRight()
+        self.corner_points.clear()
+        self.corner_points.extend([
+            self.bottom_left,
+            bbox_rect_geo.topLeft(),
+            self.upper_right,
+            bbox_rect_geo.bottomRight(),
+        ])
+        self.w = self.corner_points[3].x() - self.corner_points[1].x()
+        self.h = self.corner_points[3].y() - self.corner_points[1].y()
+
+        if self.corner_points[1].x() > 512 or self.corner_points[1].x(
+        ) + self.w > 512:
+            pass
+        if self.corner_points[1].y() > 512 or self.corner_points[1].y(
+        ) + self.h > 512:
+            pass
+        return self.corner_points
+
+    def create_lines(self):
+        pass
+
+    # ===== graphic interface to update in scene tree
+
+    def update(self):
+        l = len(self.polyline.points)
+        # print("up L:", l, " is_added:", self.is_added)
+        if l < 3:
+            if self.is_added:
+                self.remove_from_scene()
+        else:  # 大于三个点就可以更新，小于三个点删除多边形
+            if self.is_added:
+                self.add_to_scene()
+            else:
+                path_geo = QtGui.QPainterPath()
+                self.create_corners()
+                path_geo.moveTo(self.corner_points[0])
+                for i in range(4):
+                    path_geo.lineTo(self.corner_points[(i + 1) % 4])
+                self.setPath(QtGui.QPainterPath(path_geo))
+                pass
+            pass
+        pass
+
+    def add_to_scene(self):
+        # self.parentItem().scene().addItem(self)
+        self.setParentItem(self.parent)
+        self.is_added = True
+
+    def remove_from_scene(self):
+        # self.parentItem().scene().removeItem(self)
+        self.setParentItem(None)
+        self.is_added = False
+
+    # ===== annotation info
+
+    # @return : [x, y, w, h]
+    def to_array(self):
+        np_array = [
+            self._round(self.corner_points[1].x()),
+            self._round(self.corner_points[1].y()),  # topLeft
+            self._round(self.w),
+            self._round(self.h),
+        ]
+        return np_array
+
+    def _round(self, number, ind=0):
+        nint, ndec = str(number).split(".")
+        res = float(nint + "." + ndec[:ind])
+        if res <= 0:
+            res = 0.0
+        return res
+
+    def __del__(self):
+        self.corner_points.clear()

EISeg/eiseg/widget/create.py

+# 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.
+
+from qtpy.QtWidgets import QDockWidget
+from qtpy import QtCore, QtGui, QtWidgets
+from qtpy.QtCore import Qt
+
+
+## 创建文本
+def create_text(parent, text_name=None, text_text=None):
+    text = QtWidgets.QLabel(parent)
+    if text_name is not None:
+        text.setObjectName(text_name)
+    if text_text is not None:
+        text.setText(text_text)
+    return text
+
+
+## 创建可编辑文本
+def create_edit(parent, text_name=None, text_text=None):
+    edit = QtWidgets.QLineEdit(parent)
+    if text_name is not None:
+        edit.setObjectName(text_name)
+    if text_text is not None:
+        edit.setText(text_text)
+    edit.setValidator(QtGui.QIntValidator())
+    edit.setMaxLength(5)
+    return edit
+
+
+## 创建按钮
+def create_button(parent, btn_name, btn_text, ico_path=None, curt=None):
+    # 创建和设置按钮
+    sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum,
+                                       QtWidgets.QSizePolicy.Fixed)
+    min_size = QtCore.QSize(0, 40)
+    sizePolicy.setHorizontalStretch(0)
+    sizePolicy.setVerticalStretch(0)
+    btn = QtWidgets.QPushButton(parent)
+    sizePolicy.setHeightForWidth(btn.sizePolicy().hasHeightForWidth())
+    btn.setSizePolicy(sizePolicy)
+    btn.setMinimumSize(min_size)
+    btn.setObjectName(btn_name)
+    if ico_path is not None:
+        btn.setIcon(QtGui.QIcon(ico_path))
+    btn.setText(btn_text)
+    if curt is not None:
+        btn.setShortcut(curt)
+    return btn
+
+
+## 创建滑块区域
+def create_slider(parent,
+                  sld_name,
+                  text_name,
+                  text,
+                  default_value=50,
+                  max_value=100,
+                  min_value=0,
+                  text_rate=0.01,
+                  edit=False):
+    Region = QtWidgets.QHBoxLayout()
+    lab = create_text(parent, None, text)
+    Region.addWidget(lab)
+    if edit is False:
+        labShow = create_text(parent, text_name, str(default_value * text_rate))
+    else:
+        labShow = create_edit(parent, text_name, str(default_value * text_rate))
+    labShow.setMaximumWidth(100)
+    Region.addWidget(labShow)
+    Region.addStretch()
+    sld = QtWidgets.QSlider(parent)
+    sld.setMaximum(max_value)  # 好像只能整数的，这里是扩大了10倍，1 . 10
+    sld.setMinimum(min_value)
+    sld.setProperty("value", default_value)
+    sld.setOrientation(QtCore.Qt.Horizontal)
+    sld.setObjectName(sld_name)
+    sld.setStyleSheet("""
+        QSlider::sub-page:horizontal {
+            background: #9999F1
+        }
+        QSlider::handle:horizontal
+        {
+            background: #3334E3;
+            width: 12px;
+            border-radius: 4px;
+        }
+        """)
+    sld.textLab = labShow
+    return sld, labShow, Region
+
+
+class DockWidget(QDockWidget):
+    def __init__(self, parent, name, text):
+        super().__init__(parent=parent)
+        self.setObjectName(name)
+        self.setAllowedAreas(Qt.RightDockWidgetArea | Qt.LeftDockWidgetArea)
+        # 感觉不给关闭好点。可以在显示里面取消显示。有关闭的话显示里面的enable还能判断修改，累了
+        self.setFeatures(QDockWidget.DockWidgetMovable |
+                         QDockWidget.DockWidgetFloatable)
+        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Preferred,
+                                           QtWidgets.QSizePolicy.Preferred)
+        sizePolicy.setHorizontalStretch(0)
+        sizePolicy.setVerticalStretch(0)
+        sizePolicy.setHeightForWidth(self.sizePolicy().hasHeightForWidth())
+        self.setSizePolicy(sizePolicy)
+        self.setMinimumWidth(230)
+        self.setWindowTitle(text)
+        self.setStyleSheet("QDockWidget { background-color:rgb(204,204,248); }")
+        self.topLevelChanged.connect(self.changeBackColor)
+
+    def changeBackColor(self, isFloating):
+        if isFloating:
+            self.setStyleSheet(
+                "QDockWidget { background-color:rgb(255,255,255); }")
+        else:
+            self.setStyleSheet(
+                "QDockWidget { background-color:rgb(204,204,248); }")
+
+
+## 创建dock
+def creat_dock(parent, name, text, widget):
+    dock = DockWidget(parent, name, text)
+    dock.setMinimumWidth(300)  # Uniform size
+    dock.setWidget(widget)
+    return dock

EISeg/eiseg/widget/grip.py

+# 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.
+
+from PyQt5.QtCore import QPointF
+from qtpy import QtWidgets, QtGui, QtCore
+
+
+class GripItem(QtWidgets.QGraphicsPathItem):
+    maxSize = 1.5
+    minSize = 0.8
+
+    def __init__(self, annotation_item, index, color, img_size):
+        super(GripItem, self).__init__()
+        self.m_annotation_item = annotation_item
+        self.hovering = False
+        self.m_index = index
+        self.anning = True
+        color.setAlphaF(1)
+        self.color = color
+        self.img_size = img_size
+
+        self.updateSize()
+        self.setPath(self.circle)
+        self.setBrush(self.color)
+        self.setPen(QtGui.QPen(self.color, 1))
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemSendsGeometryChanges, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsFocusable, True)
+        self.setAcceptHoverEvents(True)
+        self.setZValue(12)
+        self.setCursor(QtGui.QCursor(QtCore.Qt.PointingHandCursor))
+
+    def setColor(self, color):
+        self.setBrush(color)
+        self.setPen(QtGui.QPen(color, 1))
+        self.color = color
+
+    def setAnning(self, anning=True):
+        self.anning = anning
+        self.setEnabled(anning)
+
+    # BUG: Scaling causes a crash
+    @property
+    def size(self):
+        return GripItem.minSize
+        # if not self.scene():
+        #     return GripItem.minSize
+        # else:
+        #     maxi, mini = GripItem.maxSize, GripItem.minSize
+        #     exp = 1 - mini / maxi
+        #     size = maxi * (1 - exp ** self.scene().scale)
+        #     if size > GripItem.maxSize:
+        #         size = GripItem.maxSize
+        #     if size < GripItem.minSize:
+        #         size = GripItem.minSize
+        #     return size
+
+    def updateSize(self, s=2):
+        size = self.size
+        self.circle = QtGui.QPainterPath()
+        self.circle.addEllipse(QtCore.QRectF(-size, -size, size * s, size * s))
+        self.square = QtGui.QPainterPath()
+        self.square.addRect(QtCore.QRectF(-size, -size, size * s, size * s))
+        self.setPath(self.square if self.hovering else self.circle)
+
+    def hoverEnterEvent(self, ev):
+        self.setPath(self.square)
+        self.setBrush(QtGui.QColor(0, 0, 0, 0))
+        self.m_annotation_item.item_hovering = True
+        self.hovring = True
+        super(GripItem, self).hoverEnterEvent(ev)
+
+    def hoverLeaveEvent(self, ev):
+        self.setPath(self.circle)
+        self.setBrush(self.color)
+        self.m_annotation_item.item_hovering = False
+        self.hovring = False
+        super(GripItem, self).hoverLeaveEvent(ev)
+
+    def mouseReleaseEvent(self, ev):
+        self.setSelected(False)
+        super(GripItem, self).mouseReleaseEvent(ev)
+
+    def itemChange(self, change, value):
+        tmp_val = value
+        if change == QtWidgets.QGraphicsItem.ItemPositionChange and self.isEnabled(
+        ):
+            if value.x() > self.img_size[1]:
+                x = self.img_size[1]
+            elif value.x() < 0:
+                x = 0
+            else:
+                x = value.x()
+            if value.y() > self.img_size[0]:
+                y = self.img_size[0]
+            elif value.y() < 0:
+                y = 0
+            else:
+                y = value.y()
+            tmp_val = QPointF(x, y)
+            self.m_annotation_item.movePoint(self.m_index, tmp_val)
+            self.m_annotation_item.setDirty(True)
+        return super(GripItem, self).itemChange(change, tmp_val)
+
+    def shape(self):
+        path = QtGui.QPainterPath()
+        p = self.mapFromScene(self.pos())
+        x, y = p.x(), p.y()
+        s = self.size
+        path.addEllipse(p, s + GripItem.minSize, s + GripItem.minSize)
+        return path
+
+    def mouseDoubleClickEvent(self, ev):
+        self.m_annotation_item.removeFocusPoint()

EISeg/eiseg/widget/line.py

+# 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.
+
+from qtpy import QtWidgets, QtGui, QtCore
+
+
+class LineItem(QtWidgets.QGraphicsLineItem):
+    maxWidth = 1
+    minWidth = 0.5
+
+    def __init__(self, annotation_item, idx, color):
+        super(LineItem, self).__init__()
+        self.polygon_item = annotation_item
+        self.idx = idx
+        self.color = color
+        self.anning = True
+        self.setPen(QtGui.QPen(color, self.width))
+
+        self.setZValue(11)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsFocusable, True)
+        # self.setFlag(QtWidgets.QGraphicsItem.ItemClipsToShape, True)
+        self.setAcceptHoverEvents(True)
+        self.setBoundingRegionGranularity(0.5)
+        self.updateWidth()
+
+    def setColor(self, color):
+        self.setPen(QtGui.QPen(color, self.width))
+        self.color = color
+
+    def setAnning(self, anning=True):
+        self.anning = anning
+        self.setEnabled(anning)
+        self.updateWidth()
+
+    # BUG: Scaling causes a crash
+    @property
+    def width(self):
+        return LineItem.minWidth
+        # if not self.scene():
+        #     width = LineItem.minWidth
+        # else:
+        #     maxi, mini = LineItem.maxWidth, LineItem.minWidth
+        #     exp = 1 - mini / maxi
+        #     width = maxi * (1 - exp ** self.scene().scale)
+        #     if width > LineItem.maxWidth:
+        #         width = LineItem.maxWidth
+        #     if width < LineItem.minWidth:
+        #         width = LineItem.minWidth
+        # return width
+
+    def updateWidth(self):
+        self.setPen(QtGui.QPen(self.color, self.width))
+
+    def hoverEnterEvent(self, ev):
+        self.boundingPolygon(True)
+        print("hover in")
+        if self.anning:
+            self.polygon_item.line_hovering = True
+            self.setPen(QtGui.QPen(self.color, self.width * 1.4))
+        super(LineItem, self).hoverEnterEvent(ev)
+
+    def hoverLeaveEvent(self, ev):
+        self.polygon_item.line_hovering = False
+        self.setPen(QtGui.QPen(self.color, self.width))
+        super(LineItem, self).hoverLeaveEvent(ev)
+
+    def mouseDoubleClickEvent(self, ev):
+        print("anning", self.anning)
+        if self.anning:
+            self.setPen(QtGui.QPen(self.color, self.width))
+            self.polygon_item.addPointMiddle(self.idx, ev.pos())
+        super(LineItem, self).mouseDoubleClickEvent(ev)
+
+    def shape(self):
+        path = QtGui.QPainterPath()
+        path.addPolygon(self.boundingPolygon(False))
+        return path
+
+    # def shape(self):
+    #     path = QtGui.QPainterPath()
+    #     path.moveTo(self.line().p1())
+    #     path.lineTo(self.line().p2())
+    #     path.setPen(QtGui.QPen(self.color, self.width * 3))
+    #     return path
+
+    def boundingPolygon(self, debug):
+        w = self.width * 1.5
+        w = min(w, 2)
+        s, e = self.line().p1(), self.line().p2()
+        dir = s - e
+        dx, dy = -dir.y(), dir.x()
+        norm = (dx**2 + dy**2)**(1 / 2)
+        if debug:
+            print(
+                self.width,
+                w,
+                s.x(),
+                s.y(),
+                e.x(),
+                e.y(),
+                dir.x(),
+                dir.y(),
+                dx,
+                dy,
+                norm, )
+        dx /= (norm + 1e-16)
+        dy /= (norm + 1e-16)
+        if debug:
+            print("dir", dx, dy)
+        p = QtCore.QPointF(dx * w, dy * w)
+        poly = QtGui.QPolygonF([s - p, s + p, e + p, e - p])
+        return poly

EISeg/eiseg/widget/loading.py

+# 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 os.path as osp
+
+from qtpy.QtWidgets import QWidget, QLabel, QHBoxLayout
+from qtpy.QtGui import QIcon, QMovie
+from qtpy import QtCore
+
+from eiseg import pjpath
+
+
+class LoadingWidget(QWidget):
+    def __init__(self):
+        super().__init__()
+        self.setWindowFlags(QtCore.Qt.FramelessWindowHint)
+        layout = QHBoxLayout(self)
+        self.label = QLabel()
+        layout.addWidget(self.label)
+        self.setLayout(layout)
+        self.movie = QMovie(osp.join(pjpath, "resource", "loading.gif"))
+        self.label.setMovie(self.movie)
+        self.movie.start()

EISeg/eiseg/widget/polygon.py

+# 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.
+
+from qtpy import QtWidgets, QtGui, QtCore
+
+from . import GripItem, LineItem, BBoxAnnotation
+
+
+class PolygonAnnotation(QtWidgets.QGraphicsPolygonItem):
+    def __init__(
+            self,
+            labelIndex,
+            shape,
+            delPolygon,
+            setDirty,
+            insideColor=[255, 0, 0],
+            borderColor=[0, 255, 0],
+            opacity=0.5,
+            cocoIndex=None,
+            parent=None, ):
+        super(PolygonAnnotation, self).__init__(parent)
+        self.points = []
+        self.m_items = []
+        self.m_lines = []
+        self.coco_id = cocoIndex
+        self.height, self.width = shape[:2]
+        self.delPolygon = delPolygon
+        self.setDirty = setDirty
+
+        self.labelIndex = labelIndex
+        self.item_hovering = False
+        self.polygon_hovering = False
+        self.anning = False  # 是否标注模式
+        self.line_hovering = False
+        self.noMove = False
+        self.last_focse = False  # 之前是不是焦点在
+
+        self.setZValue(10)
+        self.opacity = opacity
+        i = insideColor
+        self.insideColor = QtGui.QColor(i[0], i[1], i[2])
+        self.insideColor.setAlphaF(opacity)
+        self.halfInsideColor = QtGui.QColor(i[0], i[1], i[2])
+        self.halfInsideColor.setAlphaF(opacity / 2)
+        self.setBrush(self.halfInsideColor)
+        b = borderColor
+        self.borderColor = QtGui.QColor(b[0], b[1], b[2])
+        self.borderColor.setAlphaF(0.8)
+        self.setPen(QtGui.QPen(self.borderColor))
+        self.setAcceptHoverEvents(True)
+
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable, False)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemSendsGeometryChanges, True)
+        self.setFlag(QtWidgets.QGraphicsItem.ItemIsFocusable, True)
+
+        self.setCursor(QtGui.QCursor(QtCore.Qt.PointingHandCursor))
+
+        # persistent this bbox instance and update when needed
+        self.bbox = BBoxAnnotation(labelIndex, self, cocoIndex, self)
+        self.bbox.setParentItem(self)
+        # self.bbox.setVisible(False)
+
+    @property
+    def scnenePoints(self):
+        points = []
+        for p in self.points:
+            p = self.mapToScene(p)
+            points.append([p.x(), p.y()])
+        return points
+
+    def setAnning(self, isAnning=True):
+        if isAnning:
+            self.setAcceptHoverEvents(False)
+            self.last_focse = self.polygon_hovering
+            self.polygon_hovering = False
+            self.anning = True
+            self.setBrush(QtGui.QBrush(QtCore.Qt.NoBrush))
+            self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable, False)
+            # self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable, False)
+            self.setFlag(QtWidgets.QGraphicsItem.ItemSendsGeometryChanges,
+                         False)
+            self.setFlag(QtWidgets.QGraphicsItem.ItemIsFocusable, False)
+            self.setCursor(QtGui.QCursor(QtCore.Qt.ArrowCursor))
+            for line in self.m_lines:
+                line.setAnning(False)
+            for grip in self.m_items:
+                grip.setAnning(False)
+        else:
+            self.setAcceptHoverEvents(True)
+            self.anning = False
+            if self.last_focse:
+                self.polygon_hovering = True
+                self.setBrush(self.insideColor)
+            else:
+                self.setBrush(self.halfInsideColor)
+            self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable, True)
+            # self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable, True)
+            self.setFlag(QtWidgets.QGraphicsItem.ItemSendsGeometryChanges, True)
+            self.setFlag(QtWidgets.QGraphicsItem.ItemIsFocusable, True)
+            self.setCursor(QtGui.QCursor(QtCore.Qt.PointingHandCursor))
+            for line in self.m_lines:
+                line.setAnning(True)
+            for grip in self.m_items:
+                grip.setAnning(True)
+
+    def addPointMiddle(self, lineIdx, point):
+        gripItem = GripItem(self, lineIdx + 1, self.borderColor,
+                            (self.height, self.width))
+        gripItem.setEnabled(False)
+        gripItem.setPos(point)
+        self.scene().addItem(gripItem)
+        gripItem.updateSize()
+        gripItem.setEnabled(True)
+        for grip in self.m_items[lineIdx + 1:]:
+            grip.m_index += 1
+        self.m_items.insert(lineIdx + 1, gripItem)
+        self.points.insert(lineIdx + 1, self.mapFromScene(point))
+        self.setPolygon(QtGui.QPolygonF(self.points))
+        self.bbox.update()
+        for line in self.m_lines[lineIdx + 1:]:
+            line.idx += 1
+        line = QtCore.QLineF(self.mapToScene(self.points[lineIdx]), point)
+        self.m_lines[lineIdx].setLine(line)
+        lineItem = LineItem(self, lineIdx + 1, self.borderColor)
+        line = QtCore.QLineF(
+            point,
+            self.mapToScene(self.points[(lineIdx + 2) % len(self)]), )
+        lineItem.setLine(line)
+        self.m_lines.insert(lineIdx + 1, lineItem)
+        self.scene().addItem(lineItem)
+        lineItem.updateWidth()
+
+    def addPointLast(self, p):
+        grip = GripItem(self,
+                        len(self), self.borderColor, (self.height, self.width))
+        self.scene().addItem(grip)
+        self.m_items.append(grip)
+        grip.updateSize()
+        grip.setPos(p)
+        if len(self) == 0:
+            line = LineItem(self, len(self), self.borderColor)
+            self.scene().addItem(line)
+            self.m_lines.append(line)
+            line.setLine(QtCore.QLineF(p, p))
+        else:
+            self.m_lines[-1].setLine(QtCore.QLineF(self.points[-1], p))
+            line = LineItem(self, len(self), self.borderColor)
+            self.scene().addItem(line)
+            self.m_lines.append(line)
+            line.setLine(QtCore.QLineF(p, self.points[0]))
+
+        self.points.append(p)
+        self.setPolygon(QtGui.QPolygonF(self.points))
+        self.bbox.update()
+
+    def remove(self):
+        for grip in self.m_items:
+            self.scene().removeItem(grip)
+        for line in self.m_lines:
+            self.scene().removeItem(line)
+        while len(self.m_items) != 0:
+            self.m_items.pop()
+        while len(self.m_lines) != 0:
+            self.m_lines.pop()
+        self.scene().polygon_items.remove(self)
+        self.scene().removeItem(self)
+        self.bbox.remove_from_scene()
+        del self.bbox
+        del self
+
+    def removeFocusPoint(self):
+        focusIdx = None
+        for idx, item in enumerate(self.m_items):
+            if item.hasFocus():
+                focusIdx = idx
+                break
+        if focusIdx is not None:
+            if len(self) <= 3:
+                self.delPolygon(self)  # 调用app的删除多边形，为了同时删除coco标签
+                return
+            del self.points[focusIdx]
+            self.setPolygon(QtGui.QPolygonF(self.points))
+            self.bbox.update()
+            self.scene().removeItem(self.m_items[focusIdx])
+            del self.m_items[focusIdx]
+            for grip in self.m_items[focusIdx:]:
+                grip.m_index -= 1
+
+            self.scene().removeItem(self.m_lines[focusIdx])
+            del self.m_lines[focusIdx]
+            line = QtCore.QLineF(
+                self.mapToScene(self.points[(focusIdx - 1) % len(self)]),
+                self.mapToScene(self.points[focusIdx % len(self)]), )
+            # print((focusIdx - 1) % len(self), len(self.m_lines), len(self))
+            self.m_lines[(focusIdx - 1) % len(self)].setLine(line)
+            for line in self.m_lines[focusIdx:]:
+                line.idx -= 1
+
+    def removeLastPoint(self):
+        # TODO: 创建的时候用到，需要删line
+        if len(self.points) == 0:
+            self.points.pop()
+            self.setPolygon(QtGui.QPolygonF(self.points))
+            self.bbox.update()
+            it = self.m_items.pop()
+            self.scene().removeItem(it)
+            del it
+
+    def movePoint(self, i, p):
+        # print("Move point", i, p)
+        if 0 <= i < len(self.points):
+            p = self.mapFromScene(p)
+            self.points[i] = p
+            self.setPolygon(QtGui.QPolygonF(self.points))
+            self.bbox.update()
+            self.moveLine(i)
+
+    def moveLine(self, i):
+        # print("Moving line: ", i, self.noMove)
+        if self.noMove:
+            return
+        points = self.points
+        # line[i]
+        line = QtCore.QLineF(
+            self.mapToScene(points[i]),
+            self.mapToScene(points[(i + 1) % len(self)]))
+        self.m_lines[i].setLine(line)
+        # line[i-1]
+        line = QtCore.QLineF(
+            self.mapToScene(points[(i - 1) % len(self)]),
+            self.mapToScene(points[i]))
+        # print((i - 1) % len(self), len(self.m_lines), len(self))
+        self.m_lines[(i - 1) % len(self)].setLine(line)
+
+    def move_item(self, i, pos):
+        if 0 <= i < len(self.m_items):
+            item = self.m_items[i]
+            item.setEnabled(False)
+            item.setPos(pos)
+            item.setEnabled(True)
+            self.moveLine(i)
+
+    def itemChange(self, change, value):
+        if change == QtWidgets.QGraphicsItem.ItemPositionHasChanged:
+            for i, point in enumerate(self.points):
+                self.move_item(i, self.mapToScene(point))
+        return super(PolygonAnnotation, self).itemChange(change, value)
+
+    def hoverEnterEvent(self, ev):
+        self.polygon_hovering = True
+        self.setBrush(self.insideColor)
+        super(PolygonAnnotation, self).hoverEnterEvent(ev)
+
+    def hoverLeaveEvent(self, ev):
+        self.polygon_hovering = False
+        if not self.hasFocus():
+            self.setBrush(self.halfInsideColor)
+        super(PolygonAnnotation, self).hoverLeaveEvent(ev)
+
+    def focusInEvent(self, ev):
+        if not self.anning:
+            self.setBrush(self.insideColor)
+
+    def focusOutEvent(self, ev):
+        if not self.polygon_hovering and not self.anning:
+            self.setBrush(self.halfInsideColor)
+
+    def setColor(self, insideColor, borderColor):
+        i = insideColor
+        self.insideColor = QtGui.QColor(i[0], i[1], i[2])
+        self.insideColor.setAlphaF(self.opacity)
+        self.halfInsideColor = QtGui.QColor(i[0], i[1], i[2])
+        self.halfInsideColor.setAlphaF(self.opacity / 2)
+        self.setBrush(self.halfInsideColor)
+        b = borderColor
+        self.borderColor = QtGui.QColor(b[0], b[1], b[2])
+        self.borderColor.setAlphaF(0.8)
+        self.setPen(QtGui.QPen(self.borderColor))
+        for grip in self.m_items:
+            grip.setColor(self.borderColor)
+        for line in self.m_lines:
+            line.setColor(self.borderColor)
+
+    def __len__(self):
+        return len(self.points)

EISeg/eiseg/widget/scene.py

+# 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.
+
+from PyQt5.QtCore import QPointF
+from qtpy import QtWidgets, QtCore
+from qtpy.QtCore import Qt
+from qtpy.QtGui import QPen, QColor
+
+
+class AnnotationScene(QtWidgets.QGraphicsScene):
+    clickRequest = QtCore.Signal(int, int, bool)
+
+    def __init__(self, parent=None):
+        super(AnnotationScene, self).__init__(parent)
+        self.creating = False
+        self.polygon_items = []
+        # draw cross
+        self.coords = None
+        self.pen = QPen()
+        self.pen.setWidth(1)
+        self.pen.setColor(QColor(0, 0, 0, 127))
+
+    def setPenColor(self, color_list):
+        R, G, B, A = color_list
+        self.pen.setColor(QColor(R, G, B, A))
+
+    def updatePolygonSize(self):
+        for poly in self.polygon_items:
+            for grip in poly.m_items:
+                grip.updateSize()
+            for line in poly.m_lines:
+                line.updateWidth()
+
+    def setCreating(self, creating=True):
+        self.creating = creating
+
+    def mousePressEvent(self, ev):
+        pos = ev.scenePos()
+        if not self.creating and not self.hovering:
+            if ev.buttons() in [Qt.LeftButton, Qt.RightButton]:
+                self.clickRequest.emit(
+                    int(pos.x()), int(pos.y()), ev.buttons() == Qt.LeftButton)
+        elif self.creating:
+            self.polygon_item.removeLastPoint()
+            self.polygon_item.addPointLast(ev.scenePos())
+            # movable element
+            self.polygon_item.addPointLast(ev.scenePos())
+        super(AnnotationScene, self).mousePressEvent(ev)
+
+    def mouseMoveEvent(self, ev):
+        if self.creating:
+            self.polygon_item.movePoint(
+                # self.polygon_item.number_of_points() - 1, ev.scenePos()
+                len(self.polygon_item) - 1,
+                ev.scenePos(), )
+        super(AnnotationScene, self).mouseMoveEvent(ev)
+
+    def drawForeground(self, painter, rect):
+        if self.coords is not None and self.coords != QPointF(-1, -1):
+            painter.setClipRect(rect)
+            painter.setPen(self.pen)
+            # painter.drawLine(int(self.coords.x()), int(rect.top()),
+            #                  int(self.coords.x()), int(rect.bottom()))
+            # painter.drawLine(int(rect.left()), int(self.coords.y()),
+            #                  int(rect.right()), int(self.coords.y()))
+
+    def onMouseChanged(self, pointf):
+        self.coords = pointf
+        self.invalidate()
+
+    @property
+    def item_hovering(self):
+        for poly in self.polygon_items:
+            if poly.item_hovering:
+                return True
+        return False
+
+    @property
+    def polygon_hovering(self):
+        for poly in self.polygon_items:
+            if poly.polygon_hovering:
+                return True
+        return False
+
+    @property
+    def line_hovering(self):
+        for poly in self.polygon_items:
+            if poly.line_hovering:
+                return True
+        return False
+
+    @property
+    def hovering(self):
+        return self.item_hovering or self.polygon_hovering or self.line_hovering

EISeg/eiseg/widget/shortcut.py

+# 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 os.path as osp
+import math
+from functools import partial
+
+from PyQt5.QtCore import QPoint
+from PyQt5.QtWidgets import QDesktopWidget
+
+from qtpy import QtCore, QtWidgets
+from qtpy.QtWidgets import (
+    QWidget,
+    QLabel,
+    QPushButton,
+    QGridLayout,
+    QKeySequenceEdit,
+    QMessageBox, )
+from qtpy.QtGui import QIcon
+from qtpy import QtCore
+from qtpy.QtCore import Qt
+
+from util import save_configs
+
+
+class RecordShortcutWidget(QKeySequenceEdit):
+    def __init__(self, finishCallback, location):
+        super().__init__()
+        self.finishCallback = finishCallback
+        # 隐藏界面
+        self.setWindowFlags(Qt.FramelessWindowHint)
+        self.move(location)
+        self.show()
+        self.editingFinished.connect(lambda: finishCallback(self.keySequence()))
+
+    def keyReleaseEvent(self, ev):
+        self.finishCallback(self.keySequence())
+
+
+class ShortcutWidget(QWidget):
+    def __init__(self, actions, pjpath):
+        super().__init__()
+        self.tr = partial(QtCore.QCoreApplication.translate, "ShortcutWidget")
+        self.setWindowTitle(self.tr("编辑快捷键"))
+        self.setWindowIcon(QIcon(osp.join(pjpath, "resource/Shortcut.png")))
+        # self.setFixedSize(self.width(), self.height())
+        self.actions = actions
+        self.recorder = None
+        self.initUI()
+
+    def initUI(self):
+        grid = QGridLayout()
+        self.setLayout(grid)
+
+        actions = self.actions
+        for idx, action in enumerate(actions):
+            # 2列英文看不清
+            grid.addWidget(QLabel(action.iconText()[1:]), idx // 3, idx % 3 * 3)
+            shortcut = action.shortcut().toString()
+            if len(shortcut) == 0:
+                shortcut = self.tr("-")
+            button = QPushButton(shortcut)
+            button.setFixedWidth(150)
+            button.setFixedHeight(30)
+            button.clicked.connect(partial(self.recordShortcut, action))
+            grid.addWidget(
+                button,
+                idx // 3,
+                idx % 3 * 3 + 1, )
+
+    def refreshUi(self):
+        actions = self.actions
+        for idx, action in enumerate(actions):
+            shortcut = action.shortcut().toString()
+            if len(shortcut) == 0:
+                shortcut = self.tr("-")
+            self.layout().itemAtPosition(
+                idx // 3,
+                idx % 3 * 3 + 1, ).widget().setText(shortcut)
+
+    def recordShortcut(self, action):
+        # 打开快捷键设置的窗口时，如果之前的还在就先关闭
+        if self.recorder is not None:
+            self.recorder.close()
+        rect = self.geometry()
+        x = rect.x()
+        y = rect.y() + rect.height()
+        self.recorder = RecordShortcutWidget(self.setShortcut, QPoint(x, y))
+        self.currentAction = action
+
+    def setShortcut(self, key):
+        self.recorder.close()
+
+        for a in self.actions:
+            if a.shortcut() == key:
+                key = key.toString()
+                msg = QMessageBox()
+                msg.setIcon(QMessageBox.Warning)
+                msg.setWindowTitle(key + " " + self.tr("快捷键冲突"))
+                msg.setText(key + " " + self.tr("快捷键已被") + " " + a.data(
+                ) + " " + self.tr("使用，请设置其他快捷键或先修改") + " " + a.data() + " " +
+                            self.tr("的快捷键"))
+                msg.setStandardButtons(QMessageBox.Ok)
+                msg.exec_()
+                return
+        key = "" if key.toString() == "Esc" else key  # ESC不设置快捷键
+        self.currentAction.setShortcut(key)
+        self.refreshUi()
+        save_configs(None, None, self.actions)
+
+    def center(self):
+        qr = self.frameGeometry()
+        cp = QDesktopWidget().availableGeometry().center()
+        qr.moveCenter(cp)
+        self.move(qr.topLeft())
+
+    # 快捷键设置跟随移动
+    def moveEvent(self, event):
+        p = self.geometry()
+        x = p.x()
+        y = p.y() + p.height()
+        if self.recorder is not None:
+            self.recorder.move(x, y)
+
+    def closeEvent(self, event):
+        # 关闭时也退出快捷键设置
+        if self.recorder is not None:
+            self.recorder.close()

EISeg/eiseg/widget/table.py

+# 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.
+
+from qtpy import QtWidgets
+from qtpy.QtCore import Qt
+
+
+class TableWidget(QtWidgets.QTableWidget):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.setDragEnabled(True)
+        self.setAcceptDrops(True)
+        self.viewport().setAcceptDrops(True)
+        self.setDragDropOverwriteMode(False)
+        self.setDropIndicatorShown(True)
+        self.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
+        self.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
+        self.setDragDropMode(QtWidgets.QAbstractItemView.InternalMove)
+
+    def dropEvent(self, event):
+        if event.source() == self:
+            rows = set([mi.row() for mi in self.selectedIndexes()])
+            targetRow = self.indexAt(event.pos()).row()
+            rows.discard(targetRow)
+            rows = sorted(rows)
+            if not rows:
+                return
+            if targetRow == -1:
+                targetRow = self.rowCount()
+            for _ in range(len(rows)):
+                self.insertRow(targetRow)
+            rowMapping = dict()  # Src row to target row.
+            for idx, row in enumerate(rows):
+                if row < targetRow:
+                    rowMapping[row] = targetRow + idx
+                else:
+                    rowMapping[row + len(rows)] = targetRow + idx
+            colCount = self.columnCount()
+            for srcRow, tgtRow in sorted(rowMapping.items()):
+                for col in range(0, colCount):
+                    self.setItem(tgtRow, col, self.takeItem(srcRow, col))
+            for row in reversed(sorted(rowMapping.keys())):
+                self.removeRow(row)
+            event.accept()
+            return
+
+    def drop_on(self, event):
+        index = self.indexAt(event.pos())
+        if not index.isValid():
+            return self.rowCount()
+
+        return index.row() + 1 if self.is_below(event.pos(),
+                                                index) else index.row()
+
+    def is_below(self, pos, index):
+        rect = self.visualRect(index)
+        margin = 2
+        if pos.y() - rect.top() < margin:
+            return False
+        elif rect.bottom() - pos.y() < margin:
+            return True
+        # noinspection PyTypeChecker
+        return (rect.contains(pos, True) and
+                not (int(self.model().flags(index)) & Qt.ItemIsDropEnabled) and
+                pos.y() >= rect.center().y())

EISeg/eiseg/widget/view.py

+# 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.
+
+from qtpy import QtWidgets, QtCore, QtGui
+from qtpy.QtCore import Qt, QPointF
+
+
+class AnnotationView(QtWidgets.QGraphicsView):
+    zoomRequest = QtCore.Signal(float)
+    mousePosChanged = QtCore.Signal(QPointF)
+
+    def __init__(self, *args):
+        super(AnnotationView, self).__init__(*args)
+        self.setRenderHints(QtGui.QPainter.Antialiasing |
+                            QtGui.QPainter.SmoothPixmapTransform)
+        self.setMouseTracking(True)
+        self.setTransformationAnchor(QtWidgets.QGraphicsView.NoAnchor)
+        self.setResizeAnchor(QtWidgets.QGraphicsView.NoAnchor)
+        self.point = QtCore.QPoint(0, 0)
+        self.middle_click = False
+        self.zoom_all = 1
+        # hint mouse
+        # self.setCursor(Qt.BlankCursor)
+
+    def wheelEvent(self, ev):
+        if ev.modifiers() & QtCore.Qt.ControlModifier:
+            zoom = 1 + ev.angleDelta().y() / 2880
+            self.zoom_all *= zoom
+            oldPos = self.mapToScene(ev.pos())
+            if self.zoom_all >= 0.02 and self.zoom_all <= 50:  # 限制缩放的倍数
+                self.scale(zoom, zoom)
+            newPos = self.mapToScene(ev.pos())
+            delta = newPos - oldPos
+            self.translate(delta.x(), delta.y())
+            ev.ignore()
+            self.zoomRequest.emit(self.zoom_all)
+        else:
+            super(AnnotationView, self).wheelEvent(ev)
+
+    def mouseMoveEvent(self, ev):
+        mouse_pos = QPointF(self.mapToScene(ev.pos()))
+        self.mousePosChanged.emit(mouse_pos.toPoint())
+        if self.middle_click and (self.horizontalScrollBar().isVisible() or
+                                  self.verticalScrollBar().isVisible()):
+            # 放大到出现滚动条才允许拖动，避免出现抖动
+            self._endPos = ev.pos(
+            ) / self.zoom_all - self._startPos / self.zoom_all
+            # 这儿不写为先减后除，这样会造成速度不一致
+            self.point = self.point + self._endPos
+            self._startPos = ev.pos()
+            self.translate(self._endPos.x(), self._endPos.y())
+        super(AnnotationView, self).mouseMoveEvent(ev)
+
+    def mousePressEvent(self, ev):
+        if ev.buttons() == Qt.MiddleButton:
+            self.middle_click = True
+            self._startPos = ev.pos()
+        super(AnnotationView, self).mousePressEvent(ev)
+
+    def mouseReleaseEvent(self, ev):
+        if ev.button() == Qt.MiddleButton:
+            self.middle_click = False
+        super(AnnotationView, self).mouseReleaseEvent(ev)
+
+    def leaveEvent(self, ev):
+        self.mousePosChanged.emit(QPointF(-1, -1))
+        return super(AnnotationView, self).leaveEvent(ev)

EISeg/eiseg/widget/vtk.py

+# reference: https://blog.csdn.net/weixin_46579211/article/details/118279231
+import typing
+import numpy as np
+from qtpy.QtWidgets import QWidget, QVBoxLayout
+
+# 只有调用convert_vtk才会加载vtk相关的并把这个控件转换为真正的vtk控件
+vtk = None
+QVTKRenderWindowInteractor = None
+vtkImageImportFromArray = None
+
+
+class VTKWidget(QWidget):
+    def __init__(self, parent: typing.Optional["QWidget"]) -> None:
+        super().__init__(parent)
+        self.setObjectName("vtkWidget")
+        self.vlayer = QVBoxLayout(self)
+        # default setting
+        self.smoothing_iterations = 10
+        self.pass_band = 0.005
+        self.feature_angle = 120
+        self.import_vtk = False
+
+    def convert_vtk(self) -> bool:
+        try:
+            import vtk
+            from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
+            from vtkmodules.util.vtkImageImportFromArray import vtkImageImportFromArray
+            global vtk
+            global QVTKRenderWindowInteractor
+            global vtkImageImportFromArray
+            self.import_vtk = True
+            vtk.vtkFileOutputWindow().SetGlobalWarningDisplay(0)
+            self.init()
+        except:
+            self.import_vtk = False
+        finally:
+            return self.import_vtk
+
+    def init(self, clear: bool=True) -> None:
+        if self.import_vtk is False:
+            return
+        # remove
+        item = self.vlayer.itemAt(0)
+        self.vlayer.removeItem(item)
+        if item is not None and item.widget():
+            item.widget().deleteLater()
+        # add
+        self.renderer = vtk.vtkRenderer()
+        self.interactor = QVTKRenderWindowInteractor(self)
+        self.interactor.GetRenderWindow().AddRenderer(self.renderer)
+        self.vlayer.addWidget(self.interactor)
+        if clear:
+            # set background
+            self.renderer.SetBackground(vtk.vtkNamedColors().GetColor3d(
+                "black"))
+            self.interactor.Start()
+
+    def show_array(self,
+                   data: np.ndarray,
+                   spacing: typing.Tuple,
+                   color_map: typing.List) -> None:
+        if self.import_vtk is False:
+            return
+        print("color_map:", color_map)
+        self.num_block = len(np.unique(data))
+        print("num_block:", self.num_block)
+        self.reader = vtkImageImportFromArray()
+        self.reader.SetArray(data)
+        self.reader.SetDataSpacing(spacing)
+        mbds = vtk.vtkMultiBlockDataSet()
+        mbds.SetNumberOfBlocks(self.num_block - 1)
+        for iter in range(1, self.num_block + 1):
+            contour = self._get_mc_contour(iter)
+            smoother = self._smoothing(self.smoothing_iterations,
+                                       self.pass_band, self.feature_angle,
+                                       contour)
+            mbds.SetBlock(iter, smoother.GetOutput())
+        self._multidisplay(mbds, color_map)
+
+    def _get_mc_contour(self, setvalue: int) -> typing.Any:
+        contour = vtk.vtkDiscreteMarchingCubes()
+        contour.SetInputConnection(self.reader.GetOutputPort())
+        contour.ComputeNormalsOn()
+        contour.SetValue(0, setvalue)
+        return contour
+
+    def _smoothing(self,
+                   smoothing_iterations: int,
+                   pass_band: float,
+                   feature_angle: int,
+                   contour: typing.Any) -> typing.Any:
+        smoother = vtk.vtkWindowedSincPolyDataFilter()
+        smoother.SetInputConnection(contour.GetOutputPort())
+        smoother.SetNumberOfIterations(smoothing_iterations)
+        smoother.BoundarySmoothingOff()
+        smoother.FeatureEdgeSmoothingOff()
+        smoother.SetFeatureAngle(feature_angle)
+        smoother.SetPassBand(pass_band)
+        smoother.NonManifoldSmoothingOn()
+        smoother.NormalizeCoordinatesOn()
+        smoother.Update()
+        return smoother
+
+    def _multidisplay(self, obj: typing.Any, color_map: typing.List) -> None:
+        self.init(False)
+        mapper = vtk.vtkCompositePolyDataMapper2()
+        mapper.SetInputDataObject(obj)
+        cdsa = vtk.vtkCompositeDataDisplayAttributes()
+        mapper.SetCompositeDataDisplayAttributes(cdsa)
+        # 上色
+        color_map.insert(0, [0., 0., 0.])
+        for i in range(1, self.num_block + 1):
+            r, g, b = color_map[i - 1]
+            mapper.SetBlockColor(i, r / 255., g / 255., b / 255.)
+        actor = vtk.vtkActor()
+        actor.SetMapper(mapper)
+        actor.RotateX(180)  # 翻转一下才对
+        # Enable user interface interactor.
+        self.renderer.AddActor(actor)
+        self.interactor.Start()

EISeg/init.sh

+#!/bin/bash
+
+ROOT=`cd "$(dirname ${BASH_SOURCE[0]})" && pwd`
+
+echo "ROOT : $ROOT"
+
+export PYTHONPATH=$PYTHONPATH:$ROOT/eiseg

EISeg/med3d/CMakeLists.txt

+cmake_minimum_required(VERSION 3.16.3...3.19.7 FATAL_ERROR)
+
+project(EISeg)
+
+#-----------------------------------------------------------------------------
+# Extension meta-information
+set(EXTENSION_HOMEPAGE "https://www.slicer.org/wiki/Documentation/Nightly/Extensions/EISeg")
+set(EXTENSION_CATEGORY "Segmentation")
+set(EXTENSION_CONTRIBUTORS "Shiyu Tang (Baidu Co.), Lin Han (NYU Tandon School of Engineering)")
+set(EXTENSION_DESCRIPTION "EISeg")
+set(EXTENSION_ICONURL "https://www.example.com/Slicer/Extensions/EISeg.png")
+set(EXTENSION_SCREENSHOTURLS "https://www.example.com/Slicer/Extensions/EISeg/Screenshots/1.png")
+set(EXTENSION_DEPENDS "NA") # Specified as a list or "NA" if no dependencies
+
+#-----------------------------------------------------------------------------
+# Extension dependencies
+find_package(Slicer REQUIRED)
+include(${Slicer_USE_FILE})
+
+#-----------------------------------------------------------------------------
+# Extension modules
+## NEXT_MODULE
+
+#-----------------------------------------------------------------------------
+include(${Slicer_EXTENSION_GENERATE_CONFIG})
+include(${Slicer_EXTENSION_CPACK})

EISeg/med3d/EISegMed3D/CMakeLists.txt

+#-----------------------------------------------------------------------------
+set(MODULE_NAME placePoint)
+
+#-----------------------------------------------------------------------------
+set(MODULE_PYTHON_SCRIPTS
+  ${MODULE_NAME}.py
+  )
+
+set(MODULE_PYTHON_RESOURCES
+  Resources/Icons/${MODULE_NAME}.png
+  Resources/UI/${MODULE_NAME}.ui
+  )
+
+#-----------------------------------------------------------------------------
+slicerMacroBuildScriptedModule(
+  NAME ${MODULE_NAME}
+  SCRIPTS ${MODULE_PYTHON_SCRIPTS}
+  RESOURCES ${MODULE_PYTHON_RESOURCES}
+  WITH_GENERIC_TESTS
+  )
+
+#-----------------------------------------------------------------------------
+if(BUILD_TESTING)
+
+  # Register the unittest subclass in the main script as a ctest.
+  # Note that the test will also be available at runtime.
+  slicer_add_python_unittest(SCRIPT ${MODULE_NAME}.py)
+
+  # Additional build-time testing
+  add_subdirectory(Testing)
+endif()