自动画遮罩功能

day270010678

星级打分

• 1
• 2
• 3
• 4
• 5

平均分:NAN  参与人数:0  我的评分:未评

我刚才小幅修改了下源码，添加了个自动画遮罩的工具







原理很简单，就是加入Face Parsing人脸解析，然后把人脸直接解析出来，然后根据需求自己小幅修改，我主要是刚才画遮罩太累，就自己加了这个小功能，需要的自己修改。本质就是ONNX 人脸解析模型 + onnxruntime

安装onnxruntime库，pip安装就行了，另外下载模型，去https://hf-mirror.com/jonathandinu/face-parsing/tree/main/onnx里面下载这两个模型放到XSegEditor\models目录里面就行了

修改的文件就是XSegEditor/OcclusionDetector.py这是我自己建的ONNX 人脸解析检测器，子进程运行

1. import multiprocessing
   
2. import numpy as np
   
3. import cv2
   
4. from pathlib import Path
   
5. 
   
6. 
   
7. class OcclusionDetector:
   
8. 
   
9.     CMD_LOAD = 1
   
10.     CMD_DETECT = 2
    
11.     CMD_SHUTDOWN = 3
    
12.     RSP_LOADED = 10
    
13.     RSP_RESULT = 11
    
14.     RSP_ERROR = 12
    
15. 
    
16.     INCLUDE_LABELS = {1, 2, 4, 5, 6, 7, 11, 12}
    
17.     EXCLUDE_LABELS = {3, 8, 9, 10, 13, 14, 15, 16, 17, 18}
    
18.     HEAD_LABELS = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}
    
19. 
    
20.     MODEL_INPUT_SIZE = 512
    
21. 
    
22.     def __init__(self, model_path=None):
    
23.         self.model_path = model_path or str(
    
24.             Path(__file__).parent / 'models' / 'model_quantized.onnx')
    
25.         self._req_q = multiprocessing.Queue()
    
26.         self._rsp_q = multiprocessing.Queue()
    
27.         self._proc = None
    
28. 
    
29.     def start(self):
    
30.         if self._proc is not None and self._proc.is_alive():
    
31.             return
    
32.         self._proc = multiprocessing.Process(
    
33.             target=self._subprocess_run,
    
34.             args=(self._req_q, self._rsp_q, self.model_path),
    
35.             daemon=True,
    
36.         )
    
37.         self._proc.start()
    
38.         rsp = self._rsp_q.get(timeout=30)
    
39.         if rsp[0] == self.RSP_ERROR:
    
40.             raise RuntimeError(f"OcclusionDetector load failed: {rsp[1]}")
    
41. 
    
42.     def detect(self, img, landmarks=None):
    
43.         if self._proc is None or not self._proc.is_alive():
    
44.             self.start()
    
45.         self._req_q.put((self.CMD_DETECT, img))
    
46.         rsp = self._rsp_q.get(timeout=60)
    
47.         if rsp[0] == self.RSP_ERROR:
    
48.             raise RuntimeError(f"OcclusionDetector detect failed: {rsp[1]}")
    
49.         return rsp[1]
    
50. 
    
51.     def shutdown(self):
    
52.         if self._proc is not None and self._proc.is_alive():
    
53.             self._req_q.put((self.CMD_SHUTDOWN,))
    
54.             self._proc.join(timeout=5)
    
55.             if self._proc.is_alive():
    
56.                 self._proc.terminate()
    
57.             self._proc = None
    
58. 
    
59.     @staticmethod
    
60.     def _subprocess_run(req_q, rsp_q, model_path):
    
61.         try:
    
62.             import onnxruntime as ort
    
63.             sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
    
64.             rsp_q.put((OcclusionDetector.RSP_LOADED,))
    
65. 
    
66.             while True:
    
67.                 cmd = req_q.get()
    
68.                 if cmd[0] == OcclusionDetector.CMD_SHUTDOWN:
    
69.                     break
    
70.                 elif cmd[0] == OcclusionDetector.CMD_DETECT:
    
71.                     try:
    
72.                         img = cmd[1]
    
73.                         result = OcclusionDetector._run_detection(sess, img)
    
74.                         rsp_q.put((OcclusionDetector.RSP_RESULT, result))
    
75.                     except Exception as e:
    
76.                         rsp_q.put((OcclusionDetector.RSP_ERROR, str(e)))
    
77.         except Exception as e:
    
78.             rsp_q.put((OcclusionDetector.RSP_ERROR, str(e)))
    
79. 
    
80.     @staticmethod
    
81.     def _run_detection(sess, img):
    
82.         h, w = img.shape[:2]
    
83.         input_img = OcclusionDetector._preprocess(img)
    
84.         logits = sess.run(None, {'pixel_values': input_img})[0]
    
85.         seg_map = OcclusionDetector._logits_to_segmap(logits, h, w)
    
86.         return OcclusionDetector._seg_map_to_polys(seg_map, h, w)
    
87. 
    
88.     @staticmethod
    
89.     def _preprocess(img):
    
90.         blob = cv2.dnn.blobFromImage(img, 1.0 / 255.0,
    
91.                                       (OcclusionDetector.MODEL_INPUT_SIZE,
    
92.                                        OcclusionDetector.MODEL_INPUT_SIZE),
    
93.                                       swapRB=True, crop=False)
    
94.         return blob.astype(np.float32)
    
95. 
    
96.     @staticmethod
    
97.     def _logits_to_segmap(logits, h, w):
    
98.         seg = logits[0]
    
99.         seg = np.argmax(seg, axis=0)
    
100.         seg = seg.astype(np.uint8)
     
101.         seg = cv2.resize(seg, (w, h), interpolation=cv2.INTER_NEAREST)
     
102.         return seg
     
103. 
     
104.     @staticmethod
     
105.     def _seg_map_to_polys(seg_map, h, w):
     
106.         head_mask = np.zeros((h, w), dtype=np.uint8)
     
107.         for label in OcclusionDetector.HEAD_LABELS:
     
108.             head_mask[seg_map == label] = 255
     
109. 
     
110.         include_mask = np.zeros((h, w), dtype=np.uint8)
     
111.         for label in OcclusionDetector.INCLUDE_LABELS:
     
112.             include_mask[seg_map == label] = 255
     
113. 
     
114.         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
     
115.         include_mask = cv2.morphologyEx(include_mask, cv2.MORPH_CLOSE, kernel, iterations=1)
     
116.         include_mask = cv2.morphologyEx(include_mask, cv2.MORPH_OPEN, kernel, iterations=1)
     
117. 
     
118.         exclude_polys = []
     
119.         for label in OcclusionDetector.EXCLUDE_LABELS:
     
120.             label_mask = (seg_map == label).astype(np.uint8) * 255
     
121.             overlap = cv2.bitwise_and(label_mask, head_mask)
     
122.             if np.sum(overlap > 0) == 0:
     
123.                 continue
     
124.             overlap = cv2.morphologyEx(overlap, cv2.MORPH_CLOSE, kernel, iterations=1)
     
125.             overlap = cv2.morphologyEx(overlap, cv2.MORPH_OPEN, kernel, iterations=1)
     
126.             exclude_polys.extend(
     
127.                 OcclusionDetector.mask_to_polys(overlap, min_area=80, epsilon_factor=0.003))
     
128. 
     
129.         return {
     
130.             'include': OcclusionDetector.mask_to_polys(include_mask, min_area=300, epsilon_factor=0.002),
     
131.             'exclude': exclude_polys,
     
132.         }
     
133. 
     
134.     @staticmethod
     
135.     def mask_to_polys(mask, min_area=100, epsilon_factor=0.005):
     
136.         contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     
137.         polys = []
     
138.         for contour in contours:
     
139.             if cv2.contourArea(contour) < min_area:
     
140.                 continue
     
141.             epsilon = epsilon_factor * cv2.arcLength(contour, True)
     
142.             approx = cv2.approxPolyDP(contour, epsilon, True)
     
143.             pts = approx.reshape(-1, 2).astype(np.float32)
     
144.             if len(pts) >= 3:
     
145.                 polys.append(pts)
     
146.         return polys
     

复制代码

XSegEditor/XSegEditor.py

1. import json
   
2. import multiprocessing
   
3. import os
   
4. import pickle
   
5. import sys
   
6. import tempfile
   
7. import time
   
8. import traceback
   
9. from enum import IntEnum
   
10. from types import SimpleNamespace as sn
    
11. 
    
12. import cv2
    
13. import numpy as np
    
14. import numpy.linalg as npla
    
15. from PyQt5.QtCore import *
    
16. from PyQt5.QtGui import *
    
17. from PyQt5.QtWidgets import *
    
18. 
    
19. from core import imagelib, pathex
    
20. from core.cv2ex import *
    
21. from core.imagelib import SegIEPoly, SegIEPolys, SegIEPolyType, sd
    
22. from core.qtex import *
    
23. from DFLIMG import *
    
24. from localization import StringsDB, system_language
    
25. from samplelib import PackedFaceset
    
26. 
    
27. from .QCursorDB import QCursorDB
    
28. from .QIconDB import QIconDB
    
29. from .QStringDB import QStringDB
    
30. from .QImageDB import QImageDB
    
31. from .OcclusionDetector import OcclusionDetector
    
32. 
    
33. class OpMode(IntEnum):
    
34.     NONE = 0
    
35.     DRAW_PTS = 1
    
36.     EDIT_PTS = 2
    
37.     VIEW_BAKED = 3
    
38.     VIEW_XSEG_MASK = 4
    
39. 
    
40. class PTEditMode(IntEnum):
    
41.     MOVE = 0
    
42.     ADD_DEL = 1
    
43. 
    
44. class DragType(IntEnum):
    
45.     NONE = 0
    
46.     IMAGE_LOOK = 1
    
47.     POLY_PT = 2
    
48. 
    
49. class ViewLock(IntEnum):
    
50.     NONE = 0
    
51.     CENTER = 1
    
52. 
    
53. class QUIConfig():
    
54.     @staticmethod
    
55.     def initialize(icon_size = 48, icon_spacer_size=16, preview_bar_icon_size=64):
    
56.         QUIConfig.icon_q_size = QSize(icon_size, icon_size)
    
57.         QUIConfig.icon_spacer_q_size = QSize(icon_spacer_size, icon_spacer_size)
    
58.         QUIConfig.preview_bar_icon_q_size = QSize(preview_bar_icon_size, preview_bar_icon_size)
    
59. 
    
60. class ImagePreviewSequenceBar(QFrame):
    
61.     def __init__(self, preview_images_count, icon_size):
    
62.         super().__init__()
    
63.         self.preview_images_count = preview_images_count = max(1, preview_images_count + (preview_images_count % 2 -1) )
    
64. 
    
65.         self.icon_size = icon_size
    
66. 
    
67.         black_q_img = QImage(np.zeros( (icon_size,icon_size,3) ).data, icon_size, icon_size, 3*icon_size, QImage.Format_RGB888)
    
68.         self.black_q_pixmap = QPixmap.fromImage(black_q_img)
    
69. 
    
70.         self.image_containers = [ QLabel() for i in range(preview_images_count)]
    
71. 
    
72.         main_frame_l_cont_hl = QGridLayout()
    
73.         main_frame_l_cont_hl.setContentsMargins(0,0,0,0)
    
74.         #main_frame_l_cont_hl.setSpacing(0)
    
75. 
    
76. 
    
77. 
    
78.         for i in range(len(self.image_containers)):
    
79.             q_label = self.image_containers[i]
    
80.             q_label.setScaledContents(True)
    
81.             if i == preview_images_count//2:
    
82.                 q_label.setMinimumSize(icon_size+16, icon_size+16 )
    
83.                 q_label.setMaximumSize(icon_size+16, icon_size+16 )
    
84.             else:
    
85.                 q_label.setMinimumSize(icon_size, icon_size )
    
86.                 q_label.setMaximumSize(icon_size, icon_size )
    
87.                 opacity_effect = QGraphicsOpacityEffect()
    
88.                 opacity_effect.setOpacity(0.5)
    
89.                 q_label.setGraphicsEffect(opacity_effect)
    
90. 
    
91.             q_label.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
    
92. 
    
93.             main_frame_l_cont_hl.addWidget (q_label, 0, i)
    
94. 
    
95.         self.setLayout(main_frame_l_cont_hl)
    
96. 
    
