def process(samples,
                sample_process_options,
                output_sample_types,
                debug,
                ct_sample=None):
        SPTF = SampleProcessor.Types

        sample_rnd_seed = np.random.randint(0x80000000)

        outputs = []
        for sample in samples:
            sample_bgr = sample.load_bgr()
            ct_sample_bgr = None
            h, w, c = sample_bgr.shape

            is_face_sample = sample.landmarks is not None

            if debug and is_face_sample:
                LandmarksProcessor.draw_landmarks(sample_bgr, sample.landmarks,
                                                  (0, 1, 0))

            params = imagelib.gen_warp_params(
                sample_bgr,
                sample_process_options.random_flip,
                rotation_range=sample_process_options.rotation_range,
                scale_range=sample_process_options.scale_range,
                tx_range=sample_process_options.tx_range,
                ty_range=sample_process_options.ty_range,
                rnd_seed=sample_rnd_seed)

            outputs_sample = []
            for opts in output_sample_types:

                resolution = opts.get('resolution', 0)
                types = opts.get('types', [])

                motion_blur = opts.get('motion_blur', None)
                gaussian_blur = opts.get('gaussian_blur', None)

                ct_mode = opts.get('ct_mode', 'None')
                normalize_tanh = opts.get('normalize_tanh', False)
                data_format = opts.get('data_format', 'NHWC')

                img_type = SPTF.NONE
                target_face_type = SPTF.NONE
                mode_type = SPTF.NONE
                for t in types:
                    if t >= SPTF.IMG_TYPE_BEGIN and t < SPTF.IMG_TYPE_END:
                        img_type = t
                    elif t >= SPTF.FACE_TYPE_BEGIN and t < SPTF.FACE_TYPE_END:
                        target_face_type = t
                    elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END:
                        mode_type = t

                if mode_type == SPTF.MODE_FACE_MASK_HULL and not is_face_sample:
                    raise ValueError(
                        "MODE_FACE_MASK_HULL applicable only for face samples")
                if mode_type == SPTF.MODE_FACE_MASK_STRUCT and not is_face_sample:
                    raise ValueError(
                        "MODE_FACE_MASK_STRUCT applicable only for face samples"
                    )
                if is_face_sample:
                    if target_face_type == SPTF.NONE:
                        raise ValueError(
                            "target face type must be defined for face samples"
                        )

                can_warp = (img_type == SPTF.IMG_WARPED
                            or img_type == SPTF.IMG_WARPED_TRANSFORMED)
                can_transform = (img_type == SPTF.IMG_WARPED_TRANSFORMED
                                 or img_type == SPTF.IMG_TRANSFORMED)

                if img_type == SPTF.NONE:
                    raise ValueError('expected IMG_ type')

                if img_type == SPTF.IMG_LANDMARKS_ARRAY:
                    l = sample.landmarks
                    l = np.concatenate([
                        np.expand_dims(l[:, 0] / w, -1),
                        np.expand_dims(l[:, 1] / h, -1)
                    ], -1)
                    l = np.clip(l, 0.0, 1.0)
                    out_sample = l
                elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
                    pitch_yaw_roll = sample.get_pitch_yaw_roll()

                    if params['flip']:
                        yaw = -yaw

                    if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
                        pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
                        yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
                        roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)

                    out_sample = (pitch, yaw, roll)
                else:
                    if mode_type == SPTF.NONE:
                        raise ValueError('expected MODE_ type')

                    if mode_type == SPTF.MODE_FACE_MASK_HULL:
                        if sample.eyebrows_expand_mod is not None:
                            img = LandmarksProcessor.get_image_hull_mask(
                                sample_bgr.shape,
                                sample.landmarks,
                                eyebrows_expand_mod=sample.eyebrows_expand_mod)
                        else:
                            img = LandmarksProcessor.get_image_hull_mask(
                                sample_bgr.shape, sample.landmarks)

                        if sample.ie_polys is not None:
                            sample.ie_polys.overlay_mask(img)
                    elif mode_type == SPTF.MODE_FACE_MASK_STRUCT:
                        if sample.eyebrows_expand_mod is not None:
                            img = LandmarksProcessor.get_face_struct_mask(
                                sample_bgr.shape,
                                sample.landmarks,
                                eyebrows_expand_mod=sample.eyebrows_expand_mod)
                        else:
                            img = LandmarksProcessor.get_face_struct_mask(
                                sample_bgr.shape, sample.landmarks)
                    else:
                        img = sample_bgr
                        if motion_blur is not None:
                            chance, mb_max_size = motion_blur
                            chance = np.clip(chance, 0, 100)

                            if np.random.randint(100) < chance:
                                img = imagelib.LinearMotionBlur(
                                    img,
                                    np.random.randint(mb_max_size) + 1,
                                    np.random.randint(360))

                        if gaussian_blur is not None:
                            chance, kernel_max_size = gaussian_blur
                            chance = np.clip(chance, 0, 100)

                            if np.random.randint(100) < chance:
                                img = cv2.GaussianBlur(
                                    img,
                                    (np.random.randint(kernel_max_size) * 2 +
                                     1, ) * 2, 0)

                    if is_face_sample:
                        target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[
                            target_face_type]
                        if target_ft > sample.face_type:
                            raise Exception(
                                'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.'
                                %
                                (sample.filename, sample.face_type, target_ft))

                        if sample.face_type == FaceType.MARK_ONLY:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample.landmarks, sample.shape[0], target_ft)

                            if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
                                img = cv2.warpAffine(
                                    img,
                                    mat, (sample.shape[0], sample.shape[0]),
                                    flags=cv2.INTER_CUBIC)
                                img = imagelib.warp_by_params(
                                    params,
                                    img,
                                    can_warp,
                                    can_transform,
                                    can_flip=True,
                                    border_replicate=False)
                                img = cv2.resize(img, (resolution, resolution),
                                                 cv2.INTER_CUBIC)[..., None]
                            else:
                                img = cv2.warpAffine(
                                    img,
                                    mat, (sample.shape[0], sample.shape[0]),
                                    flags=cv2.INTER_CUBIC)
                                img = imagelib.warp_by_params(
                                    params,
                                    img,
                                    can_warp,
                                    can_transform,
                                    can_flip=True,
                                    border_replicate=True)
                                img = cv2.resize(img, (resolution, resolution),
                                                 cv2.INTER_CUBIC)

                        else:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample.landmarks, resolution, target_ft)

                            if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
                                img = imagelib.warp_by_params(
                                    params,
                                    img,
                                    can_warp,
                                    can_transform,
                                    can_flip=True,
                                    border_replicate=False)
                                img = cv2.warpAffine(
                                    img,
                                    mat, (resolution, resolution),
                                    borderMode=cv2.BORDER_CONSTANT,
                                    flags=cv2.INTER_CUBIC)[..., None]
                            else:
                                img = imagelib.warp_by_params(
                                    params,
                                    img,
                                    can_warp,
                                    can_transform,
                                    can_flip=True,
                                    border_replicate=True)
                                img = cv2.warpAffine(
                                    img,
                                    mat, (resolution, resolution),
                                    borderMode=cv2.BORDER_REPLICATE,
                                    flags=cv2.INTER_CUBIC)
                    else:
                        img = imagelib.warp_by_params(params,
                                                      img,
                                                      can_warp,
                                                      can_transform,
                                                      can_flip=True,
                                                      border_replicate=True)
                        img = cv2.resize(img, (resolution, resolution),
                                         cv2.INTER_CUBIC)

                    if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
                        out_sample = np.clip(img.astype(np.float32), 0, 1)
                    else:
                        img = np.clip(img.astype(np.float32), 0, 1)

                        if ct_mode is not None and ct_sample is not None:
                            if ct_sample_bgr is None:
                                ct_sample_bgr = ct_sample.load_bgr()
                            img = imagelib.color_transfer(
                                ct_mode, img,
                                cv2.resize(ct_sample_bgr,
                                           (resolution, resolution),
                                           cv2.INTER_LINEAR))

                        if mode_type == SPTF.MODE_BGR:
                            out_sample = img
                        elif mode_type == SPTF.MODE_BGR_SHUFFLE:
                            rnd_state = np.random.RandomState(sample_rnd_seed)
                            out_sample = np.take(img,
                                                 rnd_state.permutation(
                                                     img.shape[-1]),
                                                 axis=-1)

                        elif mode_type == SPTF.MODE_BGR_RANDOM_HSV_SHIFT:
                            rnd_state = np.random.RandomState(sample_rnd_seed)
                            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
                            h, s, v = cv2.split(hsv)
                            h = (h + rnd_state.randint(360)) % 360
                            s = np.clip(s + rnd_state.random() - 0.5, 0, 1)
                            v = np.clip(v + rnd_state.random() - 0.5, 0, 1)
                            hsv = cv2.merge([h, s, v])
                            out_sample = np.clip(
                                cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1)
                        elif mode_type == SPTF.MODE_G:
                            out_sample = cv2.cvtColor(img,
                                                      cv2.COLOR_BGR2GRAY)[...,
                                                                          None]
                        elif mode_type == SPTF.MODE_GGG:
                            out_sample = np.repeat(
                                np.expand_dims(
                                    cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
                                (3, ), -1)

                    if not debug:
                        if normalize_tanh:
                            out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
                                                 1.0)

                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))

                outputs_sample.append(out_sample)
            outputs += [outputs_sample]

        return outputs
    def process(samples,
                sample_process_options,
                output_sample_types,
                debug,
                ct_sample=None):
        SPST = SampleProcessor.SampleType
        SPCT = SampleProcessor.ChannelType
        SPFMT = SampleProcessor.FaceMaskType

        sample_rnd_seed = np.random.randint(0x80000000)

        outputs = []
        for sample in samples:
            sample_face_type = sample.face_type
            sample_bgr = sample.load_bgr()
            sample_landmarks = sample.landmarks
            ct_sample_bgr = None
            h, w, c = sample_bgr.shape

            def get_full_face_mask():
                if sample.eyebrows_expand_mod is not None:
                    full_face_mask = LandmarksProcessor.get_image_hull_mask(
                        sample_bgr.shape,
                        sample_landmarks,
                        eyebrows_expand_mod=sample.eyebrows_expand_mod)
                else:
                    full_face_mask = LandmarksProcessor.get_image_hull_mask(
                        sample_bgr.shape, sample_landmarks)
                return np.clip(full_face_mask, 0, 1)

            def get_eyes_mask():
                eyes_mask = LandmarksProcessor.get_image_eye_mask(
                    sample_bgr.shape, sample_landmarks)
                return np.clip(eyes_mask, 0, 1)

            is_face_sample = sample_landmarks is not None

            if debug and is_face_sample:
                LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks,
                                                  (0, 1, 0))

            params_per_resolution = {}
            warp_rnd_state = np.random.RandomState(sample_rnd_seed - 1)
            for opts in output_sample_types:
                resolution = opts.get('resolution', None)
                if resolution is None:
                    continue
                params_per_resolution[resolution] = imagelib.gen_warp_params(
                    resolution,
                    sample_process_options.random_flip,
                    rotation_range=sample_process_options.rotation_range,
                    scale_range=sample_process_options.scale_range,
                    tx_range=sample_process_options.tx_range,
                    ty_range=sample_process_options.ty_range,
                    rnd_state=warp_rnd_state)

            outputs_sample = []
            for opts in output_sample_types:
                sample_type = opts.get('sample_type', SPST.NONE)
                channel_type = opts.get('channel_type', SPCT.NONE)
                resolution = opts.get('resolution', 0)
                warp = opts.get('warp', False)
                transform = opts.get('transform', False)
                motion_blur = opts.get('motion_blur', None)
                gaussian_blur = opts.get('gaussian_blur', None)
                random_bilinear_resize = opts.get('random_bilinear_resize',
                                                  None)
                random_rgb_levels = opts.get('random_rgb_levels', False)
                random_hsv_shift = opts.get('random_hsv_shift', False)
                random_circle_mask = opts.get('random_circle_mask', False)
                normalize_tanh = opts.get('normalize_tanh', False)
                ct_mode = opts.get('ct_mode', None)
                data_format = opts.get('data_format', 'NHWC')

                if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE:
                    border_replicate = False
                elif sample_type == SPST.FACE_IMAGE:
                    border_replicate = True

                border_replicate = opts.get('border_replicate',
                                            border_replicate)
                borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    if not is_face_sample:
                        raise ValueError(
                            "face_samples should be provided for sample_type FACE_*"
                        )

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    face_type = opts.get('face_type', None)
                    face_mask_type = opts.get('face_mask_type', SPFMT.NONE)

                    if face_type is None:
                        raise ValueError(
                            "face_type must be defined for face samples")

                    if face_type > sample.face_type:
                        raise Exception(
                            'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.'
                            % (sample.filename, sample.face_type, face_type))

                    if sample_type == SPST.FACE_MASK:

                        if face_mask_type == SPFMT.FULL_FACE:
                            img = get_full_face_mask()
                        elif face_mask_type == SPFMT.EYES:
                            img = get_eyes_mask()
                        elif face_mask_type == SPFMT.FULL_FACE_EYES:
                            img = get_full_face_mask() + get_eyes_mask()
                        else:
                            img = np.zeros(sample_bgr.shape[0:2] + (1, ),
                                           dtype=np.float32)

                        if sample.ie_polys is not None:
                            sample.ie_polys.overlay_mask(img)

                        if sample_face_type == FaceType.MARK_ONLY:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample_landmarks, warp_resolution, face_type)
                            img = cv2.warpAffine(
                                img,
                                mat, (warp_resolution, warp_resolution),
                                flags=cv2.INTER_LINEAR)

                            img = imagelib.warp_by_params(
                                params_per_resolution[resolution],
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=border_replicate,
                                cv2_inter=cv2.INTER_LINEAR)
                            img = cv2.resize(img, (resolution, resolution),
                                             cv2.INTER_LINEAR)
                        else:
                            if face_type != sample_face_type:
                                mat = LandmarksProcessor.get_transform_mat(
                                    sample_landmarks, resolution, face_type)
                                img = cv2.warpAffine(img,
                                                     mat,
                                                     (resolution, resolution),
                                                     borderMode=borderMode,
                                                     flags=cv2.INTER_LINEAR)
                            else:
                                if w != resolution:
                                    img = cv2.resize(img,
                                                     (resolution, resolution),
                                                     cv2.INTER_CUBIC)

                            img = imagelib.warp_by_params(
                                params_per_resolution[resolution],
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=border_replicate,
                                cv2_inter=cv2.INTER_LINEAR)

                        if len(img.shape) == 2:
                            img = img[..., None]

                        if channel_type == SPCT.G:
                            out_sample = img.astype(np.float32)
                        else:
                            raise ValueError(
                                "only channel_type.G supported for the mask")

                    elif sample_type == SPST.FACE_IMAGE:
                        img = sample_bgr

                        if random_rgb_levels:
                            random_mask = sd.random_circle_faded(
                                [w, w],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed)
                            ) if random_circle_mask else None
                            img = imagelib.apply_random_rgb_levels(
                                img,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed))

                        if random_hsv_shift:
                            random_mask = sd.random_circle_faded(
                                [w, w],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    1)) if random_circle_mask else None
                            img = imagelib.apply_random_hsv_shift(
                                img,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 1))

                        if face_type != sample_face_type:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample_landmarks, resolution, face_type)
                            img = cv2.warpAffine(img,
                                                 mat, (resolution, resolution),
                                                 borderMode=borderMode,
                                                 flags=cv2.INTER_CUBIC)
                        else:
                            if w != resolution:
                                img = cv2.resize(img, (resolution, resolution),
                                                 cv2.INTER_CUBIC)

                        # Apply random color transfer
                        if ct_mode is not None and ct_sample is not None:
                            if ct_sample_bgr is None:
                                ct_sample_bgr = ct_sample.load_bgr()
                            img = imagelib.color_transfer(
                                ct_mode, img,
                                cv2.resize(ct_sample_bgr,
                                           (resolution, resolution),
                                           cv2.INTER_LINEAR))

                        img = imagelib.warp_by_params(
                            params_per_resolution[resolution],
                            img,
                            warp,
                            transform,
                            can_flip=True,
                            border_replicate=border_replicate)

                        img = np.clip(img.astype(np.float32), 0, 1)

                        if motion_blur is not None:
                            random_mask = sd.random_circle_faded(
                                [resolution, resolution],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    2)) if random_circle_mask else None
                            img = imagelib.apply_random_motion_blur(
                                img,
                                *motion_blur,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 2))

                        if gaussian_blur is not None:
                            random_mask = sd.random_circle_faded(
                                [resolution, resolution],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    3)) if random_circle_mask else None
                            img = imagelib.apply_random_gaussian_blur(
                                img,
                                *gaussian_blur,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 3))

                        if random_bilinear_resize is not None:
                            random_mask = sd.random_circle_faded(
                                [resolution, resolution],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    4)) if random_circle_mask else None
                            img = imagelib.apply_random_bilinear_resize(
                                img,
                                *random_bilinear_resize,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 4))

                        # Transform from BGR to desired channel_type
                        if channel_type == SPCT.BGR:
                            out_sample = img
                        elif channel_type == SPCT.G:
                            out_sample = cv2.cvtColor(img,
                                                      cv2.COLOR_BGR2GRAY)[...,
                                                                          None]
                        elif channel_type == SPCT.GGG:
                            out_sample = np.repeat(
                                np.expand_dims(
                                    cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
                                (3, ), -1)

                    # Final transformations
                    if not debug:
                        if normalize_tanh:
                            out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
                                                 1.0)
                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))
                elif sample_type == SPST.IMAGE:
                    img = sample_bgr
                    img = imagelib.warp_by_params(
                        params_per_resolution[resolution],
                        img,
                        warp,
                        transform,
                        can_flip=True,
                        border_replicate=True)
                    img = cv2.resize(img, (resolution, resolution),
                                     cv2.INTER_CUBIC)
                    out_sample = img

                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))

                elif sample_type == SPST.LANDMARKS_ARRAY:
                    l = sample_landmarks
                    l = np.concatenate([
                        np.expand_dims(l[:, 0] / w, -1),
                        np.expand_dims(l[:, 1] / h, -1)
                    ], -1)
                    l = np.clip(l, 0.0, 1.0)
                    out_sample = l
                elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                    pitch, yaw, roll = sample.get_pitch_yaw_roll()
                    if params_per_resolution[resolution]['flip']:
                        yaw = -yaw

                    if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                        pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
                        yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
                        roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)

                    out_sample = (pitch, yaw)
                else:
                    raise ValueError('expected sample_type')

                outputs_sample.append(out_sample)
            outputs += [outputs_sample]

        return outputs
