Ejemplo n.º 1
0
    def redraw(self):
        (h,w,c) = self.image.shape

        if not self.hide_help:
            image = cv2.addWeighted (self.image,1.0,self.text_lines_img,1.0,0)
        else:
            image = self.image.copy()

        view_rect = (np.array(self.rect) * self.view_scale).astype(np.int).tolist()
        view_landmarks  = (np.array(self.landmarks) * self.view_scale).astype(np.int).tolist()

        if self.rect_size <= 40:
            scaled_rect_size = h // 3 if w > h else w // 3

            p1 = (self.x - self.rect_size, self.y - self.rect_size)
            p2 = (self.x + self.rect_size, self.y - self.rect_size)
            p3 = (self.x - self.rect_size, self.y + self.rect_size)

            wh = h if h < w else w
            np1 = (w / 2 - wh / 4, h / 2 - wh / 4)
            np2 = (w / 2 + wh / 4, h / 2 - wh / 4)
            np3 = (w / 2 - wh / 4, h / 2 + wh / 4)

            mat = cv2.getAffineTransform( np.float32([p1,p2,p3])*self.view_scale, np.float32([np1,np2,np3]) )
            image = cv2.warpAffine(image, mat,(w,h) )
            view_landmarks = LandmarksProcessor.transform_points (view_landmarks, mat)

        landmarks_color = (255,255,0) if self.rect_locked else (0,255,0)
        LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.face_type, self.image_size, landmarks_color=landmarks_color)
        self.extract_needed = False

        io.show_image (self.wnd_name, image)
Ejemplo n.º 2
0
    def on_result(self, host_dict, data, result):
        if self.manual == True:
            filename, landmarks = result
            if landmarks is not None:
                self.landmarks = landmarks[0][1]

            (h, w, c) = self.image.shape

            if not self.hide_help:
                image = cv2.addWeighted(self.image, 1.0, self.text_lines_img,
                                        1.0, 0)
            else:
                image = self.image.copy()

            view_rect = (np.array(self.rect) * self.view_scale).astype(
                np.int).tolist()
            view_landmarks = (np.array(self.landmarks) *
                              self.view_scale).astype(np.int).tolist()

            if self.rect_size <= 40:
                scaled_rect_size = h // 3 if w > h else w // 3

                p1 = (self.x - self.rect_size, self.y - self.rect_size)
                p2 = (self.x + self.rect_size, self.y - self.rect_size)
                p3 = (self.x - self.rect_size, self.y + self.rect_size)

                wh = h if h < w else w
                np1 = (w / 2 - wh / 4, h / 2 - wh / 4)
                np2 = (w / 2 + wh / 4, h / 2 - wh / 4)
                np3 = (w / 2 - wh / 4, h / 2 + wh / 4)

                mat = cv2.getAffineTransform(
                    np.float32([p1, p2, p3]) * self.view_scale,
                    np.float32([np1, np2, np3]))
                image = cv2.warpAffine(image, mat, (w, h))
                view_landmarks = LandmarksProcessor.transform_points(
                    view_landmarks, mat)

            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)
            self.extract_needed = False

            io.show_image(self.wnd_name, image)
        else:
            if self.type == 'rects':
                self.result.append(result)
            elif self.type == 'landmarks':
                self.result.append(result)
            elif self.type == 'final':
                self.result += result

            io.progress_bar_inc(1)
Ejemplo n.º 3
0
    def onHostResult(self, host_dict, data, result):
        if self.manual == True:
            self.landmarks = result[1][0][1]

            (h, w, c) = self.original_image.shape
            image = cv2.addWeighted(self.original_image, 1.0,
                                    self.text_lines_img, 1.0, 0)
            view_rect = (np.array(self.rect) * self.view_scale).astype(
                np.int).tolist()
            view_landmarks = (np.array(self.landmarks) *
                              self.view_scale).astype(np.int).tolist()

            if self.param_rect_size <= 25:
                scaled_rect_size = h // 3 if w > h else w // 3

                p1 = (self.param_x - self.param_rect_size,
                      self.param_y - self.param_rect_size)
                p2 = (self.param_x + self.param_rect_size,
                      self.param_y - self.param_rect_size)
                p3 = (self.param_x - self.param_rect_size,
                      self.param_y + self.param_rect_size)

                np1 = (self.param_x - scaled_rect_size,
                       self.param_y - scaled_rect_size)
                np2 = (self.param_x + scaled_rect_size,
                       self.param_y - scaled_rect_size)
                np3 = (self.param_x - scaled_rect_size,
                       self.param_y + scaled_rect_size)

                mat = cv2.getAffineTransform(
                    np.float32([p1, p2, p3]) * self.view_scale,
                    np.float32([np1, np2, np3]) * self.view_scale)
                image = cv2.warpAffine(image, mat, (w, h))
                view_landmarks = LandmarksProcessor.transform_points(
                    view_landmarks, mat)

            LandmarksProcessor.draw_rect_landmarks(image, view_rect,
                                                   view_landmarks,
                                                   self.image_size,
                                                   self.face_type)

            if self.param['rect_locked']:
                LandmarksProcessor.draw_landmarks(image, view_landmarks,
                                                  (255, 255, 0))
            self.param['redraw_needed'] = False

            cv2.imshow(self.wnd_name, image)
            return 0
        else:
            if self.type == 'rects':
                self.result.append(result)
            elif self.type == 'landmarks':
                self.result.append(result)
            elif self.type == 'final':
                self.result += result

            return 1
Ejemplo n.º 4
0
    def onClientProcessData(self, data):
        filename_path = Path( data[0] )

        image = cv2_imread( str(filename_path) )
        if image is None:
            print ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) )
        else:
            if self.type == 'rects':
                rects = self.e.extract_from_bgr (image)
                return [str(filename_path), rects]

            elif self.type == 'landmarks':
                rects = data[1]   
                landmarks = self.e.extract_from_bgr (image, rects)                    
                return [str(filename_path), landmarks]

            elif self.type == 'final':     
                result = []
                faces = data[1]
                
                if self.debug:
                    debug_output_file = '{}{}'.format( str(Path(str(self.output_path) + '_debug') / filename_path.stem),  '.jpg')
                    debug_image = image.copy()
                    
                for (face_idx, face) in enumerate(faces):         
                    output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.jpg')
                    
                    rect = face[0]
                    image_landmarks = np.array(face[1])

                    if self.debug:
                        LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type)

                    if self.face_type == FaceType.MARK_ONLY:                        
                        face_image = image
                        face_image_landmarks = image_landmarks
                    else:
                        image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type)       
                        face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4)
                        face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat)
                    
                    cv2_imwrite(output_file, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), 85] )

                    DFLJPG.embed_data(output_file, face_type = FaceType.toString(self.face_type),
                                                   landmarks = face_image_landmarks.tolist(),
                                                   source_filename = filename_path.name,
                                                   source_rect=  rect,
                                                   source_landmarks = image_landmarks.tolist()
                                        )  
                        
                    result.append (output_file)
                    
                if self.debug:
                    cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
                    
                return result       
        return None
    def extract_from_bgr(self, input_image, rects, second_pass_extractor=None):
        if len(rects) == 0:
            return []

        input_image = input_image[:, :, ::-1].copy()
        (h, w, ch) = input_image.shape

        landmarks = []
        for (left, top, right, bottom) in rects:
            try:
                center = np.array([(left + right) / 2.0, (top + bottom) / 2.0])
                #center[1] -= (bottom - top) * 0.12
                scale = (right - left + bottom - top) / 195.0

                image = self.crop(input_image, center,
                                  scale).astype(np.float32)
                image = np.expand_dims(image, 0)

                predicted = self.keras_model.predict(image).transpose(
                    0, 3, 1, 2)

                pts_img = self.get_pts_from_predict(predicted[-1], center,
                                                    scale)
                pts_img = [(int(pt[0]), int(pt[1])) for pt in pts_img]
                landmarks.append(pts_img)
            except Exception as e:
                landmarks.append(None)

        if second_pass_extractor is not None:
            for i in range(len(landmarks)):
                try:
                    lmrks = landmarks[i]
                    if lmrks is None:
                        continue

                    image_to_face_mat = LandmarksProcessor.get_transform_mat(
                        lmrks, 256, FaceType.FULL)
                    face_image = cv2.warpAffine(input_image, image_to_face_mat,
                                                (256, 256), cv2.INTER_CUBIC)

                    rects2 = second_pass_extractor.extract_from_bgr(face_image)
                    if len(
                            rects2
                    ) != 1:  #dont do second pass if 1 != faces detected in cropped image
                        continue

                    lmrks2 = self.extract_from_bgr(face_image,
                                                   [rects2[0]])[0][1]
                    source_lmrks2 = LandmarksProcessor.transform_points(
                        lmrks2, image_to_face_mat, True)
                    landmarks[i] = source_lmrks2
                except:
                    continue

        return landmarks
