diff --git a/README.md b/README.md index d8106d9..95f9932 100644 --- a/README.md +++ b/README.md @@ -121,9 +121,9 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul ||bitcoin:bc1qkhh7h0gwwhxgg6h6gpllfgstkd645fefrd5s6z| |Alipay 捐款|![](doc/Alipay_donation.jpg)| ||| -|Last donations|10$ ( 朱 阳阳 )| +|Last donations|50$ ( Tomas Hajka )| +||10$ ( 朱 阳阳 )| ||24$ ( NextFace )| -||10$ ( Amien Phillips )| ||| |Collect facesets|You can collect faceset of any celebrity that can be used in DeepFaceLab and share it [in the community](https://mrdeepfakes.com/forums/forum-celebrity-facesets)| ||| diff --git a/XSegEditor/QCursorDB.py b/XSegEditor/QCursorDB.py new file mode 100644 index 0000000..0909cba --- /dev/null +++ b/XSegEditor/QCursorDB.py @@ -0,0 +1,10 @@ +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * + +class QCursorDB(): + @staticmethod + def initialize(cursor_path): + QCursorDB.cross_red = QCursor ( QPixmap ( str(cursor_path / 'cross_red.png') ) ) + QCursorDB.cross_green = QCursor ( QPixmap ( str(cursor_path / 'cross_green.png') ) ) + QCursorDB.cross_blue = QCursor ( QPixmap ( str(cursor_path / 'cross_blue.png') ) ) diff --git a/XSegEditor/QIconDB.py b/XSegEditor/QIconDB.py new file mode 100644 index 0000000..7a48cf4 --- /dev/null +++ b/XSegEditor/QIconDB.py @@ -0,0 +1,21 @@ +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * + + +class QIconDB(): + @staticmethod + def initialize(icon_path): + QIconDB.app_icon = QIcon ( str(icon_path / 'app_icon.png') ) + QIconDB.delete_poly = QIcon ( str(icon_path / 'delete_poly.png') ) + QIconDB.undo_pt = QIcon ( str(icon_path / 'undo_pt.png') ) + QIconDB.redo_pt = QIcon ( str(icon_path / 'redo_pt.png') ) + QIconDB.poly_color_red = QIcon ( str(icon_path / 'poly_color_red.png') ) + QIconDB.poly_color_green = QIcon ( str(icon_path / 'poly_color_green.png') ) + QIconDB.poly_color_blue = QIcon ( str(icon_path / 'poly_color_blue.png') ) + QIconDB.poly_type_include = QIcon ( str(icon_path / 'poly_type_include.png') ) + QIconDB.poly_type_exclude = QIcon ( str(icon_path / 'poly_type_exclude.png') ) + 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_baked = QIcon ( str(icon_path / 'view_baked.png') ) diff --git a/XSegEditor/QStringDB.py b/XSegEditor/QStringDB.py new file mode 100644 index 0000000..82b39f5 --- /dev/null +++ b/XSegEditor/QStringDB.py @@ -0,0 +1,72 @@ +from localization import system_language + + +class QStringDB(): + + @staticmethod + def initialize(): + lang = system_language + + if lang not in ['en','ru','zn']: + lang = 'en' + + QStringDB.btn_poly_color_red_tip = { 'en' : 'Poly color scheme red', + 'ru' : 'Красная цветовая схема полигонов', + 'zn' : '多边形配色方案红色', + }[lang] + + QStringDB.btn_poly_color_green_tip = { 'en' : 'Poly color scheme green', + 'ru' : 'Зелёная цветовая схема полигонов', + 'zn' : '多边形配色方案绿色', + }[lang] + + QStringDB.btn_poly_color_blue_tip = { 'en' : 'Poly color scheme blue', + 'ru' : 'Синяя цветовая схема полигонов', + 'zn' : '多边形配色方案蓝色', + }[lang] + + QStringDB.btn_view_baked_mask_tip = { 'en' : 'View baked mask', + 'ru' : 'Посмотреть запечёную маску', + 'zn' : '查看遮罩通道', + }[lang] + + QStringDB.btn_poly_type_include_tip = { 'en' : 'Poly include mode', + 'ru' : 'Режим полигонов - включение', + 'zn' : '多边形包含模式', + }[lang] + + QStringDB.btn_poly_type_exclude_tip = { 'en' : 'Poly exclude mode', + 'ru' : 'Режим полигонов - исключение', + 'zn' : '多边形排除方式', + }[lang] + + QStringDB.btn_undo_pt_tip = { 'en' : 'Undo point', + 'ru' : 'Отменить точку', + 'zn' : '撤消点', + }[lang] + + QStringDB.btn_redo_pt_tip = { 'en' : 'Redo point', + 'ru' : 'Повторить точку', + 'zn' : '重做点', + }[lang] + + QStringDB.btn_delete_poly_tip = { 'en' : 'Delete poly', + 'ru' : 'Удалить полигон', + 'zn' : '删除多边形', + }[lang] + + QStringDB.btn_pt_edit_mode_tip = { 'en' : 'Edit point mode ( HOLD CTRL )', + 'ru' : 'Режим правки точек', + 'zn' : '编辑点模式 ( 按住CTRL )', + }[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', + 'zn' : '保存和上一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n', + }[lang] + QStringDB.btn_next_image_tip = { 'en' : 'Save and Next image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n', + 'ru' : 'Сохранить и следующее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n', + 'zn' : '保存并下一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n', + }[lang] + + \ No newline at end of file diff --git a/XSegEditor/XSegEditor.py b/XSegEditor/XSegEditor.py new file mode 100644 index 0000000..0406040 --- /dev/null +++ b/XSegEditor/XSegEditor.py @@ -0,0 +1,1265 @@ +import json +import multiprocessing +import os +import pickle +import sys +import tempfile +import time +import traceback +from enum import IntEnum + +import cv2 +import numpy as np +import numpy.linalg as npla +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * + +from core import pathex +from core.cv2ex import * +from core.imagelib import SegIEPoly, SegIEPolys, SegIEPolyType, sd +from core.qtex import * +from DFLIMG import * +from localization import StringsDB, system_language + +from .QCursorDB import QCursorDB +from .QStringDB import QStringDB +from .QIconDB import QIconDB + +class OpMode(IntEnum): + NONE = 0 + DRAW_PTS = 1 + EDIT_PTS = 2 + VIEW_BAKED = 3 + +class PTEditMode(IntEnum): + MOVE = 0 + ADD_DEL = 1 + +class DragType(IntEnum): + NONE = 0 + IMAGE_LOOK = 1 + POLY_PT = 2 + +class QUIConfig(): + @staticmethod + def initialize(icon_size = 48, icon_spacer_size=16, preview_bar_icon_size=64): + QUIConfig.icon_q_size = QSize(icon_size, icon_size) + QUIConfig.icon_spacer_q_size = QSize(icon_spacer_size, icon_spacer_size) + QUIConfig.preview_bar_icon_q_size = QSize(preview_bar_icon_size, preview_bar_icon_size) + +class ImagePreviewSequenceBar(QFrame): + def __init__(self, preview_images_count, icon_size): + super().__init__() + self.preview_images_count = preview_images_count = max(1, preview_images_count + (preview_images_count % 2 -1) ) + + self.icon_size = icon_size + + black_q_img = QImage(np.zeros( (icon_size,icon_size,3) ).data, icon_size, icon_size, 3*icon_size, QImage.Format_RGB888) + self.black_q_pixmap = QPixmap.fromImage(black_q_img) + + self.image_containers = [ QLabel() for i in range(preview_images_count)] + + main_frame_l_cont_hl = QGridLayout() + main_frame_l_cont_hl.setContentsMargins(0,0,0,0) + + for i in range(len(self.image_containers)): + q_label = self.image_containers[i] + q_label.setScaledContents(True) + q_label.setMinimumSize(icon_size, icon_size ) + q_label.setSizePolicy (QSizePolicy.Ignored, QSizePolicy.Ignored) + + main_frame_l_cont_hl.addWidget (q_label, 0, i) + + self.setLayout(main_frame_l_cont_hl) + + self.prev_img_conts = self.image_containers[(preview_images_count//2) -1::-1] + self.next_img_conts = self.image_containers[preview_images_count//2:] + + self.update_images() + + def get_preview_images_count(self): + return self.preview_images_count + + def update_images(self, prev_q_imgs=None, next_q_imgs=None): + # Fix arrays + if prev_q_imgs is None: + prev_q_imgs = [] + prev_img_conts_len = len(self.prev_img_conts) + prev_q_imgs_len = len(prev_q_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) + elif prev_q_imgs_len > prev_img_conts_len: + prev_q_imgs = prev_q_imgs[:prev_img_conts_len] + + if next_q_imgs is None: + next_q_imgs = [] + next_img_conts_len = len(self.next_img_conts) + next_q_imgs_len = len(next_q_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) + elif next_q_imgs_len > next_img_conts_len: + next_q_imgs = next_q_imgs[:next_img_conts_len] + + for i,q_img in enumerate(prev_q_imgs): + if q_img is None: + self.prev_img_conts[i].setPixmap( self.black_q_pixmap ) + else: + self.prev_img_conts[i].setPixmap( QPixmap.fromImage(q_img) ) + + for i,q_img in enumerate(next_q_imgs): + if q_img is None: + self.next_img_conts[i].setPixmap( self.black_q_pixmap ) + else: + self.next_img_conts[i].setPixmap( QPixmap.fromImage(q_img) ) + +class ColorScheme(): + def __init__(self, unselected_color, selected_color, outline_color, outline_width, pt_outline_color, cross_cursor): + self.poly_unselected_brush = QBrush(unselected_color) + self.poly_selected_brush = QBrush(selected_color) + + self.poly_outline_solid_pen = QPen(outline_color, outline_width, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin) + self.poly_outline_dot_pen = QPen(outline_color, outline_width, Qt.DotLine, Qt.RoundCap, Qt.RoundJoin) + + self.pt_outline_pen = QPen(pt_outline_color) + self.cross_cursor = cross_cursor + +class CanvasConfig(): + + def __init__(self, + pt_radius=4, + pt_select_radius=8, + color_schemes=None, + **kwargs): + self.pt_radius = pt_radius + self.pt_select_radius = pt_select_radius + + if color_schemes is None: + color_schemes = [ + ColorScheme( QColor(192,0,0,alpha=0), QColor(192,0,0,alpha=72), QColor(192,0,0), 2, QColor(255,255,255), QCursorDB.cross_red ), + ColorScheme( QColor(0,192,0,alpha=0), QColor(0,192,0,alpha=72), QColor(0,192,0), 2, QColor(255,255,255), QCursorDB.cross_green ), + ColorScheme( QColor(0,0,192,alpha=0), QColor(0,0,192,alpha=72), QColor(0,0,192), 2, QColor(255,255,255), QCursorDB.cross_blue ), + ] + self.color_schemes = color_schemes + +class QCanvasControlsBar(QFrame): + + def __init__(self): + super().__init__() + #============================================== + btn_poly_color_red = QToolButton() + self.btn_poly_color_red_act = QActionEx( QIconDB.poly_color_red, QStringDB.btn_poly_color_red_tip, shortcut='1', shortcut_in_tooltip=True, is_checkable=True) + btn_poly_color_red.setDefaultAction(self.btn_poly_color_red_act) + btn_poly_color_red.setIconSize(QUIConfig.icon_q_size) + + btn_poly_color_green = QToolButton() + self.btn_poly_color_green_act = QActionEx( QIconDB.poly_color_green, QStringDB.btn_poly_color_green_tip, shortcut='2', shortcut_in_tooltip=True, is_checkable=True) + btn_poly_color_green.setDefaultAction(self.btn_poly_color_green_act) + btn_poly_color_green.setIconSize(QUIConfig.icon_q_size) + + btn_poly_color_blue = QToolButton() + self.btn_poly_color_blue_act = QActionEx( QIconDB.poly_color_blue, QStringDB.btn_poly_color_blue_tip, shortcut='3', shortcut_in_tooltip=True, is_checkable=True) + btn_poly_color_blue.setDefaultAction(self.btn_poly_color_blue_act) + btn_poly_color_blue.setIconSize(QUIConfig.icon_q_size) + + btn_view_baked_mask = QToolButton() + self.btn_view_baked_mask_act = QActionEx( QIconDB.view_baked, QStringDB.btn_view_baked_mask_tip, shortcut='4', shortcut_in_tooltip=True, is_checkable=True) + btn_view_baked_mask.setDefaultAction(self.btn_view_baked_mask_act) + btn_view_baked_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) + self.btn_poly_color_act_grp.addAction(self.btn_poly_color_green_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.setExclusive(True) + #============================================== + btn_poly_type_include = QToolButton() + self.btn_poly_type_include_act = QActionEx( QIconDB.poly_type_include, QStringDB.btn_poly_type_include_tip, shortcut='Q', shortcut_in_tooltip=True, is_checkable=True) + btn_poly_type_include.setDefaultAction(self.btn_poly_type_include_act) + btn_poly_type_include.setIconSize(QUIConfig.icon_q_size) + + btn_poly_type_exclude = QToolButton() + self.btn_poly_type_exclude_act = QActionEx( QIconDB.poly_type_exclude, QStringDB.btn_poly_type_exclude_tip, shortcut='W', shortcut_in_tooltip=True, is_checkable=True) + btn_poly_type_exclude.setDefaultAction(self.btn_poly_type_exclude_act) + btn_poly_type_exclude.setIconSize(QUIConfig.icon_q_size) + + self.btn_poly_type_act_grp = QActionGroup (self) + self.btn_poly_type_act_grp.addAction(self.btn_poly_type_include_act) + self.btn_poly_type_act_grp.addAction(self.btn_poly_type_exclude_act) + self.btn_poly_type_act_grp.setExclusive(True) + #============================================== + btn_undo_pt = QToolButton() + self.btn_undo_pt_act = QActionEx( QIconDB.undo_pt, QStringDB.btn_undo_pt_tip, shortcut='Ctrl+Z', shortcut_in_tooltip=True, is_auto_repeat=True) + btn_undo_pt.setDefaultAction(self.btn_undo_pt_act) + btn_undo_pt.setIconSize(QUIConfig.icon_q_size) + + btn_redo_pt = QToolButton() + self.btn_redo_pt_act = QActionEx( QIconDB.redo_pt, QStringDB.btn_redo_pt_tip, shortcut='Ctrl+Shift+Z', shortcut_in_tooltip=True, is_auto_repeat=True) + btn_redo_pt.setDefaultAction(self.btn_redo_pt_act) + btn_redo_pt.setIconSize(QUIConfig.icon_q_size) + + btn_delete_poly = QToolButton() + self.btn_delete_poly_act = QActionEx( QIconDB.delete_poly, QStringDB.btn_delete_poly_tip, shortcut='Delete', shortcut_in_tooltip=True) + btn_delete_poly.setDefaultAction(self.btn_delete_poly_act) + btn_delete_poly.setIconSize(QUIConfig.icon_q_size) + #============================================== + btn_pt_edit_mode = QToolButton() + 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_frame1_l = QVBoxLayout() + controls_bar_frame1_l.addWidget ( btn_poly_color_red ) + controls_bar_frame1_l.addWidget ( btn_poly_color_green ) + controls_bar_frame1_l.addWidget ( btn_poly_color_blue ) + controls_bar_frame1_l.addWidget ( btn_view_baked_mask ) + controls_bar_frame1 = QFrame() + controls_bar_frame1.setFrameShape(QFrame.StyledPanel) + 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_poly_type_include ) + controls_bar_frame2_l.addWidget ( btn_poly_type_exclude ) + 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_frame3_l = QVBoxLayout() + controls_bar_frame3_l.addWidget ( btn_undo_pt ) + controls_bar_frame3_l.addWidget ( btn_redo_pt ) + controls_bar_frame3_l.addWidget ( btn_delete_poly ) + 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_frame4_l = QVBoxLayout() + controls_bar_frame4_l.addWidget ( btn_pt_edit_mode ) + controls_bar_frame4 = QFrame() + controls_bar_frame4.setFrameShape(QFrame.StyledPanel) + controls_bar_frame4.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed) + controls_bar_frame4.setLayout(controls_bar_frame4_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) + controls_bar_l.addWidget(controls_bar_frame3) + controls_bar_l.addWidget(controls_bar_frame4) + + self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding ) + self.setLayout(controls_bar_l) + + +class QCanvasOperator(QWidget): + def __init__(self, cbar): + super().__init__() + self.cbar = cbar + + self.set_cbar_disabled(initialize=False) + + self.cbar.btn_delete_poly_act.triggered.connect ( lambda : self.action_delete_poly() ) + + self.cbar.btn_undo_pt_act.triggered.connect ( lambda : self.action_undo_pt() ) + self.cbar.btn_redo_pt_act.triggered.connect ( lambda : self.action_redo_pt() ) + + 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_blue_act.triggered.connect ( lambda : self.set_color_scheme_id(2) ) + self.cbar.btn_view_baked_mask_act.toggled.connect ( self.set_view_baked_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) ) + + 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.mouse_in_widget = False + + QXMainWindow.inst.add_keyPressEvent_listener ( self.on_keyPressEvent ) + QXMainWindow.inst.add_keyReleaseEvent_listener ( self.on_keyReleaseEvent ) + + self.qp = QPainter() + self.initialized = False + + def initialize(self, q_img, img_look_pt=None, view_scale=None, ie_polys=None, canvas_config=None ): + self.q_img = q_img + self.img_pixmap = QPixmap.fromImage(q_img) + self.img_size = QSize_to_np (self.img_pixmap.size()) + + self.img_look_pt = img_look_pt + self.view_scale = view_scale + + if ie_polys is None: + ie_polys = SegIEPolys() + self.ie_polys = ie_polys + + if canvas_config is None: + canvas_config = CanvasConfig() + self.canvas_config = canvas_config + + self.current_cursor = None + + + self.mouse_hull_poly = None + self.mouse_wire_poly = None + + self.drag_type = DragType.NONE + self.op_mode = None + self.pt_edit_mode = None + + if not hasattr(self, 'color_scheme_id' ): + self.color_scheme_id = 1 + self.set_color_scheme_id(self.color_scheme_id) + + self.set_op_mode(OpMode.NONE) + + self.set_pt_edit_mode(PTEditMode.MOVE) + self.set_view_baked_mask(False) + + self.set_cbar_disabled(initialize=True) + + if not hasattr(self, 'poly_include_type' ): + self.poly_include_type = SegIEPolyType.INCLUDE + self.set_poly_include_type(self.poly_include_type) + + + self.setMouseTracking(True) + self.update_cursor() + self.update() + self.initialized = True + + def finalize(self): + if self.initialized: + self.img_pixmap = None + self.update_cursor(is_finalize=True) + self.setMouseTracking(False) + self.setFocusPolicy(Qt.NoFocus) + self.set_cbar_disabled(initialize=False) + self.initialized = False + self.update() + + # ==================================================================================== + # ==================================================================================== + # ====================================== GETTERS ===================================== + # ==================================================================================== + # ==================================================================================== + + def is_initialized(self): + return self.initialized + + def get_ie_polys(self): + return self.ie_polys + + def get_img_look_pt(self): + img_look_pt = self.img_look_pt + if img_look_pt is None: + img_look_pt = self.img_size / 2 + return img_look_pt + + def get_view_scale(self): + view_scale = self.view_scale + if view_scale is None: + # Calc as scale to fit + min_cli_size = np.min(QSize_to_np(self.size())) + max_img_size = np.max(self.img_size) + view_scale = min_cli_size / max_img_size + + return view_scale + + def get_current_color_scheme(self): + return self.canvas_config.color_schemes[self.color_scheme_id] + + def get_poly_pt_id_under_pt(self, poly, cli_pt): + w = np.argwhere ( npla.norm ( cli_pt - self.img_to_cli_pt( poly.get_pts() ), axis=1 ) <= self.canvas_config.pt_select_radius ) + return None if len(w) == 0 else w[-1][0] + + def get_poly_edge_id_pt_under_pt(self, poly, cli_pt): + cli_pts = self.img_to_cli_pt(poly.get_pts()) + if len(cli_pts) >= 3: + edge_dists, projs = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True) + edge_id = np.argmin(edge_dists) + dist = edge_dists[edge_id] + pt = projs[edge_id] + if dist <= self.canvas_config.pt_select_radius: + return edge_id, pt + return None, None + + def get_poly_by_pt_near_wire(self, cli_pt): + pt_select_radius = self.canvas_config.pt_select_radius + + for poly in reversed(self.ie_polys.get_polys()): + pts = poly.get_pts() + if len(pts) >= 3: + cli_pts = self.img_to_cli_pt(pts) + + edge_dists, _ = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True) + + if np.min(edge_dists) <= pt_select_radius or \ + any( npla.norm ( cli_pt - cli_pts, axis=1 ) <= pt_select_radius ): + return poly + return None + + def get_poly_by_pt_in_hull(self, cli_pos): + img_pos = self.cli_to_img_pt(cli_pos) + + for poly in reversed(self.ie_polys.get_polys()): + pts = poly.get_pts() + if len(pts) >= 3: + if cv2.pointPolygonTest( pts, tuple(img_pos), False) >= 0: + return poly + + 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 + + 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() + + def img_to_cli_rect(self, rect): + tl = QPoint_to_np(rect.topLeft()) + xy = self.img_to_cli_pt(tl) + xy2 = self.img_to_cli_pt(tl + QSize_to_np(rect.size()) ) - xy + return QRect ( *xy.astype(np.int), *xy2.astype(np.int) ) + + # ==================================================================================== + # ==================================================================================== + # ====================================== SETTERS ===================================== + # ==================================================================================== + # ==================================================================================== + + def set_op_mode(self, op_mode, op_poly=None): + if op_mode != self.op_mode: + + if self.op_mode == OpMode.NONE: + self.cbar.btn_poly_type_act_grp.setDisabled(True) + elif self.op_mode == OpMode.DRAW_PTS: + self.cbar.btn_undo_pt_act.setDisabled(True) + self.cbar.btn_redo_pt_act.setDisabled(True) + + 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) + # Reset pt_edit_move when exit from EDIT_PTS + self.set_pt_edit_mode(PTEditMode.MOVE) + + self.op_mode = op_mode + + if self.op_mode == OpMode.NONE: + self.cbar.btn_poly_type_act_grp.setDisabled(False) + elif self.op_mode == OpMode.DRAW_PTS: + self.cbar.btn_undo_pt_act.setDisabled(False) + self.cbar.btn_redo_pt_act.setDisabled(False) + elif self.op_mode == OpMode.EDIT_PTS: + self.cbar.btn_pt_edit_mode_act.setDisabled(False) + self.cbar.btn_delete_poly_act.setDisabled(False) + + if self.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.set_op_poly(op_poly) + self.update_cursor() + self.update() + + def set_op_poly(self, op_poly): + self.op_poly = op_poly + if op_poly is not None: + self.update_mouse_info() + self.update() + + def set_pt_edit_mode(self, pt_edit_mode): + if self.pt_edit_mode != pt_edit_mode: + 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_cbar_disabled(self, initialize): + 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) + + if initialize: + self.cbar.btn_poly_color_act_grp.setDisabled(False) + self.cbar.btn_poly_type_act_grp.setDisabled(False) + else: + self.cbar.btn_poly_color_act_grp.setDisabled(True) + self.cbar.btn_poly_type_act_grp.setDisabled(True) + + def set_color_scheme_id(self, id): + if self.color_scheme_id != id: + self.color_scheme_id = id + self.update_cursor() + self.update() + if self.color_scheme_id == 0: + self.cbar.btn_poly_color_red_act.setChecked( True ) + elif self.color_scheme_id == 1: + self.cbar.btn_poly_color_green_act.setChecked( True ) + elif self.color_scheme_id == 2: + self.cbar.btn_poly_color_blue_act.setChecked( True ) + + def set_poly_include_type(self, poly_include_type): + if self.op_mode in [OpMode.NONE, OpMode.EDIT_PTS]: + if self.poly_include_type != poly_include_type: + self.poly_include_type = poly_include_type + self.update() + + 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) + + + + def set_view_baked_mask(self, is_checked): + if is_checked: + self.set_op_mode(OpMode.VIEW_BAKED) + + 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) + + self.cbar.btn_view_baked_mask_act.setChecked(is_checked ) + + # ==================================================================================== + # ==================================================================================== + # ====================================== METHODS ===================================== + # ==================================================================================== + # ==================================================================================== + + def update_cursor(self, is_finalize=False): + if not self.initialized: + return + + if not self.mouse_in_widget or is_finalize: + if self.current_cursor is not None: + QApplication.restoreOverrideCursor() + self.current_cursor = None + else: + color_cc = self.get_current_color_scheme().cross_cursor + nc = Qt.ArrowCursor + + if self.drag_type == DragType.IMAGE_LOOK: + nc = Qt.ClosedHandCursor + else: + + if self.op_mode == OpMode.NONE: + nc = color_cc + if self.mouse_wire_poly is not None: + nc = Qt.PointingHandCursor + + elif self.op_mode == OpMode.DRAW_PTS: + nc = color_cc + elif self.op_mode == OpMode.EDIT_PTS: + nc = Qt.ArrowCursor + + if self.mouse_op_poly_pt_id is not None: + nc = Qt.PointingHandCursor + + if self.pt_edit_mode == PTEditMode.ADD_DEL: + + if self.mouse_op_poly_edge_id is not None and \ + self.mouse_op_poly_pt_id is None: + nc = color_cc + if self.current_cursor != nc: + if self.current_cursor is None: + QApplication.setOverrideCursor(nc) + else: + QApplication.changeOverrideCursor(nc) + self.current_cursor = nc + + def update_mouse_info(self, mouse_cli_pt=None): + """ + Update selected polys/edges/points by given mouse position + """ + if mouse_cli_pt is not None: + self.mouse_cli_pt = mouse_cli_pt.astype(np.float32) + + self.mouse_img_pt = self.cli_to_img_pt(self.mouse_cli_pt) + + new_mouse_hull_poly = self.get_poly_by_pt_in_hull(self.mouse_cli_pt) + + if self.mouse_hull_poly != new_mouse_hull_poly: + self.mouse_hull_poly = new_mouse_hull_poly + self.update_cursor() + self.update() + + new_mouse_wire_poly = self.get_poly_by_pt_near_wire(self.mouse_cli_pt) + + if self.mouse_wire_poly != new_mouse_wire_poly: + self.mouse_wire_poly = new_mouse_wire_poly + self.update_cursor() + self.update() + + if self.op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]: + new_mouse_op_poly_pt_id = self.get_poly_pt_id_under_pt (self.op_poly, self.mouse_cli_pt) + if self.mouse_op_poly_pt_id != new_mouse_op_poly_pt_id: + self.mouse_op_poly_pt_id = new_mouse_op_poly_pt_id + self.update_cursor() + self.update() + + new_mouse_op_poly_edge_id,\ + new_mouse_op_poly_edge_id_pt = self.get_poly_edge_id_pt_under_pt (self.op_poly, self.mouse_cli_pt) + if self.mouse_op_poly_edge_id != new_mouse_op_poly_edge_id: + self.mouse_op_poly_edge_id = new_mouse_op_poly_edge_id + self.update_cursor() + self.update() + + if (self.mouse_op_poly_edge_id_pt.__class__ != new_mouse_op_poly_edge_id_pt.__class__) or \ + (isinstance(self.mouse_op_poly_edge_id_pt, np.ndarray) and \ + all(self.mouse_op_poly_edge_id_pt != new_mouse_op_poly_edge_id_pt)): + + self.mouse_op_poly_edge_id_pt = new_mouse_op_poly_edge_id_pt + self.update_cursor() + self.update() + + + def action_undo_pt(self): + if self.drag_type == DragType.NONE: + if self.op_mode == OpMode.DRAW_PTS: + if self.op_poly.undo() == 0: + self.ie_polys.remove_poly (self.op_poly) + self.set_op_mode(OpMode.NONE) + self.update() + + def action_redo_pt(self): + if self.drag_type == DragType.NONE: + if self.op_mode == OpMode.DRAW_PTS: + self.op_poly.redo() + self.update() + + def action_delete_poly(self): + if self.op_mode == OpMode.EDIT_PTS and \ + self.drag_type == DragType.NONE and \ + self.pt_edit_mode == PTEditMode.MOVE: + # Delete current poly + self.ie_polys.remove_poly (self.op_poly) + self.set_op_mode(OpMode.NONE) + + # ==================================================================================== + # ==================================================================================== + # ================================== OVERRIDE QT METHODS ============================= + # ==================================================================================== + # ==================================================================================== + def on_keyPressEvent(self, ev): + if not self.initialized: + return + key = ev.key() + key_mods = int(ev.modifiers()) + if 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): + if not self.initialized: + return + key = ev.key() + key_mods = int(ev.modifiers()) + if 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): + super().enterEvent(ev) + self.mouse_in_widget = True + self.update_cursor() + + def leaveEvent(self, ev): + super().leaveEvent(ev) + self.mouse_in_widget = False + self.update_cursor() + + def mousePressEvent(self, ev): + super().mousePressEvent(ev) + if not self.initialized: + return + + self.update_mouse_info(QPoint_to_np(ev.pos())) + + btn = ev.button() + + if btn == Qt.LeftButton: + if self.op_mode == OpMode.NONE: + # Clicking in NO OPERATION mode + if self.mouse_wire_poly is not None: + # Click on wire on any poly -> switch to EDIT_MODE + self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.mouse_wire_poly) + else: + # Click on empty space -> create new poly with one point + new_poly = self.ie_polys.add_poly(self.poly_include_type) + self.ie_polys.sort() + new_poly.add_pt(*self.mouse_img_pt) + self.set_op_mode(OpMode.DRAW_PTS, op_poly=new_poly ) + + elif self.op_mode == OpMode.DRAW_PTS: + # Clicking in DRAW_PTS mode + if len(self.op_poly.get_pts()) >= 3 and self.mouse_op_poly_pt_id == 0: + # Click on first point -> close poly and switch to edit mode + self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.op_poly) + else: + # Click on empty space -> add point to current poly + self.op_poly.add_pt(*self.mouse_img_pt) + self.update() + + elif self.op_mode == OpMode.EDIT_PTS: + # Clicking in EDIT_PTS mode + + if self.mouse_op_poly_pt_id is not None: + # Click on point of op_poly + if self.pt_edit_mode == PTEditMode.ADD_DEL: + # with mode -> delete point + self.op_poly.remove_pt(self.mouse_op_poly_pt_id) + if self.op_poly.get_pts_count() < 3: + # not enough points -> remove poly + self.ie_polys.remove_poly (self.op_poly) + self.set_op_mode(OpMode.NONE) + self.update() + + elif self.drag_type == DragType.NONE: + # otherwise -> start drag + self.drag_type = DragType.POLY_PT + self.drag_cli_pt = self.mouse_cli_pt + self.drag_poly_pt_id = self.mouse_op_poly_pt_id + self.drag_poly_pt = self.op_poly.get_pts()[ self.drag_poly_pt_id ] + elif self.mouse_op_poly_edge_id is not None: + # Click on edge of op_poly + if self.pt_edit_mode == PTEditMode.ADD_DEL: + # with mode -> insert new point + 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.update() + else: + # Otherwise do nothing + pass + else: + # other cases -> unselect poly + self.set_op_mode(OpMode.NONE) + + + elif btn == Qt.MiddleButton: + if self.drag_type == DragType.NONE: + # Start image drag + self.drag_type = DragType.IMAGE_LOOK + self.drag_cli_pt = self.mouse_cli_pt + self.drag_img_look_pt = self.get_img_look_pt() + self.update_cursor() + + def mouseReleaseEvent(self, ev): + super().mouseReleaseEvent(ev) + if not self.initialized: + return + + self.update_mouse_info(QPoint_to_np(ev.pos())) + + btn = ev.button() + + if btn == Qt.LeftButton: + if self.op_mode == OpMode.EDIT_PTS: + if self.drag_type == DragType.POLY_PT: + self.drag_type = DragType.NONE + self.update() + + elif btn == Qt.MiddleButton: + if self.drag_type == DragType.IMAGE_LOOK: + self.drag_type = DragType.NONE + self.update_cursor() + self.update() + + def mouseMoveEvent(self, ev): + super().mouseMoveEvent(ev) + if not self.initialized: + return + + self.update_mouse_info(QPoint_to_np(ev.pos())) + + 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 + self.update() + + 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() + + def wheelEvent(self, ev): + super().wheelEvent(ev) + + if not self.initialized: + return + + mods = int(ev.modifiers()) + delta = ev.angleDelta() + + cli_pt = QPoint_to_np(ev.pos()) + + if self.drag_type == DragType.NONE: + sign = np.sign( delta.y() ) + prev_img_pos = self.cli_to_img_pt (cli_pt) + delta_scale = sign*0.2 + sign * self.get_view_scale() / 10.0 + self.view_scale = np.clip(self.get_view_scale() + delta_scale, 1.0, 20.0) + new_img_pos = self.cli_to_img_pt (cli_pt) + if sign > 0: + self.img_look_pt = self.get_img_look_pt() + (prev_img_pos-new_img_pos)#*1.5 + else: + QCursor.setPos ( self.mapToGlobal(QPoint_from_np(self.img_to_cli_pt(prev_img_pos))) ) + self.update() + + def paintEvent(self, event): + super().paintEvent(event) + if not self.initialized: + return + + qp = self.qp + qp.begin(self) + qp.setRenderHint(QPainter.Antialiasing) + qp.setRenderHint(QPainter.HighQualityAntialiasing) + qp.setRenderHint(QPainter.SmoothPixmapTransform) + + if self.op_mode == OpMode.VIEW_BAKED: + + src_rect = QRect(0, 0, *self.img_size) + dst_rect = self.img_to_cli_rect( src_rect ) + qp.drawPixmap(dst_rect, self.img_baked_pixmap, src_rect) + else: + if self.img_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.img_pixmap, src_rect) + + polys = self.ie_polys.get_polys() + polys_len = len(polys) + + color_scheme = self.get_current_color_scheme() + + pt_rad = self.canvas_config.pt_radius + pt_rad_x2 = pt_rad*2 + + pt_select_radius = self.canvas_config.pt_select_radius + + op_mode = self.op_mode + op_poly = self.op_poly + + for i,poly in enumerate(polys): + + selected_pt_path = QPainterPath() + poly_line_path = QPainterPath() + pts_line_path = QPainterPath() + + pt_remove_cli_pt = None + poly_pts = poly.get_pts() + for pt_id, img_pt in enumerate(poly_pts): + cli_pt = self.img_to_cli_pt(img_pt) + q_cli_pt = QPoint_from_np(cli_pt) + + if pt_id == 0: + poly_line_path.moveTo(q_cli_pt) + else: + poly_line_path.lineTo(q_cli_pt) + + + if poly == op_poly: + if self.op_mode == OpMode.DRAW_PTS or \ + (self.op_mode == OpMode.EDIT_PTS and \ + (self.pt_edit_mode == PTEditMode.MOVE) or \ + (self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id == pt_id) \ + ): + pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([0,-pt_rad])) ) + pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([0,pt_rad])) ) + pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([-pt_rad,0])) ) + pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([pt_rad,0])) ) + + if (self.op_mode == OpMode.EDIT_PTS and \ + self.pt_edit_mode == PTEditMode.ADD_DEL and \ + self.mouse_op_poly_pt_id == pt_id): + pt_remove_cli_pt = cli_pt + + if self.op_mode == OpMode.DRAW_PTS and \ + len(op_poly.get_pts()) >= 3 and pt_id == 0 and self.mouse_op_poly_pt_id == pt_id: + # Circle around poly point + selected_pt_path.addEllipse(q_cli_pt, pt_rad_x2, pt_rad_x2) + + + if poly == op_poly: + if op_mode == OpMode.DRAW_PTS: + # Line from last point to mouse + poly_line_path.lineTo( QPoint_from_np(self.mouse_cli_pt) ) + + if self.mouse_op_poly_pt_id is not None: + pass + + if self.mouse_op_poly_edge_id_pt is not None: + if self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id is None: + # Ready to insert point on edge + m_cli_pt = self.mouse_op_poly_edge_id_pt + pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([0,-pt_rad])) ) + pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([0,pt_rad])) ) + pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([-pt_rad,0])) ) + pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([pt_rad,0])) ) + + if len(poly_pts) >= 2: + # Closing poly line + poly_line_path.lineTo( QPoint_from_np(self.img_to_cli_pt(poly_pts[0])) ) + + # Draw calls + qp.setPen(color_scheme.pt_outline_pen) + qp.setBrush(QBrush()) + qp.drawPath(selected_pt_path) + + qp.setPen(color_scheme.poly_outline_solid_pen) + qp.setBrush(QBrush()) + qp.drawPath(pts_line_path) + + if poly.get_type() == SegIEPolyType.INCLUDE: + qp.setPen(color_scheme.poly_outline_solid_pen) + else: + qp.setPen(color_scheme.poly_outline_dot_pen) + + qp.setBrush(color_scheme.poly_unselected_brush) + if op_mode == OpMode.NONE: + if poly == self.mouse_wire_poly: + qp.setBrush(color_scheme.poly_selected_brush) + else: + if poly == op_poly: + qp.setBrush(color_scheme.poly_selected_brush) + + qp.drawPath(poly_line_path) + + if pt_remove_cli_pt is not None: + qp.setPen(color_scheme.poly_outline_solid_pen) + qp.setBrush(QBrush()) + + qp.drawLine( *(pt_remove_cli_pt + np.float32([-pt_rad_x2,-pt_rad_x2])), *(pt_remove_cli_pt + np.float32([pt_rad_x2,pt_rad_x2])) ) + qp.drawLine( *(pt_remove_cli_pt + np.float32([-pt_rad_x2,pt_rad_x2])), *(pt_remove_cli_pt + np.float32([pt_rad_x2,-pt_rad_x2])) ) + + qp.end() + + +class QCanvas(QFrame): + def __init__(self): + super().__init__() + self.canvas_control_bar = QCanvasControlsBar() + self.op = QCanvasOperator(self.canvas_control_bar) + self.l = QHBoxLayout() + self.l.setContentsMargins(0,0,0,0) + self.l.addWidget(self.canvas_control_bar) + self.l.addWidget(self.op) + self.setLayout(self.l) + +class LoaderQSubprocessor(QSubprocessor): + def __init__(self, image_paths, q_label, q_progressbar, on_finish_func ): + + self.image_paths = image_paths + self.image_paths_len = len(image_paths) + self.idxs = [*range(self.image_paths_len)] + + self.filtered_image_paths = self.image_paths.copy() + + self.image_paths_has_ie_polys = { image_path : False for image_path in self.image_paths } + + self.q_label = q_label + self.q_progressbar = q_progressbar + self.q_progressbar.setRange(0, self.image_paths_len) + self.q_progressbar.setValue(0) + self.q_progressbar.update() + self.on_finish_func = on_finish_func + self.done_count = 0 + super().__init__('LoaderQSubprocessor', LoaderQSubprocessor.Cli, 60) + + def get_data(self, host_dict): + if len (self.idxs) > 0: + idx = self.idxs.pop(0) + image_path = self.image_paths[idx] + self.q_label.setText(f'{image_path.name}') + + return idx, image_path + + return None + + def on_clients_finalized(self): + self.on_finish_func([x for x in self.filtered_image_paths if x is not None], self.image_paths_has_ie_polys) + + def on_data_return (self, host_dict, data): + self.idxs.insert(0, data[0]) + + def on_result (self, host_dict, data, result): + idx, has_dflimg, has_ie_polys = result + + if not has_dflimg: + self.filtered_image_paths[idx] = None + self.image_paths_has_ie_polys[self.image_paths[idx]] = has_ie_polys + + self.done_count += 1 + if self.q_progressbar is not None: + self.q_progressbar.setValue(self.done_count) + + class Cli(QSubprocessor.Cli): + def process_data(self, data): + idx, filename = data + dflimg = DFLIMG.load(filename) + if dflimg is not None and dflimg.has_data(): + ie_polys = SegIEPolys.load( dflimg.get_seg_ie_polys() ) + + return idx, True, ie_polys.has_polys() + return idx, False, False + + +class MainWindow(QXMainWindow): + + def __init__(self, input_dirpath, cfg_root_path): + super().__init__() + self.input_dirpath = input_dirpath + self.cfg_root_path = cfg_root_path + + 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_has_ie_polys = {} + + self.initialize_ui() + + # Loader + self.loading_frame = QFrame(self.main_canvas_frame) + self.loading_frame.setAutoFillBackground(True) + self.loading_frame.setFrameShape(QFrame.StyledPanel) + self.loader_label = QLabel() + self.loader_progress_bar = QProgressBar() + loading_frame_l = QVBoxLayout() + loading_frame_l.addWidget (self.loader_label, alignment=Qt.AlignBottom) + loading_frame_l.addWidget (self.loader_progress_bar, alignment=Qt.AlignTop) + self.loading_frame.setLayout(loading_frame_l) + + self.loader_subprocessor = LoaderQSubprocessor( image_paths=pathex.get_image_paths(input_dirpath, return_Path_class=True), + q_label=self.loader_label, + q_progressbar=self.loader_progress_bar, + on_finish_func=self.on_loader_finish ) + + + def on_loader_finish(self, image_paths, image_paths_has_ie_polys): + self.image_paths_done = [] + self.image_paths = image_paths + self.image_paths_has_ie_polys = image_paths_has_ie_polys + self.loading_frame.hide() + self.loading_frame = None + + self.process_next_image(first_initialization=True) + + def closeEvent(self, ev): + self.cfg_dict['geometry'] = self.saveGeometry().data() + self.cfg_path.write_bytes( pickle.dumps(self.cfg_dict) ) + + + def update_cached_images (self, count=5): + d = self.cached_QImages + + for image_path in self.image_paths_done[:-count]+self.image_paths[count:]: + if image_path in d: + del d[image_path] + + for image_path in self.image_paths[:count]+self.image_paths_done[-count:]: + if image_path not in d: + img = cv2_imread(image_path) + if img is not None: + d[image_path] = QImage_from_np(img) + + def load_QImage(self, image_path): + try: + img = self.cached_QImages.get(image_path, None) + if img is None: + img = QImage_from_np(cv2_imread(image_path)) + if img is None: + raise Exception(f'Unable to load {image_path}') + except: + io.log_err(f"{traceback.format_exc()}") + + return img + + 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) + + + def canvas_initialize(self, image_path, only_has_polys=False): + if only_has_polys and not self.image_paths_has_ie_polys[image_path]: + return False + + dflimg = DFLIMG.load(image_path) + ie_polys = SegIEPolys.load( dflimg.get_seg_ie_polys() ) + q_img = self.load_QImage(image_path) + + self.canvas.op.initialize ( q_img, ie_polys=ie_polys ) + return True + + def canvas_finalize(self, image_path): + dflimg = DFLIMG.load(image_path) + + ie_polys = SegIEPolys.load( 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.dump() ) + dflimg.save() + + self.canvas.op.finalize() + + + def process_prev_image(self): + key_mods = QApplication.keyboardModifiers() + step = 5 if key_mods == Qt.ShiftModifier else 1 + only_has_polys = key_mods == Qt.ControlModifier + + if self.canvas.op.is_initialized(): + self.canvas_finalize(self.image_paths[0]) + + while True: + for _ in range(step): + if len(self.image_paths_done) != 0: + self.image_paths.insert (0, self.image_paths_done.pop(-1)) + else: + break + if len(self.image_paths) == 0: + break + if self.canvas_initialize(self.image_paths[0], only_has_polys): + break + + self.update_cached_images() + self.update_preview_bar() + + def process_next_image(self, first_initialization=False): + key_mods = QApplication.keyboardModifiers() + + step = 0 if first_initialization else 5 if key_mods == Qt.ShiftModifier else 1 + only_has_polys = False if first_initialization else key_mods == Qt.ControlModifier + + if self.canvas.op.is_initialized(): + self.canvas_finalize(self.image_paths[0]) + + while True: + for _ in range(step): + if len(self.image_paths) != 0: + self.image_paths_done.append(self.image_paths.pop(0)) + else: + break + if len(self.image_paths) == 0: + break + if self.canvas_initialize(self.image_paths[0], only_has_polys): + break + + self.update_cached_images() + self.update_preview_bar() + + def initialize_ui(self): + + self.canvas = QCanvas() + + image_bar = self.image_bar = ImagePreviewSequenceBar(preview_images_count=9, icon_size=QUIConfig.preview_bar_icon_q_size.width()) + image_bar.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed ) + + + btn_prev_image = QXIconButton(QIconDB.left, QStringDB.btn_prev_image_tip, shortcut='A', click_func=self.process_prev_image) + btn_prev_image.setIconSize(QUIConfig.preview_bar_icon_q_size) + + btn_next_image = QXIconButton(QIconDB.right, QStringDB.btn_next_image_tip, shortcut='D', click_func=self.process_next_image) + btn_next_image.setIconSize(QUIConfig.preview_bar_icon_q_size) + + + preview_image_bar_frame_l = QHBoxLayout() + preview_image_bar_frame_l.setContentsMargins(0,0,0,0) + preview_image_bar_frame_l.addWidget ( btn_prev_image, alignment=Qt.AlignCenter) + preview_image_bar_frame_l.addWidget ( image_bar) + preview_image_bar_frame_l.addWidget ( btn_next_image, alignment=Qt.AlignCenter) + + preview_image_bar_frame = QFrame() + preview_image_bar_frame.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed ) + preview_image_bar_frame.setLayout(preview_image_bar_frame_l) + + preview_image_bar_l = QHBoxLayout() + preview_image_bar_l.addWidget (preview_image_bar_frame) + + preview_image_bar = QFrame() + preview_image_bar.setFrameShape(QFrame.StyledPanel) + preview_image_bar.setSizePolicy ( QSizePolicy.Expanding, QSizePolicy.Fixed ) + preview_image_bar.setLayout(preview_image_bar_l) + + main_canvas_l = QVBoxLayout() + main_canvas_l.setContentsMargins(0,0,0,0) + main_canvas_l.addWidget (self.canvas) + main_canvas_l.addWidget (preview_image_bar) + self.main_canvas_frame = QFrame() + self.main_canvas_frame.setLayout(main_canvas_l) + + self.main_l = QHBoxLayout() + self.main_l.setContentsMargins(0,0,0,0) + self.main_l.addWidget (self.main_canvas_frame) + + self.setLayout(self.main_l) + + geometry = self.cfg_dict.get('geometry', None) + if geometry is not None: + self.restoreGeometry(geometry) + else: + self.move( QPoint(0,0)) + + def resizeEvent(self, ev): + if self.loading_frame is not None: + self.loading_frame.resize( ev.size() ) + +def start(input_dirpath): + root_path = Path(__file__).parent + cfg_root_path = Path(tempfile.gettempdir()) + + QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True) + QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True) + + app = QApplication([]) + app.setApplicationName("XSegEditor") + app.setStyle('Fusion') + + QFontDatabase.addApplicationFont( str(root_path / 'gfx' / 'fonts' / 'NotoSans-Medium.ttf') ) + + font = QFont({'en' : 'Verdana', + 'ru' : 'Verdana', + 'zn' : 'SimSun'}[system_language]) + app.setFont( QFont('NotoSans')) + + QUIConfig.initialize() + QStringDB.initialize() + + QIconDB.initialize( root_path / 'gfx' / 'icons' ) + QCursorDB.initialize( root_path / 'gfx' / 'cursors' ) + + app.setWindowIcon(QIconDB.app_icon) + app.setPalette( QDarkPalette() ) + + win = MainWindow( input_dirpath=input_dirpath, cfg_root_path=cfg_root_path) + + win.show() + win.raise_() + + app.exec_() diff --git a/XSegEditor/gfx/cursors/cross_blue.png b/XSegEditor/gfx/cursors/cross_blue.png new file mode 100644 index 0000000..8915219 Binary files /dev/null and b/XSegEditor/gfx/cursors/cross_blue.png differ diff --git a/XSegEditor/gfx/cursors/cross_green.png b/XSegEditor/gfx/cursors/cross_green.png new file mode 100644 index 0000000..3ce16f0 Binary files /dev/null and b/XSegEditor/gfx/cursors/cross_green.png differ diff --git a/XSegEditor/gfx/cursors/cross_red.png b/XSegEditor/gfx/cursors/cross_red.png new file mode 100644 index 0000000..bb851ac Binary files /dev/null and b/XSegEditor/gfx/cursors/cross_red.png differ diff --git a/XSegEditor/gfx/fonts/NotoSans-Medium.ttf b/XSegEditor/gfx/fonts/NotoSans-Medium.ttf new file mode 100644 index 0000000..25050f7 Binary files /dev/null and b/XSegEditor/gfx/fonts/NotoSans-Medium.ttf differ diff --git a/XSegEditor/gfx/icons/app_icon.png b/XSegEditor/gfx/icons/app_icon.png new file mode 100644 index 0000000..16bc03e Binary files /dev/null and b/XSegEditor/gfx/icons/app_icon.png differ diff --git a/XSegEditor/gfx/icons/delete_poly.png b/XSegEditor/gfx/icons/delete_poly.png new file mode 100644 index 0000000..afd57d1 Binary files /dev/null and b/XSegEditor/gfx/icons/delete_poly.png differ diff --git a/XSegEditor/gfx/icons/down.png b/XSegEditor/gfx/icons/down.png new file mode 100644 index 0000000..873b719 Binary files /dev/null and b/XSegEditor/gfx/icons/down.png differ diff --git a/XSegEditor/gfx/icons/left.png b/XSegEditor/gfx/icons/left.png new file mode 100644 index 0000000..2118be6 Binary files /dev/null and b/XSegEditor/gfx/icons/left.png differ diff --git a/XSegEditor/gfx/icons/poly_color.psd b/XSegEditor/gfx/icons/poly_color.psd new file mode 100644 index 0000000..9a94957 Binary files /dev/null and b/XSegEditor/gfx/icons/poly_color.psd differ diff --git a/XSegEditor/gfx/icons/poly_color_blue.png b/XSegEditor/gfx/icons/poly_color_blue.png new file mode 100644 index 0000000..80b5222 Binary files /dev/null and b/XSegEditor/gfx/icons/poly_color_blue.png differ diff --git a/XSegEditor/gfx/icons/poly_color_green.png b/XSegEditor/gfx/icons/poly_color_green.png new file mode 100644 index 0000000..2db1fbb Binary files /dev/null and b/XSegEditor/gfx/icons/poly_color_green.png differ diff --git a/XSegEditor/gfx/icons/poly_color_red.png b/XSegEditor/gfx/icons/poly_color_red.png new file mode 100644 index 0000000..d04efff Binary files /dev/null and b/XSegEditor/gfx/icons/poly_color_red.png differ diff --git a/XSegEditor/gfx/icons/poly_type_exclude.png b/XSegEditor/gfx/icons/poly_type_exclude.png new file mode 100644 index 0000000..8e36bc3 Binary files /dev/null and b/XSegEditor/gfx/icons/poly_type_exclude.png differ diff --git a/XSegEditor/gfx/icons/poly_type_include.png b/XSegEditor/gfx/icons/poly_type_include.png new file mode 100644 index 0000000..5f16c15 Binary files /dev/null and b/XSegEditor/gfx/icons/poly_type_include.png differ diff --git a/XSegEditor/gfx/icons/poly_type_source.psd b/XSegEditor/gfx/icons/poly_type_source.psd new file mode 100644 index 0000000..50943d0 Binary files /dev/null and b/XSegEditor/gfx/icons/poly_type_source.psd differ diff --git a/XSegEditor/gfx/icons/pt_edit_mode.png b/XSegEditor/gfx/icons/pt_edit_mode.png new file mode 100644 index 0000000..d385fc2 Binary files /dev/null and b/XSegEditor/gfx/icons/pt_edit_mode.png differ diff --git a/XSegEditor/gfx/icons/pt_edit_mode_source.psd b/XSegEditor/gfx/icons/pt_edit_mode_source.psd new file mode 100644 index 0000000..f73e310 Binary files /dev/null and b/XSegEditor/gfx/icons/pt_edit_mode_source.psd differ diff --git a/XSegEditor/gfx/icons/redo_pt.png b/XSegEditor/gfx/icons/redo_pt.png new file mode 100644 index 0000000..aa73329 Binary files /dev/null and b/XSegEditor/gfx/icons/redo_pt.png differ diff --git a/XSegEditor/gfx/icons/redo_pt_source.psd b/XSegEditor/gfx/icons/redo_pt_source.psd new file mode 100644 index 0000000..2771f77 Binary files /dev/null and b/XSegEditor/gfx/icons/redo_pt_source.psd differ diff --git a/XSegEditor/gfx/icons/right.png b/XSegEditor/gfx/icons/right.png new file mode 100644 index 0000000..b4ef220 Binary files /dev/null and b/XSegEditor/gfx/icons/right.png differ diff --git a/XSegEditor/gfx/icons/undo_pt.png b/XSegEditor/gfx/icons/undo_pt.png new file mode 100644 index 0000000..7a4464c Binary files /dev/null and b/XSegEditor/gfx/icons/undo_pt.png differ diff --git a/XSegEditor/gfx/icons/undo_pt_source.psd b/XSegEditor/gfx/icons/undo_pt_source.psd new file mode 100644 index 0000000..98b9d1a Binary files /dev/null and b/XSegEditor/gfx/icons/undo_pt_source.psd differ diff --git a/XSegEditor/gfx/icons/up.png b/XSegEditor/gfx/icons/up.png new file mode 100644 index 0000000..f3368b6 Binary files /dev/null and b/XSegEditor/gfx/icons/up.png differ diff --git a/XSegEditor/gfx/icons/view_baked.png b/XSegEditor/gfx/icons/view_baked.png new file mode 100644 index 0000000..3e32142 Binary files /dev/null and b/XSegEditor/gfx/icons/view_baked.png differ diff --git a/core/imagelib/SegIEPolys.py b/core/imagelib/SegIEPolys.py new file mode 100644 index 0000000..e658711 --- /dev/null +++ b/core/imagelib/SegIEPolys.py @@ -0,0 +1,152 @@ +import numpy as np +import cv2 +from enum import IntEnum + + +class SegIEPolyType(IntEnum): + EXCLUDE = 0 + INCLUDE = 1 + + + +class SegIEPoly(): + def __init__(self, type=None, pts=None, **kwargs): + self.type = type + + if pts is None: + pts = np.empty( (0,2), dtype=np.float32 ) + else: + pts = np.float32(pts) + self.pts = pts + self.n_max = self.n = len(pts) + + def dump(self): + return {'type': int(self.type), + 'pts' : self.get_pts(), + } + + def identical(self, b): + if self.n != b.n: + return False + return (self.pts[0:self.n] == b.pts[0:b.n]).all() + + def get_type(self): + return self.type + + def add_pt(self, x, y): + self.pts = np.append(self.pts[0:self.n], [ ( float(x), float(y) ) ], axis=0).astype(np.float32) + self.n_max = self.n = self.n + 1 + + def undo(self): + self.n = max(0, self.n-1) + return self.n + + def redo(self): + self.n = min(len(self.pts), self.n+1) + return self.n + + def redo_clip(self): + self.pts = self.pts[0:self.n] + self.n_max = self.n + + def insert_pt(self, n, pt): + if n < 0 or n > self.n: + raise ValueError("insert_pt out of range") + self.pts = np.concatenate( (self.pts[0:n], pt[None,...].astype(np.float32), self.pts[n:]), axis=0) + self.n_max = self.n = self.n+1 + + def remove_pt(self, n): + if n < 0 or n >= self.n: + raise ValueError("remove_pt out of range") + self.pts = np.concatenate( (self.pts[0:n], self.pts[n+1:]), axis=0) + self.n_max = self.n = self.n-1 + + def get_last_point(self): + return self.pts[self.n-1].copy() + + def get_pts(self): + return self.pts[0:self.n].copy() + + def get_pts_count(self): + return self.n + + def set_point(self, id, pt): + self.pts[id] = pt + + def set_points(self, pts): + self.pts = np.array(pts) + self.n_max = self.n = len(pts) + + + + +class SegIEPolys(): + def __init__(self): + self.polys = [] + + def identical(self, b): + polys_len = len(self.polys) + o_polys_len = len(b.polys) + if polys_len != o_polys_len: + return False + + return all ([ a_poly.identical(b_poly) for a_poly, b_poly in zip(self.polys, b.polys) ]) + + def add_poly(self, ie_poly_type): + poly = SegIEPoly(ie_poly_type) + self.polys.append (poly) + return poly + + def remove_poly(self, poly): + if poly in self.polys: + self.polys.remove(poly) + + def has_polys(self): + return len(self.polys) != 0 + + def get_poly(self, id): + return self.polys[id] + + def get_polys(self): + return self.polys + + def get_pts_count(self): + return sum([poly.get_pts_count() for poly in self.polys]) + + def sort(self): + poly_by_type = { SegIEPolyType.EXCLUDE : [], SegIEPolyType.INCLUDE : [] } + + for poly in self.polys: + poly_by_type[poly.type].append(poly) + + self.polys = poly_by_type[SegIEPolyType.INCLUDE] + poly_by_type[SegIEPolyType.EXCLUDE] + + def __iter__(self): + for poly in self.polys: + yield poly + + def overlay_mask(self, mask): + h,w,c = mask.shape + white = (1,)*c + black = (0,)*c + for poly in self.polys: + pts = poly.get_pts().astype(np.int32) + if len(pts) != 0: + cv2.fillPoly(mask, [pts], white if poly.type == SegIEPolyType.INCLUDE else black ) + + def dump(self): + return {'polys' : [ poly.dump() for poly in self.polys ] } + + @staticmethod + def load(data=None): + ie_polys = SegIEPolys() + if data is not None: + if isinstance(data, list): + # Backward comp + ie_polys.polys = [ SegIEPoly(type=type, pts=pts) for (type, pts) in data ] + elif isinstance(data, dict): + ie_polys.polys = [ SegIEPoly(**poly_cfg) for poly_cfg in data['polys'] ] + + ie_polys.sort() + + return ie_polys \ No newline at end of file diff --git a/core/imagelib/__init__.py b/core/imagelib/__init__.py index fec43f2..2a23e60 100644 --- a/core/imagelib/__init__.py +++ b/core/imagelib/__init__.py @@ -16,6 +16,7 @@ from .color_transfer import color_transfer, color_transfer_mix, color_transfer_s from .common import normalize_channels, cut_odd_image, overlay_alpha_image from .IEPolys import IEPolys +from .SegIEPolys import * from .blursharpen import LinearMotionBlur, blursharpen diff --git a/core/qtex/QSubprocessor.py b/core/qtex/QSubprocessor.py new file mode 100644 index 0000000..223f50c --- /dev/null +++ b/core/qtex/QSubprocessor.py @@ -0,0 +1,262 @@ +import multiprocessing +import sys +import time +import traceback + +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * + +from core.interact import interact as io + +from .qtex import * + +class QSubprocessor(object): + """ + + """ + + class Cli(object): + def __init__ ( self, client_dict ): + s2c = multiprocessing.Queue() + c2s = multiprocessing.Queue() + self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,s2c,c2s) ) + self.s2c = s2c + self.c2s = c2s + self.p.daemon = True + self.p.start() + + self.state = None + self.sent_time = None + self.sent_data = None + self.name = None + self.host_dict = None + + def kill(self): + self.p.terminate() + self.p.join() + + #overridable optional + def on_initialize(self, client_dict): + #initialize your subprocess here using client_dict + pass + + #overridable optional + def on_finalize(self): + #finalize your subprocess here + pass + + #overridable + def process_data(self, data): + #process 'data' given from host and return result + raise NotImplementedError + + #overridable optional + def get_data_name (self, data): + #return string identificator of your 'data' + return "undefined" + + def log_info(self, msg): self.c2s.put ( {'op': 'log_info', 'msg':msg } ) + def log_err(self, msg): self.c2s.put ( {'op': 'log_err' , 'msg':msg } ) + def progress_bar_inc(self, c): self.c2s.put ( {'op': 'progress_bar_inc' , 'c':c } ) + + def _subprocess_run(self, client_dict, s2c, c2s): + self.c2s = c2s + data = None + try: + self.on_initialize(client_dict) + c2s.put ( {'op': 'init_ok'} ) + while True: + msg = s2c.get() + op = msg.get('op','') + if op == 'data': + data = msg['data'] + result = self.process_data (data) + c2s.put ( {'op': 'success', 'data' : data, 'result' : result} ) + data = None + elif op == 'close': + break + time.sleep(0.001) + self.on_finalize() + c2s.put ( {'op': 'finalized'} ) + except Exception as e: + c2s.put ( {'op': 'error', 'data' : data} ) + if data is not None: + print ('Exception while process data [%s]: %s' % (self.get_data_name(data), traceback.format_exc()) ) + else: + print ('Exception: %s' % (traceback.format_exc()) ) + c2s.close() + s2c.close() + self.c2s = None + + # disable pickling + def __getstate__(self): + return dict() + def __setstate__(self, d): + self.__dict__.update(d) + + #overridable + def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005): + if not issubclass(SubprocessorCli_class, QSubprocessor.Cli): + raise ValueError("SubprocessorCli_class must be subclass of QSubprocessor.Cli") + + self.name = name + self.SubprocessorCli_class = SubprocessorCli_class + self.no_response_time_sec = no_response_time_sec + self.io_loop_sleep_time = io_loop_sleep_time + + self.clis = [] + + #getting info about name of subprocesses, host and client dicts, and spawning them + for name, host_dict, client_dict in self.process_info_generator(): + try: + cli = self.SubprocessorCli_class(client_dict) + cli.state = 1 + cli.sent_time = 0 + cli.sent_data = None + cli.name = name + cli.host_dict = host_dict + + self.clis.append (cli) + except: + raise Exception (f"Unable to start subprocess {name}. Error: {traceback.format_exc()}") + + if len(self.clis) == 0: + raise Exception ("Unable to start QSubprocessor '%s' " % (self.name)) + + #waiting subprocesses their success(or not) initialization + while True: + for cli in self.clis[:]: + while not cli.c2s.empty(): + obj = cli.c2s.get() + op = obj.get('op','') + if op == 'init_ok': + cli.state = 0 + elif op == 'log_info': + io.log_info(obj['msg']) + elif op == 'log_err': + io.log_err(obj['msg']) + elif op == 'error': + cli.kill() + self.clis.remove(cli) + break + if all ([cli.state == 0 for cli in self.clis]): + break + io.process_messages(0.005) + + if len(self.clis) == 0: + raise Exception ( "Unable to start subprocesses." ) + + #ok some processes survived, initialize host logic + self.on_clients_initialized() + + self.q_timer = QTimer() + self.q_timer.timeout.connect(self.tick) + self.q_timer.start(5) + + #overridable + def process_info_generator(self): + #yield per process (name, host_dict, client_dict) + for i in range(min(multiprocessing.cpu_count(), 8) ): + yield 'CPU%d' % (i), {}, {} + + #overridable optional + def on_clients_initialized(self): + #logic when all subprocesses initialized and ready + pass + + #overridable optional + def on_clients_finalized(self): + #logic when all subprocess finalized + pass + + #overridable + def get_data(self, host_dict): + #return data for processing here + raise NotImplementedError + + #overridable + def on_data_return (self, host_dict, data): + #you have to place returned 'data' back to your queue + raise NotImplementedError + + #overridable + def on_result (self, host_dict, data, result): + #your logic what to do with 'result' of 'data' + raise NotImplementedError + + def tick(self): + for cli in self.clis[:]: + while not cli.c2s.empty(): + obj = cli.c2s.get() + op = obj.get('op','') + if op == 'success': + #success processed data, return data and result to on_result + self.on_result (cli.host_dict, obj['data'], obj['result']) + self.sent_data = None + cli.state = 0 + elif op == 'error': + #some error occured while process data, returning chunk to on_data_return + if 'data' in obj.keys(): + self.on_data_return (cli.host_dict, obj['data'] ) + #and killing process + cli.kill() + self.clis.remove(cli) + elif op == 'log_info': + io.log_info(obj['msg']) + elif op == 'log_err': + io.log_err(obj['msg']) + elif op == 'progress_bar_inc': + io.progress_bar_inc(obj['c']) + + for cli in self.clis[:]: + if cli.state == 1: + if cli.sent_time != 0 and self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec: + #subprocess busy too long + io.log_info ( '%s doesnt response, terminating it.' % (cli.name) ) + self.on_data_return (cli.host_dict, cli.sent_data ) + cli.kill() + self.clis.remove(cli) + + for cli in self.clis[:]: + if cli.state == 0: + #free state of subprocess, get some data from get_data + data = self.get_data(cli.host_dict) + if data is not None: + #and send it to subprocess + cli.s2c.put ( {'op': 'data', 'data' : data} ) + cli.sent_time = time.time() + cli.sent_data = data + cli.state = 1 + + if all ([cli.state == 0 for cli in self.clis]): + #gracefully terminating subprocesses + for cli in self.clis[:]: + cli.s2c.put ( {'op': 'close'} ) + cli.sent_time = time.time() + + while True: + for cli in self.clis[:]: + terminate_it = False + while not cli.c2s.empty(): + obj = cli.c2s.get() + obj_op = obj['op'] + if obj_op == 'finalized': + terminate_it = True + break + + if (time.time() - cli.sent_time) > 30: + terminate_it = True + + if terminate_it: + cli.state = 2 + cli.kill() + + if all ([cli.state == 2 for cli in self.clis]): + break + + #finalizing host logic + self.q_timer.stop() + self.q_timer = None + self.on_clients_finalized() + diff --git a/core/qtex/QXIconButton.py b/core/qtex/QXIconButton.py new file mode 100644 index 0000000..7f9d3eb --- /dev/null +++ b/core/qtex/QXIconButton.py @@ -0,0 +1,83 @@ +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * + +from localization import StringsDB +from .QXMainWindow import * + +class QXIconButton(QPushButton): + """ + Custom Icon button that works through keyEvent system, without shortcut of QAction + works only with QXMainWindow as global window class + currently works only with one-key shortcut + """ + + def __init__(self, icon, + tooltip=None, + shortcut=None, + click_func=None, + first_repeat_delay=300, + repeat_delay=20, + ): + + super().__init__(icon, "") + + self.setIcon(icon) + + if shortcut is not None: + tooltip = f"{tooltip} ( {StringsDB['S_HOT_KEY'] }: {shortcut} )" + + self.setToolTip(tooltip) + + + self.seq = QKeySequence(shortcut) if shortcut is not None else None + + QXMainWindow.inst.add_keyPressEvent_listener ( self.on_keyPressEvent ) + QXMainWindow.inst.add_keyReleaseEvent_listener ( self.on_keyReleaseEvent ) + + self.click_func = click_func + self.first_repeat_delay = first_repeat_delay + self.repeat_delay = repeat_delay + self.repeat_timer = None + + self.op_device = None + + self.pressed.connect( lambda : self.action(is_pressed=True) ) + self.released.connect( lambda : self.action(is_pressed=False) ) + + def action(self, is_pressed=None, op_device=None): + if self.click_func is None: + return + + if is_pressed is not None: + if is_pressed: + if self.repeat_timer is None: + self.click_func() + self.repeat_timer = QTimer() + self.repeat_timer.timeout.connect(self.action) + self.repeat_timer.start(self.first_repeat_delay) + else: + if self.repeat_timer is not None: + self.repeat_timer.stop() + self.repeat_timer = None + else: + self.click_func() + if self.repeat_timer is not None: + self.repeat_timer.setInterval(self.repeat_delay) + + def on_keyPressEvent(self, ev): + key = ev.key() + if ev.isAutoRepeat(): + return + + if self.seq is not None: + if key == self.seq[0]: + self.action(is_pressed=True) + + def on_keyReleaseEvent(self, ev): + key = ev.key() + if ev.isAutoRepeat(): + return + if self.seq is not None: + if key == self.seq[0]: + self.action(is_pressed=False) diff --git a/core/qtex/QXMainWindow.py b/core/qtex/QXMainWindow.py new file mode 100644 index 0000000..a50e597 --- /dev/null +++ b/core/qtex/QXMainWindow.py @@ -0,0 +1,34 @@ +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * + +class QXMainWindow(QWidget): + """ + Custom mainwindow class that provides global single instance and event listeners + """ + inst = None + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if QXMainWindow.inst is not None: + raise Exception("QXMainWindow can only be one.") + QXMainWindow.inst = self + + self.keyPressEvent_listeners = [] + self.keyReleaseEvent_listeners = [] + self.setFocusPolicy(Qt.WheelFocus) + + def add_keyPressEvent_listener(self, func): + self.keyPressEvent_listeners.append (func) + + def add_keyReleaseEvent_listener(self, func): + self.keyReleaseEvent_listeners.append (func) + + def keyPressEvent(self, ev): + super().keyPressEvent(ev) + for func in self.keyPressEvent_listeners: + func(ev) + + def keyReleaseEvent(self, ev): + super().keyReleaseEvent(ev) + for func in self.keyReleaseEvent_listeners: + func(ev) \ No newline at end of file diff --git a/core/qtex/__init__.py b/core/qtex/__init__.py new file mode 100644 index 0000000..2cb44b5 --- /dev/null +++ b/core/qtex/__init__.py @@ -0,0 +1,3 @@ +from .qtex import * +from .QSubprocessor import * +from .QXIconButton import * \ No newline at end of file diff --git a/core/qtex/qtex.py b/core/qtex/qtex.py new file mode 100644 index 0000000..700b922 --- /dev/null +++ b/core/qtex/qtex.py @@ -0,0 +1,79 @@ +import numpy as np +from PyQt5.QtCore import * +from PyQt5.QtGui import * +from PyQt5.QtWidgets import * +from localization import StringsDB + +from .QXMainWindow import * + + +class QActionEx(QAction): + def __init__(self, icon, text, shortcut=None, trigger_func=None, shortcut_in_tooltip=False, is_checkable=False, is_auto_repeat=False ): + super().__init__(icon, text) + if shortcut is not None: + self.setShortcut(shortcut) + if shortcut_in_tooltip: + + self.setToolTip( f"{text} ( {StringsDB['S_HOT_KEY'] }: {shortcut} )") + + if trigger_func is not None: + self.triggered.connect(trigger_func) + if is_checkable: + self.setCheckable(True) + self.setAutoRepeat(is_auto_repeat) + +def QImage_from_np(img): + if img.dtype != np.uint8: + raise ValueError("img should be in np.uint8 format") + + h,w,c = img.shape + if c == 1: + fmt = QImage.Format_Grayscale8 + elif c == 3: + fmt = QImage.Format_BGR888 + elif c == 4: + fmt = QImage.Format_ARGB32 + else: + raise ValueError("unsupported channel count") + + return QImage(img.data, w, h, c*w, fmt ) + +def QImage_to_np(q_img): + q_img = q_img.convertToFormat(QImage.Format_BGR888) + + width = q_img.width() + height = q_img.height() + + b = q_img.constBits() + b.setsize(height * width * 3) + arr = np.frombuffer(b, np.uint8).reshape((height, width, 3)) + return arr#[::-1] + +def QPixmap_from_np(img): + return QPixmap.fromImage(QImage_from_np(img)) + +def QPoint_from_np(n): + return QPoint(*n.astype(np.int)) + +def QPoint_to_np(q): + return np.int32( [q.x(), q.y()] ) + +def QSize_to_np(q): + return np.int32( [q.width(), q.height()] ) + +class QDarkPalette(QPalette): + def __init__(self): + super().__init__() + self.setColor(QPalette.Window, QColor(53, 53, 53)) + self.setColor(QPalette.WindowText, Qt.white) + self.setColor(QPalette.Base, QColor(25, 25, 25)) + self.setColor(QPalette.AlternateBase, QColor(53, 53, 53)) + self.setColor(QPalette.ToolTipBase, Qt.white) + self.setColor(QPalette.ToolTipText, Qt.white) + self.setColor(QPalette.Text, Qt.white) + self.setColor(QPalette.Button, QColor(53, 53, 53)) + self.setColor(QPalette.ButtonText, Qt.white) + self.setColor(QPalette.BrightText, Qt.red) + self.setColor(QPalette.Link, QColor(42, 130, 218)) + self.setColor(QPalette.Highlight, QColor(42, 130, 218)) + self.setColor(QPalette.HighlightedText, Qt.black) \ No newline at end of file diff --git a/localization/__init__.py b/localization/__init__.py index f3bcf09..ccd8c6e 100644 --- a/localization/__init__.py +++ b/localization/__init__.py @@ -1,2 +1,2 @@ -from .localization import get_default_ttf_font_name +from .localization import StringsDB, system_language, get_default_ttf_font_name diff --git a/localization/localization.py b/localization/localization.py index a603285..eeb053d 100644 --- a/localization/localization.py +++ b/localization/localization.py @@ -4,6 +4,8 @@ import locale system_locale = locale.getdefaultlocale()[0] # system_locale may be nil system_language = system_locale[0:2] if system_locale is not None else "en" +if system_language not in ['en','ru','zn']: + system_language = 'en' windows_font_name_map = { 'en' : 'cour', @@ -28,3 +30,13 @@ def get_default_ttf_font_name(): if platform[0:3] == 'win': return windows_font_name_map.get(system_language, 'cour') elif platform == 'darwin': return darwin_font_name_map.get(system_language, 'cour') else: return linux_font_name_map.get(system_language, 'cour') + +SID_HOT_KEY = 1 + +if system_language == 'en': + StringsDB = {'S_HOT_KEY' : 'hot key'} +elif system_language == 'ru': + StringsDB = {'S_HOT_KEY' : 'горячая клавиша'} +elif system_language == 'zn': + StringsDB = {'S_HOT_KEY' : '热键'} + \ No newline at end of file diff --git a/main.py b/main.py index 75b0a55..78aafe0 100644 --- a/main.py +++ b/main.py @@ -264,26 +264,16 @@ if __name__ == "__main__": p.set_defaults (func=process_dev_test) # ========== XSeg util - xseg_parser = subparsers.add_parser( "xseg", help="XSeg utils.").add_subparsers() + p = subparsers.add_parser( "xsegeditor", help="XSegEditor.") - def process_xseg_merge(arguments): + def process_xsegeditor(arguments): osex.set_process_lowest_prio() - from mainscripts import XSegUtil - XSegUtil.merge(arguments.input_dir) - p = xseg_parser.add_parser( "merge", help="") + from XSegEditor import XSegEditor + XSegEditor.start (Path(arguments.input_dir)) p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - p.set_defaults (func=process_xseg_merge) - - def process_xseg_split(arguments): - osex.set_process_lowest_prio() - from mainscripts import XSegUtil - XSegUtil.split(arguments.input_dir) - - p = xseg_parser.add_parser( "split", help="") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - - p.set_defaults (func=process_xseg_split) + p.set_defaults (func=process_xsegeditor) + def bad_args(arguments): parser.print_help() diff --git a/requirements-cuda.txt b/requirements-cuda.txt index 905aaef..8df3478 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -7,4 +7,5 @@ scikit-image==0.14.2 scipy==1.4.1 colorama labelme==4.2.9 -tensorflow-gpu==1.13.2 \ No newline at end of file +tensorflow-gpu==1.13.2 +pyqt5 \ No newline at end of file diff --git a/samplelib/Sample.py b/samplelib/Sample.py index 1f9908d..f4165b9 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -7,7 +7,7 @@ import numpy as np from core.cv2ex import * from DFLIMG import * from facelib import LandmarksProcessor -from core.imagelib import IEPolys +from core.imagelib import IEPolys, SegIEPolys class SampleType(IntEnum): IMAGE = 0 #raw image @@ -54,7 +54,7 @@ class Sample(object): self.shape = shape self.landmarks = np.array(landmarks) if landmarks is not None else None self.ie_polys = IEPolys.load(ie_polys) - self.seg_ie_polys = IEPolys.load(seg_ie_polys) + self.seg_ie_polys = SegIEPolys.load(seg_ie_polys) self.eyebrows_expand_mod = eyebrows_expand_mod if eyebrows_expand_mod is not None else 1.0 self.source_filename = source_filename self.person_name = person_name diff --git a/samplelib/SampleGeneratorFaceXSeg.py b/samplelib/SampleGeneratorFaceXSeg.py index c744d40..2b7b061 100644 --- a/samplelib/SampleGeneratorFaceXSeg.py +++ b/samplelib/SampleGeneratorFaceXSeg.py @@ -26,8 +26,8 @@ class SampleGeneratorFaceXSeg(SampleGeneratorBase): samples = [] for path in paths: samples += SampleLoader.load (SampleType.FACE, path) - - seg_samples = [ sample for sample in samples if sample.seg_ie_polys.get_total_points() != 0] + + seg_samples = [ sample for sample in samples if sample.seg_ie_polys.get_pts_count() != 0] seg_samples_len = len(seg_samples) if seg_samples_len == 0: raise Exception(f"No segmented faces found.")