Exemple #3
0
    def process(samples,
                sample_process_options,
                output_sample_types,
                debug,
                ct_sample=None):
        SPST = SampleProcessor.SampleType
        SPCT = SampleProcessor.ChannelType
        SPFMT = SampleProcessor.FaceMaskType

        sample_rnd_seed = np.random.randint(0x80000000)

        outputs = []
        for sample in samples:
            sample_face_type = sample.face_type
            sample_bgr = sample.load_bgr()
            sample_landmarks = sample.landmarks
            ct_sample_bgr = None
            h, w, c = sample_bgr.shape

            def get_full_face_mask():
                xseg_mask = sample.get_xseg_mask()
                if xseg_mask is not None:
                    if xseg_mask.shape[0] != h or xseg_mask.shape[1] != w:
                        xseg_mask = cv2.resize(xseg_mask, (w, h),
                                               interpolation=cv2.INTER_CUBIC)
                        xseg_mask = imagelib.normalize_channels(xseg_mask, 1)
                    return np.clip(xseg_mask, 0, 1)
                else:
                    full_face_mask = LandmarksProcessor.get_image_hull_mask(
                        sample_bgr.shape,
                        sample_landmarks,
                        eyebrows_expand_mod=sample.eyebrows_expand_mod)
                    return np.clip(full_face_mask, 0, 1)

            def get_eyes_mask():
                eyes_mask = LandmarksProcessor.get_image_eye_mask(
                    sample_bgr.shape, sample_landmarks)
                # set eye masks to 1-2
                clip = np.clip(eyes_mask, 0, 1)
                clip[clip > 0.1] += 1
                return clip

            def get_mouth_mask():
                mouth_mask = LandmarksProcessor.get_image_mouth_mask(
                    sample_bgr.shape, sample_landmarks)
                # set eye masks to 2-3
                clip = np.clip(mouth_mask, 0, 1)
                clip[clip > 0.1] += 2
                return clip

            is_face_sample = sample_landmarks is not None

            if debug and is_face_sample:
                LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks,
                                                  (0, 1, 0))

            params_per_resolution = {}
            warp_rnd_state = np.random.RandomState(sample_rnd_seed - 1)
            for opts in output_sample_types:
                resolution = opts.get('resolution', None)
                if resolution is None:
                    continue
                params_per_resolution[resolution] = imagelib.gen_warp_params(
                    resolution,
                    sample_process_options.random_flip,
                    rotation_range=sample_process_options.rotation_range,
                    scale_range=sample_process_options.scale_range,
                    tx_range=sample_process_options.tx_range,
                    ty_range=sample_process_options.ty_range,
                    rnd_state=warp_rnd_state)

            outputs_sample = []
            for opts in output_sample_types:
                sample_type = opts.get('sample_type', SPST.NONE)
                channel_type = opts.get('channel_type', SPCT.NONE)
                resolution = opts.get('resolution', 0)
                nearest_resize_to = opts.get('nearest_resize_to', None)
                warp = opts.get('warp', False)
                transform = opts.get('transform', False)
                random_downsample = opts.get('random_downsample', False)
                random_noise = opts.get('random_noise', False)
                random_blur = opts.get('random_blur', False)
                random_jpeg = opts.get('random_jpeg', False)
                motion_blur = opts.get('motion_blur', None)
                gaussian_blur = opts.get('gaussian_blur', None)
                random_bilinear_resize = opts.get('random_bilinear_resize',
                                                  None)
                random_rgb_levels = opts.get('random_rgb_levels', False)
                random_hsv_shift = opts.get('random_hsv_shift', False)
                random_circle_mask = opts.get('random_circle_mask', False)
                normalize_tanh = opts.get('normalize_tanh', False)
                ct_mode = opts.get('ct_mode', None)
                data_format = opts.get('data_format', 'NHWC')

                if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE:
                    border_replicate = False
                elif sample_type == SPST.FACE_IMAGE:
                    border_replicate = True

                border_replicate = opts.get('border_replicate',
                                            border_replicate)
                borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    if not is_face_sample:
                        raise ValueError(
                            "face_samples should be provided for sample_type FACE_*"
                        )

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    face_type = opts.get('face_type', None)
                    face_mask_type = opts.get('face_mask_type', SPFMT.NONE)

                    if face_type is None:
                        raise ValueError(
                            "face_type must be defined for face samples")

                    if sample_type == SPST.FACE_MASK:
                        if face_mask_type == SPFMT.FULL_FACE:
                            img = get_full_face_mask()
                        elif face_mask_type == SPFMT.EYES:
                            img = get_eyes_mask()
                        elif face_mask_type == SPFMT.FULL_FACE_EYES:
                            # sets both eyes and mouth mask parts
                            img = get_full_face_mask()
                            mask = img.copy()
                            mask[mask != 0.0] = 1.0
                            eye_mask = get_eyes_mask() * mask
                            img = np.where(eye_mask > 1, eye_mask, img)

                            mouth_mask = get_mouth_mask() * mask
                            img = np.where(mouth_mask > 2, mouth_mask, img)
                        else:
                            img = np.zeros(sample_bgr.shape[0:2] + (1, ),
                                           dtype=np.float32)

                        if sample_face_type == FaceType.MARK_ONLY:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample_landmarks, warp_resolution, face_type)
                            img = cv2.warpAffine(
                                img,
                                mat, (warp_resolution, warp_resolution),
                                flags=cv2.INTER_LINEAR)

                            img = imagelib.warp_by_params(
                                params_per_resolution[resolution],
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=border_replicate,
                                cv2_inter=cv2.INTER_LINEAR)
                            img = cv2.resize(img, (resolution, resolution),
                                             interpolation=cv2.INTER_LINEAR)
                        else:
                            if face_type != sample_face_type:
                                mat = LandmarksProcessor.get_transform_mat(
                                    sample_landmarks, resolution, face_type)
                                img = cv2.warpAffine(img,
                                                     mat,
                                                     (resolution, resolution),
                                                     borderMode=borderMode,
                                                     flags=cv2.INTER_LINEAR)
                            else:
                                if w != resolution:
                                    img = cv2.resize(
                                        img, (resolution, resolution),
                                        interpolation=cv2.INTER_LINEAR)

                            img = imagelib.warp_by_params(
                                params_per_resolution[resolution],
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=border_replicate,
                                cv2_inter=cv2.INTER_LINEAR)

                        if len(img.shape) == 2:
                            img = img[..., None]

                        if channel_type == SPCT.G:
                            out_sample = img.astype(np.float32)
                        else:
                            raise ValueError(
                                "only channel_type.G supported for the mask")

                    elif sample_type == SPST.FACE_IMAGE:
                        img = sample_bgr

                        if random_rgb_levels:
                            random_mask = sd.random_circle_faded(
                                [w, w],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed)
                            ) if random_circle_mask else None
                            img = imagelib.apply_random_rgb_levels(
                                img,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed))

                        if random_hsv_shift:
                            random_mask = sd.random_circle_faded(
                                [w, w],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    1)) if random_circle_mask else None
                            img = imagelib.apply_random_hsv_shift(
                                img,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 1))

                        if face_type != sample_face_type:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample_landmarks, resolution, face_type)
                            img = cv2.warpAffine(img,
                                                 mat, (resolution, resolution),
                                                 borderMode=borderMode,
                                                 flags=cv2.INTER_CUBIC)
                        else:
                            if w != resolution:
                                img = cv2.resize(img, (resolution, resolution),
                                                 interpolation=cv2.INTER_CUBIC)

                        # Apply random color transfer
                        if ct_mode is not None and ct_sample is not None or ct_mode == 'fs-aug':
                            if ct_mode == 'fs-aug':
                                img = imagelib.color_augmentation(
                                    img, sample_rnd_seed)
                            else:
                                if ct_sample_bgr is None:
                                    ct_sample_bgr = ct_sample.load_bgr()
                                img = imagelib.color_transfer(
                                    ct_mode, img,
                                    cv2.resize(ct_sample_bgr,
                                               (resolution, resolution),
                                               interpolation=cv2.INTER_LINEAR))

                        randomization_order = ['blur', 'noise', 'jpeg', 'down']
                        np.random.shuffle(randomization_order)
                        for random_distortion in randomization_order:
                            # Apply random blur
                            if random_distortion == 'blur' and random_blur:
                                blur_type = np.random.choice(
                                    ['motion', 'gaussian'])

                                if blur_type == 'motion':
                                    blur_k = np.random.randint(10, 20)
                                    blur_angle = 360 * np.random.random()
                                    img = LinearMotionBlur(
                                        img, blur_k, blur_angle)
                                elif blur_type == 'gaussian':
                                    blur_sigma = 5 * np.random.random() + 3

                                    if blur_sigma < 5.0:
                                        kernel_size = 2.9 * blur_sigma  # 97% of weight
                                    else:
                                        kernel_size = 2.6 * blur_sigma  # 95% of weight
                                    kernel_size = int(kernel_size)
                                    kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size

                                    img = cv2.GaussianBlur(
                                        img, (kernel_size, kernel_size),
                                        blur_sigma)

                            # Apply random noise
                            if random_distortion == 'noise' and random_noise:
                                noise_type = np.random.choice(
                                    ['gaussian', 'laplace', 'poisson'])
                                noise_scale = (20 * np.random.random() + 20)

                                if noise_type == 'gaussian':
                                    noise = np.random.normal(scale=noise_scale,
                                                             size=img.shape)
                                    img += noise / 255.0
                                elif noise_type == 'laplace':
                                    noise = np.random.laplace(
                                        scale=noise_scale, size=img.shape)
                                    img += noise / 255.0
                                elif noise_type == 'poisson':
                                    noise_lam = (15 * np.random.random() + 15)
                                    noise = np.random.poisson(lam=noise_lam,
                                                              size=img.shape)
                                    img += noise / 255.0

                            # Apply random jpeg compression
                            if random_distortion == 'jpeg' and random_jpeg:
                                img = np.clip(img * 255, 0,
                                              255).astype(np.uint8)
                                jpeg_compression_level = np.random.randint(
                                    50, 85)
                                encode_param = [
                                    int(cv2.IMWRITE_JPEG_QUALITY),
                                    jpeg_compression_level
                                ]
                                _, enc_img = cv2.imencode(
                                    '.jpg', img, encode_param)
                                img = cv2.imdecode(
                                    enc_img, cv2.IMREAD_UNCHANGED).astype(
                                        np.float32) / 255.0

                            # Apply random downsampling
                            if random_distortion == 'down' and random_downsample:
                                down_res = np.random.randint(
                                    int(0.125 * resolution),
                                    int(0.25 * resolution))
                                img = cv2.resize(img, (down_res, down_res),
                                                 interpolation=cv2.INTER_CUBIC)
                                img = cv2.resize(img, (resolution, resolution),
                                                 interpolation=cv2.INTER_CUBIC)

                        img = imagelib.warp_by_params(
                            params_per_resolution[resolution],
                            img,
                            warp,
                            transform,
                            can_flip=True,
                            border_replicate=border_replicate)
                        img = np.clip(img.astype(np.float32), 0, 1)

                        if motion_blur is not None:
                            random_mask = sd.random_circle_faded(
                                [resolution, resolution],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    2)) if random_circle_mask else None
                            img = imagelib.apply_random_motion_blur(
                                img,
                                *motion_blur,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 2))

                        if gaussian_blur is not None:
                            random_mask = sd.random_circle_faded(
                                [resolution, resolution],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    3)) if random_circle_mask else None
                            img = imagelib.apply_random_gaussian_blur(
                                img,
                                *gaussian_blur,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 3))

                        if random_bilinear_resize is not None:
                            random_mask = sd.random_circle_faded(
                                [resolution, resolution],
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed +
                                    4)) if random_circle_mask else None
                            img = imagelib.apply_random_bilinear_resize(
                                img,
                                *random_bilinear_resize,
                                mask=random_mask,
                                rnd_state=np.random.RandomState(
                                    sample_rnd_seed + 4))

                        # Transform from BGR to desired channel_type
                        if channel_type == SPCT.BGR:
                            out_sample = img
                        elif channel_type == SPCT.LAB_RAND_TRANSFORM:
                            out_sample = random_lab_rotation(
                                img, sample_rnd_seed)
                        elif channel_type == SPCT.G:
                            out_sample = cv2.cvtColor(img,
                                                      cv2.COLOR_BGR2GRAY)[...,
                                                                          None]
                        elif channel_type == SPCT.GGG:
                            out_sample = np.repeat(
                                np.expand_dims(
                                    cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
                                (3, ), -1)

                    # Final transformations
                    if nearest_resize_to is not None:
                        out_sample = cv2_resize(
                            out_sample, (nearest_resize_to, nearest_resize_to),
                            interpolation=cv2.INTER_NEAREST)

                    if not debug:
                        if normalize_tanh:
                            out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
                                                 1.0)
                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))
                elif sample_type == SPST.IMAGE:
                    img = sample_bgr
                    img = imagelib.warp_by_params(
                        params_per_resolution[resolution],
                        img,
                        warp,
                        transform,
                        can_flip=True,
                        border_replicate=True)
                    img = cv2.resize(img, (resolution, resolution),
                                     interpolation=cv2.INTER_CUBIC)
                    out_sample = img

                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))

                elif sample_type == SPST.LANDMARKS_ARRAY:
                    l = sample_landmarks
                    l = np.concatenate([
                        np.expand_dims(l[:, 0] / w, -1),
                        np.expand_dims(l[:, 1] / h, -1)
                    ], -1)
                    l = np.clip(l, 0.0, 1.0)
                    out_sample = l
                elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                    pitch, yaw, roll = sample.get_pitch_yaw_roll()
                    if params_per_resolution[resolution]['flip']:
                        yaw = -yaw

                    if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                        pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
                        yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
                        roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)

                    out_sample = (pitch, yaw)
                else:
                    raise ValueError('expected sample_type')

                outputs_sample.append(out_sample)
            outputs += [outputs_sample]

        return outputs
    def process(samples,
                sample_process_options,
                output_sample_types,
                debug,
                ct_sample=None):
        SPST = SampleProcessor.SampleType
        SPCT = SampleProcessor.ChannelType
        SPFMT = SampleProcessor.FaceMaskType

        sample_rnd_seed = np.random.randint(0x80000000)

        outputs = []
        for sample in samples:
            sample_bgr = sample.load_bgr()
            ct_sample_bgr = None
            h, w, c = sample_bgr.shape

            is_face_sample = sample.landmarks is not None

            if debug and is_face_sample:
                LandmarksProcessor.draw_landmarks(sample_bgr, sample.landmarks,
                                                  (0, 1, 0))

            params = imagelib.gen_warp_params(
                sample_bgr,
                sample_process_options.random_flip,
                rotation_range=sample_process_options.rotation_range,
                scale_range=sample_process_options.scale_range,
                tx_range=sample_process_options.tx_range,
                ty_range=sample_process_options.ty_range)

            outputs_sample = []
            for opts in output_sample_types:
                sample_type = opts.get('sample_type', SPST.NONE)
                channel_type = opts.get('channel_type', SPCT.NONE)
                resolution = opts.get('resolution', 0)
                warp = opts.get('warp', False)
                transform = opts.get('transform', False)
                motion_blur = opts.get('motion_blur', None)
                gaussian_blur = opts.get('gaussian_blur', None)
                normalize_tanh = opts.get('normalize_tanh', False)
                ct_mode = opts.get('ct_mode', 'None')
                data_format = opts.get('data_format', 'NHWC')

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    if not is_face_sample:
                        raise ValueError(
                            "face_samples should be provided for sample_type FACE_*"
                        )

                if is_face_sample:
                    face_type = opts.get('face_type', None)
                    face_mask_type = opts.get('face_mask_type', SPFMT.NONE)

                    if face_type is None:
                        raise ValueError(
                            "face_type must be defined for face samples")

                    if face_type > sample.face_type:
                        raise Exception(
                            'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.'
                            % (sample.filename, sample.face_type, target_ft))

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:

                    if sample_type == SPST.FACE_MASK:
                        if face_mask_type == SPFMT.ALL_HULL or \
                           face_mask_type == SPFMT.EYES_HULL or \
                           face_mask_type == SPFMT.ALL_EYES_HULL:
                            if face_mask_type == SPFMT.ALL_HULL or \
                               face_mask_type == SPFMT.ALL_EYES_HULL:
                                if sample.eyebrows_expand_mod is not None:
                                    all_mask = LandmarksProcessor.get_image_hull_mask(
                                        sample_bgr.shape,
                                        sample.landmarks,
                                        eyebrows_expand_mod=sample.
                                        eyebrows_expand_mod)
                                else:
                                    all_mask = LandmarksProcessor.get_image_hull_mask(
                                        sample_bgr.shape, sample.landmarks)

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

                            if face_mask_type == SPFMT.EYES_HULL or \
                               face_mask_type == SPFMT.ALL_EYES_HULL:
                                eyes_mask = LandmarksProcessor.get_image_eye_mask(
                                    sample_bgr.shape, sample.landmarks)
                                eyes_mask = np.clip(eyes_mask, 0, 1)

                            if face_mask_type == SPFMT.ALL_HULL:
                                img = all_mask
                            elif face_mask_type == SPFMT.EYES_HULL:
                                img = eyes_mask
                            elif face_mask_type == SPFMT.ALL_EYES_HULL:
                                img = all_mask + eyes_mask
                        elif face_mask_type == SPFMT.STRUCT:
                            if sample.eyebrows_expand_mod is not None:
                                img = LandmarksProcessor.get_face_struct_mask(
                                    sample_bgr.shape,
                                    sample.landmarks,
                                    eyebrows_expand_mod=sample.
                                    eyebrows_expand_mod)
                            else:
                                img = LandmarksProcessor.get_face_struct_mask(
                                    sample_bgr.shape, sample.landmarks)

                        if sample.ie_polys is not None:
                            sample.ie_polys.overlay_mask(img)

                        if sample.face_type == FaceType.MARK_ONLY:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample.landmarks, sample.shape[0], face_type)
                            img = cv2.warpAffine(
                                img,
                                mat, (sample.shape[0], sample.shape[0]),
                                flags=cv2.INTER_LINEAR)
                            img = imagelib.warp_by_params(
                                params,
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=False,
                                cv2_inter=cv2.INTER_LINEAR)
                            img = cv2.resize(img, (resolution, resolution),
                                             cv2.INTER_LINEAR)[..., None]
                        else:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample.landmarks, resolution, face_type)
                            img = imagelib.warp_by_params(
                                params,
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=False,
                                cv2_inter=cv2.INTER_LINEAR)
                            img = cv2.warpAffine(
                                img,
                                mat, (resolution, resolution),
                                borderMode=cv2.BORDER_CONSTANT,
                                flags=cv2.INTER_LINEAR)[..., None]

                        if channel_type == SPCT.G:
                            out_sample = img.astype(np.float32)
                        else:
                            raise ValueError(
                                "only channel_type.G supported for the mask")

                    elif sample_type == SPST.FACE_IMAGE:
                        img = sample_bgr
                        if motion_blur is not None:
                            chance, mb_max_size = motion_blur
                            chance = np.clip(chance, 0, 100)

                            l_rnd_state = np.random.RandomState(
                                sample_rnd_seed)
                            mblur_rnd_chance = l_rnd_state.randint(100)
                            mblur_rnd_kernel = l_rnd_state.randint(
                                mb_max_size) + 1
                            mblur_rnd_deg = l_rnd_state.randint(360)

                            if mblur_rnd_chance < chance:
                                img = imagelib.LinearMotionBlur(
                                    img, mblur_rnd_kernel, mblur_rnd_deg)

                        if gaussian_blur is not None:
                            chance, kernel_max_size = gaussian_blur
                            chance = np.clip(chance, 0, 100)

                            l_rnd_state = np.random.RandomState(
                                sample_rnd_seed + 1)
                            gblur_rnd_chance = l_rnd_state.randint(100)
                            gblur_rnd_kernel = l_rnd_state.randint(
                                kernel_max_size) * 2 + 1

                            if gblur_rnd_chance < chance:
                                img = cv2.GaussianBlur(
                                    img, (gblur_rnd_kernel, ) * 2, 0)

                        if sample.face_type == FaceType.MARK_ONLY:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample.landmarks, sample.shape[0], face_type)
                            img = cv2.warpAffine(
                                img,
                                mat, (sample.shape[0], sample.shape[0]),
                                flags=cv2.INTER_CUBIC)
                            img = imagelib.warp_by_params(
                                params,
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=True)
                            img = cv2.resize(img, (resolution, resolution),
                                             cv2.INTER_CUBIC)
                        else:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample.landmarks, resolution, face_type)
                            img = imagelib.warp_by_params(
                                params,
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=True)
                            img = cv2.warpAffine(
                                img,
                                mat, (resolution, resolution),
                                borderMode=cv2.BORDER_REPLICATE,
                                flags=cv2.INTER_CUBIC)

                        img = np.clip(img.astype(np.float32), 0, 1)

                        # Apply random color transfer
                        if ct_mode is not None and ct_sample is not None:
                            if ct_sample_bgr is None:
                                ct_sample_bgr = ct_sample.load_bgr()
                            img = imagelib.color_transfer(
                                ct_mode, img,
                                cv2.resize(ct_sample_bgr,
                                           (resolution, resolution),
                                           cv2.INTER_LINEAR))

                        # Transform from BGR to desired channel_type
                        if channel_type == SPCT.BGR:
                            out_sample = img
                        elif channel_type == SPCT.BGR_SHUFFLE:
                            l_rnd_state = np.random.RandomState(
                                sample_rnd_seed)
                            out_sample = np.take(img,
                                                 l_rnd_state.permutation(
                                                     img.shape[-1]),
                                                 axis=-1)
                        elif channel_type == SPCT.BGR_RANDOM_HSV_SHIFT:
                            l_rnd_state = np.random.RandomState(
                                sample_rnd_seed)
                            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
                            h, s, v = cv2.split(hsv)
                            h = (h + l_rnd_state.randint(360)) % 360
                            s = np.clip(s + l_rnd_state.random() - 0.5, 0, 1)
                            v = np.clip(v + l_rnd_state.random() - 0.5, 0, 1)
                            hsv = cv2.merge([h, s, v])
                            out_sample = np.clip(
                                cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1)
                        elif channel_type == SPCT.BGR_RANDOM_RGB_LEVELS:
                            l_rnd_state = np.random.RandomState(
                                sample_rnd_seed)
                            np_rnd = l_rnd_state.rand
                            inBlack = np.array([
                                np_rnd() * 0.25,
                                np_rnd() * 0.25,
                                np_rnd() * 0.25
                            ],
                                               dtype=np.float32)
                            inWhite = np.array([
                                1.0 - np_rnd() * 0.25, 1.0 - np_rnd() * 0.25,
                                1.0 - np_rnd() * 0.25
                            ],
                                               dtype=np.float32)
                            inGamma = np.array([
                                0.5 + np_rnd(), 0.5 + np_rnd(), 0.5 + np_rnd()
                            ],
                                               dtype=np.float32)
                            outBlack = np.array([0.0, 0.0, 0.0],
                                                dtype=np.float32)
                            outWhite = np.array([1.0, 1.0, 1.0],
                                                dtype=np.float32)
                            out_sample = np.clip(
                                (img - inBlack) / (inWhite - inBlack), 0, 1)
                            out_sample = (out_sample**(1 / inGamma)) * (
                                outWhite - outBlack) + outBlack
                            out_sample = np.clip(out_sample, 0, 1)
                        elif channel_type == SPCT.G:
                            out_sample = cv2.cvtColor(img,
                                                      cv2.COLOR_BGR2GRAY)[...,
                                                                          None]
                        elif channel_type == SPCT.GGG:
                            out_sample = np.repeat(
                                np.expand_dims(
                                    cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
                                (3, ), -1)

                    # Final transformations
                    if not debug:
                        if normalize_tanh:
                            out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
                                                 1.0)
                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))
                #else:
                #    img  = imagelib.warp_by_params (params, img,  warp, transform, can_flip=True, border_replicate=True)
                #    img  = cv2.resize( img,  (resolution,resolution), cv2.INTER_CUBIC )
                elif sample_type == SPST.LANDMARKS_ARRAY:
                    l = sample.landmarks
                    l = np.concatenate([
                        np.expand_dims(l[:, 0] / w, -1),
                        np.expand_dims(l[:, 1] / h, -1)
                    ], -1)
                    l = np.clip(l, 0.0, 1.0)
                    out_sample = l
                elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                    pitch_yaw_roll = sample.get_pitch_yaw_roll()

                    if params['flip']:
                        yaw = -yaw

                    if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                        pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
                        yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
                        roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)

                    out_sample = (pitch, yaw, roll)
                else:
                    raise ValueError('expected sample_type')

                outputs_sample.append(out_sample)
            outputs += [outputs_sample]

        return outputs