Ejemplo n.º 6
0
def main(model_class_name=None,
         saved_models_path=None,
         training_data_src_path=None,
         force_model_name=None,
         input_path=None,
         output_path=None,
         aligned_path=None,
         force_gpu_idxs=None,
         cpu_only=None):
    io.log_info("Running merger.\r\n")

    try:
        if not input_path.exists():
            io.log_err('Input directory not found. Please ensure it exists.')
            return

        if not output_path.exists():
            output_path.mkdir(parents=True, exist_ok=True)

        if not saved_models_path.exists():
            io.log_err('Model directory not found. Please ensure it exists.')
            return

        is_interactive = io.input_bool("Use interactive merger?",
                                       True) if not io.is_colab() else False

        import models
        model = models.import_model(model_class_name)(
            is_training=False,
            saved_models_path=saved_models_path,
            training_data_src_path=training_data_src_path,
            force_gpu_idxs=force_gpu_idxs,
            cpu_only=cpu_only)
        merger_session_filepath = model.get_strpath_storage_for_file(
            'merger_session.dat')
        predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()

        if not is_interactive:
            cfg.ask_settings()

        input_path_image_paths = pathex.get_image_paths(input_path)

        if cfg.type == MergerConfig.TYPE_MASKED:
            if not aligned_path.exists():
                io.log_err(
                    'Aligned directory not found. Please ensure it exists.')
                return

            packed_samples = None
            try:
                packed_samples = samplelib.PackedFaceset.load(aligned_path)
            except:
                io.log_err(
                    f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
                )

            if packed_samples is not None:
                io.log_info("Using packed faceset.")

                def generator():
                    for sample in io.progress_bar_generator(
                            packed_samples, "Collecting alignments"):
                        filepath = Path(sample.filename)
                        yield DFLIMG.load(
                            filepath,
                            loader_func=lambda x: sample.read_raw_file())
            else:

                def generator():
                    for filepath in io.progress_bar_generator(
                            pathex.get_image_paths(aligned_path),
                            "Collecting alignments"):
                        filepath = Path(filepath)
                        yield DFLIMG.load(filepath)

            alignments = {}
            multiple_faces_detected = False

            for dflimg in generator():
                if dflimg is None:
                    io.log_err("%s is not a dfl image file" % (filepath.name))
                    continue

                source_filename = dflimg.get_source_filename()
                if source_filename is None or source_filename == "_":
                    continue

                source_filename = Path(source_filename)
                source_filename_stem = source_filename.stem

                if source_filename_stem not in alignments.keys():
                    alignments[source_filename_stem] = []

                alignments_ar = alignments[source_filename_stem]
                alignments_ar.append(dflimg.get_source_landmarks())
                if len(alignments_ar) > 1:
                    multiple_faces_detected = True

            if multiple_faces_detected:
                io.log_info(
                    "Warning: multiple faces detected. Strongly recommended to process them separately."
                )

            frames = [
                MergeSubprocessor.Frame(frame_info=FrameInfo(
                    filepath=Path(p),
                    landmarks_list=alignments.get(Path(p).stem, None)))
                for p in input_path_image_paths
            ]

            if multiple_faces_detected:
                io.log_info(
                    "Warning: multiple faces detected. Motion blur will not be used."
                )
            else:
                s = 256
                local_pts = [(s // 2 - 1, s // 2 - 1),
                             (s // 2 - 1, 0)]  #center+up
                frames_len = len(frames)
                for i in io.progress_bar_generator(range(len(frames)),
                                                   "Computing motion vectors"):
                    fi_prev = frames[max(0, i - 1)].frame_info
                    fi = frames[i].frame_info
                    fi_next = frames[min(i + 1, frames_len - 1)].frame_info
                    if len(fi_prev.landmarks_list) == 0 or \
                       len(fi.landmarks_list) == 0 or \
                       len(fi_next.landmarks_list) == 0:
                        continue

                    mat_prev = LandmarksProcessor.get_transform_mat(
                        fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
                    mat = LandmarksProcessor.get_transform_mat(
                        fi.landmarks_list[0], s, face_type=FaceType.FULL)
                    mat_next = LandmarksProcessor.get_transform_mat(
                        fi_next.landmarks_list[0], s, face_type=FaceType.FULL)

                    pts_prev = LandmarksProcessor.transform_points(
                        local_pts, mat_prev, True)
                    pts = LandmarksProcessor.transform_points(
                        local_pts, mat, True)
                    pts_next = LandmarksProcessor.transform_points(
                        local_pts, mat_next, True)

                    prev_vector = pts[0] - pts_prev[0]
                    next_vector = pts_next[0] - pts[0]

                    motion_vector = pts_next[0] - pts_prev[0]
                    fi.motion_power = npla.norm(motion_vector)

                    motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array(
                        [0, 0], dtype=np.float32)

                    fi.motion_deg = -math.atan2(
                        motion_vector[1], motion_vector[0]) * 180 / math.pi

        elif cfg.type == MergerConfig.TYPE_FACE_AVATAR:
            filesdata = []
            for filepath in io.progress_bar_generator(input_path_image_paths,
                                                      "Collecting info"):
                filepath = Path(filepath)

                dflimg = DFLIMG.load(filepath)
                if dflimg is None:
                    io.log_err("%s is not a dfl image file" % (filepath.name))
                    continue
                filesdata += [
                    (FrameInfo(filepath=filepath,
                               landmarks_list=[dflimg.get_landmarks()]),
                     dflimg.get_source_filename())
                ]

            filesdata = sorted(
                filesdata,
                key=operator.itemgetter(1))  #sort by source_filename
            frames = []
            filesdata_len = len(filesdata)
            for i in range(len(filesdata)):
                frame_info = filesdata[i][0]

                prev_temporal_frame_infos = []
                next_temporal_frame_infos = []

                for t in range(cfg.temporal_face_count):
                    prev_frame_info = filesdata[max(i - t, 0)][0]
                    next_frame_info = filesdata[min(i + t,
                                                    filesdata_len - 1)][0]

                    prev_temporal_frame_infos.insert(0, prev_frame_info)
                    next_temporal_frame_infos.append(next_frame_info)

                frames.append(
                    MergeSubprocessor.Frame(
                        prev_temporal_frame_infos=prev_temporal_frame_infos,
                        frame_info=frame_info,
                        next_temporal_frame_infos=next_temporal_frame_infos))

        if len(frames) == 0:
            io.log_info("No frames to merge in input_dir.")
        else:
            MergeSubprocessor(is_interactive=is_interactive,
                              merger_session_filepath=merger_session_filepath,
                              predictor_func=predictor_func,
                              predictor_input_shape=predictor_input_shape,
                              merger_config=cfg,
                              frames=frames,
                              frames_root_path=input_path,
                              output_path=output_path,
                              model_iter=model.get_iter()).run()

        model.finalize()

    except Exception as e:
        print('Error: %s' % (str(e)))
        traceback.print_exc()
Ejemplo n.º 7
0
        def process_data(self, data):
            filename_path = Path(data[0])

            filename_path_str = str(filename_path)
            if self.cached_image[0] == filename_path_str:
                image = self.cached_image[1]
            else:
                image = cv2_imread(filename_path_str)
                self.cached_image = (filename_path_str, image)

            if image is None:
                self.log_err(
                    'Failed to extract %s, reason: cv2_imread() fail.' %
                    (str(filename_path)))
            else:
                if self.type == 'rects':
                    rects = self.e.extract_from_bgr(image)
                    return [str(filename_path), rects]

                elif self.type == 'landmarks':
                    rects = data[1]
                    landmarks = self.e.extract_from_bgr(image, rects)
                    return [str(filename_path), landmarks]

                elif self.type == 'final':
                    src_dflimg = None
                    (h, w, c) = image.shape
                    if h == w:
                        #extracting from already extracted jpg image?
                        if filename_path.suffix == '.jpg':
                            src_dflimg = DFLJPG.load(str(filename_path))

                    result = []
                    faces = data[1]

                    if self.debug:
                        debug_output_file = '{}{}'.format(
                            str(
                                Path(str(self.output_path) + '_debug') /
                                filename_path.stem), '.jpg')
                        debug_image = image.copy()

                    for (face_idx, face) in enumerate(faces):
                        output_file = '{}_{}{}'.format(
                            str(self.output_path / filename_path.stem),
                            str(face_idx), '.jpg')

                        rect = face[0]
                        image_landmarks = np.array(face[1])

                        if self.debug:
                            LandmarksProcessor.draw_rect_landmarks(
                                debug_image, rect, image_landmarks,
                                self.image_size, self.face_type)

                        if self.face_type == FaceType.MARK_ONLY:
                            face_image = image
                            face_image_landmarks = image_landmarks
                        else:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat(
                                image_landmarks, self.image_size,
                                self.face_type)
                            face_image = cv2.warpAffine(
                                image, image_to_face_mat,
                                (self.image_size, self.image_size),
                                cv2.INTER_LANCZOS4)
                            face_image_landmarks = LandmarksProcessor.transform_points(
                                image_landmarks, image_to_face_mat)

                        if src_dflimg is not None:
                            #if extracting from dflimg just copy it in order not to lose quality
                            shutil.copy(str(filename_path), str(output_file))
                        else:
                            cv2_imwrite(output_file, face_image,
                                        [int(cv2.IMWRITE_JPEG_QUALITY), 85])

                        DFLJPG.embed_data(
                            output_file,
                            face_type=FaceType.toString(self.face_type),
                            landmarks=face_image_landmarks.tolist(),
                            source_filename=filename_path.name,
                            source_rect=rect,
                            source_landmarks=image_landmarks.tolist())

                        result.append(output_file)

                    if self.debug:
                        cv2_imwrite(debug_output_file, debug_image,
                                    [int(cv2.IMWRITE_JPEG_QUALITY), 50])

                    return result
            return None
Ejemplo n.º 8
0
def apply_xseg(input_path, model_path):
    if not input_path.exists():
        raise ValueError(f'{input_path} not found. Please ensure it exists.')

    if not model_path.exists():
        raise ValueError(f'{model_path} not found. Please ensure it exists.')
        
    face_type = None
    
    model_dat = model_path / 'XSeg_data.dat'
    if model_dat.exists():
        dat = pickle.loads( model_dat.read_bytes() )
        dat_options = dat.get('options', None)
        if dat_options is not None:
            face_type = dat_options.get('face_type', None)
        
        
        
    if face_type is None:
        face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower()
        if face_type == 'same':
            face_type = None
    
    if face_type is not None:
        face_type = {'h'  : FaceType.HALF,
                     'mf' : FaceType.MID_FULL,
                     'f'  : FaceType.FULL,
                     'wf' : FaceType.WHOLE_FACE,
                     'head' : FaceType.HEAD}[face_type]
                     
    io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.')

    device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
    nn.initialize(device_config)
        
    
    
    xseg = XSegNet(name='XSeg', 
                    load_weights=True,
                    weights_file_root=model_path,
                    data_format=nn.data_format,
                    raise_on_no_model_files=True)
    xseg_res = xseg.get_resolution()
              
    images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
    
    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        dflimg = DFLIMG.load(filepath)
        if dflimg is None or not dflimg.has_data():
            io.log_info(f'{filepath} is not a DFLIMG')
            continue
        
        img = cv2_imread(filepath).astype(np.float32) / 255.0
        h,w,c = img.shape
        
        img_face_type = FaceType.fromString( dflimg.get_face_type() )
        if face_type is not None and img_face_type != face_type:
            lmrks = dflimg.get_source_landmarks()
            
            fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type)
            imat = LandmarksProcessor.get_transform_mat(lmrks, w, img_face_type)
            
            g_p = LandmarksProcessor.transform_points (np.float32([(0,0),(w,0),(0,w) ]), fmat, True)
            g_p2 = LandmarksProcessor.transform_points (g_p, imat)
            
            mat = cv2.getAffineTransform( g_p2, np.float32([(0,0),(w,0),(0,w) ]) )
            
            img = cv2.warpAffine(img, mat, (w, w), cv2.INTER_LANCZOS4)
            img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
        else:
            if w != xseg_res:
                img = cv2.resize( img, (xseg_res,xseg_res), interpolation=cv2.INTER_LANCZOS4 )    
                    
        if len(img.shape) == 2:
            img = img[...,None]            
    
        mask = xseg.extract(img)
        
        if face_type is not None and img_face_type != face_type:
            mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4)
            mask = cv2.warpAffine( mask, mat, (w,w), np.zeros( (h,w,c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4)
            mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
        mask[mask < 0.5]=0
        mask[mask >= 0.5]=1    
        dflimg.set_xseg_mask(mask)
        dflimg.save()
Ejemplo n.º 9
0
        def process_data(self, data):
            filename_path = Path(data.filename)

            filename_path_str = str(filename_path)
            if self.cached_image[0] == filename_path_str:
                image = self.cached_image[
                    1]  #cached image for manual extractor
            else:
                image = cv2_imread(filename_path_str)

                if image is None:
                    self.log_err(
                        'Failed to extract %s, reason: cv2_imread() fail.' %
                        (str(filename_path)))
                    return data

                image = imagelib.normalize_channels(image, 3)
                h, w, ch = image.shape

                wm, hm = w % 2, h % 2
                if wm + hm != 0:  #fix odd image
                    image = image[0:h - hm, 0:w - wm, :]
                self.cached_image = (filename_path_str, image)

            src_dflimg = None
            h, w, ch = image.shape
            if h == w:
                #extracting from already extracted jpg image?
                if filename_path.suffix == '.png':
                    src_dflimg = DFLPNG.load(str(filename_path))
                if filename_path.suffix == '.jpg':
                    src_dflimg = DFLJPG.load(str(filename_path))

            if 'rects' in self.type:
                if min(w, h) < 128:
                    self.log_err('Image is too small %s : [%d, %d]' %
                                 (str(filename_path), w, h))
                    data.rects = []
                else:
                    for rot in ([0, 90, 270, 180]):
                        data.rects_rotation = rot
                        if rot == 0:
                            rotated_image = image
                        elif rot == 90:
                            rotated_image = image.swapaxes(0, 1)[:, ::-1, :]
                        elif rot == 180:
                            rotated_image = image[::-1, ::-1, :]
                        elif rot == 270:
                            rotated_image = image.swapaxes(0, 1)[::-1, :, :]

                        rects = data.rects = self.e.extract(rotated_image,
                                                            is_bgr=True)
                        if len(rects) != 0:
                            break

                    if self.max_faces_from_image != 0 and len(data.rects) > 1:
                        data.rects = data.rects[0:self.max_faces_from_image]

                return data

            elif self.type == 'landmarks':

                if data.rects_rotation == 0:
                    rotated_image = image
                elif data.rects_rotation == 90:
                    rotated_image = image.swapaxes(0, 1)[:, ::-1, :]
                elif data.rects_rotation == 180:
                    rotated_image = image[::-1, ::-1, :]
                elif data.rects_rotation == 270:
                    rotated_image = image.swapaxes(0, 1)[::-1, :, :]

                data.landmarks = self.e.extract(
                    rotated_image,
                    data.rects,
                    self.second_pass_e if
                    (src_dflimg is None and data.landmarks_accurate) else None,
                    is_bgr=True)
                if data.rects_rotation != 0:
                    for i, (rect,
                            lmrks) in enumerate(zip(data.rects,
                                                    data.landmarks)):
                        new_rect, new_lmrks = rect, lmrks
                        (l, t, r, b) = rect
                        if data.rects_rotation == 90:
                            new_rect = (t, h - l, b, h - r)
                            if lmrks is not None:
                                new_lmrks = lmrks[:, ::-1].copy()
                                new_lmrks[:, 1] = h - new_lmrks[:, 1]
                        elif data.rects_rotation == 180:
                            if lmrks is not None:
                                new_rect = (w - l, h - t, w - r, h - b)
                                new_lmrks = lmrks.copy()
                                new_lmrks[:, 0] = w - new_lmrks[:, 0]
                                new_lmrks[:, 1] = h - new_lmrks[:, 1]
                        elif data.rects_rotation == 270:
                            new_rect = (w - b, l, w - t, r)
                            if lmrks is not None:
                                new_lmrks = lmrks[:, ::-1].copy()
                                new_lmrks[:, 0] = w - new_lmrks[:, 0]
                        data.rects[i], data.landmarks[i] = new_rect, new_lmrks

                return data

            elif self.type == 'final':
                data.final_output_files = []
                rects = data.rects
                landmarks = data.landmarks

                if self.debug_dir is not None:
                    debug_output_file = str(
                        Path(self.debug_dir) / (filename_path.stem + '.jpg'))
                    debug_image = image.copy()

                if src_dflimg is not None and len(rects) != 1:
                    #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
                    print("src_dflimg is not None and len(rects) != 1",
                          str(filename_path))
                    output_file = str(self.final_output_path /
                                      filename_path.name)
                    if str(filename_path) != str(output_file):
                        shutil.copy(str(filename_path), str(output_file))
                    data.final_output_files.append(output_file)
                else:
                    face_idx = 0
                    for rect, image_landmarks in zip(rects, landmarks):

                        if src_dflimg is not None and face_idx > 1:
                            #cannot extract more than 1 face from dflimg
                            break

                        if image_landmarks is None:
                            continue

                        rect = np.array(rect)

                        if self.face_type == FaceType.MARK_ONLY:
                            image_to_face_mat = None
                            face_image = image
                            face_image_landmarks = image_landmarks
                        else:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat(
                                image_landmarks, self.image_size,
                                self.face_type)

                            face_image = cv2.warpAffine(
                                image, image_to_face_mat,
                                (self.image_size, self.image_size),
                                cv2.INTER_LANCZOS4)
                            face_image_landmarks = LandmarksProcessor.transform_points(
                                image_landmarks, image_to_face_mat)

                            landmarks_bbox = LandmarksProcessor.transform_points(
                                [(0, 0), (0, self.image_size - 1),
                                 (self.image_size - 1, self.image_size - 1),
                                 (self.image_size - 1, 0)], image_to_face_mat,
                                True)

                            rect_area = mathlib.polygon_area(
                                np.array(rect[[0, 2, 2, 0]]),
                                np.array(rect[[1, 1, 3, 3]]))
                            landmarks_area = mathlib.polygon_area(
                                landmarks_bbox[:, 0], landmarks_bbox[:, 1])

                            if landmarks_area > 4 * rect_area:  #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                                continue

                            if self.debug_dir is not None:
                                LandmarksProcessor.draw_rect_landmarks(
                                    debug_image,
                                    rect,
                                    image_landmarks,
                                    self.image_size,
                                    self.face_type,
                                    transparent_mask=True)

                        if src_dflimg is not None and filename_path.suffix == '.jpg':
                            #if extracting from dflimg and jpg copy it in order not to lose quality
                            output_file = str(self.final_output_path /
                                              filename_path.name)
                            if str(filename_path) != str(output_file):
                                shutil.copy(str(filename_path),
                                            str(output_file))
                        else:
                            output_file = '{}_{}{}'.format(
                                str(self.final_output_path /
                                    filename_path.stem), str(face_idx), '.jpg')
                            cv2_imwrite(output_file, face_image,
                                        [int(cv2.IMWRITE_JPEG_QUALITY), 85])

                        DFLJPG.embed_data(
                            output_file,
                            face_type=FaceType.toString(self.face_type),
                            landmarks=face_image_landmarks.tolist(),
                            source_filename=filename_path.name,
                            source_rect=rect,
                            source_landmarks=image_landmarks.tolist(),
                            image_to_face_mat=image_to_face_mat,
                            pitch_yaw_roll=data.pitch_yaw_roll)

                        data.final_output_files.append(output_file)
                        face_idx += 1
                    data.faces_detected = face_idx

                if self.debug_dir is not None:
                    cv2_imwrite(debug_output_file, debug_image,
                                [int(cv2.IMWRITE_JPEG_QUALITY), 50])

                return data

            elif self.type == 'fanseg':
                if src_dflimg is not None:
                    fanseg_mask = self.e.extract(image / 255.0)
                    src_dflimg.embed_and_set(
                        filename_path_str,
                        fanseg_mask=fanseg_mask,
                    )
Ejemplo n.º 10
0
        def final_stage(
            data,
            image,
            face_type,
            image_size,
            jpeg_quality,
            output_debug_path=None,
            final_output_path=None,
        ):
            data.final_output_files = []
            filepath = data.filepath
            rects = data.rects
            landmarks = data.landmarks

            if output_debug_path is not None:
                debug_image = image.copy()

            face_idx = 0
            for rect, image_landmarks in zip(rects, landmarks):
                if image_landmarks is None:
                    continue

                rect = np.array(rect)

                if face_type == FaceType.MARK_ONLY:
                    image_to_face_mat = None
                    face_image = image
                    face_image_landmarks = image_landmarks
                else:
                    image_to_face_mat = LandmarksProcessor.get_transform_mat(
                        image_landmarks, image_size, face_type)

                    face_image = cv2.warpAffine(image, image_to_face_mat,
                                                (image_size, image_size),
                                                cv2.INTER_LANCZOS4)
                    face_image_landmarks = LandmarksProcessor.transform_points(
                        image_landmarks, image_to_face_mat)

                    landmarks_bbox = LandmarksProcessor.transform_points(
                        [(0, 0), (0, image_size - 1),
                         (image_size - 1, image_size - 1),
                         (image_size - 1, 0)], image_to_face_mat, True)

                    rect_area = mathlib.polygon_area(
                        np.array(rect[[0, 2, 2, 0]]).astype(np.float32),
                        np.array(rect[[1, 1, 3, 3]]).astype(np.float32))
                    landmarks_area = mathlib.polygon_area(
                        landmarks_bbox[:, 0].astype(np.float32),
                        landmarks_bbox[:, 1].astype(np.float32))

                    if not data.manual and face_type <= FaceType.FULL_NO_ALIGN and landmarks_area > 4 * rect_area:  #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                        continue

                    if output_debug_path is not None:
                        LandmarksProcessor.draw_rect_landmarks(
                            debug_image,
                            rect,
                            image_landmarks,
                            face_type,
                            image_size,
                            transparent_mask=True)

                output_path = final_output_path
                if data.force_output_path is not None:
                    output_path = data.force_output_path

                output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg"
                cv2_imwrite(output_filepath, face_image,
                            [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])

                dflimg = DFLJPG.load(output_filepath)
                dflimg.set_face_type(FaceType.toString(face_type))
                dflimg.set_landmarks(face_image_landmarks.tolist())
                dflimg.set_source_filename(filepath.name)
                dflimg.set_source_rect(rect)
                dflimg.set_source_landmarks(image_landmarks.tolist())
                dflimg.set_image_to_face_mat(image_to_face_mat)
                dflimg.save()

                data.final_output_files.append(output_filepath)
                face_idx += 1
            data.faces_detected = face_idx

            if output_debug_path is not None:
                cv2_imwrite(output_debug_path / (filepath.stem + '.jpg'),
                            debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50])

            return data
Ejemplo n.º 11
0
    def extract(self,
                input_image,
                rects,
                second_pass_extractor=None,
                is_bgr=True,
                multi_sample=False):
        if len(rects) == 0:
            return []

        if is_bgr:
            input_image = input_image[:, :, ::-1]
            is_bgr = False

        (h, w, ch) = input_image.shape

        landmarks = []
        for (left, top, right, bottom) in rects:
            scale = (right - left + bottom - top) / 195.0

            center = np.array([(left + right) / 2.0, (top + bottom) / 2.0])
            centers = [center]

            if multi_sample:
                centers += [
                    center + [-1, -1],
                    center + [1, -1],
                    center + [1, 1],
                    center + [-1, 1],
                ]

            images = []
            ptss = []

            try:
                for c in centers:
                    images += [self.crop(input_image, c, scale)]

                images = np.stack(images)
                images = images.astype(np.float32) / 255.0

                predicted = []
                for i in range(len(images)):
                    predicted += [self.model.run([images[i][None, ...]])[0]]

                predicted = np.stack(predicted)

                for i, pred in enumerate(predicted):
                    ptss += [
                        self.get_pts_from_predict(pred, centers[i], scale)
                    ]
                pts_img = np.mean(np.array(ptss), 0)

                landmarks.append(pts_img)
            except:
                landmarks.append(None)

        if second_pass_extractor is not None:
            for i, lmrks in enumerate(landmarks):
                try:
                    if lmrks is not None:
                        image_to_face_mat = LandmarksProcessor.get_transform_mat(
                            lmrks, 256, FaceType.FULL)
                        face_image = cv2.warpAffine(input_image,
                                                    image_to_face_mat,
                                                    (256, 256),
                                                    cv2.INTER_CUBIC)

                        rects2 = second_pass_extractor.extract(face_image,
                                                               is_bgr=is_bgr)
                        if len(
                                rects2
                        ) == 1:  #dont do second pass if faces != 1 detected in cropped image
                            lmrks2 = self.extract(face_image, [rects2[0]],
                                                  is_bgr=is_bgr,
                                                  multi_sample=True)[0]
                            landmarks[i] = LandmarksProcessor.transform_points(
                                lmrks2, image_to_face_mat, True)
                except:
                    pass

        return landmarks
Ejemplo n.º 12
0
    def get_data(self, host_dict):
        if self.type == 'landmarks-manual':
            need_remark_face = False
            while len(self.input_data) > 0:
                data = self.input_data[0]
                filepath, data_rects, data_landmarks = data.filepath, data.rects, data.landmarks
                is_frame_done = False

                if self.image_filepath != filepath:
                    self.image_filepath = filepath
                    if self.cache_original_image[0] == filepath:
                        self.original_image = self.cache_original_image[1]
                    else:
                        self.original_image = imagelib.normalize_channels(
                            cv2_imread(filepath), 3)

                        self.cache_original_image = (filepath,
                                                     self.original_image)

                    (h, w, c) = self.original_image.shape
                    self.view_scale = 1.0 if self.manual_window_size == 0 else self.manual_window_size / (
                        h * (16.0 / 9.0))

                    if self.cache_image[0] == (h, w, c) + (self.view_scale,
                                                           filepath):
                        self.image = self.cache_image[1]
                    else:
                        self.image = cv2.resize(self.original_image, (int(
                            w * self.view_scale), int(h * self.view_scale)),
                                                interpolation=cv2.INTER_LINEAR)
                        self.cache_image = ((h, w, c) +
                                            (self.view_scale, filepath),
                                            self.image)

                    (h, w, c) = self.image.shape

                    sh = (0, 0, w, min(100, h))
                    if self.cache_text_lines_img[0] == sh:
                        self.text_lines_img = self.cache_text_lines_img[1]
                    else:
                        self.text_lines_img = (imagelib.get_draw_text_lines(
                            self.image, sh, [
                                '[L Mouse click] - lock/unlock selection. [Mouse wheel] - change rect',
                                '[R Mouse Click] - manual face rectangle',
                                '[Enter] / [Space] - confirm / skip frame',
                                '[,] [.]- prev frame, next frame. [Q] - skip remaining frames',
                                '[a] - accuracy on/off (more fps)',
                                '[h] - hide this help'
                            ], (1, 1, 1)) * 255).astype(np.uint8)

                        self.cache_text_lines_img = (sh, self.text_lines_img)

                if need_remark_face:  # need remark image from input data that already has a marked face?
                    need_remark_face = False
                    if len(
                            data_rects
                    ) != 0:  # If there was already a face then lock the rectangle to it until the mouse is clicked
                        self.rect = data_rects.pop()
                        self.landmarks = data_landmarks.pop()
                        data_rects.clear()
                        data_landmarks.clear()

                        self.rect_locked = True
                        self.rect_size = (self.rect[2] - self.rect[0]) / 2
                        self.x = (self.rect[0] + self.rect[2]) / 2
                        self.y = (self.rect[1] + self.rect[3]) / 2
                        self.redraw()

                if len(data_rects) == 0:
                    (h, w, c) = self.image.shape
                    while True:
                        io.process_messages(0.0001)

                        if not self.force_landmarks:
                            new_x = self.x
                            new_y = self.y

                        new_rect_size = self.rect_size

                        mouse_events = io.get_mouse_events(self.wnd_name)
                        for ev in mouse_events:
                            (x, y, ev, flags) = ev
                            if ev == io.EVENT_MOUSEWHEEL and not self.rect_locked:
                                mod = 1 if flags > 0 else -1
                                diff = 1 if new_rect_size <= 40 else np.clip(
                                    new_rect_size / 10, 1, 10)
                                new_rect_size = max(5,
                                                    new_rect_size + diff * mod)
                            elif ev == io.EVENT_LBUTTONDOWN:
                                if self.force_landmarks:
                                    self.x = new_x
                                    self.y = new_y
                                    self.force_landmarks = False
                                    self.rect_locked = True
                                    self.redraw()
                                else:
                                    self.rect_locked = not self.rect_locked
                                    self.extract_needed = True
                            elif ev == io.EVENT_RBUTTONDOWN:
                                self.force_landmarks = not self.force_landmarks
                                if self.force_landmarks:
                                    self.rect_locked = False
                            elif not self.rect_locked:
                                new_x = np.clip(x, 0, w - 1) / self.view_scale
                                new_y = np.clip(y, 0, h - 1) / self.view_scale

                        key_events = io.get_key_events(self.wnd_name)
                        key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[
                            -1] if len(key_events) > 0 else (0, 0, False,
                                                             False, False)

                        if key == ord('\r') or key == ord('\n'):
                            #confirm frame
                            is_frame_done = True
                            data_rects.append(self.rect)
                            data_landmarks.append(self.landmarks)
                            break
                        elif key == ord(' '):
                            #confirm skip frame
                            is_frame_done = True
                            break
                        elif key == ord(',') and len(self.result) > 0:
                            #go prev frame

                            if self.rect_locked:
                                self.rect_locked = False
                                # Only save the face if the rect is still locked
                                data_rects.append(self.rect)
                                data_landmarks.append(self.landmarks)

                            self.input_data.insert(0, self.result.pop())
                            io.progress_bar_inc(-1)
                            need_remark_face = True

                            break
                        elif key == ord('.'):
                            #go next frame

                            if self.rect_locked:
                                self.rect_locked = False
                                # Only save the face if the rect is still locked
                                data_rects.append(self.rect)
                                data_landmarks.append(self.landmarks)

                            need_remark_face = True
                            is_frame_done = True
                            break
                        elif key == ord('q'):
                            #skip remaining

                            if self.rect_locked:
                                self.rect_locked = False
                                data_rects.append(self.rect)
                                data_landmarks.append(self.landmarks)

                            while len(self.input_data) > 0:
                                self.result.append(self.input_data.pop(0))
                                io.progress_bar_inc(1)

                            break

                        elif key == ord('h'):
                            self.hide_help = not self.hide_help
                            break
                        elif key == ord('a'):
                            self.landmarks_accurate = not self.landmarks_accurate
                            break

                        if self.force_landmarks:
                            pt2 = np.float32([new_x, new_y])
                            pt1 = np.float32([self.x, self.y])

                            pt_vec_len = npla.norm(pt2 - pt1)
                            pt_vec = pt2 - pt1
                            if pt_vec_len != 0:
                                pt_vec /= pt_vec_len

                            self.rect_size = pt_vec_len
                            self.rect = (int(self.x - self.rect_size),
                                         int(self.y - self.rect_size),
                                         int(self.x + self.rect_size),
                                         int(self.y + self.rect_size))

                            if pt_vec_len > 0:
                                lmrks = np.concatenate(
                                    (np.zeros((17, 2), np.float32),
                                     LandmarksProcessor.landmarks_2D),
                                    axis=0)
                                lmrks -= lmrks[30:31, :]
                                mat = cv2.getRotationMatrix2D(
                                    (0, 0), -np.arctan2(pt_vec[1], pt_vec[0]) *
                                    180 / math.pi, pt_vec_len)
                                mat[:, 2] += (self.x, self.y)
                                self.landmarks = LandmarksProcessor.transform_points(
                                    lmrks, mat)

                            self.redraw()

                        elif self.x != new_x or \
                           self.y != new_y or \
                           self.rect_size != new_rect_size or \
                           self.extract_needed:
                            self.x = new_x
                            self.y = new_y
                            self.rect_size = new_rect_size
                            self.rect = (int(self.x - self.rect_size),
                                         int(self.y - self.rect_size),
                                         int(self.x + self.rect_size),
                                         int(self.y + self.rect_size))

                            return ExtractSubprocessor.Data(
                                filepath,
                                rects=[self.rect],
                                landmarks_accurate=self.landmarks_accurate)

                else:
                    is_frame_done = True

                if is_frame_done:
                    self.result.append(data)
                    self.input_data.pop(0)
                    io.progress_bar_inc(1)
                    self.extract_needed = True
                    self.rect_locked = False
        else:
            if len(self.input_data) > 0:
                return self.input_data.pop(0)

        return None
Ejemplo n.º 13
0
def MergeMaskedFace (predictor_func, predictor_input_shape,
                     face_enhancer_func,
                     xseg_256_extract_func,
                     cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks):

    img_size = img_bgr.shape[1], img_bgr.shape[0]
    img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks)

    input_size = predictor_input_shape[0]
    mask_subres_size = input_size*4
    output_size = input_size
    if cfg.super_resolution_power != 0:
        output_size *= 4

    face_mat        = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=cfg.face_type)
    face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=cfg.face_type, scale= 1.0 + 0.01*cfg.output_face_scale)

    if mask_subres_size == output_size:
        face_mask_output_mat = face_output_mat
    else:
        face_mask_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, mask_subres_size, face_type=cfg.face_type, scale= 1.0 + 0.01*cfg.output_face_scale)

    dst_face_bgr      = cv2.warpAffine( img_bgr        , face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
    dst_face_bgr      = np.clip(dst_face_bgr, 0, 1)

    dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
    dst_face_mask_a_0 = np.clip(dst_face_mask_a_0, 0, 1)

    predictor_input_bgr      = cv2.resize (dst_face_bgr, (input_size,input_size) )

    predicted = predictor_func (predictor_input_bgr)
    prd_face_bgr          = np.clip (predicted[0], 0, 1.0)
    prd_face_mask_a_0     = np.clip (predicted[1], 0, 1.0)
    prd_face_dst_mask_a_0 = np.clip (predicted[2], 0, 1.0)

    if cfg.super_resolution_power != 0:
        prd_face_bgr_enhanced = face_enhancer_func(prd_face_bgr, is_tanh=True, preserve_size=False)
        mod = cfg.super_resolution_power / 100.0
        prd_face_bgr = cv2.resize(prd_face_bgr, (output_size,output_size))*(1.0-mod) + prd_face_bgr_enhanced*mod
        prd_face_bgr = np.clip(prd_face_bgr, 0, 1)

    if cfg.super_resolution_power != 0:
        prd_face_mask_a_0     = cv2.resize (prd_face_mask_a_0,      (output_size, output_size), interpolation=cv2.INTER_CUBIC)
        prd_face_dst_mask_a_0 = cv2.resize (prd_face_dst_mask_a_0,  (output_size, output_size), interpolation=cv2.INTER_CUBIC)

    if cfg.mask_mode == 1: #dst
        wrk_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), interpolation=cv2.INTER_CUBIC)
    elif cfg.mask_mode == 2: #learned-prd
        wrk_face_mask_a_0 = prd_face_mask_a_0
    elif cfg.mask_mode == 3: #learned-dst
        wrk_face_mask_a_0 = prd_face_dst_mask_a_0
    elif cfg.mask_mode == 4: #learned-prd*learned-dst
        wrk_face_mask_a_0 = prd_face_mask_a_0*prd_face_dst_mask_a_0
    elif cfg.mask_mode == 5: #learned-prd+learned-dst
        wrk_face_mask_a_0 = np.clip( prd_face_mask_a_0+prd_face_dst_mask_a_0, 0, 1)
    elif cfg.mask_mode >= 6 and cfg.mask_mode <= 9:  #XSeg modes
        if cfg.mask_mode == 6 or cfg.mask_mode == 8 or cfg.mask_mode == 9:
            # obtain XSeg-prd
            prd_face_xseg_bgr = cv2.resize (prd_face_bgr, (xseg_input_size,)*2, interpolation=cv2.INTER_CUBIC)
            prd_face_xseg_mask = xseg_256_extract_func(prd_face_xseg_bgr)
            X_prd_face_mask_a_0 = cv2.resize ( prd_face_xseg_mask, (output_size, output_size), interpolation=cv2.INTER_CUBIC)

        if cfg.mask_mode >= 7 and cfg.mask_mode <= 9:
            # obtain XSeg-dst
            xseg_mat            = LandmarksProcessor.get_transform_mat (img_face_landmarks, xseg_input_size, face_type=cfg.face_type)
            dst_face_xseg_bgr   = cv2.warpAffine(img_bgr, xseg_mat, (xseg_input_size,)*2, flags=cv2.INTER_CUBIC )
            dst_face_xseg_mask  = xseg_256_extract_func(dst_face_xseg_bgr)
            X_dst_face_mask_a_0 = cv2.resize (dst_face_xseg_mask, (output_size,output_size), interpolation=cv2.INTER_CUBIC)

        if cfg.mask_mode == 6:   #'XSeg-prd'
            wrk_face_mask_a_0 = X_prd_face_mask_a_0
        elif cfg.mask_mode == 7: #'XSeg-dst'
            wrk_face_mask_a_0 = X_dst_face_mask_a_0
        elif cfg.mask_mode == 8: #'XSeg-prd*XSeg-dst'
            wrk_face_mask_a_0 = X_prd_face_mask_a_0 * X_dst_face_mask_a_0
        elif cfg.mask_mode == 9: #learned-prd*learned-dst*XSeg-prd*XSeg-dst
            wrk_face_mask_a_0 = prd_face_mask_a_0 * prd_face_dst_mask_a_0 * X_prd_face_mask_a_0 * X_dst_face_mask_a_0

    wrk_face_mask_a_0[ wrk_face_mask_a_0 < (1.0/255.0) ] = 0.0 # get rid of noise

    # resize to mask_subres_size
    if wrk_face_mask_a_0.shape[0] != mask_subres_size:
        wrk_face_mask_a_0 = cv2.resize (wrk_face_mask_a_0, (mask_subres_size, mask_subres_size), interpolation=cv2.INTER_CUBIC)

    # process mask in local predicted space
    if 'raw' not in cfg.mode:
        # add zero pad
        wrk_face_mask_a_0 = np.pad (wrk_face_mask_a_0, input_size)

        ero  = cfg.erode_mask_modifier
        blur = cfg.blur_mask_modifier

        if ero > 0:
            wrk_face_mask_a_0 = cv2.erode(wrk_face_mask_a_0, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
        elif ero < 0:
            wrk_face_mask_a_0 = cv2.dilate(wrk_face_mask_a_0, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 )

        # clip eroded/dilated mask in actual predict area
        # pad with half blur size in order to accuratelly fade to zero at the boundary
        clip_size = input_size + blur // 2

        wrk_face_mask_a_0[:clip_size,:] = 0
        wrk_face_mask_a_0[-clip_size:,:] = 0
        wrk_face_mask_a_0[:,:clip_size] = 0
        wrk_face_mask_a_0[:,-clip_size:] = 0

        if blur > 0:
            blur = blur + (1-blur % 2)
            wrk_face_mask_a_0 = cv2.GaussianBlur(wrk_face_mask_a_0, (blur, blur) , 0)

        wrk_face_mask_a_0 = wrk_face_mask_a_0[input_size:-input_size,input_size:-input_size]

        wrk_face_mask_a_0 = np.clip(wrk_face_mask_a_0, 0, 1)

    img_face_mask_a = cv2.warpAffine( wrk_face_mask_a_0, face_mask_output_mat, img_size, np.zeros(img_bgr.shape[0:2], dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC )[...,None]
    img_face_mask_a = np.clip (img_face_mask_a, 0.0, 1.0)
    img_face_mask_a [ img_face_mask_a < (1.0/255.0) ] = 0.0 # get rid of noise

    if wrk_face_mask_a_0.shape[0] != output_size:
        wrk_face_mask_a_0 = cv2.resize (wrk_face_mask_a_0, (output_size,output_size), interpolation=cv2.INTER_CUBIC)

    wrk_face_mask_a = wrk_face_mask_a_0[...,None]

    out_merging_mask_a = None
    if cfg.mode == 'original':
        return img_bgr, img_face_mask_a

    elif 'raw' in cfg.mode:
        if cfg.mode == 'raw-rgb':
            out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT )
            out_merging_mask_a = img_face_mask_a
        elif cfg.mode == 'raw-predict':
            out_img = prd_face_bgr
            out_merging_mask_a = wrk_face_mask_a
        else:
            raise ValueError(f"undefined raw type {cfg.mode}")

        out_img = np.clip (out_img, 0.0, 1.0 )
    else:
                
        # Process if the mask meets minimum size
        maxregion = np.argwhere( img_face_mask_a >= 0.1 )
        if maxregion.size != 0:
            miny,minx = maxregion.min(axis=0)[:2]
            maxy,maxx = maxregion.max(axis=0)[:2]
            lenx = maxx - minx
            leny = maxy - miny
            if min(lenx,leny) >= 4:
                wrk_face_mask_area_a = wrk_face_mask_a.copy()
                wrk_face_mask_area_a[wrk_face_mask_area_a>0] = 1.0
                
                if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0:
                    if cfg.color_transfer_mode == 1: #rct
                        prd_face_bgr = imagelib.reinhard_color_transfer ( np.clip( prd_face_bgr*wrk_face_mask_area_a*255, 0, 255).astype(np.uint8),
                                                                          np.clip( dst_face_bgr*wrk_face_mask_area_a*255, 0, 255).astype(np.uint8), )

                        prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
                    elif cfg.color_transfer_mode == 2: #lct
                        prd_face_bgr = imagelib.linear_color_transfer (prd_face_bgr, dst_face_bgr)
                    elif cfg.color_transfer_mode == 3: #mkl
                        prd_face_bgr = imagelib.color_transfer_mkl (prd_face_bgr, dst_face_bgr)
                    elif cfg.color_transfer_mode == 4: #mkl-m
                        prd_face_bgr = imagelib.color_transfer_mkl (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                    elif cfg.color_transfer_mode == 5: #idt
                        prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr, dst_face_bgr)
                    elif cfg.color_transfer_mode == 6: #idt-m
                        prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                    elif cfg.color_transfer_mode == 7: #sot-m
                        prd_face_bgr = imagelib.color_transfer_sot (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a, steps=10, batch_size=30)
                        prd_face_bgr = np.clip (prd_face_bgr, 0.0, 1.0)
                    elif cfg.color_transfer_mode == 8: #mix-m
                        prd_face_bgr = imagelib.color_transfer_mix (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)

                if cfg.mode == 'hist-match':
                    hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)

                    if cfg.masked_hist_match:
                        hist_mask_a *= wrk_face_mask_area_a

                    white =  (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)

                    hist_match_1 = prd_face_bgr*hist_mask_a + white
                    hist_match_1[ hist_match_1 > 1.0 ] = 1.0

                    hist_match_2 = dst_face_bgr*hist_mask_a + white
                    hist_match_2[ hist_match_1 > 1.0 ] = 1.0

                    prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, cfg.hist_match_threshold ).astype(dtype=np.float32)

                if 'seamless' in cfg.mode:
                    #mask used for cv2.seamlessClone
                    img_face_seamless_mask_a = None
                    for i in range(1,10):
                        a = img_face_mask_a > i / 10.0
                        if len(np.argwhere(a)) == 0:
                            continue
                        img_face_seamless_mask_a = img_face_mask_a.copy()
                        img_face_seamless_mask_a[a] = 1.0
                        img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0
                        break

                out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT )
                out_img = np.clip(out_img, 0.0, 1.0)

                if 'seamless' in cfg.mode:
                    try:
                        #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering)
                        l,t,w,h = cv2.boundingRect( (img_face_seamless_mask_a*255).astype(np.uint8) )
                        s_maskx, s_masky = int(l+w/2), int(t+h/2)
                        out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), img_bgr_uint8, (img_face_seamless_mask_a*255).astype(np.uint8), (s_maskx,s_masky) , cv2.NORMAL_CLONE )
                        out_img = out_img.astype(dtype=np.float32) / 255.0
                    except Exception as e:
                        #seamlessClone may fail in some cases
                        e_str = traceback.format_exc()

                        if 'MemoryError' in e_str:
                            raise Exception("Seamless fail: " + e_str) #reraise MemoryError in order to reprocess this data by other processes
                        else:
                            print ("Seamless fail: " + e_str)

                cfg_mp = 0.30#cfg.motion_blur_power / 100.0
                
                
                ###
                shrink_res = output_size #512
                
                shrink_prd_face_dst_mask_a_0 = cv2.resize (prd_face_dst_mask_a_0,  (shrink_res, shrink_res), interpolation=cv2.INTER_CUBIC)     
                
                shrink_blur_size = (shrink_res // 32)+1            
                shrink_blur_size += (1-shrink_blur_size % 2)
                
                # Feather the mask
                shrink_prd_face_dst_mask_a_0 = cv2.GaussianBlur(shrink_prd_face_dst_mask_a_0, (shrink_blur_size,shrink_blur_size) , 0)
                shrink_prd_face_dst_mask_a_0[shrink_prd_face_dst_mask_a_0 < 0.5] = 0.0
                shrink_prd_face_dst_mask_a_0[shrink_prd_face_dst_mask_a_0 >= 0.5] = 1.0            
                
                cnts = cv2.findContours( shrink_prd_face_dst_mask_a_0.astype(np.uint8), cv2.RETR_LIST , cv2.CHAIN_APPROX_TC89_KCOS )
                
                # Get the largest found contour
                cnt = sorted(cnts[0], key = cv2.contourArea, reverse = True)[0].squeeze()

                center = np.mean(cnt,0)
                cnt2 = cnt.copy().astype(np.float32)
                cnt2_c = center - cnt2    
                cnt2_len = npla.norm(cnt2_c, axis=1, keepdims=True)
                cnt2_vec = cnt2_c / cnt2_len
                cnt2 += cnt2_vec *  cnt2_len * cfg_mp #todo               
                cnt2 = cnt2.astype(np.int32)
                
                img_cnt = LandmarksProcessor.transform_points (cnt, face_mat, True)
                img_cnt2 = LandmarksProcessor.transform_points (cnt2, face_mat, True)
                
                 
                # Anchor perimeter
                
                
                h=img_size[1]
                w=img_size[0]
                w_pts_count = w // 16
                h_pts_count = w // 16
                perim_pts = np.concatenate (
                            (   np.concatenate ( [ np.arange(0,w+w/w_pts_count, w/w_pts_count)[...,None], np.array ( [[0]]*(h_pts_count+1) ) ], axis=-1 ),
                                np.concatenate ( [ np.arange(0,w+w/w_pts_count, w/w_pts_count)[...,None], np.array ( [[h]]*(h_pts_count+1) ) ], axis=-1 ),
                                np.concatenate ( [ np.array ( [[0]]*(w_pts_count+1) ), np.arange(0,h+h/h_pts_count, h/h_pts_count)[...,None] ], axis=-1 ),
                                np.concatenate ( [ np.array ( [[w]]*(w_pts_count+1) ), np.arange(0,h+h/h_pts_count, h/h_pts_count)[...,None] ], axis=-1 ) ), 0 ).astype(np.int32)


               
                img_cnt2 = np.concatenate ( (img_cnt2, perim_pts), 0 )
                img_cnt = np.concatenate ( (img_cnt, perim_pts), 0 )
                
                morphed_img_bgr = mls_affine_deformation_inv( img_bgr, img_cnt, img_cnt2 )
                
                while True:
                    cv2.imshow("", (img_bgr*255).astype(np.uint8) )
                    cv2.waitKey(0)
                    cv2.imshow("", (morphed_img_bgr*255).astype(np.uint8) )
                    cv2.waitKey(0)
            
                import code
                code.interact(local=dict(globals(), **locals()))

                

                out_img = img_bgr*(1-img_face_mask_a) + (out_img*img_face_mask_a)

                if ('seamless' in cfg.mode and cfg.color_transfer_mode != 0) or \
                   cfg.mode == 'seamless-hist-match' or \
                   cfg_mp != 0 or \
                   cfg.blursharpen_amount != 0 or \
                   cfg.image_denoise_power != 0 or \
                   cfg.bicubic_degrade_power != 0:

                    out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )

                    if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0:
                        if cfg.color_transfer_mode == 1:
                            out_face_bgr = imagelib.reinhard_color_transfer ( np.clip(out_face_bgr*wrk_face_mask_area_a*255, 0, 255).astype(np.uint8),
                                                                              np.clip(dst_face_bgr*wrk_face_mask_area_a*255, 0, 255).astype(np.uint8) )
                            out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
                        elif cfg.color_transfer_mode == 2: #lct
                            out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr)
                        elif cfg.color_transfer_mode == 3: #mkl
                            out_face_bgr = imagelib.color_transfer_mkl (out_face_bgr, dst_face_bgr)
                        elif cfg.color_transfer_mode == 4: #mkl-m
                            out_face_bgr = imagelib.color_transfer_mkl (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                        elif cfg.color_transfer_mode == 5: #idt
                            out_face_bgr = imagelib.color_transfer_idt (out_face_bgr, dst_face_bgr)
                        elif cfg.color_transfer_mode == 6: #idt-m
                            out_face_bgr = imagelib.color_transfer_idt (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                        elif cfg.color_transfer_mode == 7: #sot-m
                            out_face_bgr = imagelib.color_transfer_sot (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a, steps=10, batch_size=30)
                            out_face_bgr = np.clip (out_face_bgr, 0.0, 1.0)
                        elif cfg.color_transfer_mode == 8: #mix-m
                            out_face_bgr = imagelib.color_transfer_mix (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)

                    if cfg.mode == 'seamless-hist-match':
                        out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold)

                    if cfg_mp != 0:
                        k_size = int(frame_info.motion_power*cfg_mp)
                        if k_size >= 1:
                            k_size = np.clip (k_size+1, 2, 50)
                            if cfg.super_resolution_power != 0:
                                k_size *= 2
                            out_face_bgr = imagelib.LinearMotionBlur (out_face_bgr, k_size , frame_info.motion_deg)

                    if cfg.blursharpen_amount != 0:
                        out_face_bgr = imagelib.blursharpen ( out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount)

                    if cfg.image_denoise_power != 0:
                        n = cfg.image_denoise_power
                        while n > 0:
                            img_bgr_denoised = cv2.medianBlur(img_bgr, 5)
                            if int(n / 100) != 0:
                                img_bgr = img_bgr_denoised
                            else:
                                pass_power = (n % 100) / 100.0
                                img_bgr = img_bgr*(1.0-pass_power)+img_bgr_denoised*pass_power
                            n = max(n-10,0)

                    if cfg.bicubic_degrade_power != 0:
                        p = 1.0 - cfg.bicubic_degrade_power / 101.0
                        img_bgr_downscaled = cv2.resize (img_bgr, ( int(img_size[0]*p), int(img_size[1]*p ) ), interpolation=cv2.INTER_CUBIC)
                        img_bgr = cv2.resize (img_bgr_downscaled, img_size, interpolation=cv2.INTER_CUBIC)

                    new_out = cv2.warpAffine( out_face_bgr, face_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT )

                    out_img =  np.clip( img_bgr*(1-img_face_mask_a) + (new_out*img_face_mask_a) , 0, 1.0 )

                if cfg.color_degrade_power != 0:
                    out_img_reduced = imagelib.reduce_colors(out_img, 256)
                    if cfg.color_degrade_power == 100:
                        out_img = out_img_reduced
                    else:
                        alpha = cfg.color_degrade_power / 100.0
                        out_img = (out_img*(1.0-alpha) + out_img_reduced*alpha)

        out_merging_mask_a = img_face_mask_a

    return out_img, out_merging_mask_a
Ejemplo n.º 14
0
        def process_data(self, data):
            filename_path = Path(data[0])

            filename_path_str = str(filename_path)
            if self.cached_image[0] == filename_path_str:
                image = self.cached_image[
                    1]  #cached image for manual extractor
            else:
                image = cv2_imread(filename_path_str)

                if image is None:
                    self.log_err(
                        'Failed to extract %s, reason: cv2_imread() fail.' %
                        (str(filename_path)))
                    return None

                image_shape = image.shape
                if len(image_shape) == 2:
                    h, w = image.shape
                    ch = 1
                else:
                    h, w, ch = image.shape

                if ch == 1:
                    image = np.repeat(image[:, :, np.newaxis], 3, -1)
                elif ch == 4:
                    image = image[:, :, 0:3]

                wm = w % 2
                hm = h % 2
                if wm + hm != 0:  #fix odd image
                    image = image[0:h - hm, 0:w - wm, :]
                self.cached_image = (filename_path_str, image)

            src_dflimg = None
            h, w, ch = image.shape
            if h == w:
                #extracting from already extracted jpg image?
                if filename_path.suffix == '.jpg':
                    src_dflimg = DFLJPG.load(str(filename_path))

            if self.type == 'rects':
                if min(w, h) < 128:
                    self.log_err('Image is too small %s : [%d, %d]' %
                                 (str(filename_path), w, h))
                    rects = []
                else:
                    rects = self.e.extract_from_bgr(image)

                return [str(filename_path), rects]

            elif self.type == 'landmarks':
                rects = data[1]
                if rects is None:
                    landmarks = None
                else:
                    landmarks = self.e.extract_from_bgr(
                        image, rects,
                        self.second_pass_e if src_dflimg is None else None)

                return [str(filename_path), landmarks]

            elif self.type == 'final':

                result = []
                faces = data[1]

                if self.debug_dir is not None:
                    debug_output_file = str(
                        Path(self.debug_dir) / (filename_path.stem + '.jpg'))
                    debug_image = image.copy()

                if src_dflimg is not None and len(faces) != 1:
                    #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
                    print("src_dflimg is not None and len(faces) != 1",
                          str(filename_path))
                    output_file = str(self.output_path / filename_path.name)
                    if str(filename_path) != str(output_file):
                        shutil.copy(str(filename_path), str(output_file))
                    result.append(output_file)
                else:
                    face_idx = 0
                    for face in faces:
                        rect = np.array(face[0])
                        image_landmarks = face[1]
                        if image_landmarks is None:
                            continue
                        image_landmarks = np.array(image_landmarks)

                        if self.face_type == FaceType.MARK_ONLY:
                            face_image = image
                            face_image_landmarks = image_landmarks
                        else:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat(
                                image_landmarks, self.image_size,
                                self.face_type)
                            face_image = cv2.warpAffine(
                                image, image_to_face_mat,
                                (self.image_size, self.image_size),
                                cv2.INTER_LANCZOS4)
                            face_image_landmarks = LandmarksProcessor.transform_points(
                                image_landmarks, image_to_face_mat)

                            landmarks_bbox = LandmarksProcessor.transform_points(
                                [(0, 0), (0, self.image_size - 1),
                                 (self.image_size - 1, self.image_size - 1),
                                 (self.image_size - 1, 0)], image_to_face_mat,
                                True)

                            rect_area = mathlib.polygon_area(
                                np.array(rect[[0, 2, 2, 0]]),
                                np.array(rect[[1, 1, 3, 3]]))
                            landmarks_area = mathlib.polygon_area(
                                landmarks_bbox[:, 0], landmarks_bbox[:, 1])

                            if landmarks_area > 4 * rect_area:  #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                                continue

                        if self.debug_dir is not None:
                            LandmarksProcessor.draw_rect_landmarks(
                                debug_image,
                                rect,
                                image_landmarks,
                                self.image_size,
                                self.face_type,
                                transparent_mask=True)

                        if src_dflimg is not None:
                            #if extracting from dflimg copy it in order not to lose quality
                            output_file = str(self.output_path /
                                              filename_path.name)
                            if str(filename_path) != str(output_file):
                                shutil.copy(str(filename_path),
                                            str(output_file))
                        else:
                            output_file = '{}_{}{}'.format(
                                str(self.output_path / filename_path.stem),
                                str(face_idx), '.jpg')
                            cv2_imwrite(output_file, face_image,
                                        [int(cv2.IMWRITE_JPEG_QUALITY), 85])

                        DFLJPG.embed_data(
                            output_file,
                            face_type=FaceType.toString(self.face_type),
                            landmarks=face_image_landmarks.tolist(),
                            source_filename=filename_path.name,
                            source_rect=rect,
                            source_landmarks=image_landmarks.tolist(),
                            image_to_face_mat=image_to_face_mat)

                        result.append(output_file)
                        face_idx += 1

                if self.debug_dir is not None:
                    cv2_imwrite(debug_output_file, debug_image,
                                [int(cv2.IMWRITE_JPEG_QUALITY), 50])

                return result
Ejemplo n.º 15
0
def main(model_class_name=None,
         saved_models_path=None,
         training_data_src_path=None,
         force_model_name=None,
         input_path=None,
         output_path=None,
         output_mask_path=None,
         aligned_path=None,
         force_gpu_idxs=None,
         cpu_only=None):
    io.log_info("Running merger.\r\n")

    try:
        if not input_path.exists():
            io.log_err('Input directory not found. Please ensure it exists.')
            return

        if not output_path.exists():
            output_path.mkdir(parents=True, exist_ok=True)

        if not output_mask_path.exists():
            output_mask_path.mkdir(parents=True, exist_ok=True)

        if not saved_models_path.exists():
            io.log_err('Model directory not found. Please ensure it exists.')
            return

        # Initialize model
        import models
        model = models.import_model(model_class_name)(
            is_training=False,
            saved_models_path=saved_models_path,
            force_gpu_idxs=force_gpu_idxs,
            cpu_only=cpu_only)

        predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()

        # Preparing MP functions
        predictor_func = MPFunc(predictor_func)

        run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0
        xseg_256_extract_func = MPClassFuncOnDemand(
            XSegNet,
            'extract',
            name='XSeg',
            resolution=256,
            weights_file_root=saved_models_path,
            place_model_on_cpu=True,
            run_on_cpu=run_on_cpu)

        face_enhancer_func = MPClassFuncOnDemand(FaceEnhancer,
                                                 'enhance',
                                                 place_model_on_cpu=True,
                                                 run_on_cpu=run_on_cpu)

        is_interactive = io.input_bool("Use interactive merger?",
                                       True) if not io.is_colab() else False

        #         if not is_interactive:
        #             cfg.ask_settings()

        subprocess_count = multiprocessing.cpu_count()
        #         subprocess_count = io.input_int("Number of workers?", max(8, multiprocessing.cpu_count()),
        #                                         valid_range=[1, multiprocessing.cpu_count()], help_message="Specify the number of threads to process. A low value may affect performance. A high value may result in memory error. The value may not be greater than CPU cores." )

        input_path_image_paths = pathex.get_image_paths(input_path)

        if cfg.type == MergerConfig.TYPE_MASKED:
            if not aligned_path.exists():
                io.log_err(
                    'Aligned directory not found. Please ensure it exists.')
                return

            packed_samples = None
            try:
                packed_samples = samplelib.PackedFaceset.load(aligned_path)
            except:
                io.log_err(
                    f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
                )

            if packed_samples is not None:
                io.log_info("Using packed faceset.")

                def generator():
                    for sample in io.progress_bar_generator(
                            packed_samples, "Collecting alignments"):
                        filepath = Path(sample.filename)
                        yield filepath, DFLIMG.load(
                            filepath,
                            loader_func=lambda x: sample.read_raw_file())
            else:

                def generator():
                    for filepath in io.progress_bar_generator(
                            pathex.get_image_paths(aligned_path),
                            "Collecting alignments"):
                        filepath = Path(filepath)
                        yield filepath, DFLIMG.load(filepath)

            alignments = {}
            multiple_faces_detected = False

            for filepath, dflimg in generator():
                if dflimg is None or not dflimg.has_data():
                    io.log_err(f"{filepath.name} is not a dfl image file")
                    continue

                source_filename = dflimg.get_source_filename()
                if source_filename is None:
                    continue

                source_filepath = Path(source_filename)
                source_filename_stem = source_filepath.stem

                if source_filename_stem not in alignments.keys():
                    alignments[source_filename_stem] = []

                alignments_ar = alignments[source_filename_stem]
                alignments_ar.append(
                    (dflimg.get_source_landmarks(), filepath, source_filepath))

                if len(alignments_ar) > 1:
                    multiple_faces_detected = True

            if multiple_faces_detected:
                io.log_info("")
                io.log_info(
                    "Warning: multiple faces detected. Only one alignment file should refer one source file."
                )
                io.log_info("")

            for a_key in list(alignments.keys()):
                a_ar = alignments[a_key]
                if len(a_ar) > 1:
                    for _, filepath, source_filepath in a_ar:
                        io.log_info(
                            f"alignment {filepath.name} refers to {source_filepath.name} "
                        )
                    io.log_info("")

                alignments[a_key] = [a[0] for a in a_ar]

            if multiple_faces_detected:
                io.log_info(
                    "It is strongly recommended to process the faces separatelly."
                )
                io.log_info(
                    "Use 'recover original filename' to determine the exact duplicates."
                )
                io.log_info("")

            frames = [
                InteractiveMergerSubprocessor.Frame(frame_info=FrameInfo(
                    filepath=Path(p),
                    landmarks_list=alignments.get(Path(p).stem, None)))
                for p in input_path_image_paths
            ]

            if multiple_faces_detected:
                io.log_info(
                    "Warning: multiple faces detected. Motion blur will not be used."
                )
                io.log_info("")
            else:
                s = 256
                local_pts = [(s // 2 - 1, s // 2 - 1),
                             (s // 2 - 1, 0)]  #center+up
                frames_len = len(frames)
                for i in io.progress_bar_generator(range(len(frames)),
                                                   "Computing motion vectors"):
                    fi_prev = frames[max(0, i - 1)].frame_info
                    fi = frames[i].frame_info
                    fi_next = frames[min(i + 1, frames_len - 1)].frame_info
                    if len(fi_prev.landmarks_list) == 0 or \
                       len(fi.landmarks_list) == 0 or \
                       len(fi_next.landmarks_list) == 0:
                        continue

                    mat_prev = LandmarksProcessor.get_transform_mat(
                        fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
                    mat = LandmarksProcessor.get_transform_mat(
                        fi.landmarks_list[0], s, face_type=FaceType.FULL)
                    mat_next = LandmarksProcessor.get_transform_mat(
                        fi_next.landmarks_list[0], s, face_type=FaceType.FULL)

                    pts_prev = LandmarksProcessor.transform_points(
                        local_pts, mat_prev, True)
                    pts = LandmarksProcessor.transform_points(
                        local_pts, mat, True)
                    pts_next = LandmarksProcessor.transform_points(
                        local_pts, mat_next, True)

                    prev_vector = pts[0] - pts_prev[0]
                    next_vector = pts_next[0] - pts[0]

                    motion_vector = pts_next[0] - pts_prev[0]
                    fi.motion_power = npla.norm(motion_vector)

                    motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array(
                        [0, 0], dtype=np.float32)

                    fi.motion_deg = -math.atan2(
                        motion_vector[1], motion_vector[0]) * 180 / math.pi

        if len(frames) == 0:
            io.log_info("No frames to merge in input_dir.")
        else:
            if False:
                pass
            else:
                InteractiveMergerSubprocessor(
                    is_interactive=is_interactive,
                    merger_session_filepath=model.get_strpath_storage_for_file(
                        'merger_session.dat'),
                    predictor_func=predictor_func,
                    predictor_input_shape=predictor_input_shape,
                    face_enhancer_func=face_enhancer_func,
                    xseg_256_extract_func=xseg_256_extract_func,
                    merger_config=cfg,
                    frames=frames,
                    frames_root_path=input_path,
                    output_path=output_path,
                    output_mask_path=output_mask_path,
                    model_iter=model.get_iter(),
                    subprocess_count=subprocess_count,
                ).run()

        model.finalize()

    except Exception as e:
        print(traceback.format_exc())
Ejemplo n.º 16
0
def MergeMaskedFace(predictor_func, predictor_input_shape, cfg, frame_info,
                    img_bgr_uint8, img_bgr, img_face_landmarks):
    img_size = img_bgr.shape[1], img_bgr.shape[0]
    img_face_mask_a = LandmarksProcessor.get_image_hull_mask(
        img_bgr.shape, img_face_landmarks)

    if cfg.mode == 'original':
        return img_bgr, img_face_mask_a

    out_img = img_bgr.copy()
    out_merging_mask_a = None

    input_size = predictor_input_shape[0]
    mask_subres_size = input_size * 4
    output_size = input_size
    if cfg.super_resolution_power != 0:
        output_size *= 4

    face_mat = LandmarksProcessor.get_transform_mat(img_face_landmarks,
                                                    output_size,
                                                    face_type=cfg.face_type)
    face_output_mat = LandmarksProcessor.get_transform_mat(
        img_face_landmarks,
        output_size,
        face_type=cfg.face_type,
        scale=1.0 + 0.01 * cfg.output_face_scale)

    if mask_subres_size == output_size:
        face_mask_output_mat = face_output_mat
    else:
        face_mask_output_mat = LandmarksProcessor.get_transform_mat(
            img_face_landmarks,
            mask_subres_size,
            face_type=cfg.face_type,
            scale=1.0 + 0.01 * cfg.output_face_scale)

    dst_face_bgr = cv2.warpAffine(img_bgr,
                                  face_mat, (output_size, output_size),
                                  flags=cv2.INTER_CUBIC)
    dst_face_bgr = np.clip(dst_face_bgr, 0, 1)

    dst_face_mask_a_0 = cv2.warpAffine(img_face_mask_a,
                                       face_mat, (output_size, output_size),
                                       flags=cv2.INTER_CUBIC)
    dst_face_mask_a_0 = np.clip(dst_face_mask_a_0, 0, 1)

    predictor_input_bgr = cv2.resize(dst_face_bgr, (input_size, input_size))

    predicted = predictor_func(predictor_input_bgr)
    if isinstance(predicted, tuple):
        #merger return bgr,mask
        prd_face_bgr = np.clip(predicted[0], 0, 1.0)
        prd_face_mask_a_0 = np.clip(predicted[1], 0, 1.0)
        predictor_masked = True
    else:
        #merger return bgr only, using dst mask
        prd_face_bgr = np.clip(predicted, 0, 1.0)
        prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0,
                                       (input_size, input_size))
        predictor_masked = False

    if cfg.super_resolution_power != 0:
        prd_face_bgr_enhanced = cfg.superres_func(prd_face_bgr)
        mod = cfg.super_resolution_power / 100.0
        prd_face_bgr = cv2.resize(prd_face_bgr, (output_size, output_size)) * (
            1.0 - mod) + prd_face_bgr_enhanced * mod
        prd_face_bgr = np.clip(prd_face_bgr, 0, 1)

    if cfg.super_resolution_power != 0:
        if predictor_masked:
            prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0,
                                           (output_size, output_size),
                                           cv2.INTER_CUBIC)
        else:
            prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0,
                                           (output_size, output_size),
                                           cv2.INTER_CUBIC)

    if cfg.mask_mode == 2:  #dst
        prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0,
                                       (output_size, output_size),
                                       cv2.INTER_CUBIC)
    elif cfg.mask_mode >= 3 and cfg.mask_mode <= 8:

        if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6:
            prd_face_fanseg_bgr = cv2.resize(prd_face_bgr,
                                             (cfg.fanseg_input_size, ) * 2)
            prd_face_fanseg_mask = cfg.fanseg_extract_func(
                FaceType.FULL, prd_face_fanseg_bgr)
            FAN_prd_face_mask_a_0 = cv2.resize(prd_face_fanseg_mask,
                                               (output_size, output_size),
                                               cv2.INTER_CUBIC)

        if cfg.mask_mode >= 4 and cfg.mask_mode <= 7:

            full_face_fanseg_mat = LandmarksProcessor.get_transform_mat(
                img_face_landmarks,
                cfg.fanseg_input_size,
                face_type=FaceType.FULL)
            dst_face_fanseg_bgr = cv2.warpAffine(img_bgr,
                                                 full_face_fanseg_mat,
                                                 (cfg.fanseg_input_size, ) * 2,
                                                 flags=cv2.INTER_CUBIC)
            dst_face_fanseg_mask = cfg.fanseg_extract_func(
                FaceType.FULL, dst_face_fanseg_bgr)

            if cfg.face_type == FaceType.FULL:
                FAN_dst_face_mask_a_0 = cv2.resize(dst_face_fanseg_mask,
                                                   (output_size, output_size),
                                                   cv2.INTER_CUBIC)
            else:
                face_fanseg_mat = LandmarksProcessor.get_transform_mat(
                    img_face_landmarks,
                    cfg.fanseg_input_size,
                    face_type=cfg.face_type)

                fanseg_rect_corner_pts = np.array(
                    [[0, 0], [cfg.fanseg_input_size - 1, 0],
                     [0, cfg.fanseg_input_size - 1]],
                    dtype=np.float32)
                a = LandmarksProcessor.transform_points(fanseg_rect_corner_pts,
                                                        face_fanseg_mat,
                                                        invert=True)
                b = LandmarksProcessor.transform_points(
                    a, full_face_fanseg_mat)
                m = cv2.getAffineTransform(b, fanseg_rect_corner_pts)
                FAN_dst_face_mask_a_0 = cv2.warpAffine(
                    dst_face_fanseg_mask,
                    m, (cfg.fanseg_input_size, ) * 2,
                    flags=cv2.INTER_CUBIC)
                FAN_dst_face_mask_a_0 = cv2.resize(FAN_dst_face_mask_a_0,
                                                   (output_size, output_size),
                                                   cv2.INTER_CUBIC)

        if cfg.mask_mode == 3:  #FAN-prd
            prd_face_mask_a_0 = FAN_prd_face_mask_a_0
        elif cfg.mask_mode == 4:  #FAN-dst
            prd_face_mask_a_0 = FAN_dst_face_mask_a_0
        elif cfg.mask_mode == 5:
            prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
        elif cfg.mask_mode == 6:
            prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
        elif cfg.mask_mode == 7:
            prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0

    prd_face_mask_a_0[prd_face_mask_a_0 < (1.0 /
                                           255.0)] = 0.0  # get rid of noise

    # resize to mask_subres_size
    if prd_face_mask_a_0.shape[0] != mask_subres_size:
        prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0,
                                       (mask_subres_size, mask_subres_size),
                                       cv2.INTER_CUBIC)

    # process mask in local predicted space
    if 'raw' not in cfg.mode:
        # add zero pad
        prd_face_mask_a_0 = np.pad(prd_face_mask_a_0, input_size)

        ero = cfg.erode_mask_modifier
        blur = cfg.blur_mask_modifier

        if ero > 0:
            prd_face_mask_a_0 = cv2.erode(prd_face_mask_a_0,
                                          cv2.getStructuringElement(
                                              cv2.MORPH_ELLIPSE, (ero, ero)),
                                          iterations=1)
        elif ero < 0:
            prd_face_mask_a_0 = cv2.dilate(prd_face_mask_a_0,
                                           cv2.getStructuringElement(
                                               cv2.MORPH_ELLIPSE,
                                               (-ero, -ero)),
                                           iterations=1)

        # clip eroded/dilated mask in actual predict area
        # pad with half blur size in order to accuratelly fade to zero at the boundary
        clip_size = input_size + blur // 2

        prd_face_mask_a_0[:clip_size, :] = 0
        prd_face_mask_a_0[-clip_size:, :] = 0
        prd_face_mask_a_0[:, :clip_size] = 0
        prd_face_mask_a_0[:, -clip_size:] = 0

        if blur > 0:
            blur = blur + (1 - blur % 2)
            prd_face_mask_a_0 = cv2.GaussianBlur(prd_face_mask_a_0,
                                                 (blur, blur), 0)

        prd_face_mask_a_0 = prd_face_mask_a_0[input_size:-input_size,
                                              input_size:-input_size]

        prd_face_mask_a_0 = np.clip(prd_face_mask_a_0, 0, 1)

    img_face_mask_a = cv2.warpAffine(prd_face_mask_a_0,
                                     face_mask_output_mat,
                                     img_size,
                                     np.zeros(img_bgr.shape[0:2],
                                              dtype=np.float32),
                                     flags=cv2.WARP_INVERSE_MAP
                                     | cv2.INTER_CUBIC)[..., None]
    img_face_mask_a = np.clip(img_face_mask_a, 0.0, 1.0)

    img_face_mask_a[img_face_mask_a < (1.0 / 255.0)] = 0.0  # get rid of noise

    if prd_face_mask_a_0.shape[0] != output_size:
        prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0,
                                       (output_size, output_size),
                                       cv2.INTER_CUBIC)

    prd_face_mask_a = prd_face_mask_a_0[..., None]
    prd_face_mask_area_a = prd_face_mask_a.copy()
    prd_face_mask_area_a[prd_face_mask_area_a > 0] = 1.0

    if 'raw' in cfg.mode:
        if cfg.mode == 'raw-rgb':
            out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size,
                                     out_img,
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)
            out_merging_mask_a = img_face_mask_a

        out_img = np.clip(out_img, 0.0, 1.0)
    else:
        #averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask
        ar = []
        for i in range(1, 10):
            maxregion = np.argwhere(img_face_mask_a > i / 10.0)
            if maxregion.size != 0:
                miny, minx = maxregion.min(axis=0)[:2]
                maxy, maxx = maxregion.max(axis=0)[:2]
                lenx = maxx - minx
                leny = maxy - miny
                if min(lenx, leny) >= 4:
                    ar += [[lenx, leny]]

        if len(ar) > 0:

            if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0:
                if cfg.color_transfer_mode == 1:  #rct
                    prd_face_bgr = imagelib.reinhard_color_transfer(
                        np.clip(prd_face_bgr * prd_face_mask_area_a * 255, 0,
                                255).astype(np.uint8),
                        np.clip(dst_face_bgr * prd_face_mask_area_a * 255, 0,
                                255).astype(np.uint8),
                    )

                    prd_face_bgr = np.clip(
                        prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
                elif cfg.color_transfer_mode == 2:  #lct
                    prd_face_bgr = imagelib.linear_color_transfer(
                        prd_face_bgr, dst_face_bgr)
                elif cfg.color_transfer_mode == 3:  #mkl
                    prd_face_bgr = imagelib.color_transfer_mkl(
                        prd_face_bgr, dst_face_bgr)
                elif cfg.color_transfer_mode == 4:  #mkl-m
                    prd_face_bgr = imagelib.color_transfer_mkl(
                        prd_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)
                elif cfg.color_transfer_mode == 5:  #idt
                    prd_face_bgr = imagelib.color_transfer_idt(
                        prd_face_bgr, dst_face_bgr)
                elif cfg.color_transfer_mode == 6:  #idt-m
                    prd_face_bgr = imagelib.color_transfer_idt(
                        prd_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)
                elif cfg.color_transfer_mode == 7:  #sot-m
                    prd_face_bgr = imagelib.color_transfer_sot(
                        prd_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)
                    prd_face_bgr = np.clip(prd_face_bgr, 0.0, 1.0)
                elif cfg.color_transfer_mode == 8:  #mix-m
                    prd_face_bgr = imagelib.color_transfer_mix(
                        prd_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)

            if cfg.mode == 'hist-match':
                hist_mask_a = np.ones(prd_face_bgr.shape[:2] + (1, ),
                                      dtype=np.float32)

                if cfg.masked_hist_match:
                    hist_mask_a *= prd_face_mask_area_a

                white = (1.0 - hist_mask_a) * np.ones(
                    prd_face_bgr.shape[:2] + (1, ), dtype=np.float32)

                hist_match_1 = prd_face_bgr * hist_mask_a + white
                hist_match_1[hist_match_1 > 1.0] = 1.0

                hist_match_2 = dst_face_bgr * hist_mask_a + white
                hist_match_2[hist_match_1 > 1.0] = 1.0

                prd_face_bgr = imagelib.color_hist_match(
                    hist_match_1, hist_match_2,
                    cfg.hist_match_threshold).astype(dtype=np.float32)

            if 'seamless' in cfg.mode:
                #mask used for cv2.seamlessClone
                img_face_seamless_mask_a = None
                for i in range(1, 10):
                    a = img_face_mask_a > i / 10.0
                    if len(np.argwhere(a)) == 0:
                        continue
                    img_face_seamless_mask_a = img_face_mask_a.copy()
                    img_face_seamless_mask_a[a] = 1.0
                    img_face_seamless_mask_a[img_face_seamless_mask_a <= i /
                                             10.0] = 0.0
                    break

            out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size,
                                     out_img,
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)

            out_img = np.clip(out_img, 0.0, 1.0)

            if 'seamless' in cfg.mode:
                try:
                    #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering)
                    l, t, w, h = cv2.boundingRect(
                        (img_face_seamless_mask_a * 255).astype(np.uint8))
                    s_maskx, s_masky = int(l + w / 2), int(t + h / 2)
                    out_img = cv2.seamlessClone(
                        (out_img * 255).astype(np.uint8), img_bgr_uint8,
                        (img_face_seamless_mask_a * 255).astype(np.uint8),
                        (s_maskx, s_masky), cv2.NORMAL_CLONE)
                    out_img = out_img.astype(dtype=np.float32) / 255.0
                except Exception as e:
                    #seamlessClone may fail in some cases
                    e_str = traceback.format_exc()

                    if 'MemoryError' in e_str:
                        raise Exception(
                            "Seamless fail: " + e_str
                        )  #reraise MemoryError in order to reprocess this data by other processes
                    else:
                        print("Seamless fail: " + e_str)

            out_img = img_bgr * (1 - img_face_mask_a) + (out_img *
                                                         img_face_mask_a)

            out_face_bgr = cv2.warpAffine(out_img,
                                          face_mat, (output_size, output_size),
                                          flags=cv2.INTER_CUBIC)

            if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0:
                if cfg.color_transfer_mode == 1:
                    out_face_bgr = imagelib.reinhard_color_transfer(
                        np.clip(out_face_bgr * prd_face_mask_area_a * 255, 0,
                                255).astype(np.uint8),
                        np.clip(dst_face_bgr * prd_face_mask_area_a * 255, 0,
                                255).astype(np.uint8))
                    out_face_bgr = np.clip(
                        out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
                elif cfg.color_transfer_mode == 2:  #lct
                    out_face_bgr = imagelib.linear_color_transfer(
                        out_face_bgr, dst_face_bgr)
                elif cfg.color_transfer_mode == 3:  #mkl
                    out_face_bgr = imagelib.color_transfer_mkl(
                        out_face_bgr, dst_face_bgr)
                elif cfg.color_transfer_mode == 4:  #mkl-m
                    out_face_bgr = imagelib.color_transfer_mkl(
                        out_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)
                elif cfg.color_transfer_mode == 5:  #idt
                    out_face_bgr = imagelib.color_transfer_idt(
                        out_face_bgr, dst_face_bgr)
                elif cfg.color_transfer_mode == 6:  #idt-m
                    out_face_bgr = imagelib.color_transfer_idt(
                        out_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)
                elif cfg.color_transfer_mode == 7:  #sot-m
                    out_face_bgr = imagelib.color_transfer_sot(
                        out_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)
                    out_face_bgr = np.clip(out_face_bgr, 0.0, 1.0)
                elif cfg.color_transfer_mode == 8:  #mix-m
                    out_face_bgr = imagelib.color_transfer_mix(
                        out_face_bgr * prd_face_mask_area_a,
                        dst_face_bgr * prd_face_mask_area_a)

            if cfg.mode == 'seamless-hist-match':
                out_face_bgr = imagelib.color_hist_match(
                    out_face_bgr, dst_face_bgr, cfg.hist_match_threshold)

            cfg_mp = cfg.motion_blur_power / 100.0
            if cfg_mp != 0:
                k_size = int(frame_info.motion_power * cfg_mp)
                if k_size >= 1:
                    k_size = np.clip(k_size + 1, 2, 50)
                    if cfg.super_resolution_power != 0:
                        k_size *= 2
                    out_face_bgr = imagelib.LinearMotionBlur(
                        out_face_bgr, k_size, frame_info.motion_deg)

            if cfg.blursharpen_amount != 0:
                out_face_bgr = cfg.blursharpen_func(out_face_bgr,
                                                    cfg.sharpen_mode, 3,
                                                    cfg.blursharpen_amount)

            if cfg.image_denoise_power != 0:
                n = cfg.image_denoise_power
                while n > 0:
                    img_bgr_denoised = cv2.medianBlur(img_bgr, 5)
                    if int(n / 100) != 0:
                        img_bgr = img_bgr_denoised
                    else:
                        pass_power = (n % 100) / 100.0
                        img_bgr = img_bgr * (
                            1.0 - pass_power) + img_bgr_denoised * pass_power
                    n = max(n - 10, 0)

            if cfg.bicubic_degrade_power != 0:
                p = 1.0 - cfg.bicubic_degrade_power / 101.0
                img_bgr_downscaled = cv2.resize(
                    img_bgr, (int(img_size[0] * p), int(img_size[1] * p)),
                    cv2.INTER_CUBIC)
                img_bgr = cv2.resize(img_bgr_downscaled, img_size,
                                     cv2.INTER_CUBIC)

            new_out = cv2.warpAffine(out_face_bgr, face_mat, img_size,
                                     img_bgr.copy(),
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)
            out_img = np.clip(
                img_bgr * (1 - img_face_mask_a) + (new_out * img_face_mask_a),
                0, 1.0)

            if cfg.color_degrade_power != 0:
                out_img_reduced = imagelib.reduce_colors(out_img, 256)
                if cfg.color_degrade_power == 100:
                    out_img = out_img_reduced
                else:
                    alpha = cfg.color_degrade_power / 100.0
                    out_img = (out_img * (1.0 - alpha) +
                               out_img_reduced * alpha)

        out_merging_mask_a = img_face_mask_a

    return out_img, out_merging_mask_a
Ejemplo n.º 17
0
def main (args, device_args):
    io.log_info ("Running converter.\r\n")

    aligned_dir = args.get('aligned_dir', None)
    avaperator_aligned_dir = args.get('avaperator_aligned_dir', None)

    try:
        input_path = Path(args['input_dir'])
        output_path = Path(args['output_dir'])
        model_path = Path(args['model_dir'])

        if not input_path.exists():
            io.log_err('Input directory not found. Please ensure it exists.')
            return

        if output_path.exists():
            for filename in Path_utils.get_image_paths(output_path):
                Path(filename).unlink()
        else:
            output_path.mkdir(parents=True, exist_ok=True)

        if not model_path.exists():
            io.log_err('Model directory not found. Please ensure it exists.')
            return

        is_interactive = io.input_bool ("Use interactive converter? (y/n skip:y) : ", True) if not io.is_colab() else False

        import models
        model = models.import_model( args['model_name'] )(model_path, device_args=device_args)

        cfg = model.get_ConverterConfig()

        if not is_interactive:
            cfg.ask_settings()

        input_path_image_paths = Path_utils.get_image_paths(input_path)

        if cfg.type == ConverterConfig.TYPE_MASKED:
            if aligned_dir is None:
                io.log_err('Aligned directory not found. Please ensure it exists.')
                return

            aligned_path = Path(aligned_dir)
            if not aligned_path.exists():
                io.log_err('Aligned directory not found. Please ensure it exists.')
                return

            alignments = {}
            multiple_faces_detected = False
            aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
            for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"):
                filepath = Path(filepath)

                if filepath.suffix == '.png':
                    dflimg = DFLPNG.load( str(filepath) )
                elif filepath.suffix == '.jpg':
                    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

                source_filename_stem = Path( dflimg.get_source_filename() ).stem
                if source_filename_stem not in alignments.keys():
                    alignments[ source_filename_stem ] = []

                alignments_ar = alignments[ source_filename_stem ]
                alignments_ar.append (dflimg.get_source_landmarks())
                if len(alignments_ar) > 1:
                    multiple_faces_detected = True

            if multiple_faces_detected:
                io.log_info ("Warning: multiple faces detected. Strongly recommended to process them separately.")

            frames = [ ConvertSubprocessor.Frame( frame_info=FrameInfo(filename=p, landmarks_list=alignments.get(Path(p).stem, None))) for p in input_path_image_paths ]

            if multiple_faces_detected:
                io.log_info ("Warning: multiple faces detected. Motion blur will not be used.")
            else:
                s = 256
                local_pts = [ (s//2-1, s//2-1), (s//2-1,0) ] #center+up
                frames_len = len(frames)
                for i in io.progress_bar_generator( range(len(frames)) , "Computing motion vectors"):
                    fi_prev = frames[max(0, i-1)].frame_info
                    fi      = frames[i].frame_info
                    fi_next = frames[min(i+1, frames_len-1)].frame_info
                    if len(fi_prev.landmarks_list) == 0 or \
                       len(fi.landmarks_list) == 0 or \
                       len(fi_next.landmarks_list) == 0:
                            continue

                    mat_prev = LandmarksProcessor.get_transform_mat ( fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
                    mat      = LandmarksProcessor.get_transform_mat ( fi.landmarks_list[0]     , s, face_type=FaceType.FULL)
                    mat_next = LandmarksProcessor.get_transform_mat ( fi_next.landmarks_list[0], s, face_type=FaceType.FULL)

                    pts_prev = LandmarksProcessor.transform_points (local_pts, mat_prev, True)
                    pts      = LandmarksProcessor.transform_points (local_pts, mat, True)
                    pts_next = LandmarksProcessor.transform_points (local_pts, mat_next, True)

                    prev_vector = pts[0]-pts_prev[0]
                    next_vector = pts_next[0]-pts[0]

                    motion_vector = pts_next[0] - pts_prev[0]
                    fi.motion_power = npla.norm(motion_vector)

                    motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array([0,0],dtype=np.float32)

                    fi.motion_deg = -math.atan2(motion_vector[1],motion_vector[0])*180 / math.pi


        elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR:
            filesdata = []
            for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"):
                filepath = Path(filepath)

                if filepath.suffix == '.png':
                    dflimg = DFLPNG.load( str(filepath) )
                elif filepath.suffix == '.jpg':
                    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
                filesdata += [ ( FrameInfo(filename=str(filepath), landmarks_list=[dflimg.get_landmarks()] ), dflimg.get_source_filename() ) ]

            filesdata = sorted(filesdata, key=operator.itemgetter(1)) #sort by filename
            frames = []
            filesdata_len = len(filesdata)
            for i in range(len(filesdata)):
                frame_info = filesdata[i][0]

                prev_temporal_frame_infos = []
                next_temporal_frame_infos = []

                for t in range (cfg.temporal_face_count):
                    prev_frame_info = filesdata[ max(i -t, 0) ][0]
                    next_frame_info = filesdata[ min(i +t, filesdata_len-1 )][0]

                    prev_temporal_frame_infos.insert (0, prev_frame_info )
                    next_temporal_frame_infos.append (   next_frame_info )

                frames.append ( ConvertSubprocessor.Frame(prev_temporal_frame_infos=prev_temporal_frame_infos,
                                                          frame_info=frame_info,
                                                          next_temporal_frame_infos=next_temporal_frame_infos) )

        if len(frames) == 0:
            io.log_info ("No frames to convert in input_dir.")
        else:
            ConvertSubprocessor (
                        is_interactive         = is_interactive,
                        converter_config       = cfg,
                        frames                 = frames,
                        output_path            = output_path,
                    ).run()

        model.finalize()

    except Exception as e:
        print ( 'Error: %s' % (str(e)))
        traceback.print_exc()
Ejemplo n.º 18
0
def dev_segmented_extract(input_dir, output_dir ):
    # extract and merge .json labelme files within the faces

    device_config = nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(suggest_all_gpu=True) )

    input_path = Path(input_dir)
    if not input_path.exists():
        raise ValueError('input_dir not found. Please ensure it exists.')

    output_path = Path(output_dir)
    io.log_info("Performing extract segmented faces.")
    io.log_info(f'Output dir is {output_path}')

    if output_path.exists():
        output_images_paths = pathex.get_image_paths(output_path, subdirs=True)
        if len(output_images_paths) > 0:
            io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
            for filename in output_images_paths:
                Path(filename).unlink()
            shutil.rmtree(str(output_path))
    else:
        output_path.mkdir(parents=True, exist_ok=True)

    images_paths = pathex.get_image_paths(input_path, subdirs=True, return_Path_class=True)

    extract_data = []
    images_jsons = {}
    images_processed = 0


    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        json_filepath = filepath.parent / (filepath.stem+'.json')

        if json_filepath.exists():
            try:
                json_dict = json.loads(json_filepath.read_text())
                images_jsons[filepath] = json_dict

                total_points = [ [x,y] for shape in json_dict['shapes'] for x,y in shape['points'] ]
                total_points = np.array(total_points)

                if len(total_points) == 0:
                    io.log_info(f"No points found in {json_filepath}, skipping.")
                    continue

                l,r = int(total_points[:,0].min()), int(total_points[:,0].max())
                t,b = int(total_points[:,1].min()), int(total_points[:,1].max())

                force_output_path=output_path / filepath.relative_to(input_path).parent
                force_output_path.mkdir(exist_ok=True, parents=True)

                extract_data.append ( ExtractSubprocessor.Data(filepath,
                                                               rects=[ [l,t,r,b] ],
                                                               force_output_path=force_output_path ) )
                images_processed += 1
            except:
                io.log_err(f"err {filepath}, {traceback.format_exc()}")
                return
        else:
            io.log_info(f"No .json file for {filepath.relative_to(input_path)}, skipping.")
            continue

    image_size = 1024
    face_type = FaceType.HEAD
    extract_data = ExtractSubprocessor (extract_data, 'landmarks', image_size, face_type, device_config=device_config).run()
    extract_data = ExtractSubprocessor (extract_data, 'final', image_size, face_type, device_config=device_config).run()

    for data in extract_data:
        filepath = data.force_output_path / (data.filepath.stem+'_0.jpg')

        dflimg = DFLIMG.load(filepath)
        image_to_face_mat = dflimg.get_image_to_face_mat()

        json_dict = images_jsons[data.filepath]

        ie_polys = IEPolys()
        for shape in json_dict['shapes']:
            ie_poly = ie_polys.add(1)

            points = np.array( [ [x,y] for x,y in shape['points'] ] )
            points = LandmarksProcessor.transform_points(points, image_to_face_mat)

            for x,y in points:
                ie_poly.add( int(x), int(y) )

        dflimg.embed_and_set (filepath, ie_polys=ie_polys)

    io.log_info(f"Images found:     {len(images_paths)}")
    io.log_info(f"Images processed: {images_processed}")
Ejemplo n.º 19
0
def ConvertMaskedFace(cfg, frame_info, img_bgr_uint8, img_bgr,
                      img_face_landmarks):

    #if debug:
    #    debugs = [img_bgr.copy()]

    img_size = img_bgr.shape[1], img_bgr.shape[0]

    img_face_mask_a = LandmarksProcessor.get_image_hull_mask(
        img_bgr.shape, img_face_landmarks)

    if cfg.mode == 'original':
        if cfg.export_mask_alpha:
            img_bgr = np.concatenate([img_bgr, img_face_mask_a], -1)
        return img_bgr, img_face_mask_a

    out_img = img_bgr.copy()
    out_merging_mask = None

    output_size = cfg.predictor_input_shape[0]
    if cfg.super_resolution_mode != 0:
        output_size *= 2

    face_mat = LandmarksProcessor.get_transform_mat(img_face_landmarks,
                                                    output_size,
                                                    face_type=cfg.face_type)
    face_output_mat = LandmarksProcessor.get_transform_mat(
        img_face_landmarks,
        output_size,
        face_type=cfg.face_type,
        scale=1.0 + 0.01 * cfg.output_face_scale)

    dst_face_bgr = cv2.warpAffine(img_bgr,
                                  face_mat, (output_size, output_size),
                                  flags=cv2.INTER_CUBIC)
    dst_face_mask_a_0 = cv2.warpAffine(img_face_mask_a,
                                       face_mat, (output_size, output_size),
                                       flags=cv2.INTER_CUBIC)

    predictor_input_bgr = cv2.resize(dst_face_bgr,
                                     cfg.predictor_input_shape[0:2])

    if cfg.predictor_masked:
        prd_face_bgr, prd_face_mask_a_0 = cfg.predictor_func(
            predictor_input_bgr)

        prd_face_bgr = np.clip(prd_face_bgr, 0, 1.0)
        prd_face_mask_a_0 = np.clip(prd_face_mask_a_0, 0.0, 1.0)
    else:
        predicted = cfg.predictor_func(predictor_input_bgr)
        prd_face_bgr = np.clip(predicted, 0, 1.0)
        prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0,
                                       cfg.predictor_input_shape[0:2])

    if cfg.super_resolution_mode:
        #if debug:
        #    tmp = cv2.resize (prd_face_bgr,  (output_size,output_size), cv2.INTER_CUBIC)
        #    debugs += [ np.clip( cv2.warpAffine( tmp, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

        prd_face_bgr = cfg.superres_func(cfg.super_resolution_mode,
                                         prd_face_bgr)
        #if debug:
        #    debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

        if cfg.predictor_masked:
            prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0,
                                           (output_size, output_size),
                                           cv2.INTER_CUBIC)
        else:
            prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0,
                                           (output_size, output_size),
                                           cv2.INTER_CUBIC)

    if cfg.mask_mode == 2:  #dst
        prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0,
                                       (output_size, output_size),
                                       cv2.INTER_CUBIC)
    elif cfg.mask_mode >= 3 and cfg.mask_mode <= 7:

        if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6:
            prd_face_fanseg_bgr = cv2.resize(prd_face_bgr,
                                             (cfg.fanseg_input_size, ) * 2)
            prd_face_fanseg_mask = cfg.fanseg_extract_func(
                FaceType.FULL, prd_face_fanseg_bgr)
            FAN_prd_face_mask_a_0 = cv2.resize(prd_face_fanseg_mask,
                                               (output_size, output_size),
                                               cv2.INTER_CUBIC)

        if cfg.mask_mode >= 4 or cfg.mask_mode <= 7:

            full_face_fanseg_mat = LandmarksProcessor.get_transform_mat(
                img_face_landmarks,
                cfg.fanseg_input_size,
                face_type=FaceType.FULL)
            dst_face_fanseg_bgr = cv2.warpAffine(img_bgr,
                                                 full_face_fanseg_mat,
                                                 (cfg.fanseg_input_size, ) * 2,
                                                 flags=cv2.INTER_CUBIC)
            dst_face_fanseg_mask = cfg.fanseg_extract_func(
                FaceType.FULL, dst_face_fanseg_bgr)

            if cfg.face_type == FaceType.FULL:
                FAN_dst_face_mask_a_0 = cv2.resize(dst_face_fanseg_mask,
                                                   (output_size, output_size),
                                                   cv2.INTER_CUBIC)
            elif cfg.face_type == FaceType.HALF:
                half_face_fanseg_mat = LandmarksProcessor.get_transform_mat(
                    img_face_landmarks,
                    cfg.fanseg_input_size,
                    face_type=FaceType.HALF)

                fanseg_rect_corner_pts = np.array(
                    [[0, 0], [cfg.fanseg_input_size - 1, 0],
                     [0, cfg.fanseg_input_size - 1]],
                    dtype=np.float32)
                a = LandmarksProcessor.transform_points(fanseg_rect_corner_pts,
                                                        half_face_fanseg_mat,
                                                        invert=True)
                b = LandmarksProcessor.transform_points(
                    a, full_face_fanseg_mat)
                m = cv2.getAffineTransform(b, fanseg_rect_corner_pts)
                FAN_dst_face_mask_a_0 = cv2.warpAffine(
                    dst_face_fanseg_mask,
                    m, (cfg.fanseg_input_size, ) * 2,
                    flags=cv2.INTER_CUBIC)
                FAN_dst_face_mask_a_0 = cv2.resize(FAN_dst_face_mask_a_0,
                                                   (output_size, output_size),
                                                   cv2.INTER_CUBIC)
            else:
                raise ValueError("cfg.face_type unsupported")

        if cfg.mask_mode == 3:  #FAN-prd
            prd_face_mask_a_0 = FAN_prd_face_mask_a_0
        elif cfg.mask_mode == 4:  #FAN-dst
            prd_face_mask_a_0 = FAN_dst_face_mask_a_0
        elif cfg.mask_mode == 5:
            prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
        elif cfg.mask_mode == 6:
            prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
        elif cfg.mask_mode == 7:
            prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0

    prd_face_mask_a_0[prd_face_mask_a_0 < 0.001] = 0.0

    prd_face_mask_a = prd_face_mask_a_0[..., np.newaxis]
    prd_face_mask_aaa = np.repeat(prd_face_mask_a, (3, ), axis=-1)

    img_face_mask_aaa = cv2.warpAffine(prd_face_mask_aaa,
                                       face_output_mat,
                                       img_size,
                                       np.zeros(img_bgr.shape,
                                                dtype=np.float32),
                                       flags=cv2.WARP_INVERSE_MAP
                                       | cv2.INTER_CUBIC)
    img_face_mask_aaa = np.clip(img_face_mask_aaa, 0.0, 1.0)
    img_face_mask_aaa[img_face_mask_aaa <= 0.1] = 0.0  #get rid of noise

    #if debug:
    #    debugs += [img_face_mask_aaa.copy()]

    if 'raw' in cfg.mode:
        face_corner_pts = np.array(
            [[0, 0], [output_size - 1, 0], [output_size - 1, output_size - 1],
             [0, output_size - 1]],
            dtype=np.float32)
        square_mask = np.zeros(img_bgr.shape, dtype=np.float32)
        cv2.fillConvexPoly(square_mask, \
                           LandmarksProcessor.transform_points (face_corner_pts, face_output_mat, invert=True ).astype(np.int), \
                           (1,1,1) )

        if cfg.mode == 'raw-rgb':
            out_merging_mask = square_mask

        if cfg.mode == 'raw-rgb' or cfg.mode == 'raw-rgb-mask':
            out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size,
                                     out_img,
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)

        if cfg.mode == 'raw-rgb-mask':
            out_img = np.concatenate(
                [out_img,
                 np.expand_dims(img_face_mask_aaa[:, :, 0], -1)], -1)
            out_merging_mask = square_mask

        elif cfg.mode == 'raw-mask-only':
            out_img = img_face_mask_aaa
            out_merging_mask = img_face_mask_aaa
        elif cfg.mode == 'raw-predicted-only':
            out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size,
                                     np.zeros(img_bgr.shape, dtype=np.float32),
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)
            out_merging_mask = square_mask

        out_img = np.clip(out_img, 0.0, 1.0)
    else:
        #averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask
        ar = []
        for i in range(1, 10):
            maxregion = np.argwhere(img_face_mask_aaa > i / 10.0)
            if maxregion.size != 0:
                miny, minx = maxregion.min(axis=0)[:2]
                maxy, maxx = maxregion.max(axis=0)[:2]
                lenx = maxx - minx
                leny = maxy - miny
                if min(lenx, leny) >= 4:
                    ar += [[lenx, leny]]

        if len(ar) > 0:
            lenx, leny = np.mean(ar, axis=0)
            lowest_len = min(lenx, leny)
            #if debug:
            #    io.log_info ("lenx/leny:(%d/%d) " % (lenx, leny  ) )
            #    io.log_info ("lowest_len = %f" % (lowest_len) )

            if cfg.erode_mask_modifier != 0:
                ero = int(lowest_len * (0.126 - lowest_len * 0.00004551365) *
                          0.01 * cfg.erode_mask_modifier)
                #if debug:
                #    io.log_info ("erode_size = %d" % (ero) )
                if ero > 0:
                    img_face_mask_aaa = cv2.erode(img_face_mask_aaa,
                                                  cv2.getStructuringElement(
                                                      cv2.MORPH_ELLIPSE,
                                                      (ero, ero)),
                                                  iterations=1)
                elif ero < 0:
                    img_face_mask_aaa = cv2.dilate(img_face_mask_aaa,
                                                   cv2.getStructuringElement(
                                                       cv2.MORPH_ELLIPSE,
                                                       (-ero, -ero)),
                                                   iterations=1)

            if cfg.clip_hborder_mask_per > 0:  #clip hborder before blur
                prd_hborder_rect_mask_a = np.ones(prd_face_mask_a.shape,
                                                  dtype=np.float32)
                prd_border_size = int(prd_hborder_rect_mask_a.shape[1] *
                                      cfg.clip_hborder_mask_per)
                prd_hborder_rect_mask_a[:, 0:prd_border_size, :] = 0
                prd_hborder_rect_mask_a[:, -prd_border_size:, :] = 0
                prd_hborder_rect_mask_a[-prd_border_size:, :, :] = 0
                prd_hborder_rect_mask_a = np.expand_dims(
                    cv2.blur(prd_hborder_rect_mask_a,
                             (prd_border_size, prd_border_size)), -1)

                img_prd_hborder_rect_mask_a = cv2.warpAffine(
                    prd_hborder_rect_mask_a, face_output_mat, img_size,
                    np.zeros(img_bgr.shape, dtype=np.float32),
                    cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC)
                img_prd_hborder_rect_mask_a = np.expand_dims(
                    img_prd_hborder_rect_mask_a, -1)
                img_face_mask_aaa *= img_prd_hborder_rect_mask_a
                img_face_mask_aaa = np.clip(img_face_mask_aaa, 0, 1.0)

                #if debug:
                #    debugs += [img_face_mask_aaa.copy()]

            if cfg.blur_mask_modifier > 0:
                blur = int(lowest_len * 0.10 * 0.01 * cfg.blur_mask_modifier)
                #if debug:
                #    io.log_info ("blur_size = %d" % (blur) )
                if blur > 0:
                    img_face_mask_aaa = cv2.blur(img_face_mask_aaa,
                                                 (blur, blur))

            img_face_mask_aaa = np.clip(img_face_mask_aaa, 0, 1.0)

            #if debug:
            #    debugs += [img_face_mask_aaa.copy()]

            if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0:
                if cfg.color_transfer_mode == 1:
                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

                    prd_face_bgr = imagelib.reinhard_color_transfer(
                        np.clip((prd_face_bgr * 255).astype(np.uint8), 0, 255),
                        np.clip((dst_face_bgr * 255).astype(np.uint8), 0, 255),
                        source_mask=prd_face_mask_a,
                        target_mask=prd_face_mask_a)
                    prd_face_bgr = np.clip(
                        prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)

                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

                elif cfg.color_transfer_mode == 2:
                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

                    prd_face_bgr = imagelib.linear_color_transfer(
                        prd_face_bgr, dst_face_bgr)
                    prd_face_bgr = np.clip(prd_face_bgr, 0.0, 1.0)

                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

            if cfg.mode == 'hist-match-bw':
                prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY)
                prd_face_bgr = np.repeat(np.expand_dims(prd_face_bgr, -1),
                                         (3, ), -1)

            if cfg.mode == 'hist-match' or cfg.mode == 'hist-match-bw':
                #if debug:
                #    debugs += [ cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) ]

                hist_mask_a = np.ones(prd_face_bgr.shape[:2] + (1, ),
                                      dtype=np.float32)

                if cfg.masked_hist_match:
                    hist_mask_a *= prd_face_mask_a

                white = (1.0 - hist_mask_a) * np.ones(
                    prd_face_bgr.shape[:2] + (1, ), dtype=np.float32)

                hist_match_1 = prd_face_bgr * hist_mask_a + white
                hist_match_1[hist_match_1 > 1.0] = 1.0

                hist_match_2 = dst_face_bgr * hist_mask_a + white
                hist_match_2[hist_match_1 > 1.0] = 1.0

                prd_face_bgr = imagelib.color_hist_match(
                    hist_match_1, hist_match_2, cfg.hist_match_threshold)

                #if cfg.masked_hist_match:
                #    prd_face_bgr -= white

            if cfg.mode == 'hist-match-bw':
                prd_face_bgr = prd_face_bgr.astype(dtype=np.float32)

            out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size,
                                     out_img,
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)
            out_img = np.clip(out_img, 0.0, 1.0)

            #if debug:
            #    debugs += [out_img.copy()]

            if cfg.mode == 'overlay':
                pass

            if 'seamless' in cfg.mode:
                #mask used for cv2.seamlessClone
                img_face_seamless_mask_a = None
                img_face_mask_a = img_face_mask_aaa[..., 0:1]
                for i in range(1, 10):
                    a = img_face_mask_a > i / 10.0
                    if len(np.argwhere(a)) == 0:
                        continue
                    img_face_seamless_mask_a = img_face_mask_aaa[...,
                                                                 0:1].copy()
                    img_face_seamless_mask_a[a] = 1.0
                    img_face_seamless_mask_a[img_face_seamless_mask_a <= i /
                                             10.0] = 0.0
                    break

                try:
                    #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering)
                    l, t, w, h = cv2.boundingRect(
                        (img_face_seamless_mask_a * 255).astype(np.uint8))
                    s_maskx, s_masky = int(l + w / 2), int(t + h / 2)

                    out_img = cv2.seamlessClone(
                        (out_img * 255).astype(np.uint8), img_bgr_uint8,
                        (img_face_seamless_mask_a * 255).astype(np.uint8),
                        (s_maskx, s_masky), cv2.NORMAL_CLONE)
                    out_img = out_img.astype(dtype=np.float32) / 255.0
                except Exception as e:
                    #seamlessClone may fail in some cases
                    e_str = traceback.format_exc()

                    if 'MemoryError' in e_str:
                        raise Exception(
                            "Seamless fail: " + e_str
                        )  #reraise MemoryError in order to reprocess this data by other processes
                    else:
                        print("Seamless fail: " + e_str)

                #if debug:
                #    debugs += [out_img.copy()]

            out_img = img_bgr * (1 - img_face_mask_aaa) + (out_img *
                                                           img_face_mask_aaa)

            out_face_bgr = cv2.warpAffine(out_img, face_mat,
                                          (output_size, output_size))

            if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0:
                if cfg.color_transfer_mode == 1:
                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
                    face_mask_aaa = cv2.warpAffine(img_face_mask_aaa, face_mat,
                                                   (output_size, output_size))

                    out_face_bgr = imagelib.reinhard_color_transfer(
                        np.clip((out_face_bgr * 255).astype(np.uint8), 0, 255),
                        np.clip((dst_face_bgr * 255).astype(np.uint8), 0, 255),
                        source_mask=face_mask_aaa,
                        target_mask=face_mask_aaa)
                    out_face_bgr = np.clip(
                        out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)

                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

                elif cfg.color_transfer_mode == 2:
                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

                    out_face_bgr = imagelib.linear_color_transfer(
                        out_face_bgr, dst_face_bgr)
                    out_face_bgr = np.clip(out_face_bgr, 0.0, 1.0)

                    #if debug:
                    #    debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]

            if cfg.mode == 'seamless-hist-match':
                out_face_bgr = imagelib.color_hist_match(
                    out_face_bgr, dst_face_bgr, cfg.hist_match_threshold)

            cfg_mp = cfg.motion_blur_power / 100.0
            if cfg_mp != 0:
                k_size = int(frame_info.motion_power * cfg_mp)
                if k_size >= 1:
                    k_size = np.clip(k_size + 1, 2, 50)
                    if cfg.super_resolution_mode:
                        k_size *= 2
                    out_face_bgr = imagelib.LinearMotionBlur(
                        out_face_bgr, k_size, frame_info.motion_deg)

            if cfg.sharpen_mode != 0 and cfg.sharpen_amount != 0:
                out_face_bgr = cfg.sharpen_func(out_face_bgr, cfg.sharpen_mode,
                                                3, cfg.sharpen_amount)

            new_out = cv2.warpAffine(out_face_bgr, face_mat, img_size,
                                     img_bgr.copy(),
                                     cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC,
                                     cv2.BORDER_TRANSPARENT)
            out_img = np.clip(
                img_bgr * (1 - img_face_mask_aaa) +
                (new_out * img_face_mask_aaa), 0, 1.0)

            if cfg.color_degrade_power != 0:
                #if debug:
                #    debugs += [out_img.copy()]
                out_img_reduced = imagelib.reduce_colors(out_img, 256)
                if cfg.color_degrade_power == 100:
                    out_img = out_img_reduced
                else:
                    alpha = cfg.color_degrade_power / 100.0
                    out_img = (out_img * (1.0 - alpha) +
                               out_img_reduced * alpha)

            if cfg.export_mask_alpha:
                out_img = np.concatenate(
                    [out_img, img_face_mask_aaa[:, :, 0:1]], -1)
        out_merging_mask = img_face_mask_aaa

    #if debug:
    #    debugs += [out_img.copy()]

    return out_img, out_merging_mask
