fix MaskEditor, upd manual ru

This commit is contained in:
iperov 2019-04-08 14:13:36 +04:00
parent 7c930121e9
commit 8561c1a931
4 changed files with 137 additions and 134 deletions

Binary file not shown.

Binary file not shown.

View file

@ -259,6 +259,8 @@ class InteractDesktop(InteractBase):
cv2.imshow (wnd_name, img)
def on_capture_mouse (self, wnd_name):
self.last_xy = (0,0)
def onMouse(event, x, y, flags, param):
(inst, wnd_name) = param
if event == cv2.EVENT_LBUTTONDOWN: ev = InteractBase.EVENT_LBUTTONDOWN
@ -267,9 +269,12 @@ class InteractDesktop(InteractBase):
elif event == cv2.EVENT_RBUTTONUP: ev = InteractBase.EVENT_RBUTTONUP
elif event == cv2.EVENT_MBUTTONDOWN: ev = InteractBase.EVENT_MBUTTONDOWN
elif event == cv2.EVENT_MBUTTONUP: ev = InteractBase.EVENT_MBUTTONUP
elif event == cv2.EVENT_MOUSEWHEEL: ev = InteractBase.EVENT_MOUSEWHEEL
elif event == cv2.EVENT_MOUSEWHEEL:
ev = InteractBase.EVENT_MOUSEWHEEL
x,y = self.last_xy #fix opencv bug when window size more than screen size
else: ev = 0
self.last_xy = (x,y)
inst.add_mouse_event (wnd_name, x, y, ev, flags)
cv2.setMouseCallback(wnd_name, onMouse, (self,wnd_name) )

View file

