| | import json |
| | import multiprocessing |
| | import os |
| | import pickle |
| | import sys |
| | import tempfile |
| | import time |
| | import traceback |
| | from enum import IntEnum |
| | from types import SimpleNamespace as sn |
| |
|
| | import cv2 |
| | import numpy as np |
| | import numpy.linalg as npla |
| | from PyQt5.QtCore import * |
| | from PyQt5.QtGui import * |
| | from PyQt5.QtWidgets import * |
| |
|
| | from core import imagelib, pathex |
| | from core.cv2ex import * |
| | from core.imagelib import SegIEPoly, SegIEPolys, SegIEPolyType, sd |
| | from core.qtex import * |
| | from DFLIMG import * |
| | from localization import StringsDB, system_language |
| | from samplelib import PackedFaceset |
| |
|
| | from .QCursorDB import QCursorDB |
| | from .QIconDB import QIconDB |
| | from .QStringDB import QStringDB |
| | from .QImageDB import QImageDB |
| |
|
| | class OpMode(IntEnum): |
| | NONE = 0 |
| | DRAW_PTS = 1 |
| | EDIT_PTS = 2 |
| | VIEW_BAKED = 3 |
| | VIEW_XSEG_MASK = 4 |
| |
|
| | class PTEditMode(IntEnum): |
| | MOVE = 0 |
| | ADD_DEL = 1 |
| |
|
| | class DragType(IntEnum): |
| | NONE = 0 |
| | IMAGE_LOOK = 1 |
| | POLY_PT = 2 |
| |
|
| | class ViewLock(IntEnum): |
| | NONE = 0 |
| | CENTER = 1 |
| |
|
| | class QUIConfig(): |
| | @staticmethod |
| | def initialize(icon_size = 48, icon_spacer_size=16, preview_bar_icon_size=64): |
| | QUIConfig.icon_q_size = QSize(icon_size, icon_size) |
| | QUIConfig.icon_spacer_q_size = QSize(icon_spacer_size, icon_spacer_size) |
| | QUIConfig.preview_bar_icon_q_size = QSize(preview_bar_icon_size, preview_bar_icon_size) |
| |
|
| | class ImagePreviewSequenceBar(QFrame): |
| | def __init__(self, preview_images_count, icon_size): |
| | super().__init__() |
| | self.preview_images_count = preview_images_count = max(1, preview_images_count + (preview_images_count % 2 -1) ) |
| |
|
| | self.icon_size = icon_size |
| |
|
| | black_q_img = QImage(np.zeros( (icon_size,icon_size,3) ).data, icon_size, icon_size, 3*icon_size, QImage.Format_RGB888) |
| | self.black_q_pixmap = QPixmap.fromImage(black_q_img) |
| |
|
| | self.image_containers = [ QLabel() for i in range(preview_images_count)] |
| |
|
| | main_frame_l_cont_hl = QGridLayout() |
| | main_frame_l_cont_hl.setContentsMargins(0,0,0,0) |
| | |
| |
|
| |
|
| |
|
| | for i in range(len(self.image_containers)): |
| | q_label = self.image_containers[i] |
| | q_label.setScaledContents(True) |
| | if i == preview_images_count//2: |
| | q_label.setMinimumSize(icon_size+16, icon_size+16 ) |
| | q_label.setMaximumSize(icon_size+16, icon_size+16 ) |
| | else: |
| | q_label.setMinimumSize(icon_size, icon_size ) |
| | q_label.setMaximumSize(icon_size, icon_size ) |
| | opacity_effect = QGraphicsOpacityEffect() |
| | opacity_effect.setOpacity(0.5) |
| | q_label.setGraphicsEffect(opacity_effect) |
| |
|
| | q_label.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| |
|
| | main_frame_l_cont_hl.addWidget (q_label, 0, i) |
| |
|
| | self.setLayout(main_frame_l_cont_hl) |
| |
|
| | self.prev_img_conts = self.image_containers[(preview_images_count//2) -1::-1] |
| | self.next_img_conts = self.image_containers[preview_images_count//2:] |
| |
|
| | self.update_images() |
| |
|
| | def get_preview_images_count(self): |
| | return self.preview_images_count |
| |
|
| | def update_images(self, prev_imgs=None, next_imgs=None): |
| | |
| | if prev_imgs is None: |
| | prev_imgs = [] |
| | prev_img_conts_len = len(self.prev_img_conts) |
| | prev_q_imgs_len = len(prev_imgs) |
| | if prev_q_imgs_len < prev_img_conts_len: |
| | for i in range ( prev_img_conts_len - prev_q_imgs_len ): |
| | prev_imgs.append(None) |
| | elif prev_q_imgs_len > prev_img_conts_len: |
| | prev_imgs = prev_imgs[:prev_img_conts_len] |
| |
|
| | if next_imgs is None: |
| | next_imgs = [] |
| | next_img_conts_len = len(self.next_img_conts) |
| | next_q_imgs_len = len(next_imgs) |
| | if next_q_imgs_len < next_img_conts_len: |
| | for i in range ( next_img_conts_len - next_q_imgs_len ): |
| | next_imgs.append(None) |
| | elif next_q_imgs_len > next_img_conts_len: |
| | next_imgs = next_imgs[:next_img_conts_len] |
| |
|
| | for i,img in enumerate(prev_imgs): |
| | self.prev_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) if img is not None else self.black_q_pixmap ) |
| |
|
| | for i,img in enumerate(next_imgs): |
| | self.next_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) if img is not None else self.black_q_pixmap ) |
| |
|
| | class ColorScheme(): |
| | def __init__(self, unselected_color, selected_color, outline_color, outline_width, pt_outline_color, cross_cursor): |
| | self.poly_unselected_brush = QBrush(unselected_color) |
| | self.poly_selected_brush = QBrush(selected_color) |
| |
|
| | self.poly_outline_solid_pen = QPen(outline_color, outline_width, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin) |
| | self.poly_outline_dot_pen = QPen(outline_color, outline_width, Qt.DotLine, Qt.RoundCap, Qt.RoundJoin) |
| |
|
| | self.pt_outline_pen = QPen(pt_outline_color) |
| | self.cross_cursor = cross_cursor |
| |
|
| | class CanvasConfig(): |
| |
|
| | def __init__(self, |
| | pt_radius=4, |
| | pt_select_radius=8, |
| | color_schemes=None, |
| | **kwargs): |
| | self.pt_radius = pt_radius |
| | self.pt_select_radius = pt_select_radius |
| |
|
| | if color_schemes is None: |
| | color_schemes = [ |
| | ColorScheme( QColor(192,0,0,alpha=0), QColor(192,0,0,alpha=72), QColor(192,0,0), 2, QColor(255,255,255), QCursorDB.cross_red ), |
| | ColorScheme( QColor(0,192,0,alpha=0), QColor(0,192,0,alpha=72), QColor(0,192,0), 2, QColor(255,255,255), QCursorDB.cross_green ), |
| | ColorScheme( QColor(0,0,192,alpha=0), QColor(0,0,192,alpha=72), QColor(0,0,192), 2, QColor(255,255,255), QCursorDB.cross_blue ), |
| | ] |
| | self.color_schemes = color_schemes |
| |
|
| | class QCanvasControlsLeftBar(QFrame): |
| |
|
| | def __init__(self): |
| | super().__init__() |
| | |
| | btn_poly_type_include = QToolButton() |
| | self.btn_poly_type_include_act = QActionEx( QIconDB.poly_type_include, QStringDB.btn_poly_type_include_tip, shortcut='Q', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_poly_type_include.setDefaultAction(self.btn_poly_type_include_act) |
| | btn_poly_type_include.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_poly_type_exclude = QToolButton() |
| | self.btn_poly_type_exclude_act = QActionEx( QIconDB.poly_type_exclude, QStringDB.btn_poly_type_exclude_tip, shortcut='W', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_poly_type_exclude.setDefaultAction(self.btn_poly_type_exclude_act) |
| | btn_poly_type_exclude.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | self.btn_poly_type_act_grp = QActionGroup (self) |
| | self.btn_poly_type_act_grp.addAction(self.btn_poly_type_include_act) |
| | self.btn_poly_type_act_grp.addAction(self.btn_poly_type_exclude_act) |
| | self.btn_poly_type_act_grp.setExclusive(True) |
| | |
| | btn_undo_pt = QToolButton() |
| | self.btn_undo_pt_act = QActionEx( QIconDB.undo_pt, QStringDB.btn_undo_pt_tip, shortcut='Ctrl+Z', shortcut_in_tooltip=True, is_auto_repeat=True) |
| | btn_undo_pt.setDefaultAction(self.btn_undo_pt_act) |
| | btn_undo_pt.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_redo_pt = QToolButton() |
| | self.btn_redo_pt_act = QActionEx( QIconDB.redo_pt, QStringDB.btn_redo_pt_tip, shortcut='Ctrl+Shift+Z', shortcut_in_tooltip=True, is_auto_repeat=True) |
| | btn_redo_pt.setDefaultAction(self.btn_redo_pt_act) |
| | btn_redo_pt.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_delete_poly = QToolButton() |
| | self.btn_delete_poly_act = QActionEx( QIconDB.delete_poly, QStringDB.btn_delete_poly_tip, shortcut='Delete', shortcut_in_tooltip=True) |
| | btn_delete_poly.setDefaultAction(self.btn_delete_poly_act) |
| | btn_delete_poly.setIconSize(QUIConfig.icon_q_size) |
| | |
| | btn_pt_edit_mode = QToolButton() |
| | self.btn_pt_edit_mode_act = QActionEx( QIconDB.pt_edit_mode, QStringDB.btn_pt_edit_mode_tip, shortcut_in_tooltip=True, is_checkable=True) |
| | btn_pt_edit_mode.setDefaultAction(self.btn_pt_edit_mode_act) |
| | btn_pt_edit_mode.setIconSize(QUIConfig.icon_q_size) |
| | |
| |
|
| | controls_bar_frame2_l = QVBoxLayout() |
| | controls_bar_frame2_l.addWidget ( btn_poly_type_include ) |
| | controls_bar_frame2_l.addWidget ( btn_poly_type_exclude ) |
| | controls_bar_frame2 = QFrame() |
| | controls_bar_frame2.setFrameShape(QFrame.StyledPanel) |
| | controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| | controls_bar_frame2.setLayout(controls_bar_frame2_l) |
| |
|
| | controls_bar_frame3_l = QVBoxLayout() |
| | controls_bar_frame3_l.addWidget ( btn_undo_pt ) |
| | controls_bar_frame3_l.addWidget ( btn_redo_pt ) |
| | controls_bar_frame3_l.addWidget ( btn_delete_poly ) |
| | controls_bar_frame3 = QFrame() |
| | controls_bar_frame3.setFrameShape(QFrame.StyledPanel) |
| | controls_bar_frame3.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| | controls_bar_frame3.setLayout(controls_bar_frame3_l) |
| |
|
| | controls_bar_frame4_l = QVBoxLayout() |
| | controls_bar_frame4_l.addWidget ( btn_pt_edit_mode ) |
| | controls_bar_frame4 = QFrame() |
| | controls_bar_frame4.setFrameShape(QFrame.StyledPanel) |
| | controls_bar_frame4.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| | controls_bar_frame4.setLayout(controls_bar_frame4_l) |
| |
|
| | controls_bar_l = QVBoxLayout() |
| | controls_bar_l.setContentsMargins(0,0,0,0) |
| | controls_bar_l.addWidget(controls_bar_frame2) |
| | controls_bar_l.addWidget(controls_bar_frame3) |
| | controls_bar_l.addWidget(controls_bar_frame4) |
| |
|
| | self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding ) |
| | self.setLayout(controls_bar_l) |
| |
|
| | class QCanvasControlsRightBar(QFrame): |
| |
|
| | def __init__(self): |
| | super().__init__() |
| | |
| | btn_poly_color_red = QToolButton() |
| | self.btn_poly_color_red_act = QActionEx( QIconDB.poly_color_red, QStringDB.btn_poly_color_red_tip, shortcut='1', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_poly_color_red.setDefaultAction(self.btn_poly_color_red_act) |
| | btn_poly_color_red.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_poly_color_green = QToolButton() |
| | self.btn_poly_color_green_act = QActionEx( QIconDB.poly_color_green, QStringDB.btn_poly_color_green_tip, shortcut='2', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_poly_color_green.setDefaultAction(self.btn_poly_color_green_act) |
| | btn_poly_color_green.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_poly_color_blue = QToolButton() |
| | self.btn_poly_color_blue_act = QActionEx( QIconDB.poly_color_blue, QStringDB.btn_poly_color_blue_tip, shortcut='3', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_poly_color_blue.setDefaultAction(self.btn_poly_color_blue_act) |
| | btn_poly_color_blue.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_view_baked_mask = QToolButton() |
| | self.btn_view_baked_mask_act = QActionEx( QIconDB.view_baked, QStringDB.btn_view_baked_mask_tip, shortcut='4', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_view_baked_mask.setDefaultAction(self.btn_view_baked_mask_act) |
| | btn_view_baked_mask.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_view_xseg_mask = QToolButton() |
| | self.btn_view_xseg_mask_act = QActionEx( QIconDB.view_xseg, QStringDB.btn_view_xseg_mask_tip, shortcut='5', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_view_xseg_mask.setDefaultAction(self.btn_view_xseg_mask_act) |
| | btn_view_xseg_mask.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | btn_view_xseg_overlay_mask = QToolButton() |
| | self.btn_view_xseg_overlay_mask_act = QActionEx( QIconDB.view_xseg_overlay, QStringDB.btn_view_xseg_overlay_mask_tip, shortcut='`', shortcut_in_tooltip=True, is_checkable=True) |
| | btn_view_xseg_overlay_mask.setDefaultAction(self.btn_view_xseg_overlay_mask_act) |
| | btn_view_xseg_overlay_mask.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | self.btn_poly_color_act_grp = QActionGroup (self) |
| | self.btn_poly_color_act_grp.addAction(self.btn_poly_color_red_act) |
| | self.btn_poly_color_act_grp.addAction(self.btn_poly_color_green_act) |
| | self.btn_poly_color_act_grp.addAction(self.btn_poly_color_blue_act) |
| | self.btn_poly_color_act_grp.addAction(self.btn_view_baked_mask_act) |
| | self.btn_poly_color_act_grp.addAction(self.btn_view_xseg_mask_act) |
| | self.btn_poly_color_act_grp.setExclusive(True) |
| | |
| | btn_view_lock_center = QToolButton() |
| | self.btn_view_lock_center_act = QActionEx( QIconDB.view_lock_center, QStringDB.btn_view_lock_center_tip, shortcut_in_tooltip=True, is_checkable=True) |
| | btn_view_lock_center.setDefaultAction(self.btn_view_lock_center_act) |
| | btn_view_lock_center.setIconSize(QUIConfig.icon_q_size) |
| |
|
| | controls_bar_frame2_l = QVBoxLayout() |
| | controls_bar_frame2_l.addWidget ( btn_view_xseg_overlay_mask ) |
| | controls_bar_frame2 = QFrame() |
| | controls_bar_frame2.setFrameShape(QFrame.StyledPanel) |
| | controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| | controls_bar_frame2.setLayout(controls_bar_frame2_l) |
| | |
| | controls_bar_frame1_l = QVBoxLayout() |
| | controls_bar_frame1_l.addWidget ( btn_poly_color_red ) |
| | controls_bar_frame1_l.addWidget ( btn_poly_color_green ) |
| | controls_bar_frame1_l.addWidget ( btn_poly_color_blue ) |
| | controls_bar_frame1_l.addWidget ( btn_view_baked_mask ) |
| | controls_bar_frame1_l.addWidget ( btn_view_xseg_mask ) |
| | controls_bar_frame1 = QFrame() |
| | controls_bar_frame1.setFrameShape(QFrame.StyledPanel) |
| | controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| | controls_bar_frame1.setLayout(controls_bar_frame1_l) |
| | |
| | controls_bar_frame3_l = QVBoxLayout() |
| | controls_bar_frame3_l.addWidget ( btn_view_lock_center ) |
| | controls_bar_frame3 = QFrame() |
| | controls_bar_frame3.setFrameShape(QFrame.StyledPanel) |
| | controls_bar_frame3.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) |
| | controls_bar_frame3.setLayout(controls_bar_frame3_l) |
| |
|
| | controls_bar_l = QVBoxLayout() |
| | controls_bar_l.setContentsMargins(0,0,0,0) |
| | controls_bar_l.addWidget(controls_bar_frame2) |
| | controls_bar_l.addWidget(controls_bar_frame1) |
| | controls_bar_l.addWidget(controls_bar_frame3) |
| | |
| | self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding ) |
| | self.setLayout(controls_bar_l) |
| |
|
| | class QCanvasOperator(QWidget): |
| | def __init__(self, cbar): |
| | super().__init__() |
| | self.cbar = cbar |
| |
|
| | self.set_cbar_disabled() |
| |
|
| | self.cbar.btn_poly_color_red_act.triggered.connect ( lambda : self.set_color_scheme_id(0) ) |
| | self.cbar.btn_poly_color_green_act.triggered.connect ( lambda : self.set_color_scheme_id(1) ) |
| | self.cbar.btn_poly_color_blue_act.triggered.connect ( lambda : self.set_color_scheme_id(2) ) |
| | self.cbar.btn_view_baked_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_BAKED) ) |
| | self.cbar.btn_view_xseg_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_XSEG_MASK) ) |
| |
|
| | self.cbar.btn_view_xseg_overlay_mask_act.toggled.connect ( lambda is_checked: self.update() ) |
| |
|
| | self.cbar.btn_poly_type_include_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.INCLUDE) ) |
| | self.cbar.btn_poly_type_exclude_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.EXCLUDE) ) |
| |
|
| | self.cbar.btn_undo_pt_act.triggered.connect ( lambda : self.action_undo_pt() ) |
| | self.cbar.btn_redo_pt_act.triggered.connect ( lambda : self.action_redo_pt() ) |
| |
|
| | self.cbar.btn_delete_poly_act.triggered.connect ( lambda : self.action_delete_poly() ) |
| |
|
| | self.cbar.btn_pt_edit_mode_act.toggled.connect ( lambda is_checked: self.set_pt_edit_mode( PTEditMode.ADD_DEL if is_checked else PTEditMode.MOVE ) ) |
| | self.cbar.btn_view_lock_center_act.toggled.connect ( lambda is_checked: self.set_view_lock( ViewLock.CENTER if is_checked else ViewLock.NONE ) ) |
| |
|
| | self.mouse_in_widget = False |
| |
|
| | QXMainWindow.inst.add_keyPressEvent_listener ( self.on_keyPressEvent ) |
| | QXMainWindow.inst.add_keyReleaseEvent_listener ( self.on_keyReleaseEvent ) |
| |
|
| | self.qp = QPainter() |
| | self.initialized = False |
| | self.last_state = None |
| |
|
| | def initialize(self, img, img_look_pt=None, view_scale=None, ie_polys=None, xseg_mask=None, canvas_config=None ): |
| | q_img = self.q_img = QImage_from_np(img) |
| | self.img_pixmap = QPixmap.fromImage(q_img) |
| |
|
| | self.xseg_mask_pixmap = None |
| | self.xseg_overlay_mask_pixmap = None |
| | if xseg_mask is not None: |
| | h,w,c = img.shape |
| | xseg_mask = cv2.resize(xseg_mask, (w,h), interpolation=cv2.INTER_CUBIC) |
| | xseg_mask = imagelib.normalize_channels(xseg_mask, 1) |
| | xseg_img = img.astype(np.float32)/255.0 |
| | xseg_overlay_mask = xseg_img*(1-xseg_mask)*0.5 + xseg_img*xseg_mask |
| | xseg_overlay_mask = np.clip(xseg_overlay_mask*255, 0, 255).astype(np.uint8) |
| | xseg_mask = np.clip(xseg_mask*255, 0, 255).astype(np.uint8) |
| | self.xseg_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_mask)) |
| | self.xseg_overlay_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_overlay_mask)) |
| |
|
| | self.img_size = QSize_to_np (self.img_pixmap.size()) |
| |
|
| | self.img_look_pt = img_look_pt |
| | self.view_scale = view_scale |
| |
|
| | if ie_polys is None: |
| | ie_polys = SegIEPolys() |
| | self.ie_polys = ie_polys |
| |
|
| | if canvas_config is None: |
| | canvas_config = CanvasConfig() |
| | self.canvas_config = canvas_config |
| |
|
| | |
| | self.set_cbar_disabled() |
| | self.cbar.btn_poly_color_act_grp.setDisabled(False) |
| | self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(False) |
| | self.cbar.btn_poly_type_act_grp.setDisabled(False) |
| |
|
| | |
| | self.current_cursor = None |
| | self.mouse_hull_poly = None |
| | self.mouse_wire_poly = None |
| | self.drag_type = DragType.NONE |
| | self.mouse_cli_pt = np.zeros((2,), np.float32 ) |
| |
|
| | |
| | self.set_op_mode(OpMode.NONE) |
| | self.set_color_scheme_id(1) |
| | self.set_poly_include_type(SegIEPolyType.INCLUDE) |
| | self.set_pt_edit_mode(PTEditMode.MOVE) |
| | self.set_view_lock(ViewLock.NONE) |
| |
|
| | |
| | if self.last_state is not None: |
| | self.set_color_scheme_id(self.last_state.color_scheme_id) |
| | if self.last_state.op_mode is not None: |
| | self.set_op_mode(self.last_state.op_mode) |
| |
|
| | self.initialized = True |
| |
|
| | self.setMouseTracking(True) |
| | self.update_cursor() |
| | self.update() |
| |
|
| |
|
| | def finalize(self): |
| | if self.initialized: |
| | if self.op_mode == OpMode.DRAW_PTS: |
| | self.set_op_mode(OpMode.EDIT_PTS) |
| |
|
| | self.last_state = sn(op_mode = self.op_mode if self.op_mode in [OpMode.VIEW_BAKED, OpMode.VIEW_XSEG_MASK] else None, |
| | color_scheme_id = self.color_scheme_id) |
| |
|
| | self.img_pixmap = None |
| | self.update_cursor(is_finalize=True) |
| | self.setMouseTracking(False) |
| | self.setFocusPolicy(Qt.NoFocus) |
| | self.set_cbar_disabled() |
| | self.initialized = False |
| | self.update() |
| |
|
| | |
| | |
| | |
| | |
| | |
| | def is_initialized(self): |
| | return self.initialized |
| |
|
| | def get_ie_polys(self): |
| | return self.ie_polys |
| |
|
| | def get_cli_center_pt(self): |
| | return np.round(QSize_to_np(self.size())/2.0) |
| |
|
| | def get_img_look_pt(self): |
| | img_look_pt = self.img_look_pt |
| | if img_look_pt is None: |
| | img_look_pt = self.img_size / 2 |
| | return img_look_pt |
| |
|
| | def get_view_scale(self): |
| | view_scale = self.view_scale |
| | if view_scale is None: |
| | |
| | min_cli_size = np.min(QSize_to_np(self.size())) |
| | max_img_size = np.max(self.img_size) |
| | view_scale = min_cli_size / max_img_size |
| |
|
| | return view_scale |
| |
|
| | def get_current_color_scheme(self): |
| | return self.canvas_config.color_schemes[self.color_scheme_id] |
| |
|
| | def get_poly_pt_id_under_pt(self, poly, cli_pt): |
| | w = np.argwhere ( npla.norm ( cli_pt - self.img_to_cli_pt( poly.get_pts() ), axis=1 ) <= self.canvas_config.pt_select_radius ) |
| | return None if len(w) == 0 else w[-1][0] |
| |
|
| | def get_poly_edge_id_pt_under_pt(self, poly, cli_pt): |
| | cli_pts = self.img_to_cli_pt(poly.get_pts()) |
| | if len(cli_pts) >= 3: |
| | edge_dists, projs = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True) |
| | edge_id = np.argmin(edge_dists) |
| | dist = edge_dists[edge_id] |
| | pt = projs[edge_id] |
| | if dist <= self.canvas_config.pt_select_radius: |
| | return edge_id, pt |
| | return None, None |
| |
|
| | def get_poly_by_pt_near_wire(self, cli_pt): |
| | pt_select_radius = self.canvas_config.pt_select_radius |
| |
|
| | for poly in reversed(self.ie_polys.get_polys()): |
| | pts = poly.get_pts() |
| | if len(pts) >= 3: |
| | cli_pts = self.img_to_cli_pt(pts) |
| |
|
| | edge_dists, _ = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True) |
| |
|
| | if np.min(edge_dists) <= pt_select_radius or \ |
| | any( npla.norm ( cli_pt - cli_pts, axis=1 ) <= pt_select_radius ): |
| | return poly |
| | return None |
| |
|
| | def get_poly_by_pt_in_hull(self, cli_pos): |
| | img_pos = self.cli_to_img_pt(cli_pos) |
| |
|
| | for poly in reversed(self.ie_polys.get_polys()): |
| | pts = poly.get_pts() |
| | if len(pts) >= 3: |
| | if cv2.pointPolygonTest( pts, tuple(img_pos), False) >= 0: |
| | return poly |
| |
|
| | return None |
| |
|
| | def img_to_cli_pt(self, p): |
| | return (p - self.get_img_look_pt()) * self.get_view_scale() + self.get_cli_center_pt() |
| |
|
| | def cli_to_img_pt(self, p): |
| | return (p - self.get_cli_center_pt() ) / self.get_view_scale() + self.get_img_look_pt() |
| |
|
| | def img_to_cli_rect(self, rect): |
| | tl = QPoint_to_np(rect.topLeft()) |
| | xy = self.img_to_cli_pt(tl) |
| | xy2 = self.img_to_cli_pt(tl + QSize_to_np(rect.size()) ) - xy |
| | return QRect ( *xy.astype(np.int), *xy2.astype(np.int) ) |
| |
|
| | |
| | |
| | |
| | |
| | |
| | def set_op_mode(self, op_mode, op_poly=None): |
| | if not hasattr(self,'op_mode'): |
| | self.op_mode = None |
| | self.op_poly = None |
| |
|
| | if self.op_mode != op_mode: |
| | |
| | if self.op_mode == OpMode.NONE: |
| | self.cbar.btn_poly_type_act_grp.setDisabled(True) |
| | elif self.op_mode == OpMode.DRAW_PTS: |
| | self.cbar.btn_undo_pt_act.setDisabled(True) |
| | self.cbar.btn_redo_pt_act.setDisabled(True) |
| | self.cbar.btn_view_lock_center_act.setDisabled(True) |
| | |
| | self.set_view_lock(ViewLock.NONE) |
| | |
| | if self.op_poly.get_pts_count() < 3: |
| | self.ie_polys.remove_poly(self.op_poly) |
| |
|
| | elif self.op_mode == OpMode.EDIT_PTS: |
| | self.cbar.btn_pt_edit_mode_act.setDisabled(True) |
| | self.cbar.btn_delete_poly_act.setDisabled(True) |
| | |
| | self.set_pt_edit_mode(PTEditMode.MOVE) |
| | elif self.op_mode == OpMode.VIEW_BAKED: |
| | self.cbar.btn_view_baked_mask_act.setChecked(False) |
| | elif self.op_mode == OpMode.VIEW_XSEG_MASK: |
| | self.cbar.btn_view_xseg_mask_act.setChecked(False) |
| |
|
| | self.op_mode = op_mode |
| |
|
| | |
| | if op_mode == OpMode.NONE: |
| | self.cbar.btn_poly_type_act_grp.setDisabled(False) |
| | elif op_mode == OpMode.DRAW_PTS: |
| | self.cbar.btn_undo_pt_act.setDisabled(False) |
| | self.cbar.btn_redo_pt_act.setDisabled(False) |
| | self.cbar.btn_view_lock_center_act.setDisabled(False) |
| | elif op_mode == OpMode.EDIT_PTS: |
| | self.cbar.btn_pt_edit_mode_act.setDisabled(False) |
| | self.cbar.btn_delete_poly_act.setDisabled(False) |
| | elif op_mode == OpMode.VIEW_BAKED: |
| | self.cbar.btn_view_baked_mask_act.setChecked(True ) |
| | n = QImage_to_np ( self.q_img ).astype(np.float32) / 255.0 |
| | h,w,c = n.shape |
| | mask = np.zeros( (h,w,1), dtype=np.float32 ) |
| | self.ie_polys.overlay_mask(mask) |
| | n = (mask*255).astype(np.uint8) |
| | self.img_baked_pixmap = QPixmap.fromImage(QImage_from_np(n)) |
| | elif op_mode == OpMode.VIEW_XSEG_MASK: |
| | self.cbar.btn_view_xseg_mask_act.setChecked(True) |
| |
|
| | if op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]: |
| | self.mouse_op_poly_pt_id = None |
| | self.mouse_op_poly_edge_id = None |
| | self.mouse_op_poly_edge_id_pt = None |
| |
|
| | self.op_poly = op_poly |
| | if op_poly is not None: |
| | self.update_mouse_info() |
| |
|
| | self.update_cursor() |
| | self.update() |
| |
|
| | def set_pt_edit_mode(self, pt_edit_mode): |
| | if not hasattr(self, 'pt_edit_mode') or self.pt_edit_mode != pt_edit_mode: |
| | self.pt_edit_mode = pt_edit_mode |
| | self.update_cursor() |
| | self.update() |
| | self.cbar.btn_pt_edit_mode_act.setChecked( self.pt_edit_mode == PTEditMode.ADD_DEL ) |
| |
|
| | def set_view_lock(self, view_lock): |
| | if not hasattr(self, 'view_lock') or self.view_lock != view_lock: |
| | if hasattr(self, 'view_lock') and self.view_lock != view_lock: |
| | if view_lock == ViewLock.CENTER: |
| | self.img_look_pt = self.mouse_img_pt |
| | QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) )) |
| |
|
| | self.view_lock = view_lock |
| | self.update() |
| | self.cbar.btn_view_lock_center_act.setChecked( self.view_lock == ViewLock.CENTER ) |
| |
|
| | def set_cbar_disabled(self): |
| | self.cbar.btn_delete_poly_act.setDisabled(True) |
| | self.cbar.btn_undo_pt_act.setDisabled(True) |
| | self.cbar.btn_redo_pt_act.setDisabled(True) |
| | self.cbar.btn_pt_edit_mode_act.setDisabled(True) |
| | self.cbar.btn_view_lock_center_act.setDisabled(True) |
| | self.cbar.btn_poly_color_act_grp.setDisabled(True) |
| | self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(True) |
| | self.cbar.btn_poly_type_act_grp.setDisabled(True) |
| |
|
| |
|
| | def set_color_scheme_id(self, id): |
| | if self.op_mode == OpMode.VIEW_BAKED or self.op_mode == OpMode.VIEW_XSEG_MASK: |
| | self.set_op_mode(OpMode.NONE) |
| |
|
| | if not hasattr(self, 'color_scheme_id') or self.color_scheme_id != id: |
| | self.color_scheme_id = id |
| | self.update_cursor() |
| | self.update() |
| |
|
| | if self.color_scheme_id == 0: |
| | self.cbar.btn_poly_color_red_act.setChecked( True ) |
| | elif self.color_scheme_id == 1: |
| | self.cbar.btn_poly_color_green_act.setChecked( True ) |
| | elif self.color_scheme_id == 2: |
| | self.cbar.btn_poly_color_blue_act.setChecked( True ) |
| |
|
| | def set_poly_include_type(self, poly_include_type): |
| | if not hasattr(self, 'poly_include_type' ) or \ |
| | ( self.poly_include_type != poly_include_type and \ |
| | self.op_mode in [OpMode.NONE, OpMode.EDIT_PTS] ): |
| | self.poly_include_type = poly_include_type |
| | self.update() |
| | self.cbar.btn_poly_type_include_act.setChecked(self.poly_include_type == SegIEPolyType.INCLUDE) |
| | self.cbar.btn_poly_type_exclude_act.setChecked(self.poly_include_type == SegIEPolyType.EXCLUDE) |
| |
|
| | |
| | |
| | |
| | |
| | |
| |
|
| | def update_cursor(self, is_finalize=False): |
| | if not self.initialized: |
| | return |
| |
|
| | if not self.mouse_in_widget or is_finalize: |
| | if self.current_cursor is not None: |
| | QApplication.restoreOverrideCursor() |
| | self.current_cursor = None |
| | else: |
| | color_cc = self.get_current_color_scheme().cross_cursor |
| | nc = Qt.ArrowCursor |
| |
|
| | if self.drag_type == DragType.IMAGE_LOOK: |
| | nc = Qt.ClosedHandCursor |
| | else: |
| |
|
| | if self.op_mode == OpMode.NONE: |
| | nc = color_cc |
| | if self.mouse_wire_poly is not None: |
| | nc = Qt.PointingHandCursor |
| |
|
| | elif self.op_mode == OpMode.DRAW_PTS: |
| | nc = color_cc |
| | elif self.op_mode == OpMode.EDIT_PTS: |
| | nc = Qt.ArrowCursor |
| |
|
| | if self.mouse_op_poly_pt_id is not None: |
| | nc = Qt.PointingHandCursor |
| |
|
| | if self.pt_edit_mode == PTEditMode.ADD_DEL: |
| |
|
| | if self.mouse_op_poly_edge_id is not None and \ |
| | self.mouse_op_poly_pt_id is None: |
| | nc = color_cc |
| | if self.current_cursor != nc: |
| | if self.current_cursor is None: |
| | QApplication.setOverrideCursor(nc) |
| | else: |
| | QApplication.changeOverrideCursor(nc) |
| | self.current_cursor = nc |
| |
|
| | def update_mouse_info(self, mouse_cli_pt=None): |
| | """ |
| | Update selected polys/edges/points by given mouse position |
| | """ |
| | if mouse_cli_pt is not None: |
| | self.mouse_cli_pt = mouse_cli_pt.astype(np.float32) |
| |
|
| | self.mouse_img_pt = self.cli_to_img_pt(self.mouse_cli_pt) |
| |
|
| | new_mouse_hull_poly = self.get_poly_by_pt_in_hull(self.mouse_cli_pt) |
| |
|
| | if self.mouse_hull_poly != new_mouse_hull_poly: |
| | self.mouse_hull_poly = new_mouse_hull_poly |
| | self.update_cursor() |
| | self.update() |
| |
|
| | new_mouse_wire_poly = self.get_poly_by_pt_near_wire(self.mouse_cli_pt) |
| |
|
| | if self.mouse_wire_poly != new_mouse_wire_poly: |
| | self.mouse_wire_poly = new_mouse_wire_poly |
| | self.update_cursor() |
| | self.update() |
| |
|
| | if self.op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]: |
| | new_mouse_op_poly_pt_id = self.get_poly_pt_id_under_pt (self.op_poly, self.mouse_cli_pt) |
| | if self.mouse_op_poly_pt_id != new_mouse_op_poly_pt_id: |
| | self.mouse_op_poly_pt_id = new_mouse_op_poly_pt_id |
| | self.update_cursor() |
| | self.update() |
| |
|
| | new_mouse_op_poly_edge_id,\ |
| | new_mouse_op_poly_edge_id_pt = self.get_poly_edge_id_pt_under_pt (self.op_poly, self.mouse_cli_pt) |
| | if self.mouse_op_poly_edge_id != new_mouse_op_poly_edge_id: |
| | self.mouse_op_poly_edge_id = new_mouse_op_poly_edge_id |
| | self.update_cursor() |
| | self.update() |
| |
|
| | if (self.mouse_op_poly_edge_id_pt.__class__ != new_mouse_op_poly_edge_id_pt.__class__) or \ |
| | (isinstance(self.mouse_op_poly_edge_id_pt, np.ndarray) and \ |
| | all(self.mouse_op_poly_edge_id_pt != new_mouse_op_poly_edge_id_pt)): |
| |
|
| | self.mouse_op_poly_edge_id_pt = new_mouse_op_poly_edge_id_pt |
| | self.update_cursor() |
| | self.update() |
| |
|
| |
|
| | def action_undo_pt(self): |
| | if self.drag_type == DragType.NONE: |
| | if self.op_mode == OpMode.DRAW_PTS: |
| | if self.op_poly.undo() == 0: |
| | self.ie_polys.remove_poly (self.op_poly) |
| | self.set_op_mode(OpMode.NONE) |
| | self.update() |
| |
|
| | def action_redo_pt(self): |
| | if self.drag_type == DragType.NONE: |
| | if self.op_mode == OpMode.DRAW_PTS: |
| | self.op_poly.redo() |
| | self.update() |
| |
|
| | def action_delete_poly(self): |
| | if self.op_mode == OpMode.EDIT_PTS and \ |
| | self.drag_type == DragType.NONE and \ |
| | self.pt_edit_mode == PTEditMode.MOVE: |
| | |
| | self.ie_polys.remove_poly (self.op_poly) |
| | self.set_op_mode(OpMode.NONE) |
| |
|
| | |
| | |
| | |
| | |
| | |
| | def on_keyPressEvent(self, ev): |
| | if not self.initialized: |
| | return |
| | key = ev.key() |
| | key_mods = int(ev.modifiers()) |
| | if self.op_mode == OpMode.DRAW_PTS: |
| | self.set_view_lock(ViewLock.CENTER if key_mods == Qt.ShiftModifier else ViewLock.NONE ) |
| | elif self.op_mode == OpMode.EDIT_PTS: |
| | self.set_pt_edit_mode(PTEditMode.ADD_DEL if key_mods == Qt.ControlModifier else PTEditMode.MOVE ) |
| |
|
| | def on_keyReleaseEvent(self, ev): |
| | if not self.initialized: |
| | return |
| | key = ev.key() |
| | key_mods = int(ev.modifiers()) |
| | if self.op_mode == OpMode.DRAW_PTS: |
| | self.set_view_lock(ViewLock.CENTER if key_mods == Qt.ShiftModifier else ViewLock.NONE ) |
| | elif self.op_mode == OpMode.EDIT_PTS: |
| | self.set_pt_edit_mode(PTEditMode.ADD_DEL if key_mods == Qt.ControlModifier else PTEditMode.MOVE ) |
| |
|
| | def enterEvent(self, ev): |
| | super().enterEvent(ev) |
| | self.mouse_in_widget = True |
| | self.update_cursor() |
| |
|
| | def leaveEvent(self, ev): |
| | super().leaveEvent(ev) |
| | self.mouse_in_widget = False |
| | self.update_cursor() |
| |
|
| | def mousePressEvent(self, ev): |
| | super().mousePressEvent(ev) |
| | if not self.initialized: |
| | return |
| |
|
| | self.update_mouse_info(QPoint_to_np(ev.pos())) |
| |
|
| | btn = ev.button() |
| |
|
| | if btn == Qt.LeftButton: |
| | if self.op_mode == OpMode.NONE: |
| | |
| | if self.mouse_wire_poly is not None: |
| | |
| | self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.mouse_wire_poly) |
| | else: |
| | |
| | new_poly = self.ie_polys.add_poly(self.poly_include_type) |
| | self.ie_polys.sort() |
| | new_poly.add_pt(*self.mouse_img_pt) |
| | self.set_op_mode(OpMode.DRAW_PTS, op_poly=new_poly ) |
| |
|
| | elif self.op_mode == OpMode.DRAW_PTS: |
| | |
| | if len(self.op_poly.get_pts()) >= 3 and self.mouse_op_poly_pt_id == 0: |
| | |
| | self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.op_poly) |
| | else: |
| | |
| | self.op_poly.add_pt(*self.mouse_img_pt) |
| | self.update() |
| |
|
| | elif self.op_mode == OpMode.EDIT_PTS: |
| | |
| |
|
| | if self.mouse_op_poly_pt_id is not None: |
| | |
| | if self.pt_edit_mode == PTEditMode.ADD_DEL: |
| | |
| | self.op_poly.remove_pt(self.mouse_op_poly_pt_id) |
| | if self.op_poly.get_pts_count() < 3: |
| | |
| | self.ie_polys.remove_poly (self.op_poly) |
| | self.set_op_mode(OpMode.NONE) |
| | self.update() |
| |
|
| | elif self.drag_type == DragType.NONE: |
| | |
| | self.drag_type = DragType.POLY_PT |
| | self.drag_cli_pt = self.mouse_cli_pt |
| | self.drag_poly_pt_id = self.mouse_op_poly_pt_id |
| | self.drag_poly_pt = self.op_poly.get_pts()[ self.drag_poly_pt_id ] |
| | elif self.mouse_op_poly_edge_id is not None: |
| | |
| | if self.pt_edit_mode == PTEditMode.ADD_DEL: |
| | |
| | edge_img_pt = self.cli_to_img_pt(self.mouse_op_poly_edge_id_pt) |
| | self.op_poly.insert_pt (self.mouse_op_poly_edge_id+1, edge_img_pt) |
| | self.update() |
| | else: |
| | |
| | pass |
| | else: |
| | |
| | self.set_op_mode(OpMode.NONE) |
| |
|
| | elif btn == Qt.MiddleButton: |
| | if self.drag_type == DragType.NONE: |
| | |
| | self.drag_type = DragType.IMAGE_LOOK |
| | self.drag_cli_pt = self.mouse_cli_pt |
| | self.drag_img_look_pt = self.get_img_look_pt() |
| | self.update_cursor() |
| |
|
| |
|
| | def mouseReleaseEvent(self, ev): |
| | super().mouseReleaseEvent(ev) |
| | if not self.initialized: |
| | return |
| |
|
| | self.update_mouse_info(QPoint_to_np(ev.pos())) |
| |
|
| | btn = ev.button() |
| |
|
| | if btn == Qt.LeftButton: |
| | if self.op_mode == OpMode.EDIT_PTS: |
| | if self.drag_type == DragType.POLY_PT: |
| | self.drag_type = DragType.NONE |
| | self.update() |
| |
|
| | elif btn == Qt.MiddleButton: |
| | if self.drag_type == DragType.IMAGE_LOOK: |
| | self.drag_type = DragType.NONE |
| | self.update_cursor() |
| | self.update() |
| |
|
| | def mouseMoveEvent(self, ev): |
| | super().mouseMoveEvent(ev) |
| | if not self.initialized: |
| | return |
| |
|
| | prev_mouse_cli_pt = self.mouse_cli_pt |
| | self.update_mouse_info(QPoint_to_np(ev.pos())) |
| |
|
| | if self.view_lock == ViewLock.CENTER: |
| | if npla.norm(self.mouse_cli_pt - prev_mouse_cli_pt) >= 1: |
| | self.img_look_pt = self.mouse_img_pt |
| | QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) )) |
| | self.update() |
| |
|
| | if self.drag_type == DragType.IMAGE_LOOK: |
| | delta_pt = self.cli_to_img_pt(self.mouse_cli_pt) - self.cli_to_img_pt(self.drag_cli_pt) |
| | self.img_look_pt = self.drag_img_look_pt - delta_pt |
| | self.update() |
| |
|
| | if self.op_mode == OpMode.DRAW_PTS: |
| | self.update() |
| | elif self.op_mode == OpMode.EDIT_PTS: |
| | if self.drag_type == DragType.POLY_PT: |
| | delta_pt = self.cli_to_img_pt(self.mouse_cli_pt) - self.cli_to_img_pt(self.drag_cli_pt) |
| | self.op_poly.set_point(self.drag_poly_pt_id, self.drag_poly_pt + delta_pt) |
| | self.update() |
| |
|
| | def wheelEvent(self, ev): |
| | super().wheelEvent(ev) |
| |
|
| | if not self.initialized: |
| | return |
| |
|
| | mods = int(ev.modifiers()) |
| | delta = ev.angleDelta() |
| |
|
| | cli_pt = QPoint_to_np(ev.pos()) |
| |
|
| | if self.drag_type == DragType.NONE: |
| | sign = np.sign( delta.y() ) |
| | prev_img_pos = self.cli_to_img_pt (cli_pt) |
| | delta_scale = sign*0.2 + sign * self.get_view_scale() / 10.0 |
| | self.view_scale = np.clip(self.get_view_scale() + delta_scale, 1.0, 20.0) |
| | new_img_pos = self.cli_to_img_pt (cli_pt) |
| | if sign > 0: |
| | self.img_look_pt = self.get_img_look_pt() + (prev_img_pos-new_img_pos) |
| | else: |
| | QCursor.setPos ( self.mapToGlobal(QPoint_from_np(self.img_to_cli_pt(prev_img_pos))) ) |
| | self.update() |
| |
|
| | def paintEvent(self, event): |
| | super().paintEvent(event) |
| | if not self.initialized: |
| | return |
| |
|
| | qp = self.qp |
| | qp.begin(self) |
| | qp.setRenderHint(QPainter.Antialiasing) |
| | qp.setRenderHint(QPainter.HighQualityAntialiasing) |
| | qp.setRenderHint(QPainter.SmoothPixmapTransform) |
| |
|
| | src_rect = QRect(0, 0, *self.img_size) |
| | dst_rect = self.img_to_cli_rect( src_rect ) |
| |
|
| | if self.op_mode == OpMode.VIEW_BAKED: |
| | qp.drawPixmap(dst_rect, self.img_baked_pixmap, src_rect) |
| | elif self.op_mode == OpMode.VIEW_XSEG_MASK: |
| | if self.xseg_mask_pixmap is not None: |
| | qp.drawPixmap(dst_rect, self.xseg_mask_pixmap, src_rect) |
| | else: |
| | if self.cbar.btn_view_xseg_overlay_mask_act.isChecked() and \ |
| | self.xseg_overlay_mask_pixmap is not None: |
| | qp.drawPixmap(dst_rect, self.xseg_overlay_mask_pixmap, src_rect) |
| | elif self.img_pixmap is not None: |
| | qp.drawPixmap(dst_rect, self.img_pixmap, src_rect) |
| |
|
| | polys = self.ie_polys.get_polys() |
| | polys_len = len(polys) |
| |
|
| | color_scheme = self.get_current_color_scheme() |
| |
|
| | pt_rad = self.canvas_config.pt_radius |
| | pt_rad_x2 = pt_rad*2 |
| |
|
| | pt_select_radius = self.canvas_config.pt_select_radius |
| |
|
| | op_mode = self.op_mode |
| | op_poly = self.op_poly |
| |
|
| | for i,poly in enumerate(polys): |
| |
|
| | selected_pt_path = QPainterPath() |
| | poly_line_path = QPainterPath() |
| | pts_line_path = QPainterPath() |
| |
|
| | pt_remove_cli_pt = None |
| | poly_pts = poly.get_pts() |
| | for pt_id, img_pt in enumerate(poly_pts): |
| | cli_pt = self.img_to_cli_pt(img_pt) |
| | q_cli_pt = QPoint_from_np(cli_pt) |
| |
|
| | if pt_id == 0: |
| | poly_line_path.moveTo(q_cli_pt) |
| | else: |
| | poly_line_path.lineTo(q_cli_pt) |
| |
|
| |
|
| | if poly == op_poly: |
| | if self.op_mode == OpMode.DRAW_PTS or \ |
| | (self.op_mode == OpMode.EDIT_PTS and \ |
| | (self.pt_edit_mode == PTEditMode.MOVE) or \ |
| | (self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id == pt_id) \ |
| | ): |
| | pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([0,-pt_rad])) ) |
| | pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([0,pt_rad])) ) |
| | pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([-pt_rad,0])) ) |
| | pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([pt_rad,0])) ) |
| |
|
| | if (self.op_mode == OpMode.EDIT_PTS and \ |
| | self.pt_edit_mode == PTEditMode.ADD_DEL and \ |
| | self.mouse_op_poly_pt_id == pt_id): |
| | pt_remove_cli_pt = cli_pt |
| |
|
| | if self.op_mode == OpMode.DRAW_PTS and \ |
| | len(op_poly.get_pts()) >= 3 and pt_id == 0 and self.mouse_op_poly_pt_id == pt_id: |
| | |
| | selected_pt_path.addEllipse(q_cli_pt, pt_rad_x2, pt_rad_x2) |
| |
|
| |
|
| | if poly == op_poly: |
| | if op_mode == OpMode.DRAW_PTS: |
| | |
| | poly_line_path.lineTo( QPoint_from_np(self.mouse_cli_pt) ) |
| |
|
| | if self.mouse_op_poly_pt_id is not None: |
| | pass |
| |
|
| | if self.mouse_op_poly_edge_id_pt is not None: |
| | if self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id is None: |
| | |
| | m_cli_pt = self.mouse_op_poly_edge_id_pt |
| | pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([0,-pt_rad])) ) |
| | pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([0,pt_rad])) ) |
| | pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([-pt_rad,0])) ) |
| | pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([pt_rad,0])) ) |
| |
|
| | if len(poly_pts) >= 2: |
| | |
| | poly_line_path.lineTo( QPoint_from_np(self.img_to_cli_pt(poly_pts[0])) ) |
| |
|
| | |
| | qp.setPen(color_scheme.pt_outline_pen) |
| | qp.setBrush(QBrush()) |
| | qp.drawPath(selected_pt_path) |
| |
|
| | qp.setPen(color_scheme.poly_outline_solid_pen) |
| | qp.setBrush(QBrush()) |
| | qp.drawPath(pts_line_path) |
| |
|
| | if poly.get_type() == SegIEPolyType.INCLUDE: |
| | qp.setPen(color_scheme.poly_outline_solid_pen) |
| | else: |
| | qp.setPen(color_scheme.poly_outline_dot_pen) |
| |
|
| | qp.setBrush(color_scheme.poly_unselected_brush) |
| | if op_mode == OpMode.NONE: |
| | if poly == self.mouse_wire_poly: |
| | qp.setBrush(color_scheme.poly_selected_brush) |
| | |
| | |
| | |
| |
|
| | qp.drawPath(poly_line_path) |
| |
|
| | if pt_remove_cli_pt is not None: |
| | qp.setPen(color_scheme.poly_outline_solid_pen) |
| | qp.setBrush(QBrush()) |
| |
|
| | qp.drawLine( *(pt_remove_cli_pt + np.float32([-pt_rad_x2,-pt_rad_x2])), *(pt_remove_cli_pt + np.float32([pt_rad_x2,pt_rad_x2])) ) |
| | qp.drawLine( *(pt_remove_cli_pt + np.float32([-pt_rad_x2,pt_rad_x2])), *(pt_remove_cli_pt + np.float32([pt_rad_x2,-pt_rad_x2])) ) |
| |
|
| | qp.end() |
| |
|
| | class QCanvas(QFrame): |
| | def __init__(self): |
| | super().__init__() |
| |
|
| | self.canvas_control_left_bar = QCanvasControlsLeftBar() |
| | self.canvas_control_right_bar = QCanvasControlsRightBar() |
| |
|
| | cbar = sn( btn_poly_color_red_act = self.canvas_control_right_bar.btn_poly_color_red_act, |
| | btn_poly_color_green_act = self.canvas_control_right_bar.btn_poly_color_green_act, |
| | btn_poly_color_blue_act = self.canvas_control_right_bar.btn_poly_color_blue_act, |
| | btn_view_baked_mask_act = self.canvas_control_right_bar.btn_view_baked_mask_act, |
| | btn_view_xseg_mask_act = self.canvas_control_right_bar.btn_view_xseg_mask_act, |
| | btn_view_xseg_overlay_mask_act = self.canvas_control_right_bar.btn_view_xseg_overlay_mask_act, |
| | btn_poly_color_act_grp = self.canvas_control_right_bar.btn_poly_color_act_grp, |
| | btn_view_lock_center_act = self.canvas_control_right_bar.btn_view_lock_center_act, |
| |
|
| | btn_poly_type_include_act = self.canvas_control_left_bar.btn_poly_type_include_act, |
| | btn_poly_type_exclude_act = self.canvas_control_left_bar.btn_poly_type_exclude_act, |
| | btn_poly_type_act_grp = self.canvas_control_left_bar.btn_poly_type_act_grp, |
| | btn_undo_pt_act = self.canvas_control_left_bar.btn_undo_pt_act, |
| | btn_redo_pt_act = self.canvas_control_left_bar.btn_redo_pt_act, |
| | btn_delete_poly_act = self.canvas_control_left_bar.btn_delete_poly_act, |
| | btn_pt_edit_mode_act = self.canvas_control_left_bar.btn_pt_edit_mode_act ) |
| |
|
| | self.op = QCanvasOperator(cbar) |
| | self.l = QHBoxLayout() |
| | self.l.setContentsMargins(0,0,0,0) |
| | self.l.addWidget(self.canvas_control_left_bar) |
| | self.l.addWidget(self.op) |
| | self.l.addWidget(self.canvas_control_right_bar) |
| | self.setLayout(self.l) |
| |
|
| | class LoaderQSubprocessor(QSubprocessor): |
| | def __init__(self, image_paths, q_label, q_progressbar, on_finish_func ): |
| |
|
| | self.image_paths = image_paths |
| | self.image_paths_len = len(image_paths) |
| | self.idxs = [*range(self.image_paths_len)] |
| |
|
| | self.filtered_image_paths = self.image_paths.copy() |
| |
|
| | self.image_paths_has_ie_polys = { image_path : False for image_path in self.image_paths } |
| |
|
| | self.q_label = q_label |
| | self.q_progressbar = q_progressbar |
| | self.q_progressbar.setRange(0, self.image_paths_len) |
| | self.q_progressbar.setValue(0) |
| | self.q_progressbar.update() |
| | self.on_finish_func = on_finish_func |
| | self.done_count = 0 |
| | super().__init__('LoaderQSubprocessor', LoaderQSubprocessor.Cli, 60) |
| |
|
| | def get_data(self, host_dict): |
| | if len (self.idxs) > 0: |
| | idx = self.idxs.pop(0) |
| | image_path = self.image_paths[idx] |
| | self.q_label.setText(f'{QStringDB.loading_tip}... {image_path.name}') |
| |
|
| | return idx, image_path |
| |
|
| | return None |
| |
|
| | def on_clients_finalized(self): |
| | self.on_finish_func([x for x in self.filtered_image_paths if x is not None], self.image_paths_has_ie_polys) |
| |
|
| | def on_data_return (self, host_dict, data): |
| | self.idxs.insert(0, data[0]) |
| |
|
| | def on_result (self, host_dict, data, result): |
| | idx, has_dflimg, has_ie_polys = result |
| |
|
| | if not has_dflimg: |
| | self.filtered_image_paths[idx] = None |
| | self.image_paths_has_ie_polys[self.image_paths[idx]] = has_ie_polys |
| |
|
| | self.done_count += 1 |
| | if self.q_progressbar is not None: |
| | self.q_progressbar.setValue(self.done_count) |
| |
|
| | class Cli(QSubprocessor.Cli): |
| | def process_data(self, data): |
| | idx, filename = data |
| | dflimg = DFLIMG.load(filename) |
| | if dflimg is not None and dflimg.has_data(): |
| | ie_polys = dflimg.get_seg_ie_polys() |
| |
|
| | return idx, True, ie_polys.has_polys() |
| | return idx, False, False |
| |
|
| | class MainWindow(QXMainWindow): |
| |
|
| | def __init__(self, input_dirpath, cfg_root_path): |
| | self.loading_frame = None |
| | self.help_frame = None |
| |
|
| | super().__init__() |
| |
|
| | self.input_dirpath = input_dirpath |
| | self.trash_dirpath = input_dirpath.parent / (input_dirpath.name + '_trash') |
| | self.cfg_root_path = cfg_root_path |
| |
|
| | self.cfg_path = cfg_root_path / 'MainWindow_cfg.dat' |
| | self.cfg_dict = pickle.loads(self.cfg_path.read_bytes()) if self.cfg_path.exists() else {} |
| |
|
| | self.cached_images = {} |
| | self.cached_has_ie_polys = {} |
| |
|
| | self.initialize_ui() |
| |
|
| | |
| | self.loading_frame = QFrame(self.main_canvas_frame) |
| | self.loading_frame.setAutoFillBackground(True) |
| | self.loading_frame.setFrameShape(QFrame.StyledPanel) |
| | self.loader_label = QLabel() |
| | self.loader_progress_bar = QProgressBar() |
| |
|
| | intro_image = QLabel() |
| | intro_image.setPixmap( QPixmap.fromImage(QImageDB.intro) ) |
| |
|
| | intro_image_frame_l = QVBoxLayout() |
| | intro_image_frame_l.addWidget(intro_image, alignment=Qt.AlignCenter) |
| | intro_image_frame = QFrame() |
| | intro_image_frame.setSizePolicy (QSizePolicy.Expanding, QSizePolicy.Expanding) |
| | intro_image_frame.setLayout(intro_image_frame_l) |
| |
|
| | loading_frame_l = QVBoxLayout() |
| | loading_frame_l.addWidget (intro_image_frame) |
| | loading_frame_l.addWidget (self.loader_label) |
| | loading_frame_l.addWidget (self.loader_progress_bar) |
| | self.loading_frame.setLayout(loading_frame_l) |
| |
|
| | self.loader_subprocessor = LoaderQSubprocessor( image_paths=pathex.get_image_paths(input_dirpath, return_Path_class=True), |
| | q_label=self.loader_label, |
| | q_progressbar=self.loader_progress_bar, |
| | on_finish_func=self.on_loader_finish ) |
| |
|
| |
|
| | def on_loader_finish(self, image_paths, image_paths_has_ie_polys): |
| | self.image_paths_done = [] |
| | self.image_paths = image_paths |
| | self.image_paths_has_ie_polys = image_paths_has_ie_polys |
| | self.set_has_ie_polys_count ( len([ 1 for x in self.image_paths_has_ie_polys if self.image_paths_has_ie_polys[x] == True]) ) |
| | self.loading_frame.hide() |
| | self.loading_frame = None |
| |
|
| | self.process_next_image(first_initialization=True) |
| |
|
| | def closeEvent(self, ev): |
| | self.cfg_dict['geometry'] = self.saveGeometry().data() |
| | self.cfg_path.write_bytes( pickle.dumps(self.cfg_dict) ) |
| |
|
| |
|
| | def update_cached_images (self, count=5): |
| | d = self.cached_images |
| |
|
| | for image_path in self.image_paths_done[:-count]+self.image_paths[count:]: |
| | if image_path in d: |
| | del d[image_path] |
| |
|
| | for image_path in self.image_paths[:count]+self.image_paths_done[-count:]: |
| | if image_path not in d: |
| | img = cv2_imread(image_path) |
| | if img is not None: |
| | d[image_path] = img |
| |
|
| | def load_image(self, image_path): |
| | try: |
| | img = self.cached_images.get(image_path, None) |
| | if img is None: |
| | img = cv2_imread(image_path) |
| | self.cached_images[image_path] = img |
| | if img is None: |
| | io.log_err(f'Unable to load {image_path}') |
| | except: |
| | img = None |
| |
|
| | return img |
| |
|
| | def update_preview_bar(self): |
| | count = self.image_bar.get_preview_images_count() |
| | d = self.cached_images |
| | prev_imgs = [ d.get(image_path, None) for image_path in self.image_paths_done[-1:-count:-1] ] |
| | next_imgs = [ d.get(image_path, None) for image_path in self.image_paths[:count] ] |
| | self.image_bar.update_images(prev_imgs, next_imgs) |
| |
|
| |
|
| | def canvas_initialize(self, image_path, only_has_polys=False): |
| | if only_has_polys and not self.image_paths_has_ie_polys[image_path]: |
| | return False |
| |
|
| | dflimg = DFLIMG.load(image_path) |
| | if not dflimg or not dflimg.has_data(): |
| | return False |
| |
|
| | ie_polys = dflimg.get_seg_ie_polys() |
| | xseg_mask = dflimg.get_xseg_mask() |
| | img = self.load_image(image_path) |
| | if img is None: |
| | return False |
| |
|
| | self.canvas.op.initialize ( img, ie_polys=ie_polys, xseg_mask=xseg_mask ) |
| |
|
| | self.filename_label.setText(f"{image_path.name}") |
| |
|
| | return True |
| |
|
| | def canvas_finalize(self, image_path): |
| | self.canvas.op.finalize() |
| |
|
| | if image_path.exists(): |
| | dflimg = DFLIMG.load(image_path) |
| | ie_polys = dflimg.get_seg_ie_polys() |
| | new_ie_polys = self.canvas.op.get_ie_polys() |
| |
|
| | if not new_ie_polys.identical(ie_polys): |
| | prev_has_polys = self.image_paths_has_ie_polys[image_path] |
| | self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys() |
| | new_has_polys = self.image_paths_has_ie_polys[image_path] |
| |
|
| | if not prev_has_polys and new_has_polys: |
| | self.set_has_ie_polys_count ( self.get_has_ie_polys_count() +1) |
| | elif prev_has_polys and not new_has_polys: |
| | self.set_has_ie_polys_count ( self.get_has_ie_polys_count() -1) |
| |
|
| | dflimg.set_seg_ie_polys( new_ie_polys ) |
| | dflimg.save() |
| |
|
| | self.filename_label.setText(f"") |
| |
|
| | def process_prev_image(self): |
| | key_mods = QApplication.keyboardModifiers() |
| | step = 5 if key_mods == Qt.ShiftModifier else 1 |
| | only_has_polys = key_mods == Qt.ControlModifier |
| |
|
| | if self.canvas.op.is_initialized(): |
| | self.canvas_finalize(self.image_paths[0]) |
| |
|
| | while True: |
| | for _ in range(step): |
| | if len(self.image_paths_done) != 0: |
| | self.image_paths.insert (0, self.image_paths_done.pop(-1)) |
| | else: |
| | break |
| | if len(self.image_paths) == 0: |
| | break |
| |
|
| | ret = self.canvas_initialize(self.image_paths[0], len(self.image_paths_done) != 0 and only_has_polys) |
| |
|
| | if ret or len(self.image_paths_done) == 0: |
| | break |
| |
|
| | self.update_cached_images() |
| | self.update_preview_bar() |
| |
|
| | def process_next_image(self, first_initialization=False): |
| | key_mods = QApplication.keyboardModifiers() |
| |
|
| | step = 0 if first_initialization else 5 if key_mods == Qt.ShiftModifier else 1 |
| | only_has_polys = False if first_initialization else key_mods == Qt.ControlModifier |
| |
|
| | if self.canvas.op.is_initialized(): |
| | self.canvas_finalize(self.image_paths[0]) |
| |
|
| | while True: |
| | for _ in range(step): |
| | if len(self.image_paths) != 0: |
| | self.image_paths_done.append(self.image_paths.pop(0)) |
| | else: |
| | break |
| | if len(self.image_paths) == 0: |
| | break |
| | if self.canvas_initialize(self.image_paths[0], only_has_polys): |
| | break |
| |
|
| | self.update_cached_images() |
| | self.update_preview_bar() |
| | |
| | def trash_current_image(self): |
| | self.process_next_image() |
| | |
| | img_path = self.image_paths_done.pop(-1) |
| | img_path = Path(img_path) |
| | self.trash_dirpath.mkdir(parents=True, exist_ok=True) |
| | img_path.rename( self.trash_dirpath / img_path.name ) |
| | |
| | self.update_cached_images() |
| | self.update_preview_bar() |
| | |
| | def initialize_ui(self): |
| |
|
| | self.canvas = QCanvas() |
| |
|
| | image_bar = self.image_bar = ImagePreviewSequenceBar(preview_images_count=9, icon_size=QUIConfig.preview_bar_icon_q_size.width()) |
| | image_bar.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed ) |
| |
|
| |
|
| | btn_prev_image = QXIconButton(QIconDB.left, QStringDB.btn_prev_image_tip, shortcut='A', click_func=self.process_prev_image) |
| | btn_prev_image.setIconSize(QUIConfig.preview_bar_icon_q_size) |
| |
|
| | btn_next_image = QXIconButton(QIconDB.right, QStringDB.btn_next_image_tip, shortcut='D', click_func=self.process_next_image) |
| | btn_next_image.setIconSize(QUIConfig.preview_bar_icon_q_size) |
| |
|
| | btn_delete_image = QXIconButton(QIconDB.trashcan, QStringDB.btn_delete_image_tip, shortcut='X', click_func=self.trash_current_image) |
| | btn_delete_image.setIconSize(QUIConfig.preview_bar_icon_q_size) |
| | |
| | pad_image = QWidget() |
| | pad_image.setFixedSize(QUIConfig.preview_bar_icon_q_size) |
| | |
| | preview_image_bar_frame_l = QHBoxLayout() |
| | preview_image_bar_frame_l.setContentsMargins(0,0,0,0) |
| | preview_image_bar_frame_l.addWidget ( pad_image, alignment=Qt.AlignCenter) |
| | preview_image_bar_frame_l.addWidget ( btn_prev_image, alignment=Qt.AlignCenter) |
| | preview_image_bar_frame_l.addWidget ( image_bar) |
| | preview_image_bar_frame_l.addWidget ( btn_next_image, alignment=Qt.AlignCenter) |
| | |
| |
|
| | preview_image_bar_frame = QFrame() |
| | preview_image_bar_frame.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed ) |
| | preview_image_bar_frame.setLayout(preview_image_bar_frame_l) |
| |
|
| | preview_image_bar_frame2_l = QHBoxLayout() |
| | preview_image_bar_frame2_l.setContentsMargins(0,0,0,0) |
| | preview_image_bar_frame2_l.addWidget ( btn_delete_image, alignment=Qt.AlignCenter) |
| |
|
| | preview_image_bar_frame2 = QFrame() |
| | preview_image_bar_frame2.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed ) |
| | preview_image_bar_frame2.setLayout(preview_image_bar_frame2_l) |
| | |
| | preview_image_bar_l = QHBoxLayout() |
| | preview_image_bar_l.addWidget (preview_image_bar_frame, alignment=Qt.AlignCenter) |
| | preview_image_bar_l.addWidget (preview_image_bar_frame2) |
| | |
| | preview_image_bar = QFrame() |
| | preview_image_bar.setFrameShape(QFrame.StyledPanel) |
| | preview_image_bar.setSizePolicy ( QSizePolicy.Expanding, QSizePolicy.Fixed ) |
| | preview_image_bar.setLayout(preview_image_bar_l) |
| |
|
| | label_font = QFont('Courier New') |
| | self.filename_label = QLabel() |
| | self.filename_label.setFont(label_font) |
| |
|
| | self.has_ie_polys_count_label = QLabel() |
| |
|
| | status_frame_l = QHBoxLayout() |
| | status_frame_l.setContentsMargins(0,0,0,0) |
| | status_frame_l.addWidget ( QLabel(), alignment=Qt.AlignCenter) |
| | status_frame_l.addWidget (self.filename_label, alignment=Qt.AlignCenter) |
| | status_frame_l.addWidget (self.has_ie_polys_count_label, alignment=Qt.AlignCenter) |
| | status_frame = QFrame() |
| | status_frame.setLayout(status_frame_l) |
| |
|
| | main_canvas_l = QVBoxLayout() |
| | main_canvas_l.setContentsMargins(0,0,0,0) |
| | main_canvas_l.addWidget (self.canvas) |
| | main_canvas_l.addWidget (status_frame) |
| | main_canvas_l.addWidget (preview_image_bar) |
| |
|
| | self.main_canvas_frame = QFrame() |
| | self.main_canvas_frame.setLayout(main_canvas_l) |
| |
|
| | self.main_l = QHBoxLayout() |
| | self.main_l.setContentsMargins(0,0,0,0) |
| | self.main_l.addWidget (self.main_canvas_frame) |
| |
|
| | self.setLayout(self.main_l) |
| |
|
| | geometry = self.cfg_dict.get('geometry', None) |
| | if geometry is not None: |
| | self.restoreGeometry(geometry) |
| | else: |
| | self.move( QPoint(0,0)) |
| |
|
| | def get_has_ie_polys_count(self): |
| | return self.has_ie_polys_count |
| |
|
| | def set_has_ie_polys_count(self, c): |
| | self.has_ie_polys_count = c |
| | self.has_ie_polys_count_label.setText(f"{c} {QStringDB.labeled_tip}") |
| |
|
| | def resizeEvent(self, ev): |
| | if self.loading_frame is not None: |
| | self.loading_frame.resize( ev.size() ) |
| | if self.help_frame is not None: |
| | self.help_frame.resize( ev.size() ) |
| |
|
| | def start(input_dirpath): |
| | """ |
| | returns exit_code |
| | """ |
| | io.log_info("Running XSeg editor.") |
| |
|
| | if PackedFaceset.path_contains(input_dirpath): |
| | io.log_info (f'\n{input_dirpath} contains packed faceset! Unpack it first.\n') |
| | return 1 |
| |
|
| | root_path = Path(__file__).parent |
| | cfg_root_path = Path(tempfile.gettempdir()) |
| |
|
| | QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True) |
| | QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True) |
| |
|
| | app = QApplication([]) |
| | app.setApplicationName("XSegEditor") |
| | app.setStyle('Fusion') |
| |
|
| | QFontDatabase.addApplicationFont( str(root_path / 'gfx' / 'fonts' / 'NotoSans-Medium.ttf') ) |
| |
|
| | app.setFont( QFont('NotoSans')) |
| |
|
| | QUIConfig.initialize() |
| | QStringDB.initialize() |
| |
|
| | QIconDB.initialize( root_path / 'gfx' / 'icons' ) |
| | QCursorDB.initialize( root_path / 'gfx' / 'cursors' ) |
| | QImageDB.initialize( root_path / 'gfx' / 'images' ) |
| |
|
| | app.setWindowIcon(QIconDB.app_icon) |
| | app.setPalette( QDarkPalette() ) |
| |
|
| | win = MainWindow( input_dirpath=input_dirpath, cfg_root_path=cfg_root_path) |
| |
|
| | win.show() |
| | win.raise_() |
| |
|
| | app.exec_() |
| | return 0 |
| |
|