97.         self.prev_img_conts = self.image_containers[(preview_images_count//2) -1::-1]
    
98.         self.next_img_conts = self.image_containers[preview_images_count//2:]
    
99. 
    
100.         self.update_images()
     
101. 
     
102.     def get_preview_images_count(self):
     
103.         return self.preview_images_count
     
104. 
     
105.     def update_images(self, prev_imgs=None, next_imgs=None):
     
106.         # Fix arrays
     
107.         if prev_imgs is None:
     
108.             prev_imgs = []
     
109.         prev_img_conts_len = len(self.prev_img_conts)
     
110.         prev_q_imgs_len = len(prev_imgs)
     
111.         if prev_q_imgs_len < prev_img_conts_len:
     
112.             for i in range ( prev_img_conts_len - prev_q_imgs_len ):
     
113.                 prev_imgs.append(None)
     
114.         elif prev_q_imgs_len > prev_img_conts_len:
     
115.             prev_imgs = prev_imgs[:prev_img_conts_len]
     
116. 
     
117.         if next_imgs is None:
     
118.             next_imgs = []
     
119.         next_img_conts_len = len(self.next_img_conts)
     
120.         next_q_imgs_len = len(next_imgs)
     
121.         if next_q_imgs_len < next_img_conts_len:
     
122.             for i in range ( next_img_conts_len - next_q_imgs_len ):
     
123.                 next_imgs.append(None)
     
124.         elif next_q_imgs_len > next_img_conts_len:
     
125.             next_imgs = next_imgs[:next_img_conts_len]
     
126. 
     
127.         for i,img in enumerate(prev_imgs):
     
128.             self.prev_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) if img is not None else self.black_q_pixmap )
     
129. 
     
130.         for i,img in enumerate(next_imgs):
     
131.             self.next_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) if img is not None else self.black_q_pixmap )
     
132. 
     
133. class ColorScheme():
     
134.     def __init__(self, unselected_color, selected_color, outline_color, outline_width, pt_outline_color, cross_cursor):
     
135.         self.poly_unselected_brush = QBrush(unselected_color)
     
136.         self.poly_selected_brush = QBrush(selected_color)
     
137. 
     
138.         self.poly_outline_solid_pen = QPen(outline_color, outline_width, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
     
139.         self.poly_outline_dot_pen = QPen(outline_color, outline_width, Qt.DotLine, Qt.RoundCap, Qt.RoundJoin)
     
140. 
     
141.         self.pt_outline_pen = QPen(pt_outline_color)
     
142.         self.cross_cursor = cross_cursor
     
143. 
     
144. class CanvasConfig():
     
145. 
     
146.     def __init__(self,
     
147.                  pt_radius=4,
     
148.                  pt_select_radius=8,
     
149.                  color_schemes=None,
     
150.                  **kwargs):
     
151.         self.pt_radius = pt_radius
     
152.         self.pt_select_radius = pt_select_radius
     
153. 
     
154.         if color_schemes is None:
     
155.             color_schemes = [
     
156.                     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 ),
     
157.                     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 ),
     
158.                     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 ),
     
159.                     ]
     
160.         self.color_schemes = color_schemes
     
161. 
     
162. class QCanvasControlsLeftBar(QFrame):
     
163. 
     
164.     def __init__(self):
     
165.         super().__init__()
     
166.         #==============================================
     
167.         btn_poly_type_include = QToolButton()
     
168.         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)
     
169.         btn_poly_type_include.setDefaultAction(self.btn_poly_type_include_act)
     
170.         btn_poly_type_include.setIconSize(QUIConfig.icon_q_size)
     
171. 
     
172.         btn_poly_type_exclude = QToolButton()
     
173.         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)
     
174.         btn_poly_type_exclude.setDefaultAction(self.btn_poly_type_exclude_act)
     
175.         btn_poly_type_exclude.setIconSize(QUIConfig.icon_q_size)
     
176. 
     
177.         self.btn_poly_type_act_grp = QActionGroup (self)
     
178.         self.btn_poly_type_act_grp.addAction(self.btn_poly_type_include_act)
     
179.         self.btn_poly_type_act_grp.addAction(self.btn_poly_type_exclude_act)
     
180.         self.btn_poly_type_act_grp.setExclusive(True)
     
181.         #==============================================
     
182.         btn_undo_pt = QToolButton()
     
183.         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)
     
184.         btn_undo_pt.setDefaultAction(self.btn_undo_pt_act)
     
185.         btn_undo_pt.setIconSize(QUIConfig.icon_q_size)
     
186. 
     
187.         btn_redo_pt = QToolButton()
     
188.         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)
     
189.         btn_redo_pt.setDefaultAction(self.btn_redo_pt_act)
     
190.         btn_redo_pt.setIconSize(QUIConfig.icon_q_size)
     
191. 
     
192.         btn_delete_poly = QToolButton()
     
193.         self.btn_delete_poly_act = QActionEx( QIconDB.delete_poly, QStringDB.btn_delete_poly_tip, shortcut='Delete', shortcut_in_tooltip=True)
     
194.         btn_delete_poly.setDefaultAction(self.btn_delete_poly_act)
     
195.         btn_delete_poly.setIconSize(QUIConfig.icon_q_size)
     
196.         #==============================================
     
197.         btn_pt_edit_mode = QToolButton()
     
198.         self.btn_pt_edit_mode_act = QActionEx( QIconDB.pt_edit_mode, QStringDB.btn_pt_edit_mode_tip, shortcut_in_tooltip=True, is_checkable=True)
     
199.         btn_pt_edit_mode.setDefaultAction(self.btn_pt_edit_mode_act)
     
200.         btn_pt_edit_mode.setIconSize(QUIConfig.icon_q_size)
     
201.         #==============================================
     
202. 
     
203.         controls_bar_frame2_l = QVBoxLayout()
     
204.         controls_bar_frame2_l.addWidget ( btn_poly_type_include )
     
205.         controls_bar_frame2_l.addWidget ( btn_poly_type_exclude )
     
206.         controls_bar_frame2 = QFrame()
     
207.         controls_bar_frame2.setFrameShape(QFrame.StyledPanel)
     
208.         controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
209.         controls_bar_frame2.setLayout(controls_bar_frame2_l)
     
210. 
     
211.         controls_bar_frame3_l = QVBoxLayout()
     
212.         controls_bar_frame3_l.addWidget ( btn_undo_pt )
     
213.         controls_bar_frame3_l.addWidget ( btn_redo_pt )
     
214.         controls_bar_frame3_l.addWidget ( btn_delete_poly )
     
215.         controls_bar_frame3 = QFrame()
     
216.         controls_bar_frame3.setFrameShape(QFrame.StyledPanel)
     
217.         controls_bar_frame3.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
218.         controls_bar_frame3.setLayout(controls_bar_frame3_l)
     
219. 
     
220.         controls_bar_frame4_l = QVBoxLayout()
     
221.         controls_bar_frame4_l.addWidget ( btn_pt_edit_mode )
     
222.         controls_bar_frame4 = QFrame()
     
223.         controls_bar_frame4.setFrameShape(QFrame.StyledPanel)
     
224.         controls_bar_frame4.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
225.         controls_bar_frame4.setLayout(controls_bar_frame4_l)
     
226.         #==============================================
     
227.         btn_auto_detect = QToolButton()
     
228.         self.btn_auto_detect_act = QActionEx( QIconDB.auto_detect, QStringDB.btn_auto_detect_tip, shortcut='F', shortcut_in_tooltip=True)
     
229.         btn_auto_detect.setDefaultAction(self.btn_auto_detect_act)
     
230.         btn_auto_detect.setIconSize(QUIConfig.icon_q_size)
     
231. 
     
232.         controls_bar_frame5_l = QVBoxLayout()
     
233.         controls_bar_frame5_l.addWidget ( btn_auto_detect )
     
234.         controls_bar_frame5 = QFrame()
     
235.         controls_bar_frame5.setFrameShape(QFrame.StyledPanel)
     
236.         controls_bar_frame5.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
237.         controls_bar_frame5.setLayout(controls_bar_frame5_l)
     
238. 
     
239.         controls_bar_l = QVBoxLayout()
     
240.         controls_bar_l.setContentsMargins(0,0,0,0)
     
241.         controls_bar_l.addWidget(controls_bar_frame2)
     
242.         controls_bar_l.addWidget(controls_bar_frame3)
     
243.         controls_bar_l.addWidget(controls_bar_frame4)
     
244.         controls_bar_l.addWidget(controls_bar_frame5)
     
245. 
     
246.         self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding )
     
247.         self.setLayout(controls_bar_l)
     
248. 
     
249. class QCanvasControlsRightBar(QFrame):
     
250. 
     
251.     def __init__(self):
     
252.         super().__init__()
     
253.         #==============================================
     
254.         btn_poly_color_red = QToolButton()
     
255.         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)
     
256.         btn_poly_color_red.setDefaultAction(self.btn_poly_color_red_act)
     
257.         btn_poly_color_red.setIconSize(QUIConfig.icon_q_size)
     
258. 
     
259.         btn_poly_color_green = QToolButton()
     
260.         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)
     
261.         btn_poly_color_green.setDefaultAction(self.btn_poly_color_green_act)
     
262.         btn_poly_color_green.setIconSize(QUIConfig.icon_q_size)
     
263. 
     
264.         btn_poly_color_blue = QToolButton()
     
265.         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)
     
266.         btn_poly_color_blue.setDefaultAction(self.btn_poly_color_blue_act)
     
267.         btn_poly_color_blue.setIconSize(QUIConfig.icon_q_size)
     
268. 
     
269.         btn_view_baked_mask = QToolButton()
     
270.         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)
     
271.         btn_view_baked_mask.setDefaultAction(self.btn_view_baked_mask_act)
     
272.         btn_view_baked_mask.setIconSize(QUIConfig.icon_q_size)
     
273. 
     
274.         btn_view_xseg_mask = QToolButton()
     
275.         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)
     
276.         btn_view_xseg_mask.setDefaultAction(self.btn_view_xseg_mask_act)
     
277.         btn_view_xseg_mask.setIconSize(QUIConfig.icon_q_size)
     
278. 
     
279.         btn_view_xseg_overlay_mask = QToolButton()
     
280.         self.btn_view_xseg_overlay_mask_act = QActionEx( QIconDB.view_xseg_overlay, QStringDB.btn_view_xseg_overlay_mask_tip, shortcut='`', shortcut_in_tooltip=True, is_checkable=True)
     
281.         btn_view_xseg_overlay_mask.setDefaultAction(self.btn_view_xseg_overlay_mask_act)
     
282.         btn_view_xseg_overlay_mask.setIconSize(QUIConfig.icon_q_size)
     
283. 
     
284.         self.btn_poly_color_act_grp = QActionGroup (self)
     
285.         self.btn_poly_color_act_grp.addAction(self.btn_poly_color_red_act)
     
286.         self.btn_poly_color_act_grp.addAction(self.btn_poly_color_green_act)
     
287.         self.btn_poly_color_act_grp.addAction(self.btn_poly_color_blue_act)
     
288.         self.btn_poly_color_act_grp.addAction(self.btn_view_baked_mask_act)
     
289.         self.btn_poly_color_act_grp.addAction(self.btn_view_xseg_mask_act)
     
290.         self.btn_poly_color_act_grp.setExclusive(True)
     
291.         #==============================================
     
292.         btn_view_lock_center = QToolButton()
     
293.         self.btn_view_lock_center_act = QActionEx( QIconDB.view_lock_center, QStringDB.btn_view_lock_center_tip, shortcut_in_tooltip=True, is_checkable=True)
     
294.         btn_view_lock_center.setDefaultAction(self.btn_view_lock_center_act)
     
295.         btn_view_lock_center.setIconSize(QUIConfig.icon_q_size)
     
296. 
     
297.         controls_bar_frame2_l = QVBoxLayout()
     
298.         controls_bar_frame2_l.addWidget ( btn_view_xseg_overlay_mask )
     
299.         controls_bar_frame2 = QFrame()
     
300.         controls_bar_frame2.setFrameShape(QFrame.StyledPanel)
     
301.         controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
302.         controls_bar_frame2.setLayout(controls_bar_frame2_l)
     
303.         
     
304.         controls_bar_frame1_l = QVBoxLayout()
     
305.         controls_bar_frame1_l.addWidget ( btn_poly_color_red )
     
306.         controls_bar_frame1_l.addWidget ( btn_poly_color_green )
     
307.         controls_bar_frame1_l.addWidget ( btn_poly_color_blue )
     
308.         controls_bar_frame1_l.addWidget ( btn_view_baked_mask )
     
309.         controls_bar_frame1_l.addWidget ( btn_view_xseg_mask )
     
310.         controls_bar_frame1 = QFrame()
     
311.         controls_bar_frame1.setFrameShape(QFrame.StyledPanel)
     
312.         controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
313.         controls_bar_frame1.setLayout(controls_bar_frame1_l)
     
314.         
     
315.         controls_bar_frame3_l = QVBoxLayout()
     
316.         controls_bar_frame3_l.addWidget ( btn_view_lock_center )
     
317.         controls_bar_frame3 = QFrame()
     
318.         controls_bar_frame3.setFrameShape(QFrame.StyledPanel)
     
319.         controls_bar_frame3.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
     
320.         controls_bar_frame3.setLayout(controls_bar_frame3_l)
     
321. 
     
322.         controls_bar_l = QVBoxLayout()
     
