github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

XSegEditor: overlay xseg mask now works in polygon mode https://i.imgur.com/wDV6PY2.jpg

Browse source
This commit is contained in:
Colombo 2020-06-22 11:12:24 +04:00
parent f90723c9f9
commit d0ded46b51
1 changed files with 32 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

79
XSegEditor/XSegEditor.py
View file

@ -35,7 +35,6 @@ class OpMode(IntEnum):
EDIT_PTS = 2 EDIT_PTS = 2
VIEW_BAKED = 3 VIEW_BAKED = 3
VIEW_XSEG_MASK = 4 VIEW_XSEG_MASK = 4
VIEW_XSEG_OVERLAY_MASK = 5
class PTEditMode(IntEnum): class PTEditMode(IntEnum):
MOVE = 0 MOVE = 0
@ -274,7 +273,6 @@ class QCanvasControlsRightBar(QFrame):
self.btn_poly_color_act_grp.addAction(self.btn_poly_color_blue_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_baked_mask_act)
self.btn_poly_color_act_grp.addAction(self.btn_view_xseg_mask_act) self.btn_poly_color_act_grp.addAction(self.btn_view_xseg_mask_act)
self.btn_poly_color_act_grp.addAction(self.btn_view_xseg_overlay_mask_act)
self.btn_poly_color_act_grp.setExclusive(True) self.btn_poly_color_act_grp.setExclusive(True)
#============================================== #==============================================
btn_view_lock_center = QToolButton() btn_view_lock_center = QToolButton()
@ -288,24 +286,31 @@ class QCanvasControlsRightBar(QFrame):
controls_bar_frame1_l.addWidget ( btn_poly_color_blue ) 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_baked_mask )
controls_bar_frame1_l.addWidget ( btn_view_xseg_mask ) controls_bar_frame1_l.addWidget ( btn_view_xseg_mask )
controls_bar_frame1_l.addWidget ( btn_view_xseg_overlay_mask )
controls_bar_frame1 = QFrame() controls_bar_frame1 = QFrame()
controls_bar_frame1.setFrameShape(QFrame.StyledPanel) controls_bar_frame1.setFrameShape(QFrame.StyledPanel)
controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
controls_bar_frame1.setLayout(controls_bar_frame1_l) controls_bar_frame1.setLayout(controls_bar_frame1_l)
controls_bar_frame2_l = QVBoxLayout() controls_bar_frame2_l = QVBoxLayout()
controls_bar_frame2_l.addWidget ( btn_view_lock_center ) controls_bar_frame2_l.addWidget ( btn_view_xseg_overlay_mask )
controls_bar_frame2 = QFrame() controls_bar_frame2 = QFrame()
controls_bar_frame2.setFrameShape(QFrame.StyledPanel) controls_bar_frame2.setFrameShape(QFrame.StyledPanel)
controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
controls_bar_frame2.setLayout(controls_bar_frame2_l) controls_bar_frame2.setLayout(controls_bar_frame2_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 = QVBoxLayout()
controls_bar_l.setContentsMargins(0,0,0,0) controls_bar_l.setContentsMargins(0,0,0,0)
controls_bar_l.addWidget(controls_bar_frame1) controls_bar_l.addWidget(controls_bar_frame1)
controls_bar_l.addWidget(controls_bar_frame2) controls_bar_l.addWidget(controls_bar_frame2)
controls_bar_l.addWidget(controls_bar_frame3)
self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding ) self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding )
self.setLayout(controls_bar_l) self.setLayout(controls_bar_l)
@ -319,9 +324,10 @@ class QCanvasOperator(QWidget):
self.cbar.btn_poly_color_red_act.triggered.connect ( lambda : self.set_color_scheme_id(0) ) 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_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_poly_color_blue_act.triggered.connect ( lambda : self.set_color_scheme_id(2) )
self.cbar.btn_view_baked_mask_act.toggled.connect ( lambda : self.set_op_mode(OpMode.VIEW_BAKED) ) self.cbar.btn_view_baked_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_BAKED) )
self.cbar.btn_view_xseg_mask_act.toggled.connect ( self.set_view_xseg_mask ) 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 ( self.set_view_xseg_overlay_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_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_poly_type_exclude_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.EXCLUDE) )
@ -376,6 +382,7 @@ class QCanvasOperator(QWidget):
# UI init # UI init
self.set_cbar_disabled() self.set_cbar_disabled()
self.cbar.btn_poly_color_act_grp.setDisabled(False) self.cbar.btn_poly_color_act_grp.setDisabled(False)
self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(False)
self.cbar.btn_poly_type_act_grp.setDisabled(False) self.cbar.btn_poly_type_act_grp.setDisabled(False)
# Initial vars # Initial vars
@ -410,9 +417,8 @@ class QCanvasOperator(QWidget):
if self.op_mode == OpMode.DRAW_PTS: if self.op_mode == OpMode.DRAW_PTS:
self.set_op_mode(OpMode.EDIT_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, OpMode.VIEW_XSEG_OVERLAY_MASK] else None, 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, color_scheme_id = self.color_scheme_id)
)
self.img_pixmap = None self.img_pixmap = None
self.update_cursor(is_finalize=True) self.update_cursor(is_finalize=True)
@ -427,7 +433,6 @@ class QCanvasOperator(QWidget):
# ====================================== GETTERS ===================================== # ====================================== GETTERS =====================================
# ==================================================================================== # ====================================================================================
# ==================================================================================== # ====================================================================================
def is_initialized(self): def is_initialized(self):
return self.initialized return self.initialized
@ -542,8 +547,6 @@ class QCanvasOperator(QWidget):
self.cbar.btn_view_baked_mask_act.setChecked(False) self.cbar.btn_view_baked_mask_act.setChecked(False)
elif self.op_mode == OpMode.VIEW_XSEG_MASK: elif self.op_mode == OpMode.VIEW_XSEG_MASK:
self.cbar.btn_view_xseg_mask_act.setChecked(False) self.cbar.btn_view_xseg_mask_act.setChecked(False)
elif self.op_mode == OpMode.VIEW_XSEG_OVERLAY_MASK:
self.cbar.btn_view_xseg_overlay_mask_act.setChecked(False)
self.op_mode = op_mode self.op_mode = op_mode
@ -567,8 +570,6 @@ class QCanvasOperator(QWidget):
self.img_baked_pixmap = QPixmap.fromImage(QImage_from_np(n)) self.img_baked_pixmap = QPixmap.fromImage(QImage_from_np(n))
elif op_mode == OpMode.VIEW_XSEG_MASK: 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)
elif op_mode == OpMode.VIEW_XSEG_OVERLAY_MASK:
self.cbar.btn_view_xseg_overlay_mask_act.setChecked(True)
if op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]: if op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]:
self.mouse_op_poly_pt_id = None self.mouse_op_poly_pt_id = None
@ -607,10 +608,12 @@ class QCanvasOperator(QWidget):
self.cbar.btn_pt_edit_mode_act.setDisabled(True) self.cbar.btn_pt_edit_mode_act.setDisabled(True)
self.cbar.btn_view_lock_center_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_color_act_grp.setDisabled(True)
self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(True)
self.cbar.btn_poly_type_act_grp.setDisabled(True) self.cbar.btn_poly_type_act_grp.setDisabled(True)
def set_color_scheme_id(self, id): def set_color_scheme_id(self, id):
if self.op_mode == OpMode.VIEW_BAKED: if self.op_mode == OpMode.VIEW_BAKED or self.op_mode == OpMode.VIEW_XSEG_MASK:
self.set_op_mode(OpMode.NONE) self.set_op_mode(OpMode.NONE)
if not hasattr(self, 'color_scheme_id') or self.color_scheme_id != id: if not hasattr(self, 'color_scheme_id') or self.color_scheme_id != id:
@ -631,26 +634,9 @@ class QCanvasOperator(QWidget):
self.op_mode in [OpMode.NONE, OpMode.EDIT_PTS] ): self.op_mode in [OpMode.NONE, OpMode.EDIT_PTS] ):
self.poly_include_type = poly_include_type self.poly_include_type = poly_include_type
self.update() self.update()
self.cbar.btn_poly_type_include_act.setChecked(self.poly_include_type == SegIEPolyType.INCLUDE) 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) self.cbar.btn_poly_type_exclude_act.setChecked(self.poly_include_type == SegIEPolyType.EXCLUDE)
def set_view_xseg_mask(self, is_checked):
if is_checked:
self.set_op_mode(OpMode.VIEW_XSEG_MASK)
else:
self.set_op_mode(OpMode.NONE)
self.cbar.btn_view_xseg_mask_act.setChecked(is_checked )
def set_view_xseg_overlay_mask(self, is_checked):
if is_checked:
self.set_op_mode(OpMode.VIEW_XSEG_OVERLAY_MASK)
else:
self.set_op_mode(OpMode.NONE)
self.cbar.btn_view_xseg_overlay_mask_act.setChecked(is_checked )
# ==================================================================================== # ====================================================================================
# ==================================================================================== # ====================================================================================
# ====================================== METHODS ===================================== # ====================================== METHODS =====================================
@ -839,10 +825,10 @@ class QCanvasOperator(QWidget):
if self.mouse_op_poly_pt_id is not None: if self.mouse_op_poly_pt_id is not None:
# Click on point of op_poly # Click on point of op_poly
if self.pt_edit_mode == PTEditMode.ADD_DEL: if self.pt_edit_mode == PTEditMode.ADD_DEL:
# with mode -> delete point # in mode 'delete point'
self.op_poly.remove_pt(self.mouse_op_poly_pt_id) self.op_poly.remove_pt(self.mouse_op_poly_pt_id)
if self.op_poly.get_pts_count() < 3: if self.op_poly.get_pts_count() < 3:
# not enough points -> remove poly # not enough points after delete -> remove poly
self.ie_polys.remove_poly (self.op_poly) self.ie_polys.remove_poly (self.op_poly)
self.set_op_mode(OpMode.NONE) self.set_op_mode(OpMode.NONE)
self.update() self.update()
@ -856,7 +842,7 @@ class QCanvasOperator(QWidget):
elif self.mouse_op_poly_edge_id is not None: elif self.mouse_op_poly_edge_id is not None:
# Click on edge of op_poly # Click on edge of op_poly
if self.pt_edit_mode == PTEditMode.ADD_DEL: if self.pt_edit_mode == PTEditMode.ADD_DEL:
# with mode -> insert new point # in mode 'insert new point'
edge_img_pt = self.cli_to_img_pt(self.mouse_op_poly_edge_id_pt) 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.op_poly.insert_pt (self.mouse_op_poly_edge_id+1, edge_img_pt)
self.update() self.update()
@ -909,7 +895,6 @@ class QCanvasOperator(QWidget):
if npla.norm(self.mouse_cli_pt - prev_mouse_cli_pt) >= 1: if npla.norm(self.mouse_cli_pt - prev_mouse_cli_pt) >= 1:
self.img_look_pt = self.mouse_img_pt self.img_look_pt = self.mouse_img_pt
QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) )) QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) ))
self.update() self.update()
if self.drag_type == DragType.IMAGE_LOOK: if self.drag_type == DragType.IMAGE_LOOK:
@ -967,11 +952,11 @@ class QCanvasOperator(QWidget):
elif self.op_mode == OpMode.VIEW_XSEG_MASK: elif self.op_mode == OpMode.VIEW_XSEG_MASK:
if self.xseg_mask_pixmap is not None: if self.xseg_mask_pixmap is not None:
qp.drawPixmap(dst_rect, self.xseg_mask_pixmap, src_rect) qp.drawPixmap(dst_rect, self.xseg_mask_pixmap, src_rect)
elif self.op_mode == OpMode.VIEW_XSEG_OVERLAY_MASK:
if self.xseg_overlay_mask_pixmap is not None:
qp.drawPixmap(dst_rect, self.xseg_overlay_mask_pixmap, src_rect)
else: else:
if self.img_pixmap is not None: 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) qp.drawPixmap(dst_rect, self.img_pixmap, src_rect)
polys = self.ie_polys.get_polys() polys = self.ie_polys.get_polys()
@ -1294,16 +1279,16 @@ class MainWindow(QXMainWindow):
ie_polys = dflimg.get_seg_ie_polys() ie_polys = dflimg.get_seg_ie_polys()
new_ie_polys = self.canvas.op.get_ie_polys() new_ie_polys = self.canvas.op.get_ie_polys()
if not new_ie_polys.identical(ie_polys): if not new_ie_polys.identical(ie_polys):
prev_has_polys = self.image_paths_has_ie_polys[image_path] prev_has_polys = self.image_paths_has_ie_polys[image_path]
self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys() self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys()
new_has_polys = self.image_paths_has_ie_polys[image_path] new_has_polys = self.image_paths_has_ie_polys[image_path]
if not prev_has_polys and new_has_polys: if not prev_has_polys and new_has_polys:
self.set_has_ie_polys_count ( self.get_has_ie_polys_count() +1) self.set_has_ie_polys_count ( self.get_has_ie_polys_count() +1)
elif prev_has_polys and not new_has_polys: elif prev_has_polys and not new_has_polys:
self.set_has_ie_polys_count ( self.get_has_ie_polys_count() -1) self.set_has_ie_polys_count ( self.get_has_ie_polys_count() -1)
dflimg.set_seg_ie_polys( new_ie_polys ) dflimg.set_seg_ie_polys( new_ie_polys )
dflimg.save() dflimg.save()