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

refactoring

Browse source
This commit is contained in:
iperov 2019-03-16 21:29:56 +04:00
parent f009edae79
commit 4ff88865ce
1 changed files with 32 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

57
mainscripts/Extractor.py
View file

@ -113,8 +113,12 @@ class ExtractSubprocessor(Subprocessor):
return [str(filename_path), rects] return [str(filename_path), rects]
elif self.type == 'landmarks': elif self.type == 'landmarks':
rects = data[1] rects = data[1]
landmarks = self.e.extract_from_bgr (image, rects) if rects is None:
landmarks = None
else:
landmarks = self.e.extract_from_bgr (image, rects)
return [str(filename_path), landmarks] return [str(filename_path), landmarks]
elif self.type == 'final': elif self.type == 'final':
@ -291,20 +295,19 @@ class ExtractSubprocessor(Subprocessor):
return self.input_data.pop(0) return self.input_data.pop(0)
else: else:
allow_remark_faces = False need_remark_face = False
redraw_needed = False
while len (self.input_data) > 0: while len (self.input_data) > 0:
data = self.input_data[0] data = self.input_data[0]
filename, faces = data filename, faces = data
is_frame_done = False is_frame_done = False
# Can we mark an image that already has a marked face? if need_remark_face: # need remark image from input data that already has a marked face?
if allow_remark_faces: need_remark_face = False
allow_remark_faces = False if len(faces) != 0: # If there was already a face then lock the rectangle to it until the mouse is clicked
# If there was already a face then lock the rectangle to it until the mouse is clicked
if len(faces) > 0:
self.rect, self.landmarks = faces.pop() self.rect, self.landmarks = faces.pop()
faces.clear() faces.clear()
self.extract_needed = True redraw_needed = True
self.rect_locked = True self.rect_locked = True
self.rect_size = ( self.rect[2] - self.rect[0] ) / 2 self.rect_size = ( self.rect[2] - self.rect[0] ) / 2
self.x = ( self.rect[0] + self.rect[2] ) / 2 self.x = ( self.rect[0] + self.rect[2] ) / 2
@ -378,7 +381,7 @@ class ExtractSubprocessor(Subprocessor):
is_frame_done = True is_frame_done = True
break break
elif key == ord(',') and len(self.result) > 0: elif key == ord(',') and len(self.result) > 0:
#go prev frame #go prev frame
if self.rect_locked: if self.rect_locked:
# Only save the face if the rect is still locked # Only save the face if the rect is still locked
@ -386,21 +389,21 @@ class ExtractSubprocessor(Subprocessor):
self.input_data.insert(0, self.result.pop() ) self.input_data.insert(0, self.result.pop() )
io.progress_bar_inc(-1) io.progress_bar_inc(-1)
allow_remark_faces = True need_remark_face = True
self.extract_needed = True
self.rect_locked = False
break break
elif key == ord('.'): elif key == ord('.'):
#go next frame #go next frame
is_frame_done = True
allow_remark_faces = True
if self.rect_locked: if self.rect_locked:
# Only save the face if the rect is still locked # Only save the face if the rect is still locked
faces.append ( [(self.rect), self.landmarks] ) faces.append ( [(self.rect), self.landmarks] )
need_remark_face = True
is_frame_done = True
break break
elif key == ord('q'): elif key == ord('q'):
#skip remaining #skip remaining
if self.rect_locked: if self.rect_locked:
faces.append ( [(self.rect), self.landmarks] ) faces.append ( [(self.rect), self.landmarks] )
while len(self.input_data) > 0: while len(self.input_data) > 0:
@ -416,17 +419,21 @@ class ExtractSubprocessor(Subprocessor):
if self.x != new_x or \ if self.x != new_x or \
self.y != new_y or \ self.y != new_y or \
self.rect_size != new_rect_size or \ self.rect_size != new_rect_size or \
self.extract_needed: self.extract_needed or \
redraw_needed:
self.x = new_x self.x = new_x
self.y = new_y self.y = new_y
self.rect_size = new_rect_size self.rect_size = new_rect_size
self.rect = ( int(self.x-self.rect_size), self.rect = ( int(self.x-self.rect_size),
int(self.y-self.rect_size), int(self.y-self.rect_size),
int(self.x+self.rect_size), int(self.x+self.rect_size),
int(self.y+self.rect_size) ) int(self.y+self.rect_size) )
return [filename, [self.rect]] if redraw_needed:
redraw_needed = False
return [filename, None]
else:
return [filename, [self.rect]]
else: else:
is_frame_done = True is_frame_done = True
@ -448,7 +455,9 @@ class ExtractSubprocessor(Subprocessor):
#override #override
def on_result (self, host_dict, data, result): def on_result (self, host_dict, data, result):
if self.manual == True: if self.manual == True:
self.landmarks = result[1][0][1] filename, landmarks = result
if landmarks is not None:
self.landmarks = landmarks[0][1]
(h,w,c) = self.image.shape (h,w,c) = self.image.shape
@ -476,10 +485,8 @@ class ExtractSubprocessor(Subprocessor):
image = cv2.warpAffine(image, mat,(w,h) ) image = cv2.warpAffine(image, mat,(w,h) )
view_landmarks = LandmarksProcessor.transform_points (view_landmarks, mat) view_landmarks = LandmarksProcessor.transform_points (view_landmarks, mat)
LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.image_size, self.face_type) landmarks_color = (255,255,0) if self.rect_locked else (0,255,0)
LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.image_size, self.face_type, landmarks_color=landmarks_color)
if self.rect_locked:
LandmarksProcessor.draw_landmarks(image, view_landmarks, (255,255,0) )
self.extract_needed = False self.extract_needed = False
io.show_image (self.wnd_name, image) io.show_image (self.wnd_name, image)