323.         controls_bar_l.setContentsMargins(0,0,0,0)
     
324.         controls_bar_l.addWidget(controls_bar_frame2)
     
325.         controls_bar_l.addWidget(controls_bar_frame1)
     
326.         controls_bar_l.addWidget(controls_bar_frame3)
     
327.         
     
328.         self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding )
     
329.         self.setLayout(controls_bar_l)
     
330. 
     
331. class QCanvasOperator(QWidget):
     
332.     def __init__(self, cbar):
     
333.         super().__init__()
     
334.         self.cbar = cbar
     
335. 
     
336.         self.set_cbar_disabled()
     
337. 
     
338.         self.cbar.btn_poly_color_red_act.triggered.connect ( lambda : self.set_color_scheme_id(0) )
     
339.         self.cbar.btn_poly_color_green_act.triggered.connect ( lambda : self.set_color_scheme_id(1) )
     
340.         self.cbar.btn_poly_color_blue_act.triggered.connect ( lambda : self.set_color_scheme_id(2) )
     
341.         self.cbar.btn_view_baked_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_BAKED) )
     
342.         self.cbar.btn_view_xseg_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_XSEG_MASK) )
     
343. 
     
344.         self.cbar.btn_view_xseg_overlay_mask_act.toggled.connect ( lambda is_checked: self.update() )
     
345. 
     
346.         self.cbar.btn_poly_type_include_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.INCLUDE) )
     
347.         self.cbar.btn_poly_type_exclude_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.EXCLUDE) )
     
348. 
     
349.         self.cbar.btn_undo_pt_act.triggered.connect ( lambda : self.action_undo_pt() )
     
350.         self.cbar.btn_redo_pt_act.triggered.connect ( lambda : self.action_redo_pt() )
     
351. 
     
352.         self.cbar.btn_delete_poly_act.triggered.connect ( lambda : self.action_delete_poly() )
     
353. 
     
354.         self.cbar.btn_auto_detect_act.triggered.connect ( lambda : self.action_auto_detect() )
     
355. 
     
356.         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 ) )
     
357.         self.cbar.btn_view_lock_center_act.toggled.connect ( lambda is_checked: self.set_view_lock( ViewLock.CENTER if is_checked else ViewLock.NONE ) )
     
358. 
     
359.         self.mouse_in_widget = False
     
360. 
     
361.         QXMainWindow.inst.add_keyPressEvent_listener ( self.on_keyPressEvent )
     
362.         QXMainWindow.inst.add_keyReleaseEvent_listener ( self.on_keyReleaseEvent )
     
363. 
     
364.         self.qp = QPainter()
     
365.         self.initialized = False
     
366.         self.last_state = None
     
367. 
     
368.     def initialize(self, img, img_look_pt=None, view_scale=None, ie_polys=None, xseg_mask=None, canvas_config=None, landmarks=None ):
     
369.         self.q_img_copy = img.copy()
     
370.         self.landmarks = landmarks
     
371.         q_img = self.q_img = QImage_from_np(img)
     
372.         self.img_pixmap = QPixmap.fromImage(q_img)
     
373. 
     
374.         self.xseg_mask_pixmap = None
     
375.         self.xseg_overlay_mask_pixmap = None
     
376.         if xseg_mask is not None:
     
377.             h,w,c = img.shape
     
378.             xseg_mask = cv2.resize(xseg_mask, (w,h), interpolation=cv2.INTER_CUBIC)
     
379.             xseg_mask = imagelib.normalize_channels(xseg_mask, 1)
     
380.             xseg_img = img.astype(np.float32)/255.0
     
381.             xseg_overlay_mask = xseg_img*(1-xseg_mask)*0.5 + xseg_img*xseg_mask
     
382.             xseg_overlay_mask = np.clip(xseg_overlay_mask*255, 0, 255).astype(np.uint8)
     
383.             xseg_mask = np.clip(xseg_mask*255, 0, 255).astype(np.uint8)
     
384.             self.xseg_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_mask))
     
385.             self.xseg_overlay_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_overlay_mask))
     
386. 
     
387.         self.img_size = QSize_to_np (self.img_pixmap.size())
     
388. 
     
389.         self.img_look_pt = img_look_pt
     
390.         self.view_scale = view_scale
     
391. 
     
392.         if ie_polys is None:
     
393.             ie_polys = SegIEPolys()
     
394.         self.ie_polys = ie_polys
     
395. 
     
396.         if canvas_config is None:
     
397.             canvas_config = CanvasConfig()
     
398.         self.canvas_config = canvas_config
     
399. 
     
400.         # UI init
     
401.         self.set_cbar_disabled()
     
402.         self.cbar.btn_poly_color_act_grp.setDisabled(False)
     
403.         self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(False)
     
404.         self.cbar.btn_poly_type_act_grp.setDisabled(False)
     
405. 
     
406.         # Initial vars
     
407.         self.current_cursor = None
     
408.         self.mouse_hull_poly = None
     
409.         self.mouse_wire_poly = None
     
410.         self.drag_type = DragType.NONE
     
411.         self.mouse_cli_pt = np.zeros((2,), np.float32 )
     
412. 
     
413.         # Initial state
     
414.         self.set_op_mode(OpMode.NONE)
     
415.         self.set_color_scheme_id(1)
     
416.         self.set_poly_include_type(SegIEPolyType.INCLUDE)
     
417.         self.set_pt_edit_mode(PTEditMode.MOVE)
     
418.         self.set_view_lock(ViewLock.NONE)
     
419. 
     
420.         # Apply last state
     
421.         if self.last_state is not None:
     
422.             self.set_color_scheme_id(self.last_state.color_scheme_id)
     
423.             if self.last_state.op_mode is not None:
     
424.                 self.set_op_mode(self.last_state.op_mode)
     
425. 
     
426.         self.initialized = True
     
427. 
     
428.         self.setMouseTracking(True)
     
429.         self.update_cursor()
     
430.         self.update()
     
431. 
     
432. 
     
433.     def finalize(self):
     
434.         if self.initialized:
     
435.             if self.op_mode == OpMode.DRAW_PTS:
     
436.                 self.set_op_mode(OpMode.EDIT_PTS)
     
437. 
     
438.             self.last_state = sn(op_mode = self.op_mode if self.op_mode in [OpMode.VIEW_BAKED, OpMode.VIEW_XSEG_MASK] else None,
     
439.                                  color_scheme_id = self.color_scheme_id)
     
440. 
     
441.             self.img_pixmap = None
     
442.             self.update_cursor(is_finalize=True)
     
443.             self.setMouseTracking(False)
     
444.             self.setFocusPolicy(Qt.NoFocus)
     
445.             self.set_cbar_disabled()
     
446.             self.initialized = False
     
447.             self.update()
     
448. 
     
449.     # ====================================================================================
     
450.     # ====================================================================================
     
451.     # ====================================== GETTERS =====================================
     
452.     # ====================================================================================
     
453.     # ====================================================================================
     
454.     def is_initialized(self):
     
455.         return self.initialized
     
456. 
     
457.     def get_ie_polys(self):
     
458.         return self.ie_polys
     
459. 
     
460.     def get_cli_center_pt(self):
     
461.         return np.round(QSize_to_np(self.size())/2.0)
     
462. 
     
463.     def get_img_look_pt(self):
     
464.         img_look_pt = self.img_look_pt
     
465.         if img_look_pt is None:
     
466.             img_look_pt = self.img_size / 2
     
467.         return img_look_pt
     
468. 
     
469.     def get_view_scale(self):
     
470.         view_scale = self.view_scale
     
471.         if view_scale is None:
     
472.             # Calc as scale to fit
     
473.             min_cli_size = np.min(QSize_to_np(self.size()))
     
474.             max_img_size = np.max(self.img_size)
     
475.             view_scale =  min_cli_size / max_img_size
     
476. 
     
477.         return view_scale
     
478. 
     
479.     def get_current_color_scheme(self):
     
480.         return self.canvas_config.color_schemes[self.color_scheme_id]
     
481. 
     
482.     def get_poly_pt_id_under_pt(self, poly, cli_pt):
     
483.         w = np.argwhere ( npla.norm ( cli_pt - self.img_to_cli_pt( poly.get_pts() ), axis=1 )  <= self.canvas_config.pt_select_radius )
     
484.         return None if len(w) == 0 else w[-1][0]
     
485. 
     
486.     def get_poly_edge_id_pt_under_pt(self, poly, cli_pt):
     
487.         cli_pts = self.img_to_cli_pt(poly.get_pts())
     
488.         if len(cli_pts) >= 3:
     
489.             edge_dists, projs = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True)
     
490.             edge_id = np.argmin(edge_dists)
     
491.             dist = edge_dists[edge_id]
     
492.             pt = projs[edge_id]
     
493.             if dist <= self.canvas_config.pt_select_radius:
     
494.                 return edge_id, pt
     
495.         return None, None
     
496. 
     
497.     def get_poly_by_pt_near_wire(self, cli_pt):
     
498.         pt_select_radius = self.canvas_config.pt_select_radius
     
499. 
     
500.         for poly in reversed(self.ie_polys.get_polys()):
     
501.             pts = poly.get_pts()
     
502.             if len(pts) >= 3:
     
503.                 cli_pts = self.img_to_cli_pt(pts)
     
504. 
     
505.                 edge_dists, _ = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True)
     
506. 
     
507.                 if np.min(edge_dists) <= pt_select_radius or \
     
508.                    any( npla.norm ( cli_pt - cli_pts, axis=1 ) <= pt_select_radius ):
     
509.                     return poly
     
510.         return None
     
511. 
     
512.     def get_poly_by_pt_in_hull(self, cli_pos):
     
513.         img_pos = self.cli_to_img_pt(cli_pos)
     
514. 
     
515.         for poly in reversed(self.ie_polys.get_polys()):
     
516.             pts = poly.get_pts()
     
517.             if len(pts) >= 3:
     
518.                 if cv2.pointPolygonTest( pts, tuple(img_pos), False) >= 0:
     
519.                     return poly
     
520. 
     
521.         return None
     
522. 
     
523.     def img_to_cli_pt(self, p):
     
524.         return (p - self.get_img_look_pt()) * self.get_view_scale() + self.get_cli_center_pt()# QSize_to_np(self.size())/2.0
     
525. 
     
526.     def cli_to_img_pt(self, p):
     
527.         return (p - self.get_cli_center_pt() ) / self.get_view_scale() + self.get_img_look_pt()
     
528. 
     
529.     def img_to_cli_rect(self, rect):
     
530.         tl = QPoint_to_np(rect.topLeft())
     
531.         xy = self.img_to_cli_pt(tl)
     
532.         xy2 = self.img_to_cli_pt(tl + QSize_to_np(rect.size()) ) - xy
     
533.         return QRect ( *xy.astype(np.int64), *xy2.astype(np.int64) )
     
534. 
     
535.     # ====================================================================================
     
536.     # ====================================================================================
     
537.     # ====================================== SETTERS =====================================
     
538.     # ====================================================================================
     
539.     # ====================================================================================
     
540.     def set_op_mode(self, op_mode, op_poly=None):
     
541.         if not hasattr(self,'op_mode'):
     
542.             self.op_mode = None
     
543.             self.op_poly = None
     
544. 
     
545.         if self.op_mode != op_mode:
     
546.             # Finalize prev mode
     
547.             if self.op_mode == OpMode.NONE:
     
548.                 self.cbar.btn_poly_type_act_grp.setDisabled(True)
     
549.             elif self.op_mode == OpMode.DRAW_PTS:
     
550.                 self.cbar.btn_undo_pt_act.setDisabled(True)
     
551.                 self.cbar.btn_redo_pt_act.setDisabled(True)
     
552.                 self.cbar.btn_view_lock_center_act.setDisabled(True)
     
553.                 # Reset view_lock when exit from DRAW_PTS
     
554.                 self.set_view_lock(ViewLock.NONE)
     
555.                 # Remove unfinished poly
     
556.                 if self.op_poly.get_pts_count() < 3:
     
557.                     self.ie_polys.remove_poly(self.op_poly)
     
558. 
     
559.             elif self.op_mode == OpMode.EDIT_PTS:
     
560.                 self.cbar.btn_pt_edit_mode_act.setDisabled(True)
     
561.                 self.cbar.btn_delete_poly_act.setDisabled(True)
     
562.                 # Reset pt_edit_move when exit from EDIT_PTS
     
563.                 self.set_pt_edit_mode(PTEditMode.MOVE)
     
564.             elif self.op_mode == OpMode.VIEW_BAKED:
     
565.                 self.cbar.btn_view_baked_mask_act.setChecked(False)
     
566.             elif self.op_mode == OpMode.VIEW_XSEG_MASK:
     
567.                 self.cbar.btn_view_xseg_mask_act.setChecked(False)
     
568. 
     
569.             self.op_mode = op_mode
     
570. 
     
571.             # Initialize new mode
     
