# -*- 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