Exemple #5
0
    def process(samples,
                sample_process_options,
                output_sample_types,
                debug,
                ct_sample=None):
        SPST = SampleProcessor.SampleType
        SPCT = SampleProcessor.ChannelType
        SPFMT = SampleProcessor.FaceMaskType

        outputs = []
        for sample in samples:
            sample_rnd_seed = np.random.randint(0x80000000)

            sample_face_type = sample.face_type
            sample_bgr = sample.load_bgr()
            sample_landmarks = sample.landmarks
            ct_sample_bgr = None
            h, w, c = sample_bgr.shape

            def get_full_face_mask():
                xseg_mask = sample.get_xseg_mask()
                if xseg_mask is not None:
                    if xseg_mask.shape[0] != h or xseg_mask.shape[1] != w:
                        xseg_mask = cv2.resize(xseg_mask, (w, h),
                                               interpolation=cv2.INTER_CUBIC)
                        xseg_mask = imagelib.normalize_channels(xseg_mask, 1)
                    return np.clip(xseg_mask, 0, 1)
                else:
                    full_face_mask = LandmarksProcessor.get_image_hull_mask(
                        sample_bgr.shape,
                        sample_landmarks,
                        eyebrows_expand_mod=sample.eyebrows_expand_mod)
                    return np.clip(full_face_mask, 0, 1)

            def get_eyes_mask():
                eyes_mask = LandmarksProcessor.get_image_eye_mask(
                    sample_bgr.shape, sample_landmarks)
                return np.clip(eyes_mask, 0, 1)

            def get_eyes_mouth_mask():
                eyes_mask = LandmarksProcessor.get_image_eye_mask(
                    sample_bgr.shape, sample_landmarks)
                mouth_mask = LandmarksProcessor.get_image_mouth_mask(
                    sample_bgr.shape, sample_landmarks)
                mask = eyes_mask + mouth_mask
                return np.clip(mask, 0, 1)

            is_face_sample = sample_landmarks is not None

            if debug and is_face_sample:
                LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks,
                                                  (0, 1, 0))

            outputs_sample = []
            for opts in output_sample_types:
                resolution = opts.get('resolution', 0)
                sample_type = opts.get('sample_type', SPST.NONE)
                channel_type = opts.get('channel_type', SPCT.NONE)
                nearest_resize_to = opts.get('nearest_resize_to', None)
                warp = opts.get('warp', False)
                transform = opts.get('transform', False)
                random_hsv_shift_amount = opts.get('random_hsv_shift_amount',
                                                   0)
                normalize_tanh = opts.get('normalize_tanh', False)
                ct_mode = opts.get('ct_mode', None)
                data_format = opts.get('data_format', 'NHWC')

                rnd_seed_shift = opts.get('rnd_seed_shift', 0)
                warp_rnd_seed_shift = opts.get('warp_rnd_seed_shift',
                                               rnd_seed_shift)

                rnd_state = np.random.RandomState(sample_rnd_seed +
                                                  rnd_seed_shift)
                warp_rnd_state = np.random.RandomState(sample_rnd_seed +
                                                       warp_rnd_seed_shift)

                warp_params = imagelib.gen_warp_params(
                    resolution,
                    sample_process_options.random_flip,
                    rotation_range=sample_process_options.rotation_range,
                    scale_range=sample_process_options.scale_range,
                    tx_range=sample_process_options.tx_range,
                    ty_range=sample_process_options.ty_range,
                    rnd_state=rnd_state,
                    warp_rnd_state=warp_rnd_state,
                )

                if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE:
                    border_replicate = False
                elif sample_type == SPST.FACE_IMAGE:
                    border_replicate = True

                border_replicate = opts.get('border_replicate',
                                            border_replicate)
                borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    if not is_face_sample:
                        raise ValueError(
                            "face_samples should be provided for sample_type FACE_*"
                        )

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    face_type = opts.get('face_type', None)
                    face_mask_type = opts.get('face_mask_type', SPFMT.NONE)

                    if face_type is None:
                        raise ValueError(
                            "face_type must be defined for face samples")

                    if sample_type == SPST.FACE_MASK:
                        if face_mask_type == SPFMT.FULL_FACE:
                            img = get_full_face_mask()
                        elif face_mask_type == SPFMT.EYES:
                            img = get_eyes_mask()
                        elif face_mask_type == SPFMT.EYES_MOUTH:
                            mask = get_full_face_mask().copy()
                            mask[mask != 0.0] = 1.0
                            img = get_eyes_mouth_mask() * mask
                        else:
                            img = np.zeros(sample_bgr.shape[0:2] + (1, ),
                                           dtype=np.float32)

                        if sample_face_type == FaceType.MARK_ONLY:
                            raise NotImplementedError()
                            mat = LandmarksProcessor.get_transform_mat(
                                sample_landmarks, warp_resolution, face_type)
                            img = cv2.warpAffine(
                                img,
                                mat, (warp_resolution, warp_resolution),
                                flags=cv2.INTER_LINEAR)

                            img = imagelib.warp_by_params(
                                warp_params,
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=border_replicate,
                                cv2_inter=cv2.INTER_LINEAR)
                            img = cv2.resize(img, (resolution, resolution),
                                             interpolation=cv2.INTER_LINEAR)
                        else:
                            if face_type != sample_face_type:
                                mat = LandmarksProcessor.get_transform_mat(
                                    sample_landmarks, resolution, face_type)
                                img = cv2.warpAffine(img,
                                                     mat,
                                                     (resolution, resolution),
                                                     borderMode=borderMode,
                                                     flags=cv2.INTER_LINEAR)
                            else:
                                if w != resolution:
                                    img = cv2.resize(
                                        img, (resolution, resolution),
                                        interpolation=cv2.INTER_LINEAR)

                            img = imagelib.warp_by_params(
                                warp_params,
                                img,
                                warp,
                                transform,
                                can_flip=True,
                                border_replicate=border_replicate,
                                cv2_inter=cv2.INTER_LINEAR)

                        if face_mask_type == SPFMT.EYES_MOUTH:
                            div = img.max()
                            if div != 0.0:
                                img = img / div  # normalize to 1.0 after warp

                        if len(img.shape) == 2:
                            img = img[..., None]

                        if channel_type == SPCT.G:
                            out_sample = img.astype(np.float32)
                        else:
                            raise ValueError(
                                "only channel_type.G supported for the mask")

                    elif sample_type == SPST.FACE_IMAGE:
                        img = sample_bgr

                        if face_type != sample_face_type:
                            mat = LandmarksProcessor.get_transform_mat(
                                sample_landmarks, resolution, face_type)
                            img = cv2.warpAffine(img,
                                                 mat, (resolution, resolution),
                                                 borderMode=borderMode,
                                                 flags=cv2.INTER_CUBIC)
                        else:
                            if w != resolution:
                                img = cv2.resize(img, (resolution, resolution),
                                                 interpolation=cv2.INTER_CUBIC)

                        # Apply random color transfer
                        if ct_mode is not None and ct_sample is not None:
                            if ct_sample_bgr is None:
                                ct_sample_bgr = ct_sample.load_bgr()
                            img = imagelib.color_transfer(
                                ct_mode, img,
                                cv2.resize(ct_sample_bgr,
                                           (resolution, resolution),
                                           interpolation=cv2.INTER_LINEAR))

                        if random_hsv_shift_amount != 0:
                            a = random_hsv_shift_amount
                            h_amount = max(1, int(360 * a * 0.5))
                            img_h, img_s, img_v = cv2.split(
                                cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
                            img_h = (img_h + rnd_state.randint(
                                -h_amount, h_amount + 1)) % 360
                            img_s = np.clip(
                                img_s + (rnd_state.random() - 0.5) * a, 0, 1)
                            img_v = np.clip(
                                img_v + (rnd_state.random() - 0.5) * a, 0, 1)
                            img = np.clip(
                                cv2.cvtColor(cv2.merge([img_h, img_s, img_v]),
                                             cv2.COLOR_HSV2BGR), 0, 1)

                        img = imagelib.warp_by_params(
                            warp_params,
                            img,
                            warp,
                            transform,
                            can_flip=True,
                            border_replicate=border_replicate)

                        img = np.clip(img.astype(np.float32), 0, 1)

                        # Transform from BGR to desired channel_type
                        if channel_type == SPCT.BGR:
                            out_sample = img
                        elif channel_type == SPCT.G:
                            out_sample = cv2.cvtColor(img,
                                                      cv2.COLOR_BGR2GRAY)[...,
                                                                          None]
                        elif channel_type == SPCT.GGG:
                            out_sample = np.repeat(
                                np.expand_dims(
                                    cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
                                (3, ), -1)

                    # Final transformations
                    if nearest_resize_to is not None:
                        out_sample = cv2_resize(
                            out_sample, (nearest_resize_to, nearest_resize_to),
                            interpolation=cv2.INTER_NEAREST)

                    if not debug:
                        if normalize_tanh:
                            out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
                                                 1.0)
                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))
                elif sample_type == SPST.IMAGE:
                    img = sample_bgr
                    img = imagelib.warp_by_params(warp_params,
                                                  img,
                                                  warp,
                                                  transform,
                                                  can_flip=True,
                                                  border_replicate=True)
                    img = cv2.resize(img, (resolution, resolution),
                                     interpolation=cv2.INTER_CUBIC)
                    out_sample = img

                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2, 0, 1))

                elif sample_type == SPST.LANDMARKS_ARRAY:
                    l = sample_landmarks
                    l = np.concatenate([
                        np.expand_dims(l[:, 0] / w, -1),
                        np.expand_dims(l[:, 1] / h, -1)
                    ], -1)
                    l = np.clip(l, 0.0, 1.0)
                    out_sample = l
                elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                    pitch, yaw, roll = sample.get_pitch_yaw_roll()
                    if warp_params['flip']:
                        yaw = -yaw

                    if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                        pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
                        yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
                        roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)

                    out_sample = (pitch, yaw)
                else:
                    raise ValueError('expected sample_type')

                outputs_sample.append(out_sample)
            outputs += [outputs_sample]

        return outputs