572.             if op_mode == OpMode.NONE:
     
573.                 self.cbar.btn_poly_type_act_grp.setDisabled(False)
     
574.             elif op_mode == OpMode.DRAW_PTS:
     
575.                 self.cbar.btn_undo_pt_act.setDisabled(False)
     
576.                 self.cbar.btn_redo_pt_act.setDisabled(False)
     
577.                 self.cbar.btn_view_lock_center_act.setDisabled(False)
     
578.             elif op_mode == OpMode.EDIT_PTS:
     
579.                 self.cbar.btn_pt_edit_mode_act.setDisabled(False)
     
580.                 self.cbar.btn_delete_poly_act.setDisabled(False)
     
581.             elif op_mode == OpMode.VIEW_BAKED:
     
582.                 self.cbar.btn_view_baked_mask_act.setChecked(True )
     
583.                 n = QImage_to_np ( self.q_img ).astype(np.float32) / 255.0
     
584.                 h,w,c = n.shape
     
585.                 mask = np.zeros( (h,w,1), dtype=np.float32 )
     
586.                 self.ie_polys.overlay_mask(mask)
     
587.                 n = (mask*255).astype(np.uint8)
     
588.                 self.img_baked_pixmap = QPixmap.fromImage(QImage_from_np(n))
     
589.             elif op_mode == OpMode.VIEW_XSEG_MASK:
     
590.                 self.cbar.btn_view_xseg_mask_act.setChecked(True)
     
591. 
     
592.             if op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]:
     
593.                 self.mouse_op_poly_pt_id = None
     
594.                 self.mouse_op_poly_edge_id = None
     
595.                 self.mouse_op_poly_edge_id_pt = None
     
596. 
     
597.             self.op_poly = op_poly
     
598.             if op_poly is not None:
     
599.                 self.update_mouse_info()
     
600. 
     
601.             self.update_cursor()
     
602.             self.update()
     
603. 
     
604.     def set_pt_edit_mode(self, pt_edit_mode):
     
605.         if not hasattr(self, 'pt_edit_mode') or self.pt_edit_mode != pt_edit_mode:
     
606.             self.pt_edit_mode = pt_edit_mode
     
607.             self.update_cursor()
     
608.             self.update()
     
609.         self.cbar.btn_pt_edit_mode_act.setChecked( self.pt_edit_mode == PTEditMode.ADD_DEL )
     
610. 
     
611.     def set_view_lock(self, view_lock):
     
612.         if not hasattr(self, 'view_lock') or self.view_lock != view_lock:
     
613.             if hasattr(self, 'view_lock') and self.view_lock != view_lock:
     
614.                 if view_lock == ViewLock.CENTER:
     
615.                     self.img_look_pt = self.mouse_img_pt
     
616.                     QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) ))
     
617. 
     
618.             self.view_lock = view_lock
     
619.             self.update()
     
620.         self.cbar.btn_view_lock_center_act.setChecked( self.view_lock == ViewLock.CENTER )
     
621. 
     
622.     def set_cbar_disabled(self):
     
623.         self.cbar.btn_delete_poly_act.setDisabled(True)
     
624.         self.cbar.btn_undo_pt_act.setDisabled(True)
     
625.         self.cbar.btn_redo_pt_act.setDisabled(True)
     
626.         self.cbar.btn_pt_edit_mode_act.setDisabled(True)
     
627.         self.cbar.btn_view_lock_center_act.setDisabled(True)
     
628.         self.cbar.btn_poly_color_act_grp.setDisabled(True)
     
629.         self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(True)
     
630.         self.cbar.btn_poly_type_act_grp.setDisabled(True)
     
631. 
     
632. 
     
633.     def set_color_scheme_id(self, id):
     
634.         if self.op_mode == OpMode.VIEW_BAKED or self.op_mode == OpMode.VIEW_XSEG_MASK:
     
635.             self.set_op_mode(OpMode.NONE)
     
636. 
     
637.         if not hasattr(self, 'color_scheme_id') or self.color_scheme_id != id:
     
638.             self.color_scheme_id = id
     
639.             self.update_cursor()
     
640.             self.update()
     
641. 
     
642.         if self.color_scheme_id == 0:
     
643.             self.cbar.btn_poly_color_red_act.setChecked( True )
     
644.         elif self.color_scheme_id == 1:
     
645.             self.cbar.btn_poly_color_green_act.setChecked( True )
     
646.         elif self.color_scheme_id == 2:
     
647.             self.cbar.btn_poly_color_blue_act.setChecked( True )
     
648. 
     
649.     def set_poly_include_type(self, poly_include_type):
     
650.         if not hasattr(self, 'poly_include_type' ) or \
     
651.            ( self.poly_include_type != poly_include_type and \
     
652.              self.op_mode in [OpMode.NONE, OpMode.EDIT_PTS] ):
     
653.             self.poly_include_type = poly_include_type
     
654.             self.update()
     
655.         self.cbar.btn_poly_type_include_act.setChecked(self.poly_include_type == SegIEPolyType.INCLUDE)
     
656.         self.cbar.btn_poly_type_exclude_act.setChecked(self.poly_include_type == SegIEPolyType.EXCLUDE)
     
657. 
     
658.     # ====================================================================================
     
659.     # ====================================================================================
     
660.     # ====================================== METHODS =====================================
     
661.     # ====================================================================================
     
662.     # ====================================================================================
     
663. 
     
664.     def update_cursor(self, is_finalize=False):
     
665.         if not self.initialized:
     
666.             return
     
667. 
     
668.         if not self.mouse_in_widget or is_finalize:
     
669.             if self.current_cursor is not None:
     
670.                 QApplication.restoreOverrideCursor()
     
671.                 self.current_cursor = None
     
672.         else:
     
673.             color_cc = self.get_current_color_scheme().cross_cursor
     
674.             nc = Qt.ArrowCursor
     
675. 
     
676.             if self.drag_type == DragType.IMAGE_LOOK:
     
677.                 nc = Qt.ClosedHandCursor
     
678.             else:
     
679. 
     
680.                 if self.op_mode == OpMode.NONE:
     
681.                     nc = color_cc
     
682.                     if self.mouse_wire_poly is not None:
     
683.                         nc = Qt.PointingHandCursor
     
684. 
     
685.                 elif self.op_mode == OpMode.DRAW_PTS:
     
686.                     nc = color_cc
     
687.                 elif self.op_mode == OpMode.EDIT_PTS:
     
688.                     nc = Qt.ArrowCursor
     
689. 
     
690.                     if self.mouse_op_poly_pt_id is not None:
     
691.                         nc = Qt.PointingHandCursor
     
692. 
     
693.                     if self.pt_edit_mode == PTEditMode.ADD_DEL:
     
694. 
     
695.                         if self.mouse_op_poly_edge_id is not None and \
     
696.                         self.mouse_op_poly_pt_id is None:
     
697.                             nc = color_cc
     
698.             if self.current_cursor != nc:
     
699.                 if self.current_cursor is None:
     
700.                     QApplication.setOverrideCursor(nc)
     
701.                 else:
     
702.                     QApplication.changeOverrideCursor(nc)
     
703.                 self.current_cursor = nc
     
704. 
     
705.     def update_mouse_info(self, mouse_cli_pt=None):
     
706.         """
     
707.         Update selected polys/edges/points by given mouse position
     
708.         """
     
709.         if mouse_cli_pt is not None:
     
710.             self.mouse_cli_pt = mouse_cli_pt.astype(np.float32)
     
711. 
     
712.         self.mouse_img_pt = self.cli_to_img_pt(self.mouse_cli_pt)
     
713. 
     
714.         new_mouse_hull_poly = self.get_poly_by_pt_in_hull(self.mouse_cli_pt)
     
715. 
     
716.         if self.mouse_hull_poly != new_mouse_hull_poly:
     
717.             self.mouse_hull_poly = new_mouse_hull_poly
     
718.             self.update_cursor()
     
719.             self.update()
     
720. 
     
721.         new_mouse_wire_poly = self.get_poly_by_pt_near_wire(self.mouse_cli_pt)
     
722. 
     
723.         if self.mouse_wire_poly != new_mouse_wire_poly:
     
724.             self.mouse_wire_poly = new_mouse_wire_poly
     
725.             self.update_cursor()
     
726.             self.update()
     
727. 
     
728.         if self.op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]:
     
729.             new_mouse_op_poly_pt_id = self.get_poly_pt_id_under_pt (self.op_poly, self.mouse_cli_pt)
     
730.             if self.mouse_op_poly_pt_id != new_mouse_op_poly_pt_id:
     
731.                 self.mouse_op_poly_pt_id = new_mouse_op_poly_pt_id
     
732.                 self.update_cursor()
     
733.                 self.update()
     
734. 
     
735.             new_mouse_op_poly_edge_id,\
     
736.             new_mouse_op_poly_edge_id_pt = self.get_poly_edge_id_pt_under_pt (self.op_poly, self.mouse_cli_pt)
     
737.             if self.mouse_op_poly_edge_id != new_mouse_op_poly_edge_id:
     
738.                 self.mouse_op_poly_edge_id = new_mouse_op_poly_edge_id
     
739.                 self.update_cursor()
     
740.                 self.update()
     
741. 
     
742.             if (self.mouse_op_poly_edge_id_pt.__class__ != new_mouse_op_poly_edge_id_pt.__class__) or \
     
743.                (isinstance(self.mouse_op_poly_edge_id_pt, np.ndarray) and \
     
744.                 all(self.mouse_op_poly_edge_id_pt != new_mouse_op_poly_edge_id_pt)):
     
745. 
     
746.                 self.mouse_op_poly_edge_id_pt = new_mouse_op_poly_edge_id_pt
     
747.                 self.update_cursor()
     
748.                 self.update()
     
749. 
     
750. 
     
751.     def action_undo_pt(self):
     
752.         if self.drag_type == DragType.NONE:
     
753.             if self.op_mode == OpMode.DRAW_PTS:
     
754.                 if self.op_poly.undo() == 0:
     
755.                     self.ie_polys.remove_poly (self.op_poly)
     
756.                     self.set_op_mode(OpMode.NONE)
     
757.                 self.update()
     
758. 
     
759.     def action_redo_pt(self):
     
760.         if self.drag_type == DragType.NONE:
     
761.             if self.op_mode == OpMode.DRAW_PTS:
     
762.                 self.op_poly.redo()
     
763.                 self.update()
     
764. 
     
765.     def action_delete_poly(self):
     
766.         if self.op_mode == OpMode.EDIT_PTS and \
     
767.             self.drag_type == DragType.NONE and \
     
768.             self.pt_edit_mode == PTEditMode.MOVE:
     
769.             # Delete current poly
     
770.             self.ie_polys.remove_poly (self.op_poly)
     
771.             self.set_op_mode(OpMode.NONE)
     
772. 
     
773.     def action_auto_detect(self):
     
774.         if not self.initialized:
     
775.             return
     
776.         if getattr(self, '_detect_running', False):
     
777.             return
     
778.         img = self.q_img_copy if hasattr(self, 'q_img_copy') else None
     
779.         if img is None:
     
780.             return
     
781.         detector = getattr(self, '_occlusion_detector', None)
     
782.         if detector is None:
     
783.             return
     
784.         self._detect_running = True
     
785.         self.cbar.btn_auto_detect_act.setDisabled(True)
     
786. 
     
787.         class _DetectThread(QThread):
     
788.             result_ready = pyqtSignal(object)
     
789.             def __init__(self, detector, img):
     
790.                 super().__init__()
     
791.                 self.detector = detector
     
792.                 self.img = img
     
793.             def run(self):
     
794.                 try:
     
795.                     polys_data = self.detector.detect(self.img)
     
796.                     self.result_ready.emit(polys_data)
     
797.                 except Exception as e:
     
798.                     self.result_ready.emit(None)
     
799. 
     
800.         def on_result(polys_data):
     
801.             self._detect_running = False
     
802.             self.cbar.btn_auto_detect_act.setDisabled(False)
     
803.             if polys_data is not None:
     
804.                 include_polys = polys_data.get('include', [])
     
805.                 exclude_polys = polys_data.get('exclude', [])
     
806.                 for pts in include_polys:
     
807.                     poly = self.ie_polys.add_poly(SegIEPolyType.INCLUDE)
     
808.                     poly.set_points(pts)
     
809.                 for pts in exclude_polys:
     
810.                     poly = self.ie_polys.add_poly(SegIEPolyType.EXCLUDE)
     
811.                     poly.set_points(pts)
     
812.                 if len(include_polys) > 0 or len(exclude_polys) > 0:
     
813.                     self.set_op_mode(OpMode.NONE)
     
814.                     self.update()
     
815. 
     
816.         self._detect_thread = _DetectThread(detector, img)
     
817.         self._detect_thread.result_ready.connect(on_result)
     
818.         self._detect_thread.start()
     
819. 
     
820.     # ====================================================================================
     