Ejemplo n.º 20
0
        def process_data(self, filepath):
            try:
                dflimg = DFLIMG.load (filepath)
                if dflimg is None or not dflimg.has_data():
                    self.log_err (f"{filepath.name} is not a dfl image file")
                else:
                    img = cv2_imread(filepath)
                    h,w = img.shape[:2]
                    if h != w:
                        raise Exception(f'w != h in {filepath}')
                    
                    image_size = self.image_size
                    face_type = self.face_type
                    output_filepath = self.output_dirpath / filepath.name
                    
                    if face_type is not None:
                        lmrks = dflimg.get_landmarks()
                        mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type)
                        
                        img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 )
                        img = np.clip(img, 0, 255).astype(np.uint8)
                        
                        cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )

                        dfl_dict = dflimg.get_dict()
                        dflimg = DFLIMG.load (output_filepath)
                        dflimg.set_dict(dfl_dict)
                        
                        xseg_mask = dflimg.get_xseg_mask()
                        if xseg_mask is not None:
                            xseg_res = 256
                            
                            xseg_lmrks = lmrks.copy()
                            xseg_lmrks *= (xseg_res / w)
                            xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type)
                            
                            xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 )
                            xseg_mask[xseg_mask < 0.5] = 0
                            xseg_mask[xseg_mask >= 0.5] = 1

                            dflimg.set_xseg_mask(xseg_mask)
                        
                        seg_ie_polys = dflimg.get_seg_ie_polys()
                        
                        for poly in seg_ie_polys.get_polys():
                            poly_pts = poly.get_pts()
                            poly_pts = LandmarksProcessor.transform_points(poly_pts, mat)
                            poly.set_points(poly_pts)
                            
                        dflimg.set_seg_ie_polys(seg_ie_polys)
                        
                        lmrks = LandmarksProcessor.transform_points(lmrks, mat)
                        dflimg.set_landmarks(lmrks)
    
                        image_to_face_mat = dflimg.get_image_to_face_mat()
                        if image_to_face_mat is not None:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
                            dflimg.set_image_to_face_mat(image_to_face_mat)
                        dflimg.set_face_type( FaceType.toString(face_type) )
                        dflimg.save()
                        
                    else:
                        dfl_dict = dflimg.get_dict()
                         
                        scale = w / image_size
                        
                        img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)                    
                        
                        cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )

                        dflimg = DFLIMG.load (output_filepath)
                        dflimg.set_dict(dfl_dict)
                        
                        lmrks = dflimg.get_landmarks()                    
                        lmrks /= scale
                        dflimg.set_landmarks(lmrks)
                        
                        seg_ie_polys = dflimg.get_seg_ie_polys()
                        seg_ie_polys.mult_points( 1.0 / scale)
                        dflimg.set_seg_ie_polys(seg_ie_polys)
                        
                        image_to_face_mat = dflimg.get_image_to_face_mat()
    
                        if image_to_face_mat is not None:
                            face_type = FaceType.fromString ( dflimg.get_face_type() )
                            image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
                            dflimg.set_image_to_face_mat(image_to_face_mat)
                        dflimg.save()

                    return (1, filepath, output_filepath)
            except:
                self.log_err (f"Exception occured while processing file {filepath}. Error: {traceback.format_exc()}")

            return (0, filepath, None)
