mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2024-11-22 15:50:10 -08:00
1495 lines
63 KiB
Python
1495 lines
63 KiB
Python
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)
|
|
#main_frame_l_cont_hl.setSpacing(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):
|
|
# Fix arrays
|
|
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
|
|
|
|
# UI init
|
|
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)
|
|
|
|
# Initial vars
|
|
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 )
|
|
|
|
# Initial state
|
|
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)
|
|
|
|
# Apply last state
|
|
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()
|
|
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
# ====================================== GETTERS =====================================
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
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:
|
|
# Calc as scale to fit
|
|
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()# QSize_to_np(self.size())/2.0
|
|
|
|
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) )
|
|
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
# ====================================== SETTERS =====================================
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
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:
|
|
# Finalize prev 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)
|
|
# Reset view_lock when exit from DRAW_PTS
|
|
self.set_view_lock(ViewLock.NONE)
|
|
# Remove unfinished poly
|
|
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)
|
|
# Reset pt_edit_move when exit from EDIT_PTS
|
|
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
|
|
|
|
# Initialize new 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)
|
|
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
# ====================================== METHODS =====================================
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
|
|
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:
|
|
# Delete current poly
|
|
self.ie_polys.remove_poly (self.op_poly)
|
|
self.set_op_mode(OpMode.NONE)
|
|
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
# ================================== OVERRIDE QT METHODS =============================
|
|
# ====================================================================================
|
|
# ====================================================================================
|
|
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:
|
|
# Clicking in NO OPERATION mode
|
|
if self.mouse_wire_poly is not None:
|
|
# Click on wire on any poly -> switch to EDIT_MODE
|
|
self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.mouse_wire_poly)
|
|
else:
|
|
# Click on empty space -> create new poly with one point
|
|
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:
|
|
# Clicking in DRAW_PTS mode
|
|
if len(self.op_poly.get_pts()) >= 3 and self.mouse_op_poly_pt_id == 0:
|
|
# Click on first point -> close poly and switch to edit mode
|
|
self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.op_poly)
|
|
else:
|
|
# Click on empty space -> add point to current poly
|
|
self.op_poly.add_pt(*self.mouse_img_pt)
|
|
self.update()
|
|
|
|
elif self.op_mode == OpMode.EDIT_PTS:
|
|
# Clicking in EDIT_PTS mode
|
|
|
|
if self.mouse_op_poly_pt_id is not None:
|
|
# Click on point of op_poly
|
|
if self.pt_edit_mode == PTEditMode.ADD_DEL:
|
|
# in mode 'delete point'
|
|
self.op_poly.remove_pt(self.mouse_op_poly_pt_id)
|
|
if self.op_poly.get_pts_count() < 3:
|
|
# not enough points after delete -> remove poly
|
|
self.ie_polys.remove_poly (self.op_poly)
|
|
self.set_op_mode(OpMode.NONE)
|
|
self.update()
|
|
|
|
elif self.drag_type == DragType.NONE:
|
|
# otherwise -> start drag
|
|
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:
|
|
# Click on edge of op_poly
|
|
if self.pt_edit_mode == PTEditMode.ADD_DEL:
|
|
# in mode 'insert new point'
|
|
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:
|
|
# Otherwise do nothing
|
|
pass
|
|
else:
|
|
# other cases -> unselect poly
|
|
self.set_op_mode(OpMode.NONE)
|
|
|
|
elif btn == Qt.MiddleButton:
|
|
if self.drag_type == DragType.NONE:
|
|
# Start image drag
|
|
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)#*1.5
|
|
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:
|
|
# Circle around poly point
|
|
selected_pt_path.addEllipse(q_cli_pt, pt_rad_x2, pt_rad_x2)
|
|
|
|
|
|
if poly == op_poly:
|
|
if op_mode == OpMode.DRAW_PTS:
|
|
# Line from last point to mouse
|
|
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:
|
|
# Ready to insert point on edge
|
|
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:
|
|
# Closing poly line
|
|
poly_line_path.lineTo( QPoint_from_np(self.img_to_cli_pt(poly_pts[0])) )
|
|
|
|
# Draw calls
|
|
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)
|
|
#else:
|
|
# if poly == op_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()
|
|
|
|
# Loader
|
|
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_l.addWidget ( btn_delete_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
|