821.     # ====================================================================================
     
822.     # ================================== OVERRIDE QT METHODS =============================
     
823.     # ====================================================================================
     
824.     # ====================================================================================
     
825.     def on_keyPressEvent(self, ev):
     
826.         if not self.initialized:
     
827.             return
     
828.         key = ev.key()
     
829.         key_mods = int(ev.modifiers())
     
830.         if self.op_mode == OpMode.DRAW_PTS:
     
831.             self.set_view_lock(ViewLock.CENTER if key_mods == Qt.ShiftModifier else ViewLock.NONE )
     
832.         elif self.op_mode == OpMode.EDIT_PTS:
     
833.             self.set_pt_edit_mode(PTEditMode.ADD_DEL if key_mods == Qt.ControlModifier else PTEditMode.MOVE )
     
834. 
     
835.     def on_keyReleaseEvent(self, ev):
     
836.         if not self.initialized:
     
837.             return
     
838.         key = ev.key()
     
839.         key_mods = int(ev.modifiers())
     
840.         if self.op_mode == OpMode.DRAW_PTS:
     
841.             self.set_view_lock(ViewLock.CENTER if key_mods == Qt.ShiftModifier else ViewLock.NONE )
     
842.         elif self.op_mode == OpMode.EDIT_PTS:
     
843.             self.set_pt_edit_mode(PTEditMode.ADD_DEL if key_mods == Qt.ControlModifier else PTEditMode.MOVE )
     
844. 
     
845.     def enterEvent(self, ev):
     
846.         super().enterEvent(ev)
     
847.         self.mouse_in_widget = True
     
848.         self.update_cursor()
     
849. 
     
850.     def leaveEvent(self, ev):
     
851.         super().leaveEvent(ev)
     
852.         self.mouse_in_widget = False
     
853.         self.update_cursor()
     
854. 
     
855.     def mousePressEvent(self, ev):
     
856.         super().mousePressEvent(ev)
     
857.         if not self.initialized:
     
858.             return
     
859. 
     
860.         self.update_mouse_info(QPoint_to_np(ev.pos()))
     
861. 
     
862.         btn = ev.button()
     
863. 
     
864.         if btn == Qt.LeftButton:
     
865.             if self.op_mode == OpMode.NONE:
     
866.                 # Clicking in NO OPERATION mode
     
867.                 if self.mouse_wire_poly is not None:
     
868.                     # Click on wire on any poly -> switch to EDIT_MODE
     
869.                     self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.mouse_wire_poly)
     
870.                 else:
     
871.                     # Click on empty space -> create new poly with one point
     
872.                     new_poly = self.ie_polys.add_poly(self.poly_include_type)
     
873.                     self.ie_polys.sort()
     
874.                     new_poly.add_pt(*self.mouse_img_pt)
     
875.                     self.set_op_mode(OpMode.DRAW_PTS, op_poly=new_poly )
     
876. 
     
877.             elif self.op_mode == OpMode.DRAW_PTS:
     
878.                 # Clicking in DRAW_PTS mode
     
879.                 if len(self.op_poly.get_pts()) >= 3 and self.mouse_op_poly_pt_id == 0:
     
880.                     # Click on first point -> close poly and switch to edit mode
     
881.                     self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.op_poly)
     
882.                 else:
     
883.                     # Click on empty space -> add point to current poly
     
884.                     self.op_poly.add_pt(*self.mouse_img_pt)
     
885.                     self.update()
     
886. 
     
887.             elif self.op_mode == OpMode.EDIT_PTS:
     
888.                 # Clicking in EDIT_PTS mode
     
889. 
     
890.                 if self.mouse_op_poly_pt_id is not None:
     
891.                     # Click on point of op_poly
     
892.                     if self.pt_edit_mode == PTEditMode.ADD_DEL:
     
893.                         # in mode 'delete point'
     
894.                         self.op_poly.remove_pt(self.mouse_op_poly_pt_id)
     
895.                         if self.op_poly.get_pts_count() < 3:
     
896.                             # not enough points after delete -> remove poly
     
897.                             self.ie_polys.remove_poly (self.op_poly)
     
898.                             self.set_op_mode(OpMode.NONE)
     
899.                         self.update()
     
900. 
     
901.                     elif self.drag_type == DragType.NONE:
     
902.                         # otherwise -> start drag
     
903.                         self.drag_type = DragType.POLY_PT
     
904.                         self.drag_cli_pt     = self.mouse_cli_pt
     
905.                         self.drag_poly_pt_id = self.mouse_op_poly_pt_id
     
906.                         self.drag_poly_pt    = self.op_poly.get_pts()[ self.drag_poly_pt_id ]
     
907.                 elif self.mouse_op_poly_edge_id is not None:
     
908.                     # Click on edge of op_poly
     
909.                     if self.pt_edit_mode == PTEditMode.ADD_DEL:
     
910.                         # in mode 'insert new point'
     
911.                         edge_img_pt = self.cli_to_img_pt(self.mouse_op_poly_edge_id_pt)
     
912.                         self.op_poly.insert_pt (self.mouse_op_poly_edge_id+1, edge_img_pt)
     
913.                         self.update()
     
914.                     else:
     
915.                         # Otherwise do nothing
     
916.                         pass
     
917.                 else:
     
918.                     # other cases -> unselect poly
     
919.                     self.set_op_mode(OpMode.NONE)
     
920. 
     
921.         elif btn == Qt.MiddleButton:
     
922.             if self.drag_type == DragType.NONE:
     
923.                 # Start image drag
     
924.                 self.drag_type = DragType.IMAGE_LOOK
     
925.                 self.drag_cli_pt = self.mouse_cli_pt
     
926.                 self.drag_img_look_pt = self.get_img_look_pt()
     
927.                 self.update_cursor()
     
928. 
     
929. 
     
930.     def mouseReleaseEvent(self, ev):
     
931.         super().mouseReleaseEvent(ev)
     
932.         if not self.initialized:
     
933.             return
     
934. 
     
935.         self.update_mouse_info(QPoint_to_np(ev.pos()))
     
936. 
     
937.         btn = ev.button()
     
938. 
     
939.         if btn == Qt.LeftButton:
     
940.             if self.op_mode == OpMode.EDIT_PTS:
     
941.                 if self.drag_type == DragType.POLY_PT:
     
942.                     self.drag_type = DragType.NONE
     
943.                     self.update()
     
944. 
     
945.         elif btn == Qt.MiddleButton:
     
946.             if self.drag_type == DragType.IMAGE_LOOK:
     
947.                 self.drag_type = DragType.NONE
     
948.                 self.update_cursor()
     
949.                 self.update()
     
950. 
     
951.     def mouseMoveEvent(self, ev):
     
952.         super().mouseMoveEvent(ev)
     
953.         if not self.initialized:
     
954.             return
     
955. 
     
956.         prev_mouse_cli_pt = self.mouse_cli_pt
     
957.         self.update_mouse_info(QPoint_to_np(ev.pos()))
     
958. 
     
959.         if self.view_lock == ViewLock.CENTER:
     
960.             if npla.norm(self.mouse_cli_pt - prev_mouse_cli_pt) >= 1:
     
961.                 self.img_look_pt = self.mouse_img_pt
     
962.                 QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) ))
     
963.             self.update()
     
964. 
     
965.         if self.drag_type == DragType.IMAGE_LOOK:
     
966.             delta_pt = self.cli_to_img_pt(self.mouse_cli_pt) - self.cli_to_img_pt(self.drag_cli_pt)
     
967.             self.img_look_pt = self.drag_img_look_pt - delta_pt
     
968.             self.update()
     
969. 
     
970.         if self.op_mode == OpMode.DRAW_PTS:
     
971.             self.update()
     
972.         elif self.op_mode == OpMode.EDIT_PTS:
     
973.             if self.drag_type == DragType.POLY_PT:
     
974.                 delta_pt = self.cli_to_img_pt(self.mouse_cli_pt) - self.cli_to_img_pt(self.drag_cli_pt)
     
975.                 self.op_poly.set_point(self.drag_poly_pt_id, self.drag_poly_pt + delta_pt)
     
976.                 self.update()
     
977. 
     
978.     def wheelEvent(self, ev):
     
979.         super().wheelEvent(ev)
     
980. 
     
981.         if not self.initialized:
     
982.             return
     
983. 
     
984.         mods = int(ev.modifiers())
     
985.         delta = ev.angleDelta()
     
986. 
     
987.         cli_pt = QPoint_to_np(ev.pos())
     
988. 
     
989.         if self.drag_type == DragType.NONE:
     
990.             sign = np.sign( delta.y() )
     
991.             prev_img_pos = self.cli_to_img_pt (cli_pt)
     
992.             delta_scale = sign*0.2 + sign * self.get_view_scale() / 10.0
     
993.             self.view_scale = np.clip(self.get_view_scale() + delta_scale, 1.0, 20.0)
     
994.             new_img_pos = self.cli_to_img_pt (cli_pt)
     
995.             if sign > 0:
     
996.                 self.img_look_pt = self.get_img_look_pt() + (prev_img_pos-new_img_pos)#*1.5
     
997.             else:
     
998.                 QCursor.setPos ( self.mapToGlobal(QPoint_from_np(self.img_to_cli_pt(prev_img_pos))) )
     
999.             self.update()
     
1000. 
      
1001.     def paintEvent(self, event):
      
1002.         super().paintEvent(event)
      
1003.         if not self.initialized:
      
1004.             return
      
1005. 
      
1006.         qp = self.qp
      
1007.         qp.begin(self)
      
1008.         qp.setRenderHint(QPainter.Antialiasing)
      
1009.         qp.setRenderHint(QPainter.HighQualityAntialiasing)
      
1010.         qp.setRenderHint(QPainter.SmoothPixmapTransform)
      
1011. 
      
1012.         src_rect = QRect(0, 0, *self.img_size)
      
1013.         dst_rect = self.img_to_cli_rect( src_rect )
      
1014. 
      
1015.         if self.op_mode == OpMode.VIEW_BAKED:
      
1016.             qp.drawPixmap(dst_rect, self.img_baked_pixmap, src_rect)
      
1017.         elif self.op_mode == OpMode.VIEW_XSEG_MASK:
      
1018.             if self.xseg_mask_pixmap is not None:
      
1019.                 qp.drawPixmap(dst_rect, self.xseg_mask_pixmap, src_rect)
      
1020.         else:
      
1021.             if self.cbar.btn_view_xseg_overlay_mask_act.isChecked() and \
      
1022.                self.xseg_overlay_mask_pixmap is not None:
      
1023.                 qp.drawPixmap(dst_rect, self.xseg_overlay_mask_pixmap, src_rect)
      
1024.             elif self.img_pixmap is not None:
      
1025.                 qp.drawPixmap(dst_rect, self.img_pixmap, src_rect)
      
1026. 
      
1027.             polys = self.ie_polys.get_polys()
      
1028.             polys_len = len(polys)
      
1029. 
      
1030.             color_scheme = self.get_current_color_scheme()
      
1031. 
      
1032.             pt_rad = self.canvas_config.pt_radius
      
1033.             pt_rad_x2 = pt_rad*2
      
1034. 
      
1035.             pt_select_radius = self.canvas_config.pt_select_radius
      
1036. 
      
1037.             op_mode = self.op_mode
      
1038.             op_poly = self.op_poly
      
1039. 
      
1040.             for i,poly in enumerate(polys):
      
1041. 
      
1042.                 selected_pt_path = QPainterPath()
      
1043.                 poly_line_path = QPainterPath()
      
1044.                 pts_line_path = QPainterPath()
      
1045. 
      
1046.                 pt_remove_cli_pt = None
      
1047.                 poly_pts = poly.get_pts()
      
1048.                 for pt_id, img_pt in enumerate(poly_pts):
      
1049.                     cli_pt = self.img_to_cli_pt(img_pt)
      
1050.                     q_cli_pt = QPoint_from_np(cli_pt)
      
1051. 
      
1052.                     if pt_id == 0:
      
1053.                         poly_line_path.moveTo(q_cli_pt)
      
1054.                     else:
      
1055.                         poly_line_path.lineTo(q_cli_pt)
      
1056. 
      
1057. 
      
1058.                     if poly == op_poly:
      
1059.                         if self.op_mode == OpMode.DRAW_PTS or \
      
1060.                             (self.op_mode == OpMode.EDIT_PTS and \
      
1061.                             (self.pt_edit_mode == PTEditMode.MOVE) or \
      
1062.                             (self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id == pt_id) \
      
1063.                             ):
      
1064.                             pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([0,-pt_rad])) )
      
1065.                             pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([0,pt_rad])) )
      
1066.                             pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([-pt_rad,0])) )
      
1067.                             pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([pt_rad,0])) )
      
1068. 
      
1069.                         if (self.op_mode == OpMode.EDIT_PTS and \
      
1070.                             self.pt_edit_mode == PTEditMode.ADD_DEL and \
      
1071.                             self.mouse_op_poly_pt_id == pt_id):
      
