diff --git a/XSegEditor/QIconDB.py b/XSegEditor/QIconDB.py index f4be98b..a2427c3 100644 --- a/XSegEditor/QIconDB.py +++ b/XSegEditor/QIconDB.py @@ -20,4 +20,6 @@ class QIconDB(): 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 + QIconDB.view_xseg = QIcon ( str(icon_path / 'view_xseg.png') ) + QIconDB.view_xseg_overlay = QIcon ( str(icon_path / 'view_xseg_overlay.png') ) + \ No newline at end of file diff --git a/XSegEditor/QStringDB.py b/XSegEditor/QStringDB.py index c825a55..55d537a 100644 --- a/XSegEditor/QStringDB.py +++ b/XSegEditor/QStringDB.py @@ -34,6 +34,11 @@ class QStringDB(): 'ru' : 'Посмотреть тренированную XSeg маску', 'zh' : '查看导入后的XSeg遮罩', }[lang] + + QStringDB.btn_view_xseg_overlay_mask_tip = { 'en' : 'View trained XSeg mask overlay face', + 'ru' : 'Посмотреть тренированную XSeg маску поверх лица', + 'zh' : '查看导入后的XSeg遮罩于脸上方', + }[lang] QStringDB.btn_poly_type_include_tip = { 'en' : 'Poly include mode', 'ru' : 'Режим полигонов - включение', diff --git a/XSegEditor/XSegEditor.py b/XSegEditor/XSegEditor.py index 0e06b10..3e396f4 100644 --- a/XSegEditor/XSegEditor.py +++ b/XSegEditor/XSegEditor.py @@ -35,6 +35,7 @@ class OpMode(IntEnum): EDIT_PTS = 2 VIEW_BAKED = 3 VIEW_XSEG_MASK = 4 + VIEW_XSEG_OVERLAY_MASK = 5 class PTEditMode(IntEnum): MOVE = 0 @@ -89,39 +90,39 @@ class ImagePreviewSequenceBar(QFrame): def get_preview_images_count(self): return self.preview_images_count - def update_images(self, prev_q_imgs=None, next_q_imgs=None): + def update_images(self, prev_imgs=None, next_imgs=None): # Fix arrays - if prev_q_imgs is None: - prev_q_imgs = [] + if prev_imgs is None: + prev_imgs = [] prev_img_conts_len = len(self.prev_img_conts) - prev_q_imgs_len = len(prev_q_imgs) + 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_q_imgs.append(None) + prev_imgs.append(None) elif prev_q_imgs_len > prev_img_conts_len: - prev_q_imgs = prev_q_imgs[:prev_img_conts_len] + prev_imgs = prev_imgs[:prev_img_conts_len] - if next_q_imgs is None: - next_q_imgs = [] + if next_imgs is None: + next_imgs = [] next_img_conts_len = len(self.next_img_conts) - next_q_imgs_len = len(next_q_imgs) + 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_q_imgs.append(None) + next_imgs.append(None) elif next_q_imgs_len > next_img_conts_len: - next_q_imgs = next_q_imgs[:next_img_conts_len] + next_imgs = next_imgs[:next_img_conts_len] - for i,q_img in enumerate(prev_q_imgs): - if q_img is None: + for i,img in enumerate(prev_imgs): + if img is None: self.prev_img_conts[i].setPixmap( self.black_q_pixmap ) else: - self.prev_img_conts[i].setPixmap( QPixmap.fromImage(q_img) ) + self.prev_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) ) - for i,q_img in enumerate(next_q_imgs): - if q_img is None: + for i,img in enumerate(next_imgs): + if img is None: self.next_img_conts[i].setPixmap( self.black_q_pixmap ) else: - self.next_img_conts[i].setPixmap( QPixmap.fromImage(q_img) ) + self.next_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) ) class ColorScheme(): def __init__(self, unselected_color, selected_color, outline_color, outline_width, pt_outline_color, cross_cursor): @@ -255,6 +256,11 @@ class QCanvasControlsRightBar(QFrame): 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='6', 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) @@ -262,6 +268,7 @@ 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_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_overlay_mask_act) self.btn_poly_color_act_grp.setExclusive(True) #============================================== btn_view_lock_center = QToolButton() @@ -275,6 +282,7 @@ class QCanvasControlsRightBar(QFrame): 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_l.addWidget ( btn_view_xseg_overlay_mask ) controls_bar_frame1 = QFrame() controls_bar_frame1.setFrameShape(QFrame.StyledPanel) controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) @@ -307,7 +315,8 @@ class QCanvasOperator(QWidget): 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_xseg_mask_act.toggled.connect ( self.set_view_xseg_mask ) - + self.cbar.btn_view_xseg_overlay_mask_act.toggled.connect ( self.set_view_xseg_overlay_mask ) + 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) ) @@ -328,17 +337,22 @@ class QCanvasOperator(QWidget): self.initialized = False self.last_state = None - 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 + 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 if xseg_mask is not None: - w,h = QSize_to_np ( q_img.size() ) + h,w,c = img.shape 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)) - + 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 @@ -389,7 +403,7 @@ class QCanvasOperator(QWidget): 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, + 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, color_scheme_id = self.color_scheme_id, ) @@ -521,7 +535,9 @@ class QCanvasOperator(QWidget): 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) - + elif self.op_mode == OpMode.VIEW_XSEG_OVERLAY_MASK: + self.cbar.btn_view_xseg_overlay_mask_act.setChecked(False) + self.op_mode = op_mode # Initialize new mode @@ -544,6 +560,9 @@ class QCanvasOperator(QWidget): 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) + 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]: self.mouse_op_poly_pt_id = None self.mouse_op_poly_edge_id = None @@ -617,6 +636,13 @@ class QCanvasOperator(QWidget): 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 ) # ==================================================================================== # ==================================================================================== @@ -936,6 +962,11 @@ class QCanvasOperator(QWidget): src_rect = QRect(0, 0, *self.img_size) dst_rect = self.img_to_cli_rect( 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: + src_rect = QRect(0, 0, *self.img_size) + dst_rect = self.img_to_cli_rect( src_rect ) + qp.drawPixmap(dst_rect, self.xseg_overlay_mask_pixmap, src_rect) else: if self.img_pixmap is not None: src_rect = QRect(0, 0, *self.img_size) @@ -1061,6 +1092,7 @@ class QCanvas(QFrame): 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, @@ -1151,7 +1183,7 @@ class MainWindow(QXMainWindow): 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_QImages = {} + self.cached_images = {} self.cached_has_ie_polys = {} self.initialize_ui() @@ -1188,7 +1220,7 @@ class MainWindow(QXMainWindow): def update_cached_images (self, count=5): - d = self.cached_QImages + d = self.cached_images for image_path in self.image_paths_done[:-count]+self.image_paths[count:]: if image_path in d: @@ -1198,13 +1230,14 @@ class MainWindow(QXMainWindow): if image_path not in d: img = cv2_imread(image_path) if img is not None: - d[image_path] = QImage_from_np(img) + d[image_path] = img - def load_QImage(self, image_path): + def load_image(self, image_path): try: - img = self.cached_QImages.get(image_path, None) + img = self.cached_images.get(image_path, None) if img is None: - img = QImage_from_np(cv2_imread(image_path)) + 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: @@ -1214,10 +1247,10 @@ class MainWindow(QXMainWindow): def update_preview_bar(self): count = self.image_bar.get_preview_images_count() - d = self.cached_QImages - prev_q_imgs = [ d.get(image_path, None) for image_path in self.image_paths_done[-1:-count:-1] ] - next_q_imgs = [ d.get(image_path, None) for image_path in self.image_paths[:count] ] - self.image_bar.update_images(prev_q_imgs, next_q_imgs) + 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): @@ -1230,11 +1263,11 @@ class MainWindow(QXMainWindow): 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: + img = self.load_image(image_path) + if img is None: return False - self.canvas.op.initialize ( q_img, ie_polys=ie_polys, xseg_mask=xseg_mask ) + self.canvas.op.initialize ( img, ie_polys=ie_polys, xseg_mask=xseg_mask ) self.filename_label.setText(str(image_path.name)) diff --git a/XSegEditor/gfx/icons/view_xseg_overlay.png b/XSegEditor/gfx/icons/view_xseg_overlay.png new file mode 100644 index 0000000..d188285 Binary files /dev/null and b/XSegEditor/gfx/icons/view_xseg_overlay.png differ