Ejemplo n.º 21
0
        def final_stage(
            data,
            image,
            face_type,
            image_size,
            extract_from_dflimg=False,
            output_debug_path=None,
            final_output_path=None,
        ):
            data.final_output_files = []
            filepath = data.filepath
            rects = data.rects
            landmarks = data.landmarks

            if output_debug_path is not None:
                debug_image = image.copy()

            if extract_from_dflimg and len(rects) != 1:
                #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
                print("extract_from_dflimg and len(rects) != 1", filepath)
                output_filepath = final_output_path / filepath.name
                if filepath != str(output_file):
                    shutil.copy(str(filepath), str(output_filepath))
                data.final_output_files.append(output_filepath)
            else:
                face_idx = 0
                for rect, image_landmarks in zip(rects, landmarks):

                    if extract_from_dflimg and face_idx > 1:
                        #cannot extract more than 1 face from dflimg
                        break

                    if image_landmarks is None:
                        continue

                    rect = np.array(rect)

                    if face_type == FaceType.MARK_ONLY:
                        image_to_face_mat = None
                        face_image = image
                        face_image_landmarks = image_landmarks
                    else:
                        image_to_face_mat = LandmarksProcessor.get_transform_mat(
                            image_landmarks, image_size, face_type)

                        face_image = cv2.warpAffine(image, image_to_face_mat,
                                                    (image_size, image_size),
                                                    cv2.INTER_LANCZOS4)
                        face_image_landmarks = LandmarksProcessor.transform_points(
                            image_landmarks, image_to_face_mat)

                        landmarks_bbox = LandmarksProcessor.transform_points(
                            [(0, 0), (0, image_size - 1),
                             (image_size - 1, image_size - 1),
                             (image_size - 1, 0)], image_to_face_mat, True)

                        rect_area = mathlib.polygon_area(
                            np.array(rect[[0, 2, 2, 0]]),
                            np.array(rect[[1, 1, 3, 3]]))
                        landmarks_area = mathlib.polygon_area(
                            landmarks_bbox[:, 0], landmarks_bbox[:, 1])

                        if not data.manual and face_type <= FaceType.FULL_NO_ALIGN and landmarks_area > 4 * rect_area:  #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                            continue

                        if output_debug_path is not None:
                            LandmarksProcessor.draw_rect_landmarks(
                                debug_image,
                                rect,
                                image_landmarks,
                                image_size,
                                face_type,
                                transparent_mask=True)

                    output_path = final_output_path
                    if data.force_output_path is not None:
                        output_path = data.force_output_path

                    if extract_from_dflimg and filepath.suffix == '.jpg':
                        #if extracting from dflimg and jpg copy it in order not to lose quality
                        output_filepath = output_path / filepath.name
                        if filepath != output_filepath:
                            shutil.copy(str(filepath), str(output_filepath))
                    else:
                        output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg"
                        cv2_imwrite(output_filepath, face_image,
                                    [int(cv2.IMWRITE_JPEG_QUALITY), 100])

                    DFLJPG.embed_data(
                        output_filepath,
                        face_type=FaceType.toString(face_type),
                        landmarks=face_image_landmarks.tolist(),
                        source_filename=filepath.name,
                        source_rect=rect,
                        source_landmarks=image_landmarks.tolist(),
                        image_to_face_mat=image_to_face_mat)

                    data.final_output_files.append(output_filepath)
                    face_idx += 1
                data.faces_detected = face_idx

            if output_debug_path is not None:
                cv2_imwrite(output_debug_path / (filepath.stem + '.jpg'),
                            debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50])

            return data
