Blame for clean/video/fakestormer/datasets/face_forensic_hm.py

# -*- coding: utf-8 -*-
import random

import cv2
import numpy as np
import torch
from imgaug import augmenters as iaa
from package_utils.bi_online_generation import (
    blendImages,
    colorTransfer,
    random_erode_dilate,
    random_get_hull,
)
from package_utils.image_utils import load_image
from package_utils.transform import (
    get_affine_transform,
    get_center_scale,
)
from package_utils.utils import draw_landmarks, vis_heatmap
from PIL import Image
from skimage import transform as sktransform

from .builder import DATASETS, PIPELINES, build_pipeline
from .master import MasterDataset


@DATASETS.register_module()
class HeatmapFaceForensic(MasterDataset):
    def __init__(self, config, split, **kwargs):
        """
        @params:
        config: Dataset config
        split: train/val/test which directs to the split folders
        """
        self.split = split
        super().__init__(config, **kwargs)
        if kwargs is not None:
            for k, v in kwargs.items():
                if v is None:
                    raise ValueError(f"{k}:{v} retrieve a None value!")
                self.__setattr__(k, v)
        self.rot = 0
        self.pixel_std = 200
        self.target_w = self._cfg.IMAGE_SIZE[1]
        self.target_h = self._cfg.IMAGE_SIZE[0]
        self.aspect_ratio = self.target_w * 1.0 / self.target_h
        self.sigma = self._cfg.SIGMA
        self.heatmap_type = self._cfg.HEATMAP_TYPE
        self.debug = self._cfg.DEBUG
        # self.train = self._cfg.TRAIN
        self.dynamic_fxray = self._cfg.DYNAMIC_FXRAY

        # Load data
        self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data(
            split
        )

        # predefine mask distortion
        self.distortion = iaa.Sequential([iaa.PiecewiseAffine(scale=(0.01, 0.15))])

        # Calling transform methods for inputs
        self.geo_transform = build_pipeline(config.TRANSFORM.geometry, PIPELINES)
        self.colorjitter_transform = build_pipeline(config.TRANSFORM.color, PIPELINES)

    def _load_data(self, split):
        from_file = self._cfg.DATA[self.split.upper()].FROM_FILE

        if not from_file:
            image_paths, labels, mask_paths, ot_props = self._load_from_path(split)
        else:
            image_paths, labels, mask_paths, ot_props = self._load_from_file(split)

        if self.sampler_active and self.train:
            print("Running sampler...")
            params = dict(mask_paths=mask_paths, ot_props=ot_props)
            data_sampler = self._sampler(image_paths, labels, **params)
            image_paths, labels = data_sampler["image_paths"], data_sampler["labels"]
            if len(mask_paths):
                mask_paths = data_sampler["mask_paths"]
            if len(ot_props):
                ot_props = data_sampler["ot_props"]
            print(f"n samples after running sampling --- {len(image_paths)}")

        assert (
            len(image_paths) != 0
        ), "Image paths have not been loaded! Please check image directory!"
        assert (
            len(labels) != 0
        ), "Labels have not been loaded! Please check annotation file!"
        # if not self.dynamic_fxray or self.split == 'val':
        if from_file and not (self.dynamic_fxray):
            assert (
                len(mask_paths) != 0
            ), "Mask paths have not been loaded! Please check mask directory!"
        return image_paths, labels, mask_paths, ot_props

    def _reload_data(self):
        self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data(
            self.split
        )

    def _gen_BI(
        self, background_face, background_landmark, foreground_face_path, idx=None
    ):
        foreground_face = load_image(foreground_face_path)

        # down sample before blending
        aug_size = random.randint(128, 317)
        background_landmark = background_landmark * (aug_size / 317)
        foreground_face = sktransform.resize(
            foreground_face, (aug_size, aug_size), preserve_range=True
        ).astype(np.uint8)
        background_face = sktransform.resize(
            background_face, (aug_size, aug_size), preserve_range=True
        ).astype(np.uint8)

        # get random type of initial blending mask
        mask = random_get_hull(background_landmark, background_face)

        if self.debug:
            Image.fromarray((mask * 255).astype(np.uint8)).save(
                f"samples/debugs/orig_CH_{idx}.jpg"
            )

        #  random deform mask
        mask = self.distortion.augment_image(mask)
        mask = random_erode_dilate(mask)

        if self.debug:
            Image.fromarray((mask * 255).astype(np.uint8)).save(
                f"samples/debugs/deformed_CH_{idx}.jpg"
            )

        # filte empty mask after deformation
        if np.sum(mask) == 0:
            raise NotImplementedError

        # apply color transfer
        foreground_face = colorTransfer(background_face, foreground_face, mask * 255)

        # blend two face
        blended_face, mask = blendImages(foreground_face, background_face, mask * 255)
        blended_face = blended_face.astype(np.uint8)

        # resize back to default resolution
        blended_face = sktransform.resize(
            blended_face, (317, 317), preserve_range=True
        ).astype(np.uint8)
        mask = sktransform.resize(mask, (317, 317), preserve_range=True)
        mask = mask[:, :, 0:1]
        return blended_face, mask

    def _gen_target(
        self, background_face, background_landmark, foreground_face_path, idx=None
    ):
        data_type = "real" if random.randint(0, 1) else "fake"

        if not background_landmark.any():
            data_type = "real"

        if data_type == "fake":
            face_img, mask = self._gen_BI(
                background_face, background_landmark, foreground_face_path, idx=idx
            )
            mask = (1 - mask) * mask * 4
        else:
            face_img = background_face
            mask = np.zeros((317, 317, 1))

        face_img = Image.fromarray(face_img)
        # randomly downsample after BI pipeline
        if random.randint(0, 1):
            aug_size = random.randint(64, 317)
            if random.randint(0, 1):
                face_img = face_img.resize((aug_size, aug_size), Image.BILINEAR)
            else:
                face_img = face_img.resize((aug_size, aug_size), Image.NEAREST)
        face_img = face_img.resize((317, 317), Image.BILINEAR)
        face_img = np.array(face_img)

        face_img = face_img[60:(317), 30:(287), :]
        mask = mask[60:(317), 30:(287), :]
        mask = np.repeat(mask, 3, 2)
        mask = (mask * 255).astype(np.uint8)
        return face_img, mask, int(data_type == "fake")

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        flag = True
        while flag:
            try:
                img_path = self.image_paths[idx]
                label = self.labels[idx]
                vid_id = img_path.split("/")[-2]
                img = load_image(img_path)
                mask = None

                if self.debug:
                    Image.fromarray(img).save(f"samples/debugs/orig_{idx}.jpg")

                # Applying color transform to inputs
                if self.split == "train":
                    color_transfomed = self.colorjitter_transform(img)
                    img = color_transfomed["image"]

                # if not self.dynamic_fxray or self.split == 'val':
                if not self.dynamic_fxray:
                    if bool(self.mask_paths):
                        mask_path = self.mask_paths[idx]
                        mask = load_image(mask_path)
                else:
                    if self.train:
                        best_match = (
                            self.ot_props[idx]["best_match"]
                            if len(self.ot_props[idx]["best_match"])
                            else []
                        )
                        if len(self.ot_props[idx]["aligned_lms"]):
                            f_lms = self.ot_props[idx]["aligned_lms"]
                        elif len(self.ot_props[idx]["orig_lms"]):
                            f_lms = self.ot_props[idx]["orig_lms"]
                        else:
                            f_lms = []

                        if self.debug:
                            img_lms_draw = draw_landmarks(img, f_lms)
                            Image.fromarray(img_lms_draw).save(
                                f"samples/debugs/orig_{idx}_lms.jpg"
                            )

                        if len(best_match):
                            best_match_idx = random.randint(0, len(best_match) - 1)
                            best_match_path = best_match[best_match_idx]
                            img, mask, label = self._gen_target(
                                img, f_lms, best_match_path, idx=idx
                            )
                        else:
                            img, mask, label = self._gen_target(img, np.array([]), "")
                    else:
                        img = cv2.resize(img, (317, 317))
                        # Best croppings from 18-299 for testing
                        img = img[18:(299), 18:(299), :]
                target = None

                if mask is not None:
                    assert (
                        mask.shape[:2] == img.shape[:2]
                    ), "Color Image and Mask must have the same shape!"

                # Applying geo transform to inputs and masks
                if self.split == "train":
                    geo_transfomed = self.geo_transform(img, mask=mask)
                    img = geo_transfomed["image"]
                    mask = geo_transfomed["mask"]

                # Applying affine transform
                c, s = get_center_scale(
                    img.shape[:2], self.aspect_ratio, pixel_std=self.pixel_std
                )
                trans = get_affine_transform(c, s, self.rot, self._cfg.IMAGE_SIZE)
                trans_heatmap = get_affine_transform(
                    c, s, self.rot, self._cfg.HEATMAP_SIZE
                )

                input = cv2.warpAffine(
                    img,
                    trans,
                    (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])),
                    flags=cv2.INTER_LINEAR,
                )

                if mask is not None:
                    target = cv2.warpAffine(
                        mask,
                        trans_heatmap,
                        (
                            int(self._cfg.HEATMAP_SIZE[0]),
                            int(self._cfg.HEATMAP_SIZE[1]),
                        ),
                        flags=cv2.INTER_LINEAR,
                    )

                # Target encoding
                # 0 for original, 1 for FXRay, 2 for NoFXRay. If 2, comment: mask = (1 - mask) * mask * 4
                heatmap, cstency_hm = (
                    self.select_encode_method(version=1)(target)
                    if (
                        target is not None
                        and self.heatmap_type == "gaussian"
                        and self.train
                    )
                    else (None, None)
                )

                if self.debug:
                    Image.fromarray(input).save(f"samples/debugs/affine_{idx}.jpg")
                    Image.fromarray(target).save(
                        f"samples/debugs/mask_affine_{idx}.jpg"
                    )
                    vis_heatmap(
                        input,
                        cstency_hm / 255,
                        f"samples/debugs/cstency_mask_{idx}.jpg",
                    )
                    vis_heatmap(input, heatmap, f"samples/debugs/hm_{idx}.jpg")

                if self.train:
                    if self.split_image:
                        patch_img_trans = []
                        patch_heatmap = []
                        patch_cstency_hm = []
                        patch_target = []
                        patch_label = np.expand_dims(np.tile(label, len(heatmap)), -1)

                        for i, (k, l) in enumerate(
                            [[0, 0], [1 / 2, 0], [0, 1 / 2], [1 / 2, 1 / 2]]
                        ):
                            input_ = input[
                                int(self.target_h * k) : int(
                                    self.target_h * (k + 1 / 2)
                                ),
                                int(self.target_w * l) : int(
                                    self.target_w * (l + 1 / 2)
                                ),
                                :,
                            ]
                            heatmap_ = heatmap[i][
                                int(self.heatmap_h * k) : int(
                                    self.heatmap_h * (k + 1 / 2)
                                ),
                                int(self.heatmap_w * l) : int(
                                    self.heatmap_w * (l + 1 / 2)
                                ),
                            ]
                            cstency_ = cstency_hm[i][
                                int(self.heatmap_h * k) : int(
                                    self.heatmap_h * (k + 1 / 2)
                                ),
                                int(self.heatmap_w * l) : int(
                                    self.heatmap_w * (l + 1 / 2)
                                ),
                            ]
                            target_ = target[..., 0][
                                int(self.heatmap_h * k) : int(
                                    self.heatmap_h * (k + 1 / 2)
                                ),
                                int(self.heatmap_w * l) : int(
                                    self.heatmap_w * (l + 1 / 2)
                                ),
                            ]

                            # Normalise + Convert numpy array to tensor
                            input_ = input_ / 255
                            patch_img_trans.append(self.final_transforms(input_))

                            patch_heatmap.append(heatmap_)
                            patch_cstency_hm.append(cstency_ / 255)
                            patch_target.append(target_ / 255)
                    else:
                        patch_img_trans = self.final_transforms(input / 255)
                        patch_heatmap = heatmap
                        patch_cstency_hm = cstency_hm / 255
                        patch_target = target / 255
                        patch_label = np.expand_dims(label, axis=-1)
                        # patch_label = label
                else:
                    # Normalise + Convert numpy array to tensor
                    img_trans = input / 255
                    img_trans = self.final_transforms(img_trans)
                    label = np.expand_dims(label, axis=-1)
                flag = False
            except Exception as e:
                print("There is an exception during loading data, please check --- ", e)
                idx = torch.randint(low=0, high=self.__len__(), size=(1,)).item()

        if self.train:
            return (
                patch_img_trans,
                patch_label,
                patch_target,
                patch_heatmap,
                patch_cstency_hm,
            )
        else:
            meta = {"vid_id": vid_id, "img_path": img_path}
            return img_trans, label, meta

    def train_collate_fn(self, batch):
        batch_data = {}
        img, label, target, hm, cstency_hm = zip(*batch)

        # Collating data in case of using spliting images into patches
        if self.split_image:
            hm_H, hm_W = hm[0][0].shape

            img = np.reshape(img, (-1))
            hm = np.reshape(hm, (-1, 1, hm_H, hm_W))
            cstency_hm = np.reshape(cstency_hm, (-1, 1, hm_H, hm_W))
            target = np.reshape(target, (-1, 1, hm_H, hm_W))
            label = np.reshape(label, (-1, 1))

        img = torch.tensor([it.numpy() for it in img])
        heatmap = torch.tensor(hm).float()
        cstency_heatmap = torch.tensor(cstency_hm).float()
        target = torch.tensor(target).float()
        label = torch.tensor(label)

        batch_data["img"] = img
        batch_data["label"] = label
        batch_data["target"] = target
        batch_data["heatmap"] = heatmap
        batch_data["cstency"] = cstency_heatmap

        return batch_data


if __name__ == "__main__":
    # from datasets import *
    from configs.get_config import load_config
    from pipelines.color_transform import ColorJitterTransform
    from pipelines.geo_transform import GeometryTransform
    from torch.utils.data import DataLoader

    PIPELINES.register_module(module=GeometryTransform)
    PIPELINES.register_module(module=ColorJitterTransform)

    config = load_config("configs/efn4_fpn_hm_adv.yaml")
    hm_ff = DATASETS.build(
        cfg=config.DATASET, default_args=dict(split="train", config=config.DATASET)
    )
    hm_ff_loader = DataLoader(
        hm_ff, batch_size=10, shuffle=False, collate_fn=hm_ff.train_collate_fn
    )
    for b, batch_data in enumerate(hm_ff_loader):
        inputs, labels, targets, heatmaps, cstency_heatmap = (
            batch_data["img"],
            batch_data["label"],
            batch_data["target"],
            batch_data["heatmap"],
            batch_data["cstency_heatmap"],
        )
        print(
            f"X.shape - {inputs.shape}, y shape - {labels.shape}, heatmaps - {heatmaps.shape}, consistency -- {cstency_heatmap.shape}"
        )
        break