1072.                             pt_remove_cli_pt = cli_pt
      
1073. 
      
1074.                         if self.op_mode == OpMode.DRAW_PTS and \
      
1075.                             len(op_poly.get_pts()) >= 3 and pt_id == 0 and self.mouse_op_poly_pt_id == pt_id:
      
1076.                             # Circle around poly point
      
1077.                             selected_pt_path.addEllipse(q_cli_pt, pt_rad_x2, pt_rad_x2)
      
1078. 
      
1079. 
      
1080.                 if poly == op_poly:
      
1081.                     if op_mode == OpMode.DRAW_PTS:
      
1082.                         # Line from last point to mouse
      
1083.                         poly_line_path.lineTo( QPoint_from_np(self.mouse_cli_pt) )
      
1084. 
      
1085.                     if self.mouse_op_poly_pt_id is not None:
      
1086.                         pass
      
1087. 
      
1088.                     if self.mouse_op_poly_edge_id_pt is not None:
      
1089.                         if self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id is None:
      
1090.                             # Ready to insert point on edge
      
1091.                             m_cli_pt = self.mouse_op_poly_edge_id_pt
      
1092.                             pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([0,-pt_rad])) )
      
1093.                             pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([0,pt_rad])) )
      
1094.                             pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([-pt_rad,0])) )
      
1095.                             pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([pt_rad,0])) )
      
1096. 
      
1097.                 if len(poly_pts) >= 2:
      
1098.                     # Closing poly line
      
1099.                     poly_line_path.lineTo( QPoint_from_np(self.img_to_cli_pt(poly_pts[0])) )
      
1100. 
      
1101.                 # Draw calls
      
1102.                 qp.setPen(color_scheme.pt_outline_pen)
      
1103.                 qp.setBrush(QBrush())
      
1104.                 qp.drawPath(selected_pt_path)
      
1105. 
      
1106.                 qp.setPen(color_scheme.poly_outline_solid_pen)
      
1107.                 qp.setBrush(QBrush())
      
1108.                 qp.drawPath(pts_line_path)
      
1109. 
      
1110.                 if poly.get_type() == SegIEPolyType.INCLUDE:
      
1111.                     qp.setPen(color_scheme.poly_outline_solid_pen)
      
1112.                 else:
      
1113.                     qp.setPen(color_scheme.poly_outline_dot_pen)
      
1114. 
      
1115.                 qp.setBrush(color_scheme.poly_unselected_brush)
      
1116.                 if op_mode == OpMode.NONE:
      
1117.                     if poly == self.mouse_wire_poly:
      
1118.                         qp.setBrush(color_scheme.poly_selected_brush)
      
1119.                 #else:
      
1120.                 #    if poly == op_poly:
      
1121.                 #        qp.setBrush(color_scheme.poly_selected_brush)
      
1122. 
      
1123.                 qp.drawPath(poly_line_path)
      
1124. 
      
1125.                 if pt_remove_cli_pt is not None:
      
1126.                     qp.setPen(color_scheme.poly_outline_solid_pen)
      
1127.                     qp.setBrush(QBrush())
      
1128. 
      
1129.                     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])) )
      
1130.                     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])) )
      
1131. 
      
1132.         qp.end()
      
1133. 
      
1134. class QCanvas(QFrame):
      
1135.     def __init__(self):
      
1136.         super().__init__()
      
1137. 
      
1138.         self.canvas_control_left_bar = QCanvasControlsLeftBar()
      
1139.         self.canvas_control_right_bar = QCanvasControlsRightBar()
      
1140. 
      
1141.         cbar = sn( btn_poly_color_red_act   = self.canvas_control_right_bar.btn_poly_color_red_act,
      
1142.                    btn_poly_color_green_act = self.canvas_control_right_bar.btn_poly_color_green_act,
      
1143.                    btn_poly_color_blue_act  = self.canvas_control_right_bar.btn_poly_color_blue_act,
      
1144.                    btn_view_baked_mask_act  = self.canvas_control_right_bar.btn_view_baked_mask_act,
      
1145.                    btn_view_xseg_mask_act  = self.canvas_control_right_bar.btn_view_xseg_mask_act,
      
1146.                    btn_view_xseg_overlay_mask_act  = self.canvas_control_right_bar.btn_view_xseg_overlay_mask_act,
      
1147.                    btn_poly_color_act_grp = self.canvas_control_right_bar.btn_poly_color_act_grp,
      
1148.                    btn_view_lock_center_act = self.canvas_control_right_bar.btn_view_lock_center_act,
      
1149. 
      
1150.                    btn_poly_type_include_act = self.canvas_control_left_bar.btn_poly_type_include_act,
      
1151.                    btn_poly_type_exclude_act = self.canvas_control_left_bar.btn_poly_type_exclude_act,
      
1152.                    btn_poly_type_act_grp = self.canvas_control_left_bar.btn_poly_type_act_grp,
      
1153.                    btn_undo_pt_act = self.canvas_control_left_bar.btn_undo_pt_act,
      
1154.                    btn_redo_pt_act = self.canvas_control_left_bar.btn_redo_pt_act,
      
1155.                    btn_delete_poly_act = self.canvas_control_left_bar.btn_delete_poly_act,
      
1156.                    btn_pt_edit_mode_act = self.canvas_control_left_bar.btn_pt_edit_mode_act,
      
1157.                    btn_auto_detect_act = self.canvas_control_left_bar.btn_auto_detect_act )
      
1158. 
      
1159.         self.op = QCanvasOperator(cbar)
      
1160.         self.l = QHBoxLayout()
      
1161.         self.l.setContentsMargins(0,0,0,0)
      
1162.         self.l.addWidget(self.canvas_control_left_bar)
      
1163.         self.l.addWidget(self.op)
      
1164.         self.l.addWidget(self.canvas_control_right_bar)
      
1165.         self.setLayout(self.l)
      
1166. 
      
1167. class LoaderQSubprocessor(QSubprocessor):
      
1168.     def __init__(self, image_paths, q_label, q_progressbar, on_finish_func ):
      
1169. 
      
1170.         self.image_paths = image_paths
      
1171.         self.image_paths_len = len(image_paths)
      
1172.         self.idxs = [*range(self.image_paths_len)]
      
1173. 
      
1174.         self.filtered_image_paths = self.image_paths.copy()
      
1175. 
      
1176.         self.image_paths_has_ie_polys = { image_path : False for image_path in self.image_paths }
      
1177. 
      
1178.         self.q_label = q_label
      
1179.         self.q_progressbar = q_progressbar
      
1180.         self.q_progressbar.setRange(0, self.image_paths_len)
      
1181.         self.q_progressbar.setValue(0)
      
1182.         self.q_progressbar.update()
      
1183.         self.on_finish_func = on_finish_func
      
1184.         self.done_count = 0
      
1185.         super().__init__('LoaderQSubprocessor', LoaderQSubprocessor.Cli, 60)
      
1186. 
      
1187.     def get_data(self, host_dict):
      
1188.         if len (self.idxs) > 0:
      
1189.             idx = self.idxs.pop(0)
      
1190.             image_path = self.image_paths[idx]
      
1191.             self.q_label.setText(f'{QStringDB.loading_tip}... {image_path.name}')
      
1192. 
      
1193.             return idx, image_path
      
1194. 
      
1195.         return None
      
1196. 
      
1197.     def on_clients_finalized(self):
      
1198.         self.on_finish_func([x for x in self.filtered_image_paths if x is not None], self.image_paths_has_ie_polys)
      
1199. 
      
1200.     def on_data_return (self, host_dict, data):
      
1201.         self.idxs.insert(0, data[0])
      
1202. 
      
1203.     def on_result (self, host_dict, data, result):
      
1204.         idx, has_dflimg, has_ie_polys = result
      
1205. 
      
1206.         if not has_dflimg:
      
1207.             self.filtered_image_paths[idx] = None
      
1208.         self.image_paths_has_ie_polys[self.image_paths[idx]] = has_ie_polys
      
1209. 
      
1210.         self.done_count += 1
      
1211.         if self.q_progressbar is not None:
      
1212.             self.q_progressbar.setValue(self.done_count)
      
1213. 
      
1214.     class Cli(QSubprocessor.Cli):
      
1215.         def process_data(self, data):
      
1216.             idx, filename = data
      
1217.             dflimg = DFLIMG.load(filename)
      
1218.             if dflimg is not None and dflimg.has_data():
      
1219.                 ie_polys = dflimg.get_seg_ie_polys()
      
1220. 
      
1221.                 return idx, True, ie_polys.has_polys()
      
1222.             return idx, False, False
      
1223. 
      
1224. class MainWindow(QXMainWindow):
      
1225. 
      
1226.     def __init__(self, input_dirpath, cfg_root_path):
      
1227.         self.loading_frame = None
      
1228.         self.help_frame = None
      
1229. 
      
1230.         super().__init__()
      
1231. 
      
1232.         self.input_dirpath = input_dirpath
      
1233.         self.trash_dirpath = input_dirpath.parent / (input_dirpath.name + '_trash')
      
1234.         self.cfg_root_path = cfg_root_path
      
1235. 
      
1236.         self.cfg_path = cfg_root_path / 'MainWindow_cfg.dat'
      
1237.         self.cfg_dict = pickle.loads(self.cfg_path.read_bytes()) if self.cfg_path.exists() else {}
      
1238. 
      
1239.         self.cached_images = {}
      
1240.         self.cached_has_ie_polys = {}
      
1241. 
      
1242.         self._occlusion_detector = OcclusionDetector()
      
1243. 
      
1244.         self.initialize_ui()
      
1245. 
      
1246.         # Loader
      
1247.         self.loading_frame = QFrame(self.main_canvas_frame)
      
1248.         self.loading_frame.setAutoFillBackground(True)
      
1249.         self.loading_frame.setFrameShape(QFrame.StyledPanel)
      
1250.         self.loader_label = QLabel()
      
1251.         self.loader_progress_bar = QProgressBar()
      
1252. 
      
1253.         intro_image = QLabel()
      
1254.         intro_image.setPixmap( QPixmap.fromImage(QImageDB.intro) )
      
1255. 
      
1256.         intro_image_frame_l = QVBoxLayout()
      
1257.         intro_image_frame_l.addWidget(intro_image, alignment=Qt.AlignCenter)
      
1258.         intro_image_frame = QFrame()
      
1259.         intro_image_frame.setSizePolicy (QSizePolicy.Expanding, QSizePolicy.Expanding)
      
1260.         intro_image_frame.setLayout(intro_image_frame_l)
      
1261. 
      
1262.         loading_frame_l = QVBoxLayout()
      
1263.         loading_frame_l.addWidget (intro_image_frame)
      
1264.         loading_frame_l.addWidget (self.loader_label)
      
1265.         loading_frame_l.addWidget (self.loader_progress_bar)
      
1266.         self.loading_frame.setLayout(loading_frame_l)
      
1267. 
      
1268.         self.loader_subprocessor = LoaderQSubprocessor( image_paths=pathex.get_image_paths(input_dirpath, return_Path_class=True),
      
1269.                                                         q_label=self.loader_label,
      
1270.                                                         q_progressbar=self.loader_progress_bar,
      
1271.                                                         on_finish_func=self.on_loader_finish )
      
1272. 
      
1273. 
      
1274.     def on_loader_finish(self, image_paths, image_paths_has_ie_polys):
      
1275.         self.image_paths_done = []
      
1276.         self.image_paths = image_paths
      
1277.         self.image_paths_has_ie_polys = image_paths_has_ie_polys
      
1278.         self.set_has_ie_polys_count ( len([ 1 for x in self.image_paths_has_ie_polys if self.image_paths_has_ie_polys[x] == True]) )
      
1279.         self.loading_frame.hide()
      
1280.         self.loading_frame = None
      
1281. 
      
1282.         self.process_next_image(first_initialization=True)
      
1283. 
      
1284.     def closeEvent(self, ev):
      
1285.         self._occlusion_detector.shutdown()
      
1286.         self.cfg_dict['geometry'] = self.saveGeometry().data()
      
1287.         self.cfg_path.write_bytes( pickle.dumps(self.cfg_dict) )
      
1288. 
      
1289. 
      
1290.     def update_cached_images (self, count=5):
      
1291.         d = self.cached_images
      
1292. 
      
1293.         for image_path in self.image_paths_done[:-count]+self.image_paths[count:]:
      
1294.             if image_path in d:
      
1295.                 del d[image_path]
      
1296. 
      
1297.         for image_path in self.image_paths[:count]+self.image_paths_done[-count:]:
      
1298.             if image_path not in d:
      
1299.                 img = cv2_imread(image_path)
      
1300.                 if img is not None:
      
1301.                     d[image_path] = img
      
1302. 
      
1303.     def load_image(self, image_path):
      