@ -20,28 +20,28 @@ from utils.DFLPNG import DFLPNG
class MaskEditor:
STATE_NONE=0
STATE_MASKING=1
def __init__(self, img, prev_images, next_images, mask=None, ie_polys=None, get_status_lines_func=None):
self.img = imagelib.normalize_channels (img,3)
self.img = imagelib.normalize_channels (img,3)
h, w, c = img.shape
if h != w and w != 256:
#to support any square res, scale img,mask and ie_polys to 256, then scale ie_polys back on .get_ie_polys()
raise Exception ("MaskEditor does not support image size != 256x256")
ph, pw = h // 4, w // 4 #pad wh
self.prev_images = prev_images
self.next_images = next_images
if mask is not None:
self.mask = imagelib.normalize_channels (mask,3)
else:
self.mask = np.zeros ( (h,w,3) )
self.get_status_lines_func = get_status_lines_func
self.state_prop = self.STATE_NONE
self.w, self.h = w, h
self.pw, self.ph = pw, ph
self.pwh = np.array([self.pw, self.ph])
@ -52,18 +52,18 @@ class MaskEditor:
if ie_polys is None:
ie_polys = IEPolys()
self.ie_polys = ie_polys
self.polys_mask = None
self.preview_images = None
self.mouse_x = self.mouse_y = 9999
self.screen_status_block = None
self.screen_status_block_dirty = True
self.screen_changed = True
def set_state(self, state):
self.state = state
@property
def state(self):
return self.state_prop
@ -73,62 +73,61 @@ class MaskEditor:
self.state_prop = value
if value == self.STATE_MASKING:
self.ie_polys.dirty = True
def get_mask(self):
def get_mask(self):
if self.ie_polys.switch_dirty():
self.screen_status_block_dirty = True
self.ie_mask = img = self.mask.copy()
self.ie_polys.overlay_mask(img)
self.ie_polys.overlay_mask(img)
return img
return self.ie_mask
def get_screen_overlay(self):
img = np.zeros ( (self.sh, self.sw, 3) )
if self.state == self.STATE_MASKING:
mouse_xy = self.mouse_xy.copy() + self.pwh
l = self.ie_polys.n_list()
if l.n > 0:
if l.n > 0:
p = l.cur_point().copy() + self.pwh
color = (0,1,0) if l.type == 1 else (0,0,1)
cv2.line(img, tuple(p), tuple(mouse_xy), color )
return img
def undo_to_begin_point(self):
while not self.undo_point():
pass
def undo_point(self):
self.screen_changed = True
if self.state == self.STATE_NONE:
if self.ie_polys.n > 0:
if self.ie_polys.n > 0:
self.state = self.STATE_MASKING
if self.state == self.STATE_MASKING:
if self.ie_polys.n_list().n_dec() == 0 and \
self.ie_polys.n_dec() == 0:
self.state = self.STATE_NONE
else:
return False
return True
def redo_to_end_point(self):
while not self.redo_point():
pass
def redo_point(self):
self.screen_changed = True
if self.state == self.STATE_NONE:
if self.ie_polys.n_max > 0:
if self.ie_polys.n_max > 0:
self.state = self.STATE_MASKING
if self.ie_polys.n == 0:
self.ie_polys.n_inc()
if self.state == self.STATE_MASKING:
while True:
l = self.ie_polys.n_list()
@ -138,30 +137,30 @@ class MaskEditor:
self.ie_polys.n_inc()
else:
return False
return True
def combine_screens(self, screens):
screens_len = len(screens)
new_screens = []
for screen, padded_overlay in screens:
screen_img = np.zeros( (self.sh, self.sw, 3), dtype=np.float32 )
screen = imagelib.normalize_channels (screen, 3)
h,w,c = screen.shape
screen_img[self.ph:-self.ph, self.pw:-self.pw, :] = screen
if padded_overlay is not None:
screen_img = screen_img + padded_overlay
screen_img = np.clip(screen_img*255, 0, 255).astype(np.uint8)
new_screens.append(screen_img)
return np.concatenate (new_screens, axis=1)
def get_screen_status_block(self, w, c):
if self.screen_status_block_dirty:
self.screen_status_block_dirty = False
@ -170,86 +169,86 @@ class MaskEditor:
]
if self.get_status_lines_func is not None:
lines += self.get_status_lines_func()
lines_count = len(lines)
h_line = 21
h = lines_count * h_line
img = np.ones ( (h,w,c) ) * 0.1
for i in range(lines_count):
img[ i*h_line:(i+1)*h_line, 0:w] += \
imagelib.get_text_image ( (h_line,w,c), lines[i], color=[0.8]*c )
self.screen_status_block = np.clip(img*255, 0, 255).astype(np.uint8)
return self.screen_status_block
def set_screen_status_block_dirty(self):
self.screen_status_block_dirty = True
def set_screen_changed(self):
self.screen_changed = True
def switch_screen_changed(self):
result = self.screen_changed
self.screen_changed = False
return result
def make_screen(self):
screen_overlay = self.get_screen_overlay()
final_mask = self.get_mask()
masked_img = self.img*final_mask*0.5 + self.img*(1-final_mask)
pink = np.full ( (self.h, self.w, 3), (1,0,1) )
pink = np.full ( (self.h, self.w, 3), (1,0,1) )
pink_masked_img = self.img*final_mask + pink*(1-final_mask)
screens = [ (self.img, screen_overlay),
(masked_img, screen_overlay),
(pink_masked_img, screen_overlay),
]
(pink_masked_img, screen_overlay),
]
screens = self.combine_screens(screens)
if self.preview_images is None:
sh,sw,sc = screens.shape
prh, prw = self.prwh, self.prwh
total_w = sum ([ img.shape[1] for (t,img) in self.prev_images ]) + \
sum ([ img.shape[1] for (t,img) in self.next_images ])
total_images_len = len(self.prev_images) + len(self.next_images)
max_hor_images_count = sw // prw
max_side_images_count = (max_hor_images_count - 1) // 2
prev_images = self.prev_images[-max_side_images_count:]
next_images = self.next_images[:max_side_images_count]
border = 2
max_wh_bordered = (prw-border*2, prh-border*2)
prev_images = [ (t, cv2.resize( imagelib.normalize_channels(img, 3), max_wh_bordered )) for t,img in prev_images ]
next_images = [ (t, cv2.resize( imagelib.normalize_channels(img, 3), max_wh_bordered )) for t,img in next_images ]
for images in [prev_images, next_images]:
for i, (t, img) in enumerate(images):
new_img = np.zeros ( (prh,prw, sc) )
new_img[border:-border,border:-border] = img
if t == 2:
cv2.line (new_img, ( prw//2, int(prh//1.5) ), (int(prw/1.5), prh ) , (0,1,0), thickness=2 )
cv2.line (new_img, ( int(prw/1.5), prh ), ( prw, prh // 2 ) , (0,1,0), thickness=2 )
elif t == 1:
cv2.line (new_img, ( prw//2, prh//2 ), ( prw, prh ) , (0,0,1), thickness=2 )
cv2.line (new_img, ( prw//2, prh ), ( prw, prh // 2 ) , (0,0,1), thickness=2 )
images[i] = new_img
@ -261,27 +260,27 @@ class MaskEditor:
img[border:-border,border:-border] = cv2.resize( self.img, max_wh_bordered )
preview_images += [ img ]
if len(next_images) > 0:
preview_images += [ np.concatenate (next_images, axis=1) ]
preview_images = np.concatenate ( preview_images, axis=1 )
left_pad = sw // 2 - len(prev_images) * prw - prw // 2
left_pad = sw // 2 - len(prev_images) * prw - prw // 2
right_pad = sw // 2 - len(next_images) * prw - prw // 2
preview_images = np.concatenate ([np.zeros ( (preview_images.shape[0], left_pad, preview_images.shape[2]) ),
preview_images,
np.zeros ( (preview_images.shape[0], right_pad, preview_images.shape[2]) )
np.zeros ( (preview_images.shape[0], right_pad, preview_images.shape[2]) )
], axis=1)
self.preview_images = np.clip(preview_images * 255, 0, 255 ).astype(np.uint8)
status_img = self.get_screen_status_block( screens.shape[1], screens.shape[2] )
result = np.concatenate ( [self.preview_images, screens, status_img], axis=0 )
return result
def mask_finish(self, n_clip=True):
if self.state == self.STATE_MASKING:
self.screen_changed = True
@ -290,16 +289,16 @@ class MaskEditor:
self.state = self.STATE_NONE
if n_clip:
self.ie_polys.n_clip()
def set_mouse_pos(self,x,y):
if self.preview_images is not None:
y -= self.preview_images.shape[0]
mouse_x = x % (self.sw) - self.pw
mouse_y = y % (self.sh) - self.ph
if mouse_x != self.mouse_x or mouse_y != self.mouse_y:
self.mouse_xy = np.array( [mouse_x, mouse_y] )
self.mouse_x, self.mouse_y = self.mouse_xy
@ -310,20 +309,20 @@ class MaskEditor:
if self.state == self.STATE_MASKING and \
self.ie_polys.n_list().type != type:
self.mask_finish()
elif self.state == self.STATE_NONE:
self.state = self.STATE_MASKING
self.state = self.STATE_MASKING
self.ie_polys.add(type)
if self.state == self.STATE_MASKING:
self.ie_polys.n_list().add (self.mouse_x, self.mouse_y)
self.ie_polys.n_list().add (self.mouse_x, self.mouse_y)
def get_ie_polys(self):
return self.ie_polys
def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
input_path = Path(input_dir)
confirmed_path = Path(confirmed_dir)
skipped_path = Path(skipped_dir)
@ -342,28 +341,28 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
io.capture_keys(wnd_name)
cached_images = {}
image_paths = [ Path(x) for x in Path_utils.get_image_paths(input_path)]
done_paths = []
done_paths = []
done_images_types = {}
image_paths_total = len(image_paths)
zoom_factor = 1.0
preview_images_count = 9
target_wh = 256
target_wh = 256
do_prev_count = 0
do_save_move_count = 0
do_save_count = 0
do_skip_move_count = 0
do_skip_count = 0
def jobs_count():
return do_prev_count + do_save_move_count + do_save_count + do_skip_move_count + do_skip_count
is_exit = False
while not is_exit:
if len(image_paths) > 0:
filepath = image_paths.pop(0)
else:
@ -373,12 +372,12 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
next_image_paths_names = [ path.name for path in next_image_paths ]
prev_image_paths = done_paths[-preview_images_count:]
prev_image_paths_names = [ path.name for path in prev_image_paths ]
for key in list( cached_images.keys() ):
if key not in prev_image_paths_names and \
key not in next_image_paths_names:
cached_images.pop(key)
for paths in [prev_image_paths, next_image_paths]:
for path in paths:
if path.name not in cached_images:
@ -391,19 +390,19 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
dflimg = DFLJPG.load ( str(filepath) )
else:
dflimg = None
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
else:
lmrks = dflimg.get_landmarks()
ie_polys = dflimg.get_ie_polys()
if filepath.name in cached_images:
img = cached_images[filepath.name]
else:
img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks)
else:
img = np.zeros ( (target_wh,target_wh,3) )
@ -411,7 +410,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
ie_polys = None
def get_status_lines_func():
return ['Progress: %d / %d . Current file: %s' % (len(done_paths), image_paths_total, str(filepath.name) if filepath is not None else "end" ),
return ['Progress: %d / %d . Current file: %s' % (len(done_paths), image_paths_total, str(filepath.name) if filepath is not None else "end" ),
'[Left mouse button] - mark include mask.',
'[Right mouse button] - mark exclude mask.',
'[Middle mouse button] - finish current poly.',
@ -421,25 +420,24 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
'hold [shift] - speed up the frame counter by 10.',
'[-/+] - window zoom [esc] - quit',
]
try:
ed = MaskEditor(img,
ed = MaskEditor(img,
[ (done_images_types[name], cached_images[name]) for name in prev_image_paths_names ],
[ (0, cached_images[name]) for name in next_image_paths_names ],
mask, ie_polys, get_status_lines_func)
except Exception as e:
print(e)
continue
next = False
while not next:
io.process_messages(0.005)
if jobs_count() == 0:
for (x,y,ev,flags) in io.get_mouse_events(wnd_name):
x, y = int (x / zoom_factor), int(y / zoom_factor)
ed.set_mouse_pos(x, y)
ed.set_mouse_pos(x, y)
if filepath is not None:
if ev == io.EVENT_LBUTTONDOWN:
ed.mask_point(1)
@ -460,7 +458,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
ed.redo_point()
for key, chr_key, ctrl_pressed, alt_pressed, shift_pressed in io.get_key_events(wnd_name):
if chr_key == 'q' or chr_key == 'z':
do_prev_count = 1 if not shift_pressed else 10
elif chr_key == '-':
@ -482,64 +480,64 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
do_skip_move_count = 1 if not shift_pressed else 10
elif chr_key == 'x':
do_skip_count = 1 if not shift_pressed else 10
if do_prev_count > 0:
do_prev_count -= 1
if len(done_paths) > 0:
if filepath is not None:
image_paths.insert(0, filepath)
filepath = done_paths.pop(-1)
done_images_types[filepath.name] = 0
if filepath.parent != input_path:
new_filename_path = input_path / filepath.name
filepath.rename ( new_filename_path )
image_paths.insert(0, new_filename_path)
filepath.rename ( new_filename_path )
image_paths.insert(0, new_filename_path)
else:
image_paths.insert(0, filepath)
image_paths.insert(0, filepath)
next = True
elif filepath is not None:
if do_save_move_count > 0:
do_save_move_count -= 1
ed.mask_finish()
dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() )
done_paths += [ confirmed_path / filepath.name ]
done_paths += [ confirmed_path / filepath.name ]
done_images_types[filepath.name] = 2
filepath.rename(done_paths[-1])
next = True
next = True
elif do_save_count > 0:
do_save_count -= 1
ed.mask_finish()
dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() )
done_paths += [ filepath ]
done_images_types[filepath.name] = 2
next = True
next = True
elif do_skip_move_count > 0:
do_skip_move_count -= 1
done_paths += [ skipped_path / filepath.name ]
done_paths += [ skipped_path / filepath.name ]
done_images_types[filepath.name] = 1
filepath.rename(done_paths[-1])
next = True
next = True
elif do_skip_count > 0:
do_skip_count -= 1
done_paths += [ filepath ]
done_images_types[filepath.name] = 1
next = True
else:
do_save_move_count = do_save_count = do_skip_move_count = do_skip_count = 0
if jobs_count() == 0:
if ed.switch_screen_changed():
screen = ed.make_screen()
@ -548,8 +546,8 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
screen = cv2.resize ( screen, ( int(w*zoom_factor), int(h*zoom_factor) ) )
io.show_image (wnd_name, screen )
io.process_messages(0.005)
io.destroy_all_windows()
io.process_messages(0.005)
io.destroy_all_windows()