StarMirror/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2024-03-22 13:10:55 +08:00
Code Issues Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Colombo 2020-04-15 22:12:54 +04:00
parent 6bbc607312
commit ea607edfc9
6 changed files with 116 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
XSegEditor/QIconDB.py
View File

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

12
XSegEditor/QStringDB.py
View File

@ -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',

140
XSegEditor/XSegEditor.py
View File

@ -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

BIN
XSegEditor/gfx/icons/view_lock_center.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 4.0 KiB

11
main.py
View File

@ -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

5
samplelib/PackedFaceset.py
View File

@ -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