1304.         try:
      
1305.             img = self.cached_images.get(image_path, None)
      
1306.             if img is None:
      
1307.                 img = cv2_imread(image_path)
      
1308.                 self.cached_images[image_path] = img
      
1309.             if img is None:
      
1310.                 io.log_err(f'Unable to load {image_path}')
      
1311.         except:
      
1312.             img = None
      
1313. 
      
1314.         return img
      
1315. 
      
1316.     def update_preview_bar(self):
      
1317.         count = self.image_bar.get_preview_images_count()
      
1318.         d = self.cached_images
      
1319.         prev_imgs = [ d.get(image_path, None) for image_path in self.image_paths_done[-1:-count:-1] ]
      
1320.         next_imgs = [ d.get(image_path, None) for image_path in self.image_paths[:count] ]
      
1321.         self.image_bar.update_images(prev_imgs, next_imgs)
      
1322. 
      
1323. 
      
1324.     def canvas_initialize(self, image_path, only_has_polys=False):
      
1325.         if only_has_polys and not self.image_paths_has_ie_polys[image_path]:
      
1326.             return False
      
1327. 
      
1328.         dflimg = DFLIMG.load(image_path)
      
1329.         if not dflimg or not dflimg.has_data():
      
1330.             return False
      
1331. 
      
1332.         ie_polys = dflimg.get_seg_ie_polys()
      
1333.         xseg_mask = dflimg.get_xseg_mask()
      
1334.         landmarks = dflimg.get_landmarks()
      
1335.         img = self.load_image(image_path)
      
1336.         if img is None:
      
1337.             return False
      
1338. 
      
1339.         self.canvas.op.initialize ( img,  ie_polys=ie_polys, xseg_mask=xseg_mask, landmarks=landmarks )
      
1340. 
      
1341.         self.filename_label.setText(f"{image_path.name}")
      
1342. 
      
1343.         return True
      
1344. 
      
1345.     def canvas_finalize(self, image_path):
      
1346.         self.canvas.op.finalize()
      
1347. 
      
1348.         if image_path.exists():
      
1349.             dflimg = DFLIMG.load(image_path)
      
1350.             ie_polys = dflimg.get_seg_ie_polys()
      
1351.             new_ie_polys = self.canvas.op.get_ie_polys()
      
1352. 
      
1353.             if not new_ie_polys.identical(ie_polys):
      
1354.                 prev_has_polys = self.image_paths_has_ie_polys[image_path]
      
1355.                 self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys()
      
1356.                 new_has_polys = self.image_paths_has_ie_polys[image_path]
      
1357. 
      
1358.                 if not prev_has_polys and new_has_polys:
      
1359.                     self.set_has_ie_polys_count ( self.get_has_ie_polys_count() +1)
      
1360.                 elif prev_has_polys and not new_has_polys:
      
1361.                     self.set_has_ie_polys_count ( self.get_has_ie_polys_count() -1)
      
1362. 
      
1363.                 dflimg.set_seg_ie_polys( new_ie_polys )
      
1364.                 dflimg.save()
      
1365. 
      
1366.         self.filename_label.setText(f"")
      
1367. 
      
1368.     def process_prev_image(self):
      
1369.         key_mods = QApplication.keyboardModifiers()
      
1370.         step = 5 if key_mods == Qt.ShiftModifier else 1
      
1371.         only_has_polys = key_mods == Qt.ControlModifier
      
1372. 
      
1373.         if self.canvas.op.is_initialized():
      
1374.             self.canvas_finalize(self.image_paths[0])
      
1375. 
      
1376.         while True:
      
1377.             for _ in range(step):
      
1378.                 if len(self.image_paths_done) != 0:
      
1379.                     self.image_paths.insert (0, self.image_paths_done.pop(-1))
      
1380.                 else:
      
1381.                     break
      
1382.             if len(self.image_paths) == 0:
      
1383.                 break
      
1384. 
      
1385.             ret = self.canvas_initialize(self.image_paths[0], len(self.image_paths_done) != 0 and only_has_polys)
      
1386. 
      
1387.             if ret or len(self.image_paths_done) == 0:
      
1388.                 break
      
1389. 
      
1390.         self.update_cached_images()
      
1391.         self.update_preview_bar()
      
1392. 
      
1393.     def process_next_image(self, first_initialization=False):
      
1394.         key_mods = QApplication.keyboardModifiers()
      
1395. 
      
1396.         step = 0 if first_initialization else 5 if key_mods == Qt.ShiftModifier else 1
      
1397.         only_has_polys = False if first_initialization else key_mods == Qt.ControlModifier
      
1398. 
      
1399.         if self.canvas.op.is_initialized():
      
1400.             self.canvas_finalize(self.image_paths[0])
      
1401. 
      
1402.         while True:
      
1403.             for _ in range(step):
      
1404.                 if len(self.image_paths) != 0:
      
1405.                     self.image_paths_done.append(self.image_paths.pop(0))
      
1406.                 else:
      
1407.                     break
      
1408.             if len(self.image_paths) == 0:
      
1409.                 break
      
1410.             if self.canvas_initialize(self.image_paths[0], only_has_polys):
      
1411.                 break
      
1412. 
      
1413.         self.update_cached_images()
      
1414.         self.update_preview_bar()
      
1415.         
      
1416.     def trash_current_image(self):
      
1417.         self.process_next_image()
      
1418.         
      
1419.         img_path = self.image_paths_done.pop(-1)
      
1420.         img_path = Path(img_path)
      
1421.         self.trash_dirpath.mkdir(parents=True, exist_ok=True)
      
1422.         img_path.rename( self.trash_dirpath / img_path.name )
      
1423.         
      
1424.         self.update_cached_images()
      
1425.         self.update_preview_bar()
      
1426.         
      
1427.     def initialize_ui(self):
      
1428. 
      
1429.         self.canvas = QCanvas()
      
1430.         self.canvas.op._occlusion_detector = self._occlusion_detector
      
1431. 
      
1432.         image_bar = self.image_bar = ImagePreviewSequenceBar(preview_images_count=9, icon_size=QUIConfig.preview_bar_icon_q_size.width())
      
1433.         image_bar.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed )
      
1434. 
      
1435. 
      
1436.         btn_prev_image = QXIconButton(QIconDB.left, QStringDB.btn_prev_image_tip, shortcut='A', click_func=self.process_prev_image)
      
1437.         btn_prev_image.setIconSize(QUIConfig.preview_bar_icon_q_size)
      
1438. 
      
1439.         btn_next_image = QXIconButton(QIconDB.right, QStringDB.btn_next_image_tip, shortcut='D', click_func=self.process_next_image)
      
1440.         btn_next_image.setIconSize(QUIConfig.preview_bar_icon_q_size)
      
1441. 
      
1442.         btn_delete_image = QXIconButton(QIconDB.trashcan, QStringDB.btn_delete_image_tip, shortcut='X', click_func=self.trash_current_image)
      
1443.         btn_delete_image.setIconSize(QUIConfig.preview_bar_icon_q_size)
      
1444.    
      
1445.         pad_image = QWidget()
      
1446.         pad_image.setFixedSize(QUIConfig.preview_bar_icon_q_size)
      
1447.         
      
1448.         preview_image_bar_frame_l = QHBoxLayout()
      
1449.         preview_image_bar_frame_l.setContentsMargins(0,0,0,0)
      
1450.         preview_image_bar_frame_l.addWidget ( pad_image, alignment=Qt.AlignCenter)
      
1451.         preview_image_bar_frame_l.addWidget ( btn_prev_image, alignment=Qt.AlignCenter)
      
1452.         preview_image_bar_frame_l.addWidget ( image_bar)
      
1453.         preview_image_bar_frame_l.addWidget ( btn_next_image, alignment=Qt.AlignCenter)
      
1454.         #preview_image_bar_frame_l.addWidget ( btn_delete_image, alignment=Qt.AlignCenter)
      
1455. 
      
1456.         preview_image_bar_frame = QFrame()
      
1457.         preview_image_bar_frame.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed )
      
1458.         preview_image_bar_frame.setLayout(preview_image_bar_frame_l)
      
1459. 
      
1460.         preview_image_bar_frame2_l = QHBoxLayout()
      
1461.         preview_image_bar_frame2_l.setContentsMargins(0,0,0,0)
      
1462.         preview_image_bar_frame2_l.addWidget ( btn_delete_image, alignment=Qt.AlignCenter)
      
1463. 
      
1464.         preview_image_bar_frame2 = QFrame()
      
1465.         preview_image_bar_frame2.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed )
      
1466.         preview_image_bar_frame2.setLayout(preview_image_bar_frame2_l)
      
1467.         
      
1468.         preview_image_bar_l = QHBoxLayout()
      
1469.         preview_image_bar_l.addWidget (preview_image_bar_frame, alignment=Qt.AlignCenter)
      
1470.         preview_image_bar_l.addWidget (preview_image_bar_frame2)
      
1471.         
      
1472.         preview_image_bar = QFrame()
      
1473.         preview_image_bar.setFrameShape(QFrame.StyledPanel)
      
1474.         preview_image_bar.setSizePolicy ( QSizePolicy.Expanding, QSizePolicy.Fixed )
      
1475.         preview_image_bar.setLayout(preview_image_bar_l)
      
1476. 
      
1477.         label_font = QFont('Courier New')
      
1478.         self.filename_label = QLabel()
      
1479.         self.filename_label.setFont(label_font)
      
1480. 
      
1481.         self.has_ie_polys_count_label = QLabel()
      
1482. 
      
1483.         status_frame_l = QHBoxLayout()
      
1484.         status_frame_l.setContentsMargins(0,0,0,0)
      
1485.         status_frame_l.addWidget ( QLabel(), alignment=Qt.AlignCenter)
      
1486.         status_frame_l.addWidget (self.filename_label, alignment=Qt.AlignCenter)
      
1487.         status_frame_l.addWidget (self.has_ie_polys_count_label, alignment=Qt.AlignCenter)
      
1488.         status_frame = QFrame()
      
1489.         status_frame.setLayout(status_frame_l)
      
1490. 
      
1491.         main_canvas_l = QVBoxLayout()
      
1492.         main_canvas_l.setContentsMargins(0,0,0,0)
      
1493.         main_canvas_l.addWidget (self.canvas)
      
1494.         main_canvas_l.addWidget (status_frame)
      
1495.         main_canvas_l.addWidget (preview_image_bar)
      
1496. 
      
1497.         self.main_canvas_frame = QFrame()
      
1498.         self.main_canvas_frame.setLayout(main_canvas_l)
      
1499. 
      
1500.         self.main_l = QHBoxLayout()
      
1501.         self.main_l.setContentsMargins(0,0,0,0)
      
1502.         self.main_l.addWidget (self.main_canvas_frame)
      
1503. 
      
1504.         self.setLayout(self.main_l)
      
1505. 
      
1506.         geometry = self.cfg_dict.get('geometry', None)
      
1507.         if geometry is not None:
      
1508.             self.restoreGeometry(geometry)
      
1509.         else:
      
1510.             self.move( QPoint(0,0))
      
1511. 
      
1512.     def get_has_ie_polys_count(self):
      
1513.         return self.has_ie_polys_count
      
1514. 
      
1515.     def set_has_ie_polys_count(self, c):
      
1516.         self.has_ie_polys_count = c
      
1517.         self.has_ie_polys_count_label.setText(f"{c} {QStringDB.labeled_tip}")
      
1518. 
      
1519.     def resizeEvent(self, ev):
      
1520.         if self.loading_frame is not None:
      
1521.             self.loading_frame.resize( ev.size() )
      
1522.         if self.help_frame is not None:
      
1523.             self.help_frame.resize( ev.size() )
      
1524. 
      
1525. def start(input_dirpath):
      
1526.     """
      
1527.     returns exit_code
      
1528.     """
      
1529.     io.log_info("Running XSeg editor.")
      
1530. 
      
1531.     if PackedFaceset.path_contains(input_dirpath):
      
1532.         io.log_info (f'\n{input_dirpath} contains packed faceset! Unpack it first.\n')
      
1533.         return 1
      
1534. 
      
1535.     root_path = Path(__file__).parent
      
1536.     cfg_root_path = Path(tempfile.gettempdir())
      
1537. 
      
1538.     QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
      
1539.     QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
      
1540. 
      
1541.     app = QApplication([])
      
1542.     app.setApplicationName("XSegEditor")
      
1543.     app.setStyle('Fusion')
      
1544. 
      
1545.     QFontDatabase.addApplicationFont( str(root_path / 'gfx' / 'fonts' / 'NotoSans-Medium.ttf') )
      
1546. 
      
1547.     app.setFont( QFont('NotoSans'))
      
1548. 
      
1549.     QUIConfig.initialize()
      
1550.     QStringDB.initialize()
      
1551. 
      
1552.     QIconDB.initialize( root_path / 'gfx' / 'icons' )
      