Ejemplo n.º 22
0
        def process_data(self, data):
            filename_path = Path(data[0])

            filename_path_str = str(filename_path)
            if self.cached_image[0] == filename_path_str:
                image = self.cached_image[1]
            else:
                image = cv2_imread(filename_path_str)
                self.cached_image = (filename_path_str, image)

            if image is None:
                self.log_err(
                    'Failed to extract %s, reason: cv2_imread() fail.' %
                    (str(filename_path)))
            else:
                if self.type == 'rects':
                    rects = self.e.extract_from_bgr(image)
                    return [str(filename_path), rects]

                elif self.type == 'landmarks':
                    rects = data[1]
                    landmarks = self.e.extract_from_bgr(image, rects)
                    return [str(filename_path), landmarks]

                elif self.type == 'final':
                    src_dflimg = None
                    (h, w, c) = image.shape
                    if h == w:
                        #extracting from already extracted jpg image?
                        if filename_path.suffix == '.jpg':
                            src_dflimg = DFLJPG.load(str(filename_path))

                    result = []
                    faces = data[1]

                    if self.debug:
                        debug_output_file = '{}{}'.format(
                            str(
                                Path(str(self.output_path) + '_debug') /
                                filename_path.stem), '.jpg')
                        debug_image = image.copy()

                    face_idx = 0
                    for face in faces:
                        rect = np.array(face[0])
                        image_landmarks = np.array(face[1])

                        if self.face_type == FaceType.MARK_ONLY:
                            face_image = image
                            face_image_landmarks = image_landmarks
                        else:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat(
                                image_landmarks, self.image_size,
                                self.face_type)
                            face_image = cv2.warpAffine(
                                image, image_to_face_mat,
                                (self.image_size, self.image_size),
                                cv2.INTER_LANCZOS4)
                            face_image_landmarks = LandmarksProcessor.transform_points(
                                image_landmarks, image_to_face_mat)

                            landmarks_bbox = LandmarksProcessor.transform_points(
                                [(0, 0), (0, self.image_size - 1),
                                 (self.image_size - 1, self.image_size - 1),
                                 (self.image_size - 1, 0)], image_to_face_mat,
                                True)

                            rect_area = mathlib.polygon_area(
                                np.array(rect[[0, 2, 2, 0]]),
                                np.array(rect[[1, 1, 3, 3]]))
                            landmarks_area = mathlib.polygon_area(
                                landmarks_bbox[:, 0], landmarks_bbox[:, 1])

                            if landmarks_area > 4 * rect_area:  #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                                continue

                        if self.debug:
                            LandmarksProcessor.draw_rect_landmarks(
                                debug_image, rect, image_landmarks,
                                self.image_size, self.face_type)

                        output_file = '{}_{}{}'.format(
                            str(self.output_path / filename_path.stem),
                            str(face_idx), '.jpg')
                        face_idx += 1

                        if src_dflimg is not None:
                            #if extracting from dflimg just copy it in order not to lose quality
                            shutil.copy(str(filename_path), str(output_file))
                        else:
                            cv2_imwrite(output_file, face_image,
                                        [int(cv2.IMWRITE_JPEG_QUALITY), 85])

                        DFLJPG.embed_data(
                            output_file,
                            face_type=FaceType.toString(self.face_type),
                            landmarks=face_image_landmarks.tolist(),
                            source_filename=filename_path.name,
                            source_rect=rect,
                            source_landmarks=image_landmarks.tolist())

                        result.append(output_file)

                    if self.debug:
                        cv2_imwrite(debug_output_file, debug_image,
                                    [int(cv2.IMWRITE_JPEG_QUALITY), 50])

                    return result
            return None