|
import os
|
|
import sys
|
|
current_file_path = os.path.abspath(__file__)
|
|
parent_dir = os.path.dirname(os.path.dirname(current_file_path))
|
|
project_root_dir = os.path.dirname(parent_dir)
|
|
sys.path.append(parent_dir)
|
|
sys.path.append(project_root_dir)
|
|
|
|
import cv2
|
|
import random
|
|
import yaml
|
|
import torch
|
|
import numpy as np
|
|
from copy import deepcopy
|
|
import albumentations as A
|
|
from .abstract_dataset import DeepfakeAbstractBaseDataset
|
|
from PIL import Image
|
|
|
|
c=0
|
|
|
|
class LRLDataset(DeepfakeAbstractBaseDataset):
|
|
def __init__(self, config=None, mode='train'):
|
|
super().__init__(config, mode)
|
|
global c
|
|
c=config
|
|
|
|
def multi_pass_filter(self, img, r1=0.33, r2=0.66):
|
|
rows, cols = img.shape
|
|
k = cols / rows
|
|
|
|
mask = np.zeros((rows, cols), np.uint8)
|
|
x, y = np.ogrid[:rows, :cols]
|
|
mask_area = (k * x + y < r1 * cols)
|
|
mask[mask_area] = 1
|
|
low_mask = mask
|
|
|
|
mask = np.ones((rows, cols), np.uint8)
|
|
x, y = np.ogrid[:rows, :cols]
|
|
mask_area = (k * x + y < r2 * cols)
|
|
mask[mask_area] = 0
|
|
high_mask = mask
|
|
|
|
mask1 = np.zeros((rows, cols), np.uint8)
|
|
mask1[low_mask == 0] = 1
|
|
mask2 = np.zeros((rows, cols), np.uint8)
|
|
mask2[high_mask == 0] = 1
|
|
mid_mask = mask1 * mask2
|
|
|
|
return low_mask, mid_mask, high_mask
|
|
|
|
def image2dct(self,img):
|
|
img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
img_gray = np.float32(img_gray)
|
|
img_dct = cv2.dct(img_gray)
|
|
|
|
|
|
low_mask, mid_mask, high_mask = self.multi_pass_filter(img_dct, r1=0.33, r2=0.33)
|
|
img_dct_filterd = high_mask * img_dct
|
|
img_idct = cv2.idct(img_dct_filterd)
|
|
|
|
return img_idct
|
|
|
|
def __getitem__(self, index):
|
|
image_trans, label, landmark_tensors, mask_trans = super().__getitem__(index, no_norm=True)
|
|
|
|
img_idct = self.image2dct(image_trans)
|
|
|
|
img_idct = (img_idct / 255 - 0.5) / 0.5
|
|
|
|
|
|
|
|
image_trans = self.normalize(self.to_tensor(image_trans))
|
|
img_idct_trans = self.to_tensor(img_idct)
|
|
mask_trans = torch.from_numpy(mask_trans)
|
|
mask_trans = mask_trans.squeeze(2).permute(2, 0, 1)
|
|
mask_trans = torch.mean(mask_trans, dim=0, keepdim=True)
|
|
return image_trans, label, img_idct_trans, mask_trans
|
|
|
|
def __len__(self):
|
|
return len(self.image_list)
|
|
|
|
|
|
@staticmethod
|
|
def collate_fn(batch):
|
|
"""
|
|
Collate a batch of data points.
|
|
|
|
Args:
|
|
batch (list): A list of tuples containing the image tensor and label tensor.
|
|
|
|
Returns:
|
|
A tuple containing the image tensor, the label tensor, the landmark tensor,
|
|
and the mask tensor.
|
|
"""
|
|
global c
|
|
images, labels, img_idct_trans, masks = zip(*batch)
|
|
|
|
images = torch.stack(images, dim=0)
|
|
labels = torch.LongTensor(labels)
|
|
masks = torch.stack(masks, dim=0)
|
|
img_idct_trans = torch.stack(img_idct_trans, dim=0)
|
|
|
|
data_dict = {
|
|
'image': images,
|
|
'label': labels,
|
|
'landmark': None,
|
|
'idct': img_idct_trans,
|
|
'mask': masks,
|
|
}
|
|
return data_dict
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
with open(r'H:\code\DeepfakeBench\training\config\detector\lrl_effnb4.yaml', 'r') as f:
|
|
config = yaml.safe_load(f)
|
|
with open(r'H:\code\DeepfakeBench\training\config\train_config.yaml', 'r') as f:
|
|
config2 = yaml.safe_load(f)
|
|
random.seed(config['manualSeed'])
|
|
torch.manual_seed(config['manualSeed'])
|
|
if config['cuda']:
|
|
torch.cuda.manual_seed_all(config['manualSeed'])
|
|
config2['data_manner'] = 'lmdb'
|
|
config['dataset_json_folder'] = 'preprocessing/dataset_json_v3'
|
|
config.update(config2)
|
|
train_set = LRLDataset(config=config, mode='train')
|
|
train_data_loader = \
|
|
torch.utils.data.DataLoader(
|
|
dataset=train_set,
|
|
batch_size=4,
|
|
shuffle=True,
|
|
num_workers=0,
|
|
collate_fn=train_set.collate_fn,
|
|
)
|
|
from tqdm import tqdm
|
|
for iteration, batch in enumerate(tqdm(train_data_loader)):
|
|
print(iteration)
|
|
if iteration > 10:
|
|
break |