1553.     QCursorDB.initialize( root_path / 'gfx' / 'cursors' )
      
1554.     QImageDB.initialize( root_path / 'gfx' / 'images' )
      
1555. 
      
1556.     app.setWindowIcon(QIconDB.app_icon)
      
1557.     app.setPalette( QDarkPalette() )
      
1558. 
      
1559.     win = MainWindow( input_dirpath=input_dirpath, cfg_root_path=cfg_root_path)
      
1560. 
      
1561.     win.show()
      
1562.     win.raise_()
      
1563. 
      
1564.     app.exec_()
      
1565.     return 0
      

复制代码

XSegEditor/QIconDB.py

1. from PyQt5.QtCore import *
   
2. from PyQt5.QtGui import *
   
3. from PyQt5.QtWidgets import *
   
4. 
   
5. 
   
6. class QIconDB():
   
7.     @staticmethod
   
8.     def initialize(icon_path):
   
9.         QIconDB.app_icon         = QIcon ( str(icon_path / 'app_icon.png') )
   
10.         QIconDB.delete_poly      = QIcon ( str(icon_path / 'delete_poly.png') )
    
11.         QIconDB.undo_pt          = QIcon ( str(icon_path / 'undo_pt.png') )
    
12.         QIconDB.redo_pt          = QIcon ( str(icon_path / 'redo_pt.png') )
    
13.         QIconDB.poly_color_red   = QIcon ( str(icon_path / 'poly_color_red.png') )
    
14.         QIconDB.poly_color_green = QIcon ( str(icon_path / 'poly_color_green.png') )
    
15.         QIconDB.poly_color_blue  = QIcon ( str(icon_path / 'poly_color_blue.png') )
    
16.         QIconDB.poly_type_include = QIcon ( str(icon_path / 'poly_type_include.png') )
    
17.         QIconDB.poly_type_exclude = QIcon ( str(icon_path / 'poly_type_exclude.png') )
    
18.         QIconDB.left  = QIcon ( str(icon_path / 'left.png') )
    
19.         QIconDB.right = QIcon ( str(icon_path / 'right.png') )
    
20.         QIconDB.trashcan = QIcon ( str(icon_path / 'trashcan.png') )
    
21.         QIconDB.pt_edit_mode = QIcon ( str(icon_path / 'pt_edit_mode.png') )
    
22.         QIconDB.view_lock_center = QIcon ( str(icon_path / 'view_lock_center.png') )
    
23.         QIconDB.view_baked = QIcon ( str(icon_path / 'view_baked.png') )
    
24.         QIconDB.view_xseg = QIcon ( str(icon_path / 'view_xseg.png') )
    
25.         QIconDB.view_xseg_overlay = QIcon ( str(icon_path / 'view_xseg_overlay.png') )
    
26.         QIconDB.auto_detect = QIcon ( str(icon_path / 'auto_detect.png') )
    
27.         

复制代码

XSegEditor/QStringDB.py

1. from localization import system_language
   
2. 
   
3. 
   
4. class QStringDB():
   
5.    
   
6.     @staticmethod
   
7.     def initialize():
   
8.         lang = system_language
   
9.         
   
10.         if lang not in ['en','ru','zh']:
    
11.             lang = 'en'
    
12.         
    
13.         QStringDB.btn_poly_color_red_tip = {    'en' : 'Poly color scheme red',
    
14.                                                 'ru' : 'Красная цветовая схема полигонов',
    
15.                                                 'zh' : '选区配色方案红色',
    
16.                                            }[lang]
    
17.                                           
    
18.         QStringDB.btn_poly_color_green_tip = {  'en' : 'Poly color scheme green',
    
19.                                                 'ru' : 'Зелёная цветовая схема полигонов',
    
20.                                                 'zh' : '选区配色方案绿色',
    
21.                                            }[lang]
    
22.                                           
    
23.         QStringDB.btn_poly_color_blue_tip = {   'en' : 'Poly color scheme blue',
    
24.                                                 'ru' : 'Синяя цветовая схема полигонов',
    
25.                                                 'zh' : '选区配色方案蓝色',
    
26.                                            }[lang]
    
27.                                           
    
28.         QStringDB.btn_view_baked_mask_tip = {   'en' : 'View baked mask',
    
29.                                                 'ru' : 'Посмотреть запечёную маску',
    
30.                                                 'zh' : '查看遮罩通道',
    
31.                                            }[lang]
    
32.                                           
    
33.         QStringDB.btn_view_xseg_mask_tip =  {   'en' : 'View trained XSeg mask',
    
34.                                                 'ru' : 'Посмотреть тренированную XSeg маску',
    
35.                                                 'zh' : '查看导入后的XSeg遮罩',
    
36.                                             }[lang]
    
37.                                             
    
38.         QStringDB.btn_view_xseg_overlay_mask_tip =  {   'en' : 'View trained XSeg mask overlay face',
    
39.                                                         'ru' : 'Посмотреть тренированную XSeg маску поверх лица',
    
40.                                                         'zh' : '查看导入后的XSeg遮罩于脸上方',
    
41.                                                     }[lang]
    
42.                                           
    
43.         QStringDB.btn_poly_type_include_tip = { 'en' : 'Poly include mode',
    
44.                                                 'ru' : 'Режим полигонов - включение',
    
45.                                                 'zh' : '包含选区模式',
    
46.                                            }[lang]
    
47.                                           
    
48.         QStringDB.btn_poly_type_exclude_tip = { 'en' : 'Poly exclude mode',
    
49.                                                 'ru' : 'Режим полигонов - исключение',
    
50.                                                 'zh' : '排除选区模式',
    
51.                                            }[lang]        
    
52.                                           
    
53.         QStringDB.btn_undo_pt_tip = {   'en' : 'Undo point',
    
54.                                         'ru' : 'Отменить точку',
    
55.                                         'zh' : '撤消点',
    
56.                                     }[lang]      
    
57.                                     
    
58.         QStringDB.btn_redo_pt_tip = {   'en' : 'Redo point',
    
59.                                         'ru' : 'Повторить точку',
    
60.                                         'zh' : '重做点',
    
61.                                      }[lang]      
    
62.                                       
    
63.         QStringDB.btn_delete_poly_tip = {   'en' : 'Delete poly',
    
64.                                             'ru' : 'Удалить полигон',
    
65.                                             'zh' : '删除选区',
    
66.                                            }[lang]     
    
67.                                              
    
68.         QStringDB.btn_pt_edit_mode_tip = {  'en' : 'Add/delete point mode ( HOLD CTRL )',
    
69.                                             'ru' : 'Режим добавления/удаления точек ( удерживайте CTRL )',
    
70.                                             'zh' : '点加/删除模式 ( 按住CTRL )',
    
71.                                            }[lang]   
    
72.                                           
    
73.         QStringDB.btn_view_lock_center_tip = { 'en' : 'Lock cursor at the center ( HOLD SHIFT )',
    
74.                                                'ru' : 'Заблокировать курсор в центре ( удерживайте SHIFT )',
    
75.                                                'zh' : '将光标锁定在中心 ( 按住SHIFT )',
    
76.                                              }[lang]                    
    
77.                                        
    
78.                                           
    
79.         QStringDB.btn_prev_image_tip = {    'en' : 'Save and Prev image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n',
    
80.                                             'ru' : 'Сохранить и предыдущее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n',
    
81.                                             'zh' : '保存并转到上一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n',
    
82.                                            }[lang]   
    
83.         QStringDB.btn_next_image_tip = {    'en' : 'Save and Next image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n',
    
84.                                             'ru' : 'Сохранить и следующее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n',
    
85.                                             'zh' : '保存并转到下一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n',
    
86.                                            }[lang]  
    
87.         
    
88.         QStringDB.btn_delete_image_tip = {  'en' : 'Move to _trash and Next image\n',
    
89.                                             'ru' : 'Переместить в _trash и следующее изображение\n',
    
90.                                             'zh' : '移至_trash，转到下一张图片 ',
    
91.                                            }[lang]  
    
92.                                           
    
93.         QStringDB.loading_tip = {'en' : 'Loading',
    
94.                                  'ru' : 'Загрузка',
    
95.                                  'zh' : '正在载入',
    
96.                                 }[lang]   
    
97.                                 
    
98.         QStringDB.labeled_tip = {'en' : 'labeled',
    
99.                                  'ru' : 'размечено',
    
100.                                  'zh' : '标记的',
     
101.                                 }[lang]
     
102. 
     
103.         QStringDB.btn_auto_detect_tip = {'en' : 'Auto detect occlusion',
     
104.                                          'ru' : 'Автообнаружение окклюзии',
     
105.                                          'zh' : '自动检测遮挡物',
     
106.                                         }[lang]   
     
107.                                           
     

复制代码

如果不会修改就用我网盘里面存的文件覆盖你原始的文件就行了，

通过网盘分享的文件：
链接: https://pan.baidu.com/s/12lH1u11mEnDDYnV75_ZK5g?pwd=t8i8 提取码: t8i8 复制这段内容后打开百度网盘手机App，操作更方便哦

你看看你的python里面有没有这个onnxruntime库，正确的安装方法是打开cmd ，然后用python安装，比方说我的python路径是D:\AI\DeepFaceLab\_internal\python\,所以在cmd中运行“ D:/AI/DeepFaceLab/_internal/python/python.exe -m pip install onnxruntime -i http://mirrors.aliyun.com/pypi/simple --trusted-host mirrors.aliyun.com”，记住了是你dfl的python环境，不是电脑的python环境，不知道 的看看你setenv.bat里面的python目录


在这里我说明一点：假如要实现更高效的更高级的，更准确的遮罩，就必须用pytorch了，因为现有的所有的高级人脸解析都是建立在pytorch基础上的，这也是为什么很多人希望把pytorch替换dfl的tf框架。

pasanonic

day270010678 发表于 2026-7-15 22:56
是按照到你dfl的python环境里面，setenv.bat里面的python环境。不是你自己电脑的python ...

大佬牛B  好了,自动画了  就是线条不太圆润,头部意外的有些多余,不过我要求不高,没那么追求精细

pasanonic

day270010678 发表于 2026-7-15 22:35
你装的是那个版本？是gpu版本，还是就是我上面说的？尽量不要用gpu版本，我当时就是怕出现问题，所以没有 ...

pip install onnxruntime

pasanonic

day270010678 发表于 2026-7-15 22:27
你安装onnxruntime库了吗？

安装了  本来就有 后来还升级了 1.26  不过我是amd卡, 装的cpu推理  

mzh920

day270010678 发表于 2026-7-15 21:09
我把嘴唇区域排除用虚线标出来了，主要是为了防止嘴巴含东西，你们不需要的话，直接删除就行了，还有一个小 ...

图3图4有部分头发没圈出去，影响会很大，所以一般遇到这样的我都是手动画，直接从眉毛上打横线不走头发轮廓，楼主可以改成这样的吗？如果能改就好了
图12无遮挡部分那种其实万能遮罩应用即可

pasanonic

所有步骤都做了 按了左下角图标没反应

day270010678

我把嘴唇区域排除用虚线标出来了，主要是为了防止嘴巴含东西，你们不需要的话，直接删除就行了，还有一个小图标，就是这个放到XSegEditor/gfx/icons/auto_detect.png目录里面就行了

day270010678

pasanonic 发表于 2026-7-15 22:17
所有步骤都做了 按了左下角图标没反应

你安装onnxruntime库了吗？

day270010678

mzh920 发表于 2026-7-15 22:21
图3图4有部分头发没圈出去，影响会很大，所以一般遇到这样的我都是手动画，直接从眉毛上打横线不走头发轮 ...

有些东西需要精修的，你说的那些手动修改下就行了，理论上是像素级别的，所以只需要手动修改部分就行了。总比手一个个画要省事的多。等我空了最多我把嘴巴模型升级下。其他的不升级了，就靠官方的模型就行了

day270010678

pasanonic 发表于 2026-7-15 22:33
安装了  本来就有 后来还升级了

你装的是那个版本？是gpu版本，还是就是我上面说的？尽量不要用gpu版本，我当时就是怕出现问题，所以没有用gpu版本。

day270010678

pasanonic 发表于 2026-7-15 22:36
pip install onnxruntime

通过网盘分享的文件：
链接: https://pan.baidu.com/s/12lH1u11mEnDDYnV75_ZK5g?pwd=t8i8 提取码: t8i8 复制这段内容后打开百度网盘手机App，操作更方便哦

你下载我修改的测试下，我觉得你是按照库导致的问题， 你看看你的python里面有没有这个onnxruntime库，正确的安装方法是打开cmd ，然后用python安装，比方说我的python路径是D:\AI\DeepFaceLab\_internal\python\,所以在cmd中运行“ D:/AI/DeepFaceLab/_internal/python/python.exe -m pip install onnxruntime -i http://mirrors.aliyun.com/pypi/simple --trusted-host mirrors.aliyun.com”