diff --git a/XSegEditor/QIconDB.py b/XSegEditor/QIconDB.py index 064efad..f4be98b 100644 --- a/XSegEditor/QIconDB.py +++ b/XSegEditor/QIconDB.py @@ -18,5 +18,6 @@ class QIconDB(): QIconDB.left = QIcon ( str(icon_path / 'left.png') ) QIconDB.right = QIcon ( str(icon_path / 'right.png') ) QIconDB.pt_edit_mode = QIcon ( str(icon_path / 'pt_edit_mode.png') ) + QIconDB.view_lock_center = QIcon ( str(icon_path / 'view_lock_center.png') ) QIconDB.view_baked = QIcon ( str(icon_path / 'view_baked.png') ) QIconDB.view_xseg = QIcon ( str(icon_path / 'view_xseg.png') ) \ No newline at end of file diff --git a/XSegEditor/QStringDB.py b/XSegEditor/QStringDB.py index 02617d0..c825a55 100644 --- a/XSegEditor/QStringDB.py +++ b/XSegEditor/QStringDB.py @@ -60,11 +60,17 @@ class QStringDB(): 'zh' : '删除选区', }[lang] - QStringDB.btn_pt_edit_mode_tip = { 'en' : 'Edit point mode ( HOLD CTRL )', - 'ru' : 'Режим правки точек', - 'zh' : '编辑点模式 ( 按住CTRL )', + QStringDB.btn_pt_edit_mode_tip = { 'en' : 'Add/delete point mode ( HOLD CTRL )', + 'ru' : 'Режим добавления/удаления точек ( удерживайте CTRL )', + 'zh' : '点加/删除模式 ( 按住CTRL )', }[lang] + QStringDB.btn_view_lock_center_tip = { 'en' : 'Lock cursor at the center ( HOLD SHIFT )', + 'ru' : 'Заблокировать курсор в центре ( удерживайте SHIFT )', + 'zh' : '将光标锁定在中心 ( 按住SHIFT )', + }[lang] + + QStringDB.btn_prev_image_tip = { 'en' : 'Save and Prev image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n', 'ru' : 'Сохранить и предыдущее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n', 'zh' : '保存并转到上一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n', diff --git a/XSegEditor/XSegEditor.py b/XSegEditor/XSegEditor.py index 025eb36..d40f4dc 100644 --- a/XSegEditor/XSegEditor.py +++ b/XSegEditor/XSegEditor.py @@ -16,13 +16,13 @@ from PyQt5.QtCore import * from PyQt5.QtGui import * from PyQt5.QtWidgets import * -from core import pathex +from core import imagelib, pathex from core.cv2ex import * -from core import imagelib 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 @@ -45,6 +45,10 @@ class DragType(IntEnum): 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): @@ -187,6 +191,7 @@ class QCanvasControlsLeftBar(QFrame): 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 ) @@ -259,6 +264,10 @@ class QCanvasControlsRightBar(QFrame): 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_frame1_l = QVBoxLayout() controls_bar_frame1_l.addWidget ( btn_poly_color_red ) @@ -271,9 +280,17 @@ class QCanvasControlsRightBar(QFrame): controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) controls_bar_frame1.setLayout(controls_bar_frame1_l) + controls_bar_frame2_l = QVBoxLayout() + controls_bar_frame2_l.addWidget ( btn_view_lock_center ) + 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_l = QVBoxLayout() controls_bar_l.setContentsMargins(0,0,0,0) controls_bar_l.addWidget(controls_bar_frame1) + controls_bar_l.addWidget(controls_bar_frame2) self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding ) self.setLayout(controls_bar_l) @@ -300,6 +317,7 @@ class QCanvasOperator(QWidget): 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 @@ -313,14 +331,14 @@ class QCanvasOperator(QWidget): def initialize(self, q_img, img_look_pt=None, view_scale=None, ie_polys=None, xseg_mask=None, canvas_config=None ): self.q_img = q_img self.img_pixmap = QPixmap.fromImage(q_img) - + self.xseg_mask_pixmap = None if xseg_mask is not None: w,h = QSize_to_np ( q_img.size() ) xseg_mask = cv2.resize(xseg_mask, (w,h), cv2.INTER_CUBIC) xseg_mask = (imagelib.normalize_channels(xseg_mask, 1) * 255).astype(np.uint8) self.xseg_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_mask)) - + self.img_size = QSize_to_np (self.img_pixmap.size()) self.img_look_pt = img_look_pt @@ -333,46 +351,48 @@ class QCanvasOperator(QWidget): 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_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_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) @@ -393,6 +413,9 @@ class QCanvasOperator(QWidget): 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: @@ -454,10 +477,10 @@ class QCanvasOperator(QWidget): return None def img_to_cli_pt(self, p): - return (p - self.get_img_look_pt()) * self.get_view_scale() + QSize_to_np(self.size())/2.0 + 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 - QSize_to_np(self.size())/2.0 ) / self.get_view_scale() + self.get_img_look_pt() + 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()) @@ -474,7 +497,7 @@ class QCanvasOperator(QWidget): 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: @@ -482,9 +505,13 @@ class QCanvasOperator(QWidget): 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: - # Remove unfinished poly 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) @@ -496,18 +523,19 @@ class QCanvasOperator(QWidget): 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 ) + 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 ) @@ -515,12 +543,12 @@ class QCanvasOperator(QWidget): 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) + 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() @@ -533,14 +561,25 @@ class QCanvasOperator(QWidget): 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_poly_type_act_grp.setDisabled(True) @@ -573,16 +612,6 @@ class QCanvasOperator(QWidget): def set_view_xseg_mask(self, is_checked): if is_checked: self.set_op_mode(OpMode.VIEW_XSEG_MASK) - - #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)) else: self.set_op_mode(OpMode.NONE) @@ -714,7 +743,9 @@ class QCanvasOperator(QWidget): return key = ev.key() key_mods = int(ev.modifiers()) - if self.op_mode == OpMode.EDIT_PTS: + 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): @@ -722,7 +753,9 @@ class QCanvasOperator(QWidget): return key = ev.key() key_mods = int(ev.modifiers()) - if self.op_mode == OpMode.EDIT_PTS: + 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): @@ -836,8 +869,16 @@ class QCanvasOperator(QWidget): 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 @@ -846,9 +887,7 @@ class QCanvasOperator(QWidget): 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() @@ -1010,7 +1049,6 @@ class QCanvasOperator(QWidget): qp.end() - class QCanvas(QFrame): def __init__(self): super().__init__() @@ -1024,16 +1062,14 @@ class QCanvas(QFrame): 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_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) @@ -1188,13 +1224,13 @@ class MainWindow(QXMainWindow): 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() q_img = self.load_QImage(image_path) if q_img is None: return False - + self.canvas.op.initialize ( q_img, ie_polys=ie_polys, xseg_mask=xseg_mask ) self.filename_label.setText(str(image_path.name)) @@ -1203,17 +1239,17 @@ class MainWindow(QXMainWindow): 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() + 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): self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys() dflimg.set_seg_ie_polys( new_ie_polys ) dflimg.save() - + self.filename_label.setText("") def process_prev_image(self): @@ -1234,7 +1270,7 @@ class MainWindow(QXMainWindow): 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 @@ -1327,8 +1363,15 @@ class MainWindow(QXMainWindow): self.loading_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()) @@ -1358,3 +1401,4 @@ def start(input_dirpath): win.raise_() app.exec_() + return 0 diff --git a/XSegEditor/gfx/icons/view_lock_center.png b/XSegEditor/gfx/icons/view_lock_center.png new file mode 100644 index 0000000..2a10408 Binary files /dev/null and b/XSegEditor/gfx/icons/view_lock_center.png differ diff --git a/main.py b/main.py index 922726d..36e79b1 100644 --- a/main.py +++ b/main.py @@ -22,6 +22,8 @@ if __name__ == "__main__": def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values))) + exit_code = 0 + parser = argparse.ArgumentParser() subparsers = parser.add_subparsers() @@ -262,7 +264,9 @@ if __name__ == "__main__": def process_xsegeditor(arguments): osex.set_process_lowest_prio() from XSegEditor import XSegEditor - XSegEditor.start (Path(arguments.input_dir)) + global exit_code + exit_code = XSegEditor.start (Path(arguments.input_dir)) + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.set_defaults (func=process_xsegeditor) @@ -313,7 +317,10 @@ if __name__ == "__main__": arguments = parser.parse_args() arguments.func(arguments) - print ("Done.") + if exit_code == 0: + print ("Done.") + + exit(exit_code) ''' import code diff --git a/samplelib/PackedFaceset.py b/samplelib/PackedFaceset.py index 867fcd9..2ea01d9 100644 --- a/samplelib/PackedFaceset.py +++ b/samplelib/PackedFaceset.py @@ -120,6 +120,11 @@ class PackedFaceset(): samples_dat_path.unlink() + @staticmethod + def path_contains(samples_path): + samples_dat_path = samples_path / packed_faceset_filename + return samples_dat_path.exists() + @staticmethod def load(samples_path): samples_dat_path = samples_path / packed_faceset_filename