XSegEditor: added view lock at the center by holding shift.

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') )

XSegEditor/QStringDB.py

@@ -60,11 +60,17 @@ class QStringDB():
                                             'zh' : '删除选区',
                                            }[lang]     
                                               
-        QStringDB.btn_pt_edit_mode_tip = {  'en' : 'Edit point mode ( HOLD CTRL )',
+        QStringDB.btn_pt_edit_mode_tip = {  'en' : 'Add/delete point mode ( HOLD CTRL )',
-                                            'ru' : 'Режим правки точек',
+                                            'ru' : 'Режим добавления/удаления точек ( удерживайте CTRL )',
-                                            'zh' : '编辑点模式 ( 按住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',

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_op_mode(OpMode.NONE)
-        self.set_color_scheme_id(1)        
+        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)        
+            dflimg = DFLIMG.load(image_path)
-            ie_polys = dflimg.get_seg_ie_polys()            
+            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

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

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