├── runfiles ├── no-noise.sh ├── quantization.sh ├── crop-quantization.sh └── combined-noise.sh ├── scripts_adv_loss ├── val_jpeg.sh ├── val_quant.sh ├── val_webp.sh ├── run_jpeg.sh ├── run_webp.sh ├── run_quant.sh ├── val_diff_jpeg.sh ├── val_jpeg2000.sh ├── run_jpeg2000.sh ├── val_jpeg2.sh ├── run_diff_jpeg.sh ├── run_mpeg4.sh ├── val_jpeg_plus_quant.sh ├── run_jpeg_quant.sh ├── val_diff_qf_jpeg2.sh ├── val_diff_corruptions.sh ├── run_h264_video.sh ├── run_xvid_video.sh ├── run_mpeg4_video.sh ├── run_diff_corruptions.sh ├── run_jpeg2.sh └── run_diff_qf_jpeg2.sh ├── noise_layers ├── identity.py ├── resize.py ├── dropout.py ├── cropout.py ├── crop2.py ├── noiser.py ├── diff_jpeg.py ├── gaussian.py ├── h264_compression.py ├── xvid_compression.py ├── mpeg4_compression.py ├── webp.py ├── jpeg_compression2.py ├── combined2.py ├── jpeg_compression2000.py ├── dct_filters.py ├── diff_quality_jpeg_compression2.py ├── quantization.py ├── crop.py ├── diff_corruptions.py └── jpeg_compression.py ├── average_meter.py ├── README.md ├── model ├── conv_bn_relu.py ├── discriminator.py ├── decoder.py ├── encoder_universal.py ├── encoder.py ├── encoder_decoder.py └── hidden.py ├── .gitignore ├── vgg_loss.py ├── tensorboard_logger.py ├── options.py ├── oops_dataset.py ├── test_model.py ├── validate-trained-models.py ├── validate-trained-models_diff_jpeg_levels.py ├── validate-trained-models_diff_webp_levels.py ├── validate-trained-models_diff_jpeg2000_levels.py ├── validate-trained-models_diff_corruptions.py ├── noise_argparser.py ├── modules ├── decompression.py └── compression.py ├── main.py ├── utils.py └── train.py /runfiles/no-noise.sh: -------------------------------------------------------------------------------- 1 | nohup python -u main.py new -d /data/coco/10K -e 200 -b 32 --name no-noise & 2 | sleep 1 3 | tail -f nohup.out 4 | -------------------------------------------------------------------------------- /runfiles/quantization.sh: -------------------------------------------------------------------------------- 1 | nohup python -u main.py new -d /data/coco/10K -e 200 -b 32 --noise 'quant()' --name quantization & 2 | sleep 1 3 | tail -f nohup.out 4 | -------------------------------------------------------------------------------- /runfiles/crop-quantization.sh: -------------------------------------------------------------------------------- 1 | nohup python -u main.py new -d /data/coco/10K -e 200 -b 32 --noise 'crop((0.4,0.55),(0.4,0.55))+quant()' --name crop_quantization & 2 | sleep 1 3 | tail -f nohup.out 4 | -------------------------------------------------------------------------------- /runfiles/combined-noise.sh: -------------------------------------------------------------------------------- 1 | nohup python -u main.py new -d /data/coco/10K -e 400 -b 32 --noise "crop((0.4,0.55),(0.4,0.55))+cropout((0.25,0.35),(0.25,0.35))+dropout(0.25,0.35)+resize(0.4,0.6)+jpeg()" --name combined-noise & 2 | sleep 1 3 | tail -f nohup.out 4 | -------------------------------------------------------------------------------- /scripts_adv_loss/val_jpeg.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/jpeg_cover_dependent 2020.11.03--05-56-39/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_quant.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/quant_cover_dependent 2020.11.03--05-57-27/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_webp.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_webp_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/webp_cover_dependent 2020.11.03--05-58-01/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_jpeg.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'jpeg_cover_dependent' \ 5 | --noise 'jpeg' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_webp.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'webp_cover_dependent' \ 5 | --noise 'webp' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_quant.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'quant_cover_dependent' \ 5 | --noise 'quant' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_diff_jpeg.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/diff_jpeg_cover_dependent 2020.11.03--05-56-03/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_jpeg2000.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg2000_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/jpeg2000_cover_dependent 2020.11.03--05-57-11/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_jpeg2000.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'jpeg2000_cover_dependent' \ 5 | --noise 'jpeg2000' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_jpeg2.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/jpeg2_cover_dependent_adv_loss_1e-2 2020.11.02--12-49-11/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_diff_jpeg.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'diff_jpeg_cover_dependent' \ 5 | --noise 'diff_jpeg' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_mpeg4.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'mpeg4_cover_dependent_adv_loss_1e-1' \ 5 | --noise 'mpeg4' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-1 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_jpeg_plus_quant.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/jpeg_plus_quant_cover_dependent 2020.11.03--14-58-56/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_jpeg_quant.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'jpeg_plus_quant_cover_dependent' \ 5 | --noise 'jpeg+quant' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_diff_qf_jpeg2.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_jpeg_levels.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/diff_qf_jpeg2_cover_dependent_adv_loss_1e-2 2020.12.05--03-24-46/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/val_diff_corruptions.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 validate-trained-models_diff_corruptions.py \ 3 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 4 | --runs_root './runs_adv_loss/1e-2/diff_corruptions_cover_dependent_adv_loss_1e-2 2020.12.05--04-02-01/' \ 5 | --batch-size 32 \ -------------------------------------------------------------------------------- /scripts_adv_loss/run_h264_video.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'h264_cover_dependent_video_trained' \ 5 | --noise 'h264' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --adv_loss 1e-2 \ 9 | --video_dataset 1 \ 10 | --save-dir 'runs_adv_loss/1e-2' -------------------------------------------------------------------------------- /scripts_adv_loss/run_xvid_video.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'xvid_cover_dependent_video_trained' \ 5 | --noise 'xvid' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --adv_loss 1e-2 \ 9 | --video_dataset 1 \ 10 | --save-dir 'runs_adv_loss/1e-2' -------------------------------------------------------------------------------- /noise_layers/identity.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class Identity(nn.Module): 5 | """ 6 | Identity-mapping noise layer. Does not change the image 7 | """ 8 | def __init__(self): 9 | super(Identity, self).__init__() 10 | 11 | def forward(self, noised_and_cover): 12 | return noised_and_cover 13 | -------------------------------------------------------------------------------- /scripts_adv_loss/run_mpeg4_video.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'mpeg4_cover_dependent_video_trained' \ 5 | --noise 'mpeg4' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --adv_loss 1e-2 \ 9 | --video_dataset 1 \ 10 | --save-dir 'runs_adv_loss/1e-2' -------------------------------------------------------------------------------- /scripts_adv_loss/run_diff_corruptions.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'different_corruptions_cover_dependent_adv_loss_1e-2' \ 5 | --noise 'diff_corruptions' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 48 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ 10 | --save-dir 'runs_adv_loss/1e-2' -------------------------------------------------------------------------------- /scripts_adv_loss/run_jpeg2.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'jpeg2_cover_dependent_adv_loss_100' \ 5 | --noise 'jpeg2' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 32 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-2 \ 10 | --save-dir 'runs_adv_loss/1e-2' 11 | # --save-dir 'runs_adv_loss/debug' -------------------------------------------------------------------------------- /scripts_adv_loss/run_diff_qf_jpeg2.sh: -------------------------------------------------------------------------------- 1 | #### NON-ROBUST DATASET #### 2 | python3 main.py \ 3 | 'new' \ 4 | --name 'different_qf_jpeg2_cover_dependent_adv_loss_1e-3' \ 5 | --noise 'diff_qf_jpeg2' \ 6 | --data-dir '/media/user/SSD1TB-1/ImageNet/' \ 7 | --batch-size 50 \ 8 | --cover-dependent 1 \ 9 | --adv_loss 1e-3 \ 10 | --save-dir 'runs_adv_loss/1e-3' 11 | # --save-dir 'runs_adv_loss/debug' -------------------------------------------------------------------------------- /average_meter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class AverageMeter(object): 4 | """Computes and stores the average and current value""" 5 | def __init__(self): 6 | self.reset() 7 | 8 | def reset(self): 9 | self.val = 0 10 | self.avg = 0 11 | self.sum = 0 12 | self.count = 0 13 | 14 | def update(self, val, n=1): 15 | if val != np.nan and val != np.inf: 16 | self.val = val 17 | self.sum += val * n 18 | self.count += n 19 | self.avg = self.sum / self.count -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Towards Robust Data Hiding Against (JPEG) Compression: A Pseudo-Differentiable Deep Learning Approach 2 | 3 | Official pytorch code for the paper "Towards Robust Data Hiding Against (JPEG) Compression: A Pseudo-Differentiable Deep Learning Approach" by Chaoning Zhang*, Adil Karjauv*, Philipp Benz*, and In So Kweon (*Equal Contribution). 4 | 5 | ## Experiments 6 | 7 | In order to run the experiments, just use the scripts from the "scripts_adv_loss" folder. 8 | For example: 9 | ``` 10 | bash run_jpeg2.sh 11 | ``` 12 | 13 | ## Note 14 | 15 | More detailed instructions about the code are coming soon. -------------------------------------------------------------------------------- /model/conv_bn_relu.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class ConvBNRelu(nn.Module): 4 | """ 5 | Building block used in HiDDeN network. Is a sequence of Convolution, Batch Normalization, and ReLU activation 6 | """ 7 | def __init__(self, channels_in, channels_out, stride=1): 8 | 9 | super(ConvBNRelu, self).__init__() 10 | 11 | self.layers = nn.Sequential( 12 | nn.Conv2d(channels_in, channels_out, 3, stride, padding=1), 13 | nn.BatchNorm2d(channels_out), 14 | nn.ReLU(inplace=True) 15 | ) 16 | 17 | def forward(self, x): 18 | return self.layers(x) 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Prerequisites 2 | *.d 3 | 4 | # Compiled Object files 5 | *.slo 6 | *.lo 7 | *.o 8 | *.obj 9 | 10 | # Precompiled Headers 11 | *.gch 12 | *.pch 13 | 14 | # Compiled Dynamic libraries 15 | *.so 16 | *.dylib 17 | *.dll 18 | 19 | # Fortran module files 20 | *.mod 21 | *.smod 22 | 23 | # Compiled Static libraries 24 | *.lai 25 | *.la 26 | *.a 27 | *.lib 28 | 29 | # Executables 30 | *.exe 31 | *.out 32 | *.app 33 | 34 | # Folders 35 | __pycache__/ 36 | .idea/ 37 | .ipynb_checkpoints/ 38 | .vscode/ 39 | data/ 40 | runs/ 41 | runs_adv_loss/ 42 | runs_cn/ 43 | runs_new/ 44 | runs_res/ 45 | runs_test/ 46 | scripts/ 47 | scripts_cn/ 48 | scripts_res/ 49 | experiments/ 50 | oops_dataset/ 51 | 52 | # Files 53 | validate-trained-models_video.py 54 | README_old.md 55 | LICENSE 56 | -------------------------------------------------------------------------------- /noise_layers/resize.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | import numpy as np 4 | from noise_layers.crop import random_float 5 | 6 | class Resize(nn.Module): 7 | """ 8 | Resize the image. The target size is original size * resize_ratio 9 | """ 10 | def __init__(self, resize_ratio_range, interpolation_method='nearest'): 11 | super(Resize, self).__init__() 12 | self.resize_ratio_min = resize_ratio_range[0] 13 | self.resize_ratio_max = resize_ratio_range[1] 14 | self.interpolation_method = interpolation_method 15 | 16 | 17 | def forward(self, noised_and_cover): 18 | 19 | resize_ratio = random_float(self.resize_ratio_min, self.resize_ratio_max) 20 | noised_image = noised_and_cover[0] 21 | noised_and_cover[0] = F.interpolate( 22 | noised_image, 23 | scale_factor=(resize_ratio, resize_ratio), 24 | mode=self.interpolation_method) 25 | 26 | return noised_and_cover 27 | -------------------------------------------------------------------------------- /noise_layers/dropout.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | 5 | class Dropout(nn.Module): 6 | """ 7 | Drops random pixels from the noised image and substitues them with the pixels from the cover image 8 | """ 9 | def __init__(self, keep_ratio_range): 10 | super(Dropout, self).__init__() 11 | self.keep_min = keep_ratio_range[0] 12 | self.keep_max = keep_ratio_range[1] 13 | 14 | 15 | def forward(self, noised_and_cover): 16 | 17 | noised_image = noised_and_cover[0] 18 | cover_image = noised_and_cover[1] 19 | 20 | mask_percent = np.random.uniform(self.keep_min, self.keep_max) 21 | 22 | mask = np.random.choice([0.0, 1.0], noised_image.shape[2:], p=[1 - mask_percent, mask_percent]) 23 | mask_tensor = torch.tensor(mask, device=noised_image.device, dtype=torch.float) 24 | # mask_tensor.unsqueeze_(0) 25 | # mask_tensor.unsqueeze_(0) 26 | mask_tensor = mask_tensor.expand_as(noised_image) 27 | noised_image = noised_image * mask_tensor + cover_image * (1-mask_tensor) 28 | return [noised_image, cover_image] 29 | 30 | 31 | -------------------------------------------------------------------------------- /vgg_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torchvision 4 | 5 | class VGGLoss(nn.Module): 6 | """ 7 | Part of pre-trained VGG16. This is used in case we want perceptual loss instead of Mean Square Error loss. 8 | See for instance https://arxiv.org/abs/1603.08155 9 | """ 10 | def __init__(self, block_no: int, layer_within_block: int, use_batch_norm_vgg: bool): 11 | super(VGGLoss, self).__init__() 12 | if use_batch_norm_vgg: 13 | vgg16 = torchvision.models.vgg16_bn(pretrained=True) 14 | else: 15 | vgg16 = torchvision.models.vgg16(pretrained=True) 16 | curr_block = 1 17 | curr_layer = 1 18 | layers = [] 19 | for layer in vgg16.features.children(): 20 | layers.append(layer) 21 | if curr_block == block_no and curr_layer == layer_within_block: 22 | break 23 | if isinstance(layer, nn.MaxPool2d): 24 | curr_block += 1 25 | curr_layer = 1 26 | else: 27 | curr_layer += 1 28 | 29 | self.vgg_loss = nn.Sequential(*layers) 30 | 31 | def forward(self, img): 32 | return self.vgg_loss(img) 33 | -------------------------------------------------------------------------------- /tensorboard_logger.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorboardX 3 | 4 | 5 | class TensorBoardLogger: 6 | """ 7 | Wrapper class for easy TensorboardX logging 8 | """ 9 | def __init__(self, log_dir): 10 | self.grads = {} 11 | self.tensors = {} 12 | self.writer = tensorboardX.SummaryWriter(log_dir) 13 | 14 | def grad_hook_by_name(self, grad_name): 15 | def backprop_hook(grad): 16 | self.grads[grad_name] = grad 17 | return backprop_hook 18 | 19 | def save_losses(self, losses_accu: dict, epoch: int): 20 | for loss_name, loss_value in losses_accu.items(): 21 | self.writer.add_scalar('losses/{}'.format(loss_name.strip()), loss_value.avg, global_step=epoch) 22 | 23 | def save_grads(self, epoch: int): 24 | for grad_name, grad_values in self.grads.items(): 25 | self.writer.add_histogram(grad_name, grad_values, global_step=epoch) 26 | 27 | def add_tensor(self, name: str, tensor): 28 | self.tensors[name] = tensor 29 | 30 | def save_tensors(self, epoch: int): 31 | for tensor_name, tensor_value in self.tensors.items(): 32 | self.writer.add_histogram('tensor/{}'.format(tensor_name), tensor_value, global_step=epoch) 33 | -------------------------------------------------------------------------------- /model/discriminator.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from options import HiDDenConfiguration 3 | from model.conv_bn_relu import ConvBNRelu 4 | 5 | class Discriminator(nn.Module): 6 | """ 7 | Discriminator network. Receives an image and has to figure out whether it has a watermark inserted into it, or not. 8 | """ 9 | def __init__(self, config: HiDDenConfiguration): 10 | super(Discriminator, self).__init__() 11 | 12 | layers = [ConvBNRelu(3, config.discriminator_channels)] 13 | for _ in range(config.discriminator_blocks-1): 14 | layers.append(ConvBNRelu(config.discriminator_channels, config.discriminator_channels)) 15 | 16 | layers.append(nn.AdaptiveAvgPool2d(output_size=(1, 1))) 17 | self.before_linear = nn.Sequential(*layers) 18 | self.linear = nn.Linear(config.discriminator_channels, 1) 19 | 20 | def forward(self, image): 21 | X = self.before_linear(image) 22 | # the output is of shape b x c x 1 x 1, and we want to squeeze out the last two dummy dimensions and make 23 | # the tensor of shape b x c. If we just call squeeze_() it will also squeeze the batch dimension when b=1. 24 | X.squeeze_(3).squeeze_(2) 25 | X = self.linear(X) 26 | # X = torch.sigmoid(X) 27 | return X -------------------------------------------------------------------------------- /noise_layers/cropout.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from noise_layers.crop import get_random_rectangle_inside 4 | 5 | 6 | class Cropout(nn.Module): 7 | """ 8 | Combines the noised and cover images into a single image, as follows: Takes a crop of the noised image, and takes the rest from 9 | the cover image. The resulting image has the same size as the original and the noised images. 10 | """ 11 | def __init__(self, height_ratio_range, width_ratio_range): 12 | super(Cropout, self).__init__() 13 | self.height_ratio_range = height_ratio_range 14 | self.width_ratio_range = width_ratio_range 15 | 16 | def forward(self, noised_and_cover): 17 | noised_image = noised_and_cover[0] 18 | cover_image = noised_and_cover[1] 19 | assert noised_image.shape == cover_image.shape 20 | 21 | cropout_mask = torch.zeros_like(noised_image) 22 | h_start, h_end, w_start, w_end = get_random_rectangle_inside(image=noised_image, 23 | height_ratio_range=self.height_ratio_range, 24 | width_ratio_range=self.width_ratio_range) 25 | cropout_mask[:, :, h_start:h_end, w_start:w_end] = 1 26 | 27 | noised_and_cover[0] = noised_image * cropout_mask + cover_image * (1-cropout_mask) 28 | return noised_and_cover -------------------------------------------------------------------------------- /noise_layers/crop2.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from noise_layers.crop import get_random_rectangle_inside 4 | 5 | 6 | class Crop2(nn.Module): 7 | """ 8 | Combines the noised and cover images into a single image, as follows: Takes a crop of the noised image, and takes the rest from 9 | the cover image. The resulting image has the same size as the original and the noised images. 10 | """ 11 | def __init__(self, height_ratio_range, width_ratio_range): 12 | super(Crop2, self).__init__() 13 | self.height_ratio_range = height_ratio_range 14 | self.width_ratio_range = width_ratio_range 15 | 16 | def forward(self, noised_and_cover): 17 | noised_image = noised_and_cover[0] 18 | cover_image = noised_and_cover[1] 19 | assert noised_image.shape == cover_image.shape 20 | 21 | cropout_mask = torch.zeros_like(noised_image) 22 | h_start, h_end, w_start, w_end = get_random_rectangle_inside(image=noised_image, 23 | height_ratio_range=self.height_ratio_range, 24 | width_ratio_range=self.width_ratio_range) 25 | cropout_mask[:, :, h_start:h_end, w_start:w_end] = 1 26 | 27 | noised_and_cover[0] = noised_image * cropout_mask + torch.ones_like(cover_image) * (1-cropout_mask) 28 | return noised_and_cover -------------------------------------------------------------------------------- /model/decoder.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from options import HiDDenConfiguration 3 | from model.conv_bn_relu import ConvBNRelu 4 | 5 | 6 | class Decoder(nn.Module): 7 | """ 8 | Decoder module. Receives a watermarked image and extracts the watermark. 9 | The input image may have various kinds of noise applied to it, 10 | such as Crop, JpegCompression, and so on. See Noise layers for more. 11 | """ 12 | def __init__(self, config: HiDDenConfiguration): 13 | 14 | super(Decoder, self).__init__() 15 | self.channels = config.decoder_channels 16 | 17 | layers = [ConvBNRelu(3, self.channels)] 18 | for _ in range(config.decoder_blocks - 1): 19 | layers.append(ConvBNRelu(self.channels, self.channels)) 20 | 21 | # layers.append(block_builder(self.channels, config.message_length)) 22 | layers.append(ConvBNRelu(self.channels, config.message_length)) 23 | 24 | layers.append(nn.AdaptiveAvgPool2d(output_size=(1, 1))) 25 | self.layers = nn.Sequential(*layers) 26 | 27 | self.linear = nn.Linear(config.message_length, config.message_length) 28 | 29 | def forward(self, image_with_wm): 30 | x = self.layers(image_with_wm) 31 | # the output is of shape b x c x 1 x 1, and we want to squeeze out the last two dummy dimensions and make 32 | # the tensor of shape b x c. If we just call squeeze_() it will also squeeze the batch dimension when b=1. 33 | x.squeeze_(3).squeeze_(2) 34 | x = self.linear(x) 35 | return x 36 | -------------------------------------------------------------------------------- /noise_layers/noiser.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch.nn as nn 3 | from noise_layers.identity import Identity 4 | from noise_layers.jpeg_compression import JpegCompression 5 | from noise_layers.quantization import Quantization 6 | from noise_layers.combined2 import Combined2 7 | 8 | 9 | 10 | class Noiser(nn.Module): 11 | """ 12 | This module allows to combine different noise layers into a sequential noise module. The 13 | configuration and the sequence of the noise layers is controlled by the noise_config parameter. 14 | """ 15 | def __init__(self, noise_layers: list, device, jpeg_type): 16 | super(Noiser, self).__init__() 17 | self.noise_layers = [] 18 | for layer in noise_layers: 19 | if type(layer) is str: 20 | if layer == 'JpegPlaceholder': 21 | self.noise_layers.append(JpegCompression(device)) 22 | elif layer == 'QuantizationPlaceholder': 23 | self.noise_layers.append(Quantization(device)) 24 | elif layer == 'Combined2Placeholder': 25 | self.noise_layers.append(Combined2(device, jpeg_type)) 26 | else: 27 | raise ValueError(f'Wrong layer placeholder string in Noiser.__init__().' 28 | f' Expected "JpegPlaceholder" or "QuantizationPlaceholder" but got {layer} instead') 29 | else: 30 | self.noise_layers.append(layer) 31 | # self.noise_layers = nn.Sequential(*noise_layers) 32 | 33 | def forward(self, encoded_and_cover): 34 | random_noise_layer = np.random.choice(self.noise_layers, 1)[0] 35 | return random_noise_layer(encoded_and_cover) 36 | 37 | -------------------------------------------------------------------------------- /noise_layers/diff_jpeg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | # Local 4 | from modules.compression import compress_jpeg 5 | from modules.decompression import decompress_jpeg 6 | 7 | def diff_round(x): 8 | """ Differentiable rounding function 9 | Input: 10 | x(tensor) 11 | Output: 12 | x(tensor) 13 | """ 14 | return torch.round(x) + (x - torch.round(x))**3 15 | 16 | def quality_to_factor(quality): 17 | """ Calculate factor corresponding to quality 18 | Input: 19 | quality(float): Quality for jpeg compression 20 | Output: 21 | factor(float): Compression factor 22 | """ 23 | if quality < 50: 24 | quality = 5000. / quality 25 | else: 26 | quality = 200. - quality*2 27 | return quality / 100. 28 | 29 | class DiffJPEG(nn.Module): 30 | def __init__(self, height=128, width=128, differentiable=True, quality=50): 31 | ''' Initialize the DiffJPEG layer 32 | Inputs: 33 | height(int): Original image height 34 | width(int): Original image width 35 | differentiable(bool): If true uses custom differentiable 36 | rounding function, if false uses standrard torch.round 37 | quality(float): Quality factor for jpeg compression scheme. 38 | ''' 39 | super(DiffJPEG, self).__init__() 40 | if differentiable: 41 | rounding = diff_round 42 | else: 43 | rounding = torch.round 44 | factor = quality_to_factor(quality) 45 | self.compress = compress_jpeg(rounding=rounding, factor=factor) 46 | self.decompress = decompress_jpeg(height, width, rounding=rounding, 47 | factor=factor) 48 | 49 | def forward(self, x): 50 | ''' 51 | ''' 52 | x0, x1 = x 53 | y, cb, cr = self.compress(x0) 54 | recovered = self.decompress(y, cb, cr) 55 | return [recovered, x1] -------------------------------------------------------------------------------- /model/encoder_universal.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from options import HiDDenConfiguration 4 | from model.conv_bn_relu import ConvBNRelu 5 | 6 | 7 | class Encoder(nn.Module): 8 | """ 9 | Inserts a watermark into an image. 10 | """ 11 | def __init__(self, config: HiDDenConfiguration): 12 | super(Encoder, self).__init__() 13 | self.H = config.H 14 | self.W = config.W 15 | self.conv_channels = config.encoder_channels 16 | self.num_blocks = config.encoder_blocks 17 | 18 | layers = [ConvBNRelu(30, self.conv_channels)] 19 | 20 | for _ in range(config.encoder_blocks-1): 21 | layer = ConvBNRelu(self.conv_channels, self.conv_channels) 22 | layers.append(layer) 23 | 24 | self.conv_layers = nn.Sequential(*layers) 25 | self.after_concat_layer = ConvBNRelu(self.conv_channels + 3 + config.message_length, 26 | self.conv_channels) 27 | 28 | self.final_layer = nn.Conv2d(self.conv_channels, 3, kernel_size=1) 29 | 30 | def forward(self, image, message): 31 | 32 | # First, add two dummy dimensions in the end of the message. 33 | # This is required for the .expand to work correctly 34 | expanded_message = message.unsqueeze(-1) 35 | expanded_message.unsqueeze_(-1) 36 | 37 | expanded_message = expanded_message.expand(-1,-1, self.H, self.W) 38 | if self.cover_dependent: 39 | encoded_image = self.conv_layers(image) 40 | # concatenate expanded message and image 41 | concat = torch.cat([expanded_message, encoded_image, image], dim=1) 42 | im_w = self.after_concat_layer(concat) 43 | im_w = self.final_layer(im_w) 44 | else: 45 | encoded_image = self.conv_layers(message) 46 | # concatenate expanded message and image 47 | concat = expanded_message # torch.cat([expanded_message, encoded_image, image], dim=1) 48 | im_w = self.after_concat_layer(concat) 49 | im_w = self.final_layer(im_w) + image 50 | return im_w 51 | -------------------------------------------------------------------------------- /noise_layers/gaussian.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import math 4 | 5 | def gaussian_kernel(kernel_size=3, sigma=2, channels=3): 6 | # Create a x, y coordinate grid of shape (kernel_size, kernel_size, 2) 7 | x_coord = torch.arange(kernel_size) 8 | x_grid = x_coord.repeat(kernel_size).view(kernel_size, kernel_size) 9 | y_grid = x_grid.t() 10 | xy_grid = torch.stack([x_grid, y_grid], dim=-1).float() 11 | 12 | mean = (kernel_size - 1)/2. 13 | variance = sigma**2. 14 | 15 | # Calculate the 2-dimensional gaussian kernel which is 16 | # the product of two gaussian distributions for two different 17 | # variables (in this case called x and y) 18 | gaussian_kernel = (1./(2.*math.pi*variance)) *\ 19 | torch.exp( 20 | -torch.sum((xy_grid - mean)**2., dim=-1) /\ 21 | (2*variance) 22 | ) 23 | 24 | # Make sure sum of values in gaussian kernel equals 1. 25 | gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel) 26 | 27 | # Reshape to 2d depthwise convolutional weight 28 | gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size, kernel_size) 29 | gaussian_kernel = gaussian_kernel.repeat(channels, 1, 1, 1) 30 | 31 | gaussian_filter = nn.Conv2d(in_channels=channels, out_channels=channels, 32 | kernel_size=kernel_size, padding=1, groups=channels, bias=False) 33 | 34 | gaussian_filter.weight.data = gaussian_kernel 35 | gaussian_filter.weight.requires_grad = False 36 | 37 | return gaussian_filter.cuda() 38 | 39 | class Gaussian(nn.Module): 40 | """ 41 | Gaussian kernel. 42 | """ 43 | def __init__(self, kernel_size=3, sigma=2, channels=3): 44 | super(Gaussian, self).__init__() 45 | self.gaussian = gaussian_kernel(kernel_size=kernel_size, sigma=sigma, channels=channels) 46 | 47 | def forward(self, noised_and_cover): 48 | noised_image = noised_and_cover[0] 49 | cover_image = noised_and_cover[1] 50 | noised_image = self.gaussian(noised_image) 51 | return [noised_image, cover_image] 52 | -------------------------------------------------------------------------------- /options.py: -------------------------------------------------------------------------------- 1 | #### TEST 2 | class TrainingOptions: 3 | """ 4 | Configuration options for the training 5 | """ 6 | 7 | def __init__(self, 8 | batch_size: int, 9 | number_of_epochs: int, 10 | train_folder: str, validation_folder: str, runs_folder: str, 11 | start_epoch: int, experiment_name: str, video_dataset: int): 12 | self.batch_size = batch_size 13 | self.number_of_epochs = number_of_epochs 14 | self.train_folder = train_folder 15 | self.validation_folder = validation_folder 16 | self.runs_folder = runs_folder 17 | self.start_epoch = start_epoch 18 | self.experiment_name = experiment_name 19 | self.video_dataset = video_dataset 20 | 21 | 22 | class HiDDenConfiguration(): 23 | """ 24 | The HiDDeN network configuration. 25 | """ 26 | 27 | def __init__(self, H: int, W: int, message_length: int, 28 | encoder_blocks: int, encoder_channels: int, 29 | decoder_blocks: int, decoder_channels: int, 30 | use_discriminator: bool, 31 | use_vgg: bool, 32 | discriminator_blocks: int, discriminator_channels: int, 33 | decoder_loss: float, 34 | encoder_loss: float, 35 | adversarial_loss: float, 36 | cover_dependent: int, 37 | residual: int = 0, 38 | enable_fp16: bool = False): 39 | self.H = H 40 | self.W = W 41 | self.message_length = message_length 42 | self.encoder_blocks = encoder_blocks 43 | self.encoder_channels = encoder_channels 44 | self.use_discriminator = use_discriminator 45 | self.use_vgg = use_vgg 46 | self.decoder_blocks = decoder_blocks 47 | self.decoder_channels = decoder_channels 48 | self.discriminator_blocks = discriminator_blocks 49 | self.discriminator_channels = discriminator_channels 50 | self.decoder_loss = decoder_loss 51 | self.encoder_loss = encoder_loss 52 | self.adversarial_loss = adversarial_loss 53 | self.cover_dependent = cover_dependent 54 | self.residual = residual 55 | self.enable_fp16 = enable_fp16 56 | -------------------------------------------------------------------------------- /noise_layers/h264_compression.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class H264(nn.Module): 15 | def __init__(self): 16 | super(H264, self).__init__() 17 | self.key = generate_random_key() 18 | 19 | def forward(self, noised_and_cover): 20 | video_folder_path = "./h264/" + self.key 21 | if not os.path.exists(video_folder_path): 22 | os.makedirs(video_folder_path) 23 | fourcc = cv2.VideoWriter_fourcc(*'hvc1') 24 | video = cv2.VideoWriter(video_folder_path + '/video.mov', fourcc, 30, (128, 128)) 25 | 26 | noised_image = noised_and_cover[0] 27 | 28 | container_img_copy = noised_image.clone() 29 | containers_ori = container_img_copy.detach().cpu().numpy() 30 | 31 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 32 | N, _, _, _ = containers.shape 33 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 34 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 35 | 36 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 37 | 38 | for i in range(N): 39 | img = containers[i] 40 | video.write(img) 41 | 42 | cv2.destroyAllWindows() 43 | video.release() 44 | 45 | containers_loaded = np.copy(containers) 46 | video_saved = cv2.VideoCapture(video_folder_path + '/video.mov') 47 | 48 | for i in range(N): 49 | _, img = video_saved.read() 50 | containers_loaded[i] = img 51 | 52 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 53 | 54 | containers_loaded = containers_loaded.astype(np.float32) / 255 55 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 56 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 57 | 58 | container_gap = containers_loaded - containers_ori 59 | container_gap = torch.from_numpy(container_gap).float().cuda() 60 | 61 | container_img_noised_mpeg = noised_image + container_gap 62 | 63 | noised_and_cover[0] = container_img_noised_mpeg 64 | 65 | return noised_and_cover -------------------------------------------------------------------------------- /noise_layers/xvid_compression.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class XVID(nn.Module): 15 | def __init__(self): 16 | super(XVID, self).__init__() 17 | self.key = generate_random_key() 18 | 19 | def forward(self, noised_and_cover): 20 | video_folder_path = "./xvid/" + self.key 21 | if not os.path.exists(video_folder_path): 22 | os.makedirs(video_folder_path) 23 | fourcc = cv2.VideoWriter_fourcc(*'XVID') 24 | video = cv2.VideoWriter(video_folder_path + '/video.avi', fourcc, 30, (128, 128)) 25 | 26 | noised_image = noised_and_cover[0] 27 | 28 | container_img_copy = noised_image.clone() 29 | containers_ori = container_img_copy.detach().cpu().numpy() 30 | 31 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 32 | N, _, _, _ = containers.shape 33 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 34 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 35 | 36 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 37 | 38 | for i in range(N): 39 | img = containers[i] 40 | video.write(img) 41 | 42 | cv2.destroyAllWindows() 43 | video.release() 44 | 45 | containers_loaded = np.copy(containers) 46 | video_saved = cv2.VideoCapture(video_folder_path + '/video.avi') 47 | 48 | for i in range(N): 49 | _, img = video_saved.read() 50 | containers_loaded[i] = img 51 | 52 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 53 | 54 | containers_loaded = containers_loaded.astype(np.float32) / 255 55 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 56 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 57 | 58 | container_gap = containers_loaded - containers_ori 59 | container_gap = torch.from_numpy(container_gap).float().cuda() 60 | 61 | container_img_noised_mpeg = noised_image + container_gap 62 | 63 | noised_and_cover[0] = container_img_noised_mpeg 64 | 65 | return noised_and_cover -------------------------------------------------------------------------------- /noise_layers/mpeg4_compression.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class MPEG4(nn.Module): 15 | def __init__(self): 16 | super(MPEG4, self).__init__() 17 | self.key = generate_random_key() 18 | 19 | def forward(self, noised_and_cover): 20 | video_folder_path = "./mpeg4/" + self.key 21 | if not os.path.exists(video_folder_path): 22 | os.makedirs(video_folder_path) 23 | fourcc = cv2.VideoWriter_fourcc(*'mp4v') 24 | video = cv2.VideoWriter(video_folder_path + '/video.mp4', fourcc, 30, (128, 128)) 25 | 26 | noised_image = noised_and_cover[0] 27 | 28 | container_img_copy = noised_image.clone() 29 | containers_ori = container_img_copy.detach().cpu().numpy() 30 | 31 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 32 | N, _, _, _ = containers.shape 33 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 34 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 35 | 36 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 37 | 38 | for i in range(N): 39 | img = containers[i] 40 | video.write(img) 41 | 42 | cv2.destroyAllWindows() 43 | video.release() 44 | 45 | containers_loaded = np.copy(containers) 46 | video_saved = cv2.VideoCapture(video_folder_path + '/video.mp4') 47 | 48 | for i in range(N): 49 | _, img = video_saved.read() 50 | containers_loaded[i] = img 51 | 52 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 53 | 54 | containers_loaded = containers_loaded.astype(np.float32) / 255 55 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 56 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 57 | 58 | container_gap = containers_loaded - containers_ori 59 | container_gap = torch.from_numpy(container_gap).float().cuda() 60 | 61 | container_img_noised_mpeg = noised_image + container_gap 62 | 63 | noised_and_cover[0] = container_img_noised_mpeg 64 | 65 | return noised_and_cover -------------------------------------------------------------------------------- /noise_layers/webp.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class WebP(nn.Module): 15 | def __init__(self, quality=50): 16 | super(WebP, self).__init__() 17 | self.quality = quality 18 | self.key = generate_random_key() 19 | 20 | def forward(self, noised_and_cover): 21 | jpeg_folder_path = "./webp_" + str(self.quality) + "/" + self.key 22 | if not os.path.exists(jpeg_folder_path): 23 | os.makedirs(jpeg_folder_path) 24 | 25 | noised_image = noised_and_cover[0] 26 | 27 | container_img_copy = noised_image.clone() 28 | containers_ori = container_img_copy.detach().cpu().numpy() 29 | 30 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 31 | N, _, _, _ = containers.shape 32 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 33 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 34 | 35 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 36 | 37 | for i in range(N): 38 | img = cv2.cvtColor(containers[i], cv2.COLOR_RGB2BGR) 39 | folder_imgs = jpeg_folder_path + "/webp_" + str(i).zfill(2) + ".webp" 40 | cv2.imwrite(folder_imgs, img, [int(cv2.IMWRITE_WEBP_QUALITY), self.quality]) 41 | 42 | containers_loaded = np.copy(containers) 43 | 44 | for i in range(N): 45 | folder_imgs = jpeg_folder_path + "/webp_" + str(i).zfill(2) + ".webp" 46 | img = cv2.imread(folder_imgs) 47 | containers_loaded[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 48 | 49 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 50 | 51 | containers_loaded = containers_loaded.astype(np.float32) / 255 52 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 53 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 54 | 55 | container_gap = containers_loaded - containers_ori 56 | container_gap = torch.from_numpy(container_gap).float().cuda() 57 | 58 | container_img_noised_jpeg = noised_image + container_gap 59 | 60 | noised_and_cover[0] = container_img_noised_jpeg 61 | 62 | return noised_and_cover -------------------------------------------------------------------------------- /noise_layers/jpeg_compression2.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class JpegCompression2(nn.Module): 15 | def __init__(self, quality=50): 16 | super(JpegCompression2, self).__init__() 17 | self.quality = quality 18 | self.key = generate_random_key() 19 | 20 | def forward(self, noised_and_cover): 21 | jpeg_folder_path = "./jpeg_" + str(self.quality) + "/" + self.key 22 | if not os.path.exists(jpeg_folder_path): 23 | os.makedirs(jpeg_folder_path) 24 | 25 | noised_image = noised_and_cover[0] 26 | container_img_copy = noised_image.clone() 27 | containers_ori = container_img_copy.detach().cpu().numpy() 28 | 29 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 30 | N, _, _, _ = containers.shape 31 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 32 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 33 | 34 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 35 | 36 | for i in range(N): 37 | img = cv2.cvtColor(containers[i], cv2.COLOR_RGB2BGR) 38 | folder_imgs = jpeg_folder_path + "/jpg_" + str(i).zfill(2) + ".jpg" 39 | cv2.imwrite(folder_imgs, img, [int(cv2.IMWRITE_JPEG_QUALITY), self.quality]) 40 | 41 | containers_loaded = np.copy(containers) 42 | 43 | for i in range(N): 44 | folder_imgs = jpeg_folder_path + "/jpg_" + str(i).zfill(2) + ".jpg" 45 | img = cv2.imread(folder_imgs) 46 | containers_loaded[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 47 | 48 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 49 | 50 | containers_loaded = containers_loaded.astype(np.float32) / 255 51 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 52 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 53 | 54 | container_gap = containers_loaded - containers_ori 55 | container_gap = torch.from_numpy(container_gap).float().cuda() 56 | 57 | container_img_noised_jpeg = noised_image + container_gap 58 | 59 | noised_and_cover[0] = container_img_noised_jpeg 60 | 61 | return noised_and_cover 62 | -------------------------------------------------------------------------------- /noise_layers/combined2.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | import math 5 | from noise_layers.identity import Identity 6 | from noise_layers.jpeg_compression2 import JpegCompression2 7 | from noise_layers.jpeg_compression import JpegCompression 8 | from noise_layers.quantization import Quantization 9 | from noise_layers.diff_jpeg import DiffJPEG 10 | from noise_layers.dropout import Dropout 11 | from noise_layers.crop2 import Crop2 12 | from noise_layers.cropout import Cropout 13 | from noise_layers.gaussian import Gaussian 14 | 15 | class Combined2(nn.Module): 16 | """ 17 | Combined noise 18 | """ 19 | def __init__(self, device, jpeg_type='jpeg2'): 20 | super(Combined2, self).__init__() 21 | self.identity = Identity() 22 | if jpeg_type == 'jpeg2': 23 | self.jpeg = JpegCompression2() 24 | elif jpeg_type == 'jpeg': 25 | self.jpeg = JpegCompression(device) 26 | elif jpeg_type == 'diff_jpeg': 27 | self.jpeg = DiffJPEG() 28 | else: 29 | self.jpeg = Quantization() 30 | self.dropout = Dropout([0.3, 0.3]) 31 | self.gaussian = Gaussian() 32 | self.crop2 = Crop2([0.187, 0.187], [0.187, 0.187]) # Crop2([0.547, 0.547], [0.547, 0.547]) 33 | self.cropout = Cropout([0.547, 0.547], [0.547, 0.547]) 34 | 35 | def forward(self, noised_and_cover): 36 | noised_image = noised_and_cover[0] 37 | 38 | N, _, _, _ = noised_image.shape 39 | 40 | n0, c0 = noised_and_cover[0][:N//6], noised_and_cover[1][:N//6] 41 | n1, c1 = noised_and_cover[0][N//6:2*N//6], noised_and_cover[1][N//6:2*N//6] 42 | n2, c2 = noised_and_cover[0][2*N//6:3*N//6], noised_and_cover[1][2*N//6:3*N//6] 43 | n3, c3 = noised_and_cover[0][3*N//6:4*N//6], noised_and_cover[1][3*N//6:4*N//6] 44 | n4, c4 = noised_and_cover[0][4*N//6:5*N//6], noised_and_cover[1][4*N//6:5*N//6] 45 | n5, c5 = noised_and_cover[0][5*N//6:], noised_and_cover[1][5*N//6:] 46 | 47 | nc0 = self.identity([n0, c0]) 48 | nc1 = self.jpeg([n1, c1]) 49 | nc2 = self.dropout([n2, c2]) 50 | nc3 = self.gaussian([n3, c3]) 51 | nc4 = self.crop2([n4, c4]) 52 | nc5 = self.cropout([n5, c5]) 53 | 54 | # noised_new = torch.cat((nc0[0], nc1[0], nc2[0], nc3[0]), 0) 55 | # cover_new = torch.cat((nc0[1], nc1[1], nc2[1], nc3[1]), 0) 56 | noised_new = torch.cat((nc0[0], nc1[0], nc2[0], nc3[0], nc4[0], nc5[0]), 0) 57 | cover_new = torch.cat((nc0[1], nc1[1], nc2[1], nc3[1], nc4[1], nc5[1]), 0) 58 | 59 | return [noised_new, cover_new] -------------------------------------------------------------------------------- /noise_layers/jpeg_compression2000.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from os import environ 6 | environ["OPENCV_IO_ENABLE_JASPER"] = "true" 7 | import cv2 8 | import os 9 | import random 10 | import string 11 | 12 | def generate_random_key(l=30): 13 | s=string.ascii_lowercase+string.digits 14 | return ''.join(random.sample(s,l)) 15 | 16 | class JpegCompression2000(nn.Module): 17 | def __init__(self, quality=500): 18 | super(JpegCompression2000, self).__init__() 19 | self.quality = quality 20 | self.key = generate_random_key() 21 | 22 | def forward(self, noised_and_cover): 23 | jpeg_folder_path = "./jpeg2000_" + str(self.quality) + "/" + self.key 24 | if not os.path.exists(jpeg_folder_path): 25 | os.makedirs(jpeg_folder_path) 26 | 27 | noised_image = noised_and_cover[0] 28 | 29 | container_img_copy = noised_image.clone() 30 | containers_ori = container_img_copy.detach().cpu().numpy() 31 | 32 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 33 | N, _, _, _ = containers.shape 34 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 35 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 36 | 37 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 38 | 39 | for i in range(N): 40 | img = cv2.cvtColor(containers[i], cv2.COLOR_RGB2BGR) 41 | folder_imgs = jpeg_folder_path + "/jpg2_" + str(i).zfill(2) + ".jp2" 42 | cv2.imwrite(folder_imgs, img, [int(cv2.IMWRITE_JPEG2000_COMPRESSION_X1000), self.quality]) 43 | 44 | containers_loaded = np.copy(containers) 45 | 46 | for i in range(N): 47 | folder_imgs = jpeg_folder_path + "/jpg2_" + str(i).zfill(2) + ".jp2" 48 | img = cv2.imread(folder_imgs) 49 | containers_loaded[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 50 | 51 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 52 | 53 | containers_loaded = containers_loaded.astype(np.float32) / 255 54 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 55 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 56 | 57 | container_gap = containers_loaded - containers_ori 58 | container_gap = torch.from_numpy(container_gap).float().cuda() 59 | 60 | container_img_noised_jpeg = noised_image + container_gap 61 | 62 | noised_and_cover[0] = container_img_noised_jpeg 63 | 64 | return noised_and_cover -------------------------------------------------------------------------------- /oops_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | from moviepy.editor import * 4 | import numpy as np 5 | import torch 6 | from torch.utils.data import Dataset 7 | import warnings 8 | warnings.filterwarnings("ignore") 9 | 10 | def read_video(video_path): 11 | # video_path = "/workspace/smb_ssd_data1/Data/oops_dataset/oops_video/train/27 Hilarious Cooking Fail Nominees - FailArmy Hall of Fame (July 2017)11.mp4" 12 | 13 | # sample = cv2.VideoCapture(video_path) 14 | # length = int(sample.get(cv2.CAP_PROP_FRAME_COUNT)) 15 | 16 | sample = VideoFileClip(video_path) 17 | length = sample.reader.nframes 18 | 19 | frames = np.zeros((length, 128, 128, 3), dtype=np.uint8) 20 | 21 | # for i in range(length): 22 | # success, fr = sample.read() 23 | # H, W, _ = fr.shape 24 | # S = min(H, W) 25 | # fr = cv2.resize(fr[H//2-S//2:H//2+S//2, W//2-S//2:W//2+S//2], (128, 128)) 26 | # frames[i] = fr 27 | 28 | for i, fr in enumerate(sample.iter_frames()): 29 | H, W, _ = fr.shape 30 | S = min(H, W) 31 | fr = cv2.resize(fr[H//2-S//2:H//2+S//2, W//2-S//2:W//2+S//2], (128, 128)) 32 | frames[i] = fr 33 | 34 | # sample.release() 35 | # cv2.destroyAllWindows() 36 | 37 | return frames, length 38 | 39 | class OOPSDataset(Dataset): 40 | def __init__(self, data_path, transform=None): 41 | self.data_path = data_path 42 | self.dataset = os.listdir(data_path) 43 | self.dataset.sort() 44 | 45 | # # Filter problematic entries 46 | # bad_entries = [] 47 | # print(len(self.dataset)) 48 | # for i in range(len(self.dataset)): 49 | # if i % 1000 == 0: 50 | # print(i) 51 | # s = self.dataset[i] 52 | # if 'ata Party' in s: 53 | # bad_entries.append(s) 54 | # else: 55 | # try: 56 | # sample = VideoFileClip(self.data_path + "/" + s) 57 | # except OSError: 58 | # print("problematic video") 59 | # bad_entries.append(s) 60 | 61 | # for entry in bad_entries: 62 | # os.remove(self.data_path + "/" + entry) 63 | 64 | # for s in self.dataset: 65 | # if 'ata Party' in s: 66 | # print(s) 67 | # import pdb; pdb.set_trace() 68 | 69 | self.transform = transform 70 | 71 | def __getitem__(self, idx): 72 | frames, l = read_video(self.data_path + "/" + self.dataset[idx]) 73 | 74 | if (l < 32): 75 | frames = frames.transpose(0, 3, 1, 2).astype(np.float32) 76 | frames = torch.from_numpy(frames / 255) 77 | else: 78 | start_ind = np.random.randint(l - 31) 79 | frames = frames[start_ind:start_ind + 32] 80 | frames = frames.transpose(0, 3, 1, 2).astype(np.float32) 81 | frames = torch.from_numpy(frames / 255) 82 | 83 | return (frames, 0) 84 | 85 | def __len__(self): 86 | return len(self.dataset) -------------------------------------------------------------------------------- /noise_layers/dct_filters.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy as np 3 | 4 | 5 | def delta(a: int, b: int) -> int: 6 | if a == b: 7 | return 1 8 | else: 9 | return 0 10 | 11 | 12 | def dct_coefficient(n, k, N): 13 | """Discrete cosine coefficient 14 | Coefficient for dct 15 | :param n: 16 | :param k: 17 | :param N: 18 | :return: 19 | """ 20 | return math.cos(np.pi / N * (n + 1. / 2.) * k) 21 | 22 | 23 | def idct_coefficient(n, k, N): 24 | return (delta(0, n) * (- 1 / 2) + math.cos( 25 | np.pi / N * (k + 1. / 2.) * n)) * math.sqrt(1 / (2. * N)) 26 | 27 | 28 | class DctFilterGenerator: 29 | """ 30 | Generate dct filters 31 | """ 32 | def __init__(self, tile_size_x: int = 8, tile_size_y: int = 8, channels: int = 3): 33 | self.tile_size_x = tile_size_x 34 | self.tile_size_y = tile_size_y 35 | self.channels = channels 36 | 37 | 38 | def generate_per_channel_filter(self, size_x: int, size_y: int, function: callable) -> np.ndarray: 39 | filters = np.zeros((size_x, size_y, size_x * size_y)) 40 | for k_y in range(size_y): 41 | for m_x in range(size_x): 42 | for n_y in range(size_y): 43 | for l_x in range(size_x): 44 | filters[n_y, l_x, k_y * self.tile_size_x + m_x] = function(n_y, k_y, size_y) * function(l_x, 45 | m_x, 46 | size_x) 47 | return filters 48 | 49 | 50 | def get_dct_filters(self) -> np.ndarray: 51 | 52 | filters = self.generate_per_channel_filter(self.tile_size_x, self.tile_size_y, dct_coefficient) 53 | result = np.zeros( 54 | (self.channels, self.tile_size_x, self.tile_size_y, self.channels, self.tile_size_x * self.tile_size_y), 55 | dtype=np.float32) 56 | for channel in range(self.channels): 57 | result[channel, :, :, channel, :] = filters[:, :, :] 58 | return result 59 | 60 | def get_idct_filters(self) -> np.ndarray: 61 | 62 | filters = self.generate_per_channel_filter(self.tile_size_x, self.tile_size_y, idct_coefficient) 63 | result = np.zeros( 64 | (self.channels, self.tile_size_x, self.tile_size_y, self.channels, self.tile_size_x * self.tile_size_x), 65 | dtype=np.float32) 66 | for channel in range(self.channels): 67 | result[channel, :, :, channel, :] = filters[:, :, :] 68 | return result 69 | 70 | 71 | def get_jpeg_yuv_filter_mask(self, shape: tuple, N: int, count: int): 72 | mask = np.zeros((N, N), dtype=np.uint8) 73 | 74 | index_order = sorted(((x, y) for x in range(N) for y in range(N)), 75 | key=lambda p: (p[0] + p[1], -p[1] if (p[0] + p[1]) % 2 else p[1])) 76 | 77 | for i, j in index_order[0:count]: 78 | mask[i, j] = 1 79 | 80 | return np.tile(mask, (int(np.ceil(shape[0] / N)), int(np.ceil(shape[1] / N))))[0: shape[0], 0: shape[1]] -------------------------------------------------------------------------------- /test_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn 3 | import argparse 4 | import os 5 | import numpy as np 6 | from options import HiDDenConfiguration 7 | 8 | import utils 9 | from model.hidden import * 10 | from noise_layers.noiser import Noiser 11 | from PIL import Image 12 | import torchvision.transforms.functional as TF 13 | 14 | 15 | def randomCrop(img, height, width): 16 | assert img.shape[0] >= height 17 | assert img.shape[1] >= width 18 | x = np.random.randint(0, img.shape[1] - width) 19 | y = np.random.randint(0, img.shape[0] - height) 20 | img = img[y:y+height, x:x+width] 21 | return img 22 | 23 | 24 | def main(): 25 | if torch.cuda.is_available(): 26 | device = torch.device('cuda') 27 | else: 28 | device = torch.device('cpu') 29 | 30 | parser = argparse.ArgumentParser(description='Test trained models') 31 | parser.add_argument('--options-file', '-o', default='options-and-config.pickle', type=str, 32 | help='The file where the simulation options are stored.') 33 | parser.add_argument('--checkpoint-file', '-c', required=True, type=str, help='Model checkpoint file') 34 | parser.add_argument('--batch-size', '-b', default=12, type=int, help='The batch size.') 35 | parser.add_argument('--source-image', '-s', required=True, type=str, 36 | help='The image to watermark') 37 | # parser.add_argument('--times', '-t', default=10, type=int, 38 | # help='Number iterations (insert watermark->extract).') 39 | 40 | args = parser.parse_args() 41 | 42 | train_options, hidden_config, noise_config = utils.load_options(args.options_file) 43 | noiser = Noiser(noise_config) 44 | 45 | checkpoint = torch.load(args.checkpoint_file) 46 | hidden_net = Hidden(hidden_config, device, noiser, None) 47 | utils.model_from_checkpoint(hidden_net, checkpoint) 48 | 49 | 50 | image_pil = Image.open(args.source_image) 51 | image = randomCrop(np.array(image_pil), hidden_config.H, hidden_config.W) 52 | image_tensor = TF.to_tensor(image).to(device) 53 | image_tensor = image_tensor * 2 - 1 # transform from [0, 1] to [-1, 1] 54 | image_tensor.unsqueeze_(0) 55 | 56 | # for t in range(args.times): 57 | message = torch.Tensor(np.random.choice([0, 1], (image_tensor.shape[0], 58 | hidden_config.message_length))).to(device) 59 | losses, (encoded_images, noised_images, decoded_messages) = hidden_net.validate_on_batch([image_tensor, message]) 60 | decoded_rounded = decoded_messages.detach().cpu().numpy().round().clip(0, 1) 61 | message_detached = message.detach().cpu().numpy() 62 | print('original: {}'.format(message_detached)) 63 | print('decoded : {}'.format(decoded_rounded)) 64 | print('error : {:.3f}'.format(np.mean(np.abs(decoded_rounded - message_detached)))) 65 | utils.save_images(image_tensor.cpu(), encoded_images.cpu(), 'test', '.', resize_to=(256, 256)) 66 | 67 | # bitwise_avg_err = np.sum(np.abs(decoded_rounded - message.detach().cpu().numpy()))/(image_tensor.shape[0] * messages.shape[1]) 68 | 69 | 70 | 71 | if __name__ == '__main__': 72 | main() 73 | -------------------------------------------------------------------------------- /noise_layers/diff_quality_jpeg_compression2.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class DiffQFJpegCompression2(nn.Module): 15 | def __init__(self): 16 | super(DiffQFJpegCompression2, self).__init__() 17 | self.qualities = [10, 25, 50, 75, 90] 18 | self.key = generate_random_key() 19 | 20 | def forward(self, noised_and_cover): 21 | 22 | noised_image = noised_and_cover[0] 23 | res_noised_image = torch.zeros_like(noised_image) 24 | batch_size = noised_image.shape[0] 25 | 26 | for q in range(len(self.qualities)): 27 | quality = self.qualities[q] 28 | 29 | jpeg_folder_path = "./jpeg_" + str(quality) + "/" + self.key 30 | if not os.path.exists(jpeg_folder_path): 31 | os.makedirs(jpeg_folder_path) 32 | 33 | noised_image_q = noised_image[q*batch_size//len(self.qualities):(q+1)*batch_size//len(self.qualities)] 34 | container_img_copy = noised_image_q.clone() 35 | containers_ori = container_img_copy.detach().cpu().numpy() 36 | 37 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 38 | N, _, _, _ = containers.shape 39 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 40 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 41 | 42 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 43 | 44 | for i in range(N): 45 | img = cv2.cvtColor(containers[i], cv2.COLOR_RGB2BGR) 46 | folder_imgs = jpeg_folder_path + "/jpg_" + str(i).zfill(2) + ".jpg" 47 | cv2.imwrite(folder_imgs, img, [int(cv2.IMWRITE_JPEG_QUALITY), quality]) 48 | 49 | containers_loaded = np.copy(containers) 50 | 51 | for i in range(N): 52 | folder_imgs = jpeg_folder_path + "/jpg_" + str(i).zfill(2) + ".jpg" 53 | img = cv2.imread(folder_imgs) 54 | containers_loaded[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 55 | 56 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 57 | 58 | containers_loaded = containers_loaded.astype(np.float32) / 255 59 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 60 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 61 | 62 | container_gap = containers_loaded - containers_ori 63 | container_gap = torch.from_numpy(container_gap).float().cuda() 64 | 65 | container_img_noised_jpeg = noised_image_q + container_gap 66 | 67 | res_noised_image[q*batch_size//len(self.qualities):(q+1)*batch_size//len(self.qualities)] = container_img_noised_jpeg 68 | 69 | noised_and_cover[0] = res_noised_image 70 | 71 | return noised_and_cover 72 | -------------------------------------------------------------------------------- /noise_layers/quantization.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | 5 | # def transform_clamp(tensor, target_range): 6 | # tensor = torch.clamp(tensor, 0., 1.) 7 | # source_min = tensor.min() 8 | # source_max = tensor.max() 9 | 10 | # # normalize to [0, 1] 11 | # tensor_target = (tensor - source_min)/(source_max - source_min) 12 | # # move to target range 13 | # tensor_target = tensor_target * (target_range[1] - target_range[0]) + target_range[0] 14 | # return tensor_target 15 | 16 | def transform(tensor, target_range): 17 | source_min = tensor.min() 18 | source_max = tensor.max() 19 | 20 | # normalize to [0, 1] 21 | tensor_target = (tensor - source_min)/(source_max - source_min) 22 | # move to target range 23 | tensor_target = tensor_target * (target_range[1] - target_range[0]) + target_range[0] 24 | return tensor_target 25 | 26 | 27 | class Quantization(nn.Module): 28 | def __init__(self, device=None): 29 | super(Quantization, self).__init__() 30 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 31 | 32 | self.min_value = 0.0 33 | self.max_value = 255.0 34 | self.N = 10 35 | self.weights = torch.tensor([((-1) ** (n + 1)) / (np.pi * (n + 1)) for n in range(self.N)]).to(device) 36 | self.scales = torch.tensor([2 * np.pi * (n + 1) for n in range(self.N)]).to(device) 37 | for _ in range(4): 38 | self.weights.unsqueeze_(-1) 39 | self.scales.unsqueeze_(-1) 40 | 41 | 42 | def fourier_rounding(self, tensor): 43 | shape = tensor.shape 44 | z = torch.mul(self.weights, torch.sin(torch.mul(tensor, self.scales))) 45 | z = torch.sum(z, dim=0) 46 | return tensor + z 47 | 48 | 49 | def forward(self, noised_and_cover): 50 | noised_image = noised_and_cover[0] 51 | noised_image = transform(noised_image, (0, 255)) 52 | # noised_image = noised_image.clamp(self.min_value, self.max_value).round() 53 | noised_image = self.fourier_rounding(noised_image.clamp(self.min_value, self.max_value)) 54 | noised_image = transform(noised_image, (noised_and_cover[0].min(), noised_and_cover[0].max())) 55 | return [noised_image, noised_and_cover[1]] 56 | 57 | 58 | # def main(): 59 | # import numpy as np 60 | # qq = Quantization() 61 | # min_value = -1.4566678979 62 | # max_value = 2.334567325678 63 | # # min_value = 0.0 64 | # # max_value = 255.0 65 | # data = torch.rand((12, 3, 64, 64)) 66 | # data = (max_value - min_value) * data + min_value 67 | # print(f'data.min(): {data.min()}') 68 | # print(f'data.max(): {data.max()}') 69 | # print(f'data.avg(): {data.mean()}') 70 | # data_tf, _ = qq.forward([data, None]) 71 | # print('perform quantization') 72 | # print(f'data.min(): {data_tf.min()}') 73 | # print(f'data.max(): {data_tf.max()}') 74 | # print(f'data.avg(): {data_tf.mean()}') 75 | # 76 | # print(f'mse diff: {np.mean((data.numpy()-data_tf.numpy())**2)}') 77 | # print(f'mabs diff: {np.mean(np.abs(data.numpy()-data_tf.numpy()))}') 78 | # 79 | # if __name__ == '__main__': 80 | # main() -------------------------------------------------------------------------------- /noise_layers/crop.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import numpy as np 3 | 4 | 5 | def random_float(min, max): 6 | """ 7 | Return a random number 8 | :param min: 9 | :param max: 10 | :return: 11 | """ 12 | return np.random.rand() * (max - min) + min 13 | 14 | 15 | def get_random_rectangle_inside(image, height_ratio_range, width_ratio_range): 16 | """ 17 | Returns a random rectangle inside the image, where the size is random and is controlled by height_ratio_range and width_ratio_range. 18 | This is analogous to a random crop. For example, if height_ratio_range is (0.7, 0.9), then a random number in that range will be chosen 19 | (say it is 0.75 for illustration), and the image will be cropped such that the remaining height equals 0.75. In fact, 20 | a random 'starting' position rs will be chosen from (0, 0.25), and the crop will start at rs and end at rs + 0.75. This ensures 21 | that we crop from top/bottom with equal probability. 22 | The same logic applies to the width of the image, where width_ratio_range controls the width crop range. 23 | :param image: The image we want to crop 24 | :param height_ratio_range: The range of remaining height ratio 25 | :param width_ratio_range: The range of remaining width ratio. 26 | :return: "Cropped" rectange with width and height drawn randomly height_ratio_range and width_ratio_range 27 | """ 28 | image_height = image.shape[2] 29 | image_width = image.shape[3] 30 | 31 | remaining_height = int(np.rint(random_float(height_ratio_range[0], height_ratio_range[1]) * image_height)) 32 | remaining_width = int(np.rint(random_float(width_ratio_range[0], width_ratio_range[0]) * image_width)) 33 | 34 | if remaining_height == image_height: 35 | height_start = 0 36 | else: 37 | height_start = np.random.randint(0, image_height - remaining_height) 38 | 39 | if remaining_width == image_width: 40 | width_start = 0 41 | else: 42 | width_start = np.random.randint(0, image_width - remaining_width) 43 | 44 | return height_start, height_start+remaining_height, width_start, width_start+remaining_width 45 | 46 | 47 | class Crop(nn.Module): 48 | """ 49 | Randomly crops the image from top/bottom and left/right. The amount to crop is controlled by parameters 50 | heigth_ratio_range and width_ratio_range 51 | """ 52 | def __init__(self, height_ratio_range, width_ratio_range): 53 | """ 54 | 55 | :param height_ratio_range: 56 | :param width_ratio_range: 57 | """ 58 | super(Crop, self).__init__() 59 | self.height_ratio_range = height_ratio_range 60 | self.width_ratio_range = width_ratio_range 61 | 62 | 63 | def forward(self, noised_and_cover): 64 | noised_image = noised_and_cover[0] 65 | # crop_rectangle is in form (from, to) where @from and @to are 2D points -- (height, width) 66 | 67 | h_start, h_end, w_start, w_end = get_random_rectangle_inside(noised_image, self.height_ratio_range, self.width_ratio_range) 68 | 69 | noised_and_cover[0] = noised_image[ 70 | :, 71 | :, 72 | h_start: h_end, 73 | w_start: w_end].clone() 74 | 75 | return noised_and_cover -------------------------------------------------------------------------------- /model/encoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from options import HiDDenConfiguration 4 | from model.conv_bn_relu import ConvBNRelu 5 | 6 | 7 | class Encoder(nn.Module): 8 | """ 9 | Inserts a watermark into an image. 10 | """ 11 | def __init__(self, config: HiDDenConfiguration): 12 | super(Encoder, self).__init__() 13 | self.H = config.H 14 | self.W = config.W 15 | self.conv_channels = config.encoder_channels 16 | self.num_blocks = config.encoder_blocks 17 | self.cover_dependent = config.cover_dependent 18 | # self.residual = config.residual 19 | 20 | if self.cover_dependent == 1: 21 | layers = [ConvBNRelu(3, self.conv_channels)] 22 | else: 23 | layers = [ConvBNRelu(config.message_length, self.conv_channels)] 24 | 25 | for _ in range(config.encoder_blocks-1): 26 | layer = ConvBNRelu(self.conv_channels, self.conv_channels) 27 | layers.append(layer) 28 | 29 | self.conv_layers = nn.Sequential(*layers) 30 | if self.cover_dependent == 1: 31 | self.after_concat_layer = ConvBNRelu(self.conv_channels + 3 + config.message_length, 32 | self.conv_channels) 33 | else: 34 | self.after_concat_layer = ConvBNRelu(config.message_length, 35 | self.conv_channels) 36 | 37 | self.final_layer = nn.Conv2d(self.conv_channels, 3, kernel_size=1) 38 | self.final_tanh = nn.Tanh() 39 | self.factor = 10/255 40 | 41 | def forward(self, image, message): 42 | 43 | # # First, add two dummy dimensions in the end of the message. 44 | # # This is required for the .expand to work correctly 45 | # expanded_message = message.unsqueeze(-1) 46 | # expanded_message.unsqueeze_(-1) 47 | 48 | # expanded_message = expanded_message.expand(-1,-1, self.H, self.W) 49 | # encoded_image = self.conv_layers(image) 50 | # # concatenate expanded message and image 51 | # concat = torch.cat([expanded_message, encoded_image, image], dim=1) 52 | # im_w = self.after_concat_layer(concat) 53 | # im_w = self.final_layer(im_w) 54 | 55 | expanded_message = message.unsqueeze(-1) 56 | expanded_message.unsqueeze_(-1) 57 | 58 | expanded_message = expanded_message.expand(-1,-1, self.H, self.W) 59 | if self.cover_dependent: 60 | encoded_image = self.conv_layers(image) 61 | # concatenate expanded message and image 62 | concat = torch.cat([expanded_message, encoded_image, image], dim=1) 63 | im_w = self.after_concat_layer(concat) 64 | im_w = self.final_layer(im_w) 65 | # if self.residual: 66 | # im_w = self.factor * self.final_tanh(im_w) + image 67 | else: 68 | #import pdb; pdb.set_trace() 69 | #encoded_message = self.conv_layers(expanded_message) 70 | # concatenate expanded message and image 71 | #concat = encoded_message # torch.cat([expanded_message, encoded_image, image], dim=1) 72 | im_w = self.after_concat_layer(expanded_message) 73 | im_w = self.final_layer(im_w) + image 74 | return im_w 75 | -------------------------------------------------------------------------------- /noise_layers/diff_corruptions.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import cv2 6 | import os 7 | import random 8 | import string 9 | 10 | def generate_random_key(l=30): 11 | s=string.ascii_lowercase+string.digits 12 | return ''.join(random.sample(s,l)) 13 | 14 | class DiffCorruptions(nn.Module): 15 | def __init__(self): 16 | super(DiffCorruptions, self).__init__() 17 | self.corruptions = ["jpeg", "jpeg2000", "webp"] 18 | self.file_formats = ["jpg", "jp2", "webp"] 19 | self.flags = [int(cv2.IMWRITE_JPEG_QUALITY), int(cv2.IMWRITE_JPEG2000_COMPRESSION_X1000), int(cv2.IMWRITE_WEBP_QUALITY)] 20 | self.qualities = [50, 500, 50] 21 | self.key = generate_random_key() 22 | 23 | def forward(self, noised_and_cover): 24 | 25 | noised_image = noised_and_cover[0] 26 | res_noised_image = torch.zeros_like(noised_image) 27 | batch_size = noised_image.shape[0] 28 | 29 | for q in range(len(self.corruptions)): 30 | corruption = self.corruptions[q] 31 | file_format = self.file_formats[q] 32 | flag = self.flags[q] 33 | quality = self.qualities[q] 34 | 35 | jpeg_folder_path = "./" + corruption + "_" + str(quality) + "/" + self.key 36 | if not os.path.exists(jpeg_folder_path): 37 | os.makedirs(jpeg_folder_path) 38 | 39 | noised_image_q = noised_image[q*batch_size//len(self.corruptions):(q+1)*batch_size//len(self.corruptions)] 40 | 41 | container_img_copy = noised_image_q.clone() 42 | containers_ori = container_img_copy.detach().cpu().numpy() 43 | 44 | containers = np.transpose(containers_ori, (0, 2, 3, 1)) 45 | N, _, _, _ = containers.shape 46 | # containers = (containers + 1) / 2 # transform range of containers from [-1, 1] to [0, 1] 47 | containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 48 | 49 | # containers = (np.clip(containers, 0.0, 1.0)*255).astype(np.uint8) 50 | 51 | for i in range(N): 52 | img = cv2.cvtColor(containers[i], cv2.COLOR_RGB2BGR) 53 | folder_imgs = jpeg_folder_path + "/" + file_format + "_" + str(i).zfill(2) + "." + file_format 54 | cv2.imwrite(folder_imgs, img, [flag, quality]) 55 | 56 | containers_loaded = np.copy(containers) 57 | 58 | for i in range(N): 59 | folder_imgs = jpeg_folder_path + "/" + file_format + "_" + str(i).zfill(2) + "." + file_format 60 | img = cv2.imread(folder_imgs) 61 | containers_loaded[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 62 | 63 | # containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)).astype(np.float32) / 255 64 | 65 | containers_loaded = containers_loaded.astype(np.float32) / 255 66 | # containers_loaded = containers_loaded * 2 - 1 # transform range of containers from [0, 1] to [-1, 1] 67 | containers_loaded = np.transpose(containers_loaded, (0, 3, 1, 2)) 68 | 69 | container_gap = containers_loaded - containers_ori 70 | container_gap = torch.from_numpy(container_gap).float().cuda() 71 | 72 | container_img_noised_jpeg = noised_image_q + container_gap 73 | 74 | res_noised_image[q*batch_size//len(self.qualities):(q+1)*batch_size//len(self.qualities)] = container_img_noised_jpeg 75 | 76 | noised_and_cover[0] = res_noised_image 77 | 78 | return noised_and_cover 79 | -------------------------------------------------------------------------------- /validate-trained-models.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import pprint 4 | import argparse 5 | import torch 6 | import numpy as np 7 | import pickle 8 | import utils 9 | import csv 10 | 11 | from model.hidden import Hidden 12 | from noise_layers.noiser import Noiser 13 | from average_meter import AverageMeter 14 | 15 | 16 | def write_validation_loss(file_name, losses_accu, experiment_name, epoch, write_header=False): 17 | with open(file_name, 'a', newline='') as csvfile: 18 | writer = csv.writer(csvfile) 19 | if write_header: 20 | row_to_write = ['experiment_name', 'epoch'] + [loss_name.strip() for loss_name in losses_accu.keys()] 21 | writer.writerow(row_to_write) 22 | row_to_write = [experiment_name, epoch] + ['{:.4f}'.format(loss_avg.avg) for loss_avg in losses_accu.values()] 23 | writer.writerow(row_to_write) 24 | 25 | 26 | def main(): 27 | # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 28 | device = torch.device('cpu') 29 | 30 | parser = argparse.ArgumentParser(description='Training of HiDDeN nets') 31 | # parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).') 32 | parser.add_argument('--data-dir', '-d', required=True, type=str, help='The directory where the data is stored.') 33 | parser.add_argument('--runs_root', '-r', default=os.path.join('.', 'experiments'), type=str, 34 | help='The root folder where data about experiments are stored.') 35 | parser.add_argument('--batch-size', '-b', default=1, type=int, help='Validation batch size.') 36 | 37 | args = parser.parse_args() 38 | print_each = 25 39 | 40 | completed_runs = [o for o in os.listdir(args.runs_root) 41 | if os.path.isdir(os.path.join(args.runs_root, o)) and o != 'no-noise-defaults'] 42 | 43 | print(completed_runs) 44 | 45 | write_csv_header = True 46 | for run_name in completed_runs: 47 | current_run = os.path.join(args.runs_root, run_name) 48 | print(f'Run folder: {current_run}') 49 | options_file = os.path.join(current_run, 'options-and-config.pickle') 50 | train_options, hidden_config, noise_config = utils.load_options(options_file) 51 | train_options.train_folder = os.path.join(args.data_dir, 'val') 52 | train_options.validation_folder = os.path.join(args.data_dir, 'val') 53 | train_options.batch_size = args.batch_size 54 | checkpoint, chpt_file_name = utils.load_last_checkpoint(os.path.join(current_run, 'checkpoints')) 55 | print(f'Loaded checkpoint from file {chpt_file_name}') 56 | 57 | noiser = Noiser(noise_config) 58 | model = Hidden(hidden_config, device, noiser, tb_logger=None) 59 | utils.model_from_checkpoint(model, checkpoint) 60 | 61 | print('Model loaded successfully. Starting validation run...') 62 | _, val_data = utils.get_data_loaders(hidden_config, train_options) 63 | file_count = len(val_data.dataset) 64 | if file_count % train_options.batch_size == 0: 65 | steps_in_epoch = file_count // train_options.batch_size 66 | else: 67 | steps_in_epoch = file_count // train_options.batch_size + 1 68 | 69 | losses_accu = {} 70 | step = 0 71 | for image, _ in val_data: 72 | step += 1 73 | image = image.to(device) 74 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 75 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch([image, message], 76 | set_eval_mode=True) 77 | if not losses_accu: # dict is empty, initialize 78 | for name in losses: 79 | losses_accu[name] = AverageMeter() 80 | for name, loss in losses.items(): 81 | losses_accu[name].update(loss) 82 | if step % print_each == 0 or step == steps_in_epoch: 83 | print(f'Step {step}/{steps_in_epoch}') 84 | utils.print_progress(losses_accu) 85 | print('-' * 40) 86 | 87 | # utils.print_progress(losses_accu) 88 | write_validation_loss(os.path.join(args.runs_root, 'validation_run.csv'), losses_accu, run_name, 89 | checkpoint['epoch'], 90 | write_header=write_csv_header) 91 | write_csv_header = False 92 | 93 | # train(model, device, hidden_config, train_options, this_run_folder, tb_logger) 94 | 95 | 96 | if __name__ == '__main__': 97 | main() -------------------------------------------------------------------------------- /validate-trained-models_diff_jpeg_levels.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import pprint 4 | import argparse 5 | import torch 6 | import numpy as np 7 | import pickle 8 | import utils 9 | import csv 10 | import socket 11 | 12 | from model.hidden import Hidden 13 | from noise_layers.noiser import Noiser 14 | from average_meter import AverageMeter 15 | 16 | def write_validation_loss(file_name, losses_accu, experiment_name, epoch, write_header=False): 17 | with open(file_name, 'a', newline='') as csvfile: 18 | writer = csv.writer(csvfile) 19 | if write_header: 20 | row_to_write = ['experiment_name', 'epoch'] + [loss_name.strip() for loss_name in losses_accu.keys()] 21 | writer.writerow(row_to_write) 22 | row_to_write = [experiment_name, epoch] + ['{:.4f}'.format(loss_avg.avg) for loss_avg in losses_accu.values()] 23 | writer.writerow(row_to_write) 24 | 25 | def main(): 26 | # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 27 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 28 | 29 | parser = argparse.ArgumentParser(description='Training of HiDDeN nets') 30 | parser.add_argument('--hostname', default=socket.gethostname(), help='the host name of the running server') 31 | # parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).') 32 | parser.add_argument('--data-dir', '-d', required=True, type=str, help='The directory where the data is stored.') 33 | parser.add_argument('--runs_root', '-r', default=os.path.join('.', 'experiments'), type=str, 34 | help='The root folder where data about experiments are stored.') 35 | parser.add_argument('--batch-size', '-b', default=1, type=int, help='Validation batch size.') 36 | 37 | args = parser.parse_args() 38 | 39 | if args.hostname == 'ee898-System-Product-Name': 40 | args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet' 41 | args.hostname = 'ee898' 42 | elif args.hostname == 'DL178': 43 | args.data_dir = '/media/user/SSD1TB-2/ImageNet' 44 | else: 45 | args.data_dir = '/workspace/data_local/imagenet_pytorch' 46 | assert args.data_dir 47 | 48 | print_each = 25 49 | 50 | completed_runs = [o for o in os.listdir(args.runs_root) 51 | if os.path.isdir(os.path.join(args.runs_root, o)) and o != 'no-noise-defaults'] 52 | 53 | print(completed_runs) 54 | 55 | write_csv_header = True 56 | current_run = args.runs_root 57 | print(f'Run folder: {current_run}') 58 | options_file = os.path.join(current_run, 'options-and-config.pickle') 59 | train_options, hidden_config, noise_config = utils.load_options(options_file) 60 | train_options.train_folder = os.path.join(args.data_dir, 'val') 61 | train_options.validation_folder = os.path.join(args.data_dir, 'val') 62 | train_options.batch_size = args.batch_size 63 | checkpoint, chpt_file_name = utils.load_last_checkpoint(os.path.join(current_run, 'checkpoints')) 64 | print(f'Loaded checkpoint from file {chpt_file_name}') 65 | 66 | noiser = Noiser(noise_config, device, 'jpeg') 67 | model = Hidden(hidden_config, device, noiser, tb_logger=None) 68 | utils.model_from_checkpoint(model, checkpoint) 69 | 70 | print('Model loaded successfully. Starting validation run...') 71 | _, val_data = utils.get_data_loaders(hidden_config, train_options) 72 | file_count = len(val_data.dataset) 73 | if file_count % train_options.batch_size == 0: 74 | steps_in_epoch = file_count // train_options.batch_size 75 | else: 76 | steps_in_epoch = file_count // train_options.batch_size + 1 77 | 78 | with torch.no_grad(): 79 | noises = ['jpeg_10', 'jpeg_25', 'jpeg_50', 'jpeg_75', 'jpeg_90'] 80 | for noise in noises: 81 | losses_accu = {} 82 | step = 0 83 | for image, _ in val_data: 84 | step += 1 85 | image = image.to(device) 86 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 87 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch_specific_noise([image, message], noise=noise) 88 | if not losses_accu: # dict is empty, initialize 89 | for name in losses: 90 | losses_accu[name] = AverageMeter() 91 | for name, loss in losses.items(): 92 | losses_accu[name].update(loss) 93 | if step % print_each == 0 or step == steps_in_epoch: 94 | print(f'Step {step}/{steps_in_epoch}') 95 | utils.print_progress(losses_accu) 96 | print('-' * 40) 97 | 98 | # utils.print_progress(losses_accu) 99 | write_validation_loss(os.path.join(args.runs_root, 'validation_run.csv'), losses_accu, noise, 100 | checkpoint['epoch'], 101 | write_header=write_csv_header) 102 | write_csv_header = False 103 | 104 | if __name__ == '__main__': 105 | main() -------------------------------------------------------------------------------- /validate-trained-models_diff_webp_levels.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import pprint 4 | import argparse 5 | import torch 6 | import numpy as np 7 | import pickle 8 | import utils 9 | import csv 10 | import socket 11 | 12 | from model.hidden import Hidden 13 | from noise_layers.noiser import Noiser 14 | from average_meter import AverageMeter 15 | 16 | def write_validation_loss(file_name, losses_accu, experiment_name, epoch, write_header=False): 17 | with open(file_name, 'a', newline='') as csvfile: 18 | writer = csv.writer(csvfile) 19 | if write_header: 20 | row_to_write = ['experiment_name', 'epoch'] + [loss_name.strip() for loss_name in losses_accu.keys()] 21 | writer.writerow(row_to_write) 22 | row_to_write = [experiment_name, epoch] + ['{:.4f}'.format(loss_avg.avg) for loss_avg in losses_accu.values()] 23 | writer.writerow(row_to_write) 24 | 25 | def main(): 26 | # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 27 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 28 | 29 | parser = argparse.ArgumentParser(description='Training of HiDDeN nets') 30 | parser.add_argument('--hostname', default=socket.gethostname(), help='the host name of the running server') 31 | # parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).') 32 | parser.add_argument('--data-dir', '-d', required=True, type=str, help='The directory where the data is stored.') 33 | parser.add_argument('--runs_root', '-r', default=os.path.join('.', 'experiments'), type=str, 34 | help='The root folder where data about experiments are stored.') 35 | parser.add_argument('--batch-size', '-b', default=1, type=int, help='Validation batch size.') 36 | 37 | args = parser.parse_args() 38 | 39 | if args.hostname == 'ee898-System-Product-Name': 40 | args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet' 41 | args.hostname = 'ee898' 42 | elif args.hostname == 'DL178': 43 | args.data_dir = '/media/user/SSD1TB-2/ImageNet' 44 | else: 45 | args.data_dir = '/workspace/data_local/imagenet_pytorch' 46 | assert args.data_dir 47 | 48 | print_each = 25 49 | 50 | completed_runs = [o for o in os.listdir(args.runs_root) 51 | if os.path.isdir(os.path.join(args.runs_root, o)) and o != 'no-noise-defaults'] 52 | 53 | print(completed_runs) 54 | 55 | write_csv_header = True 56 | current_run = args.runs_root 57 | print(f'Run folder: {current_run}') 58 | options_file = os.path.join(current_run, 'options-and-config.pickle') 59 | train_options, hidden_config, noise_config = utils.load_options(options_file) 60 | train_options.train_folder = os.path.join(args.data_dir, 'val') 61 | train_options.validation_folder = os.path.join(args.data_dir, 'val') 62 | train_options.batch_size = args.batch_size 63 | checkpoint, chpt_file_name = utils.load_last_checkpoint(os.path.join(current_run, 'checkpoints')) 64 | print(f'Loaded checkpoint from file {chpt_file_name}') 65 | 66 | noiser = Noiser(noise_config, device, 'jpeg') 67 | model = Hidden(hidden_config, device, noiser, tb_logger=None) 68 | utils.model_from_checkpoint(model, checkpoint) 69 | 70 | print('Model loaded successfully. Starting validation run...') 71 | _, val_data = utils.get_data_loaders(hidden_config, train_options) 72 | file_count = len(val_data.dataset) 73 | if file_count % train_options.batch_size == 0: 74 | steps_in_epoch = file_count // train_options.batch_size 75 | else: 76 | steps_in_epoch = file_count // train_options.batch_size + 1 77 | 78 | with torch.no_grad(): 79 | noises = ['webp_10', 'webp_25', 'webp_50', 'webp_75', 'webp_90'] 80 | for noise in noises: 81 | losses_accu = {} 82 | step = 0 83 | for image, _ in val_data: 84 | step += 1 85 | image = image.to(device) 86 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 87 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch_specific_noise([image, message], noise=noise) 88 | if not losses_accu: # dict is empty, initialize 89 | for name in losses: 90 | losses_accu[name] = AverageMeter() 91 | for name, loss in losses.items(): 92 | losses_accu[name].update(loss) 93 | if step % print_each == 0 or step == steps_in_epoch: 94 | print(f'Step {step}/{steps_in_epoch}') 95 | utils.print_progress(losses_accu) 96 | print('-' * 40) 97 | 98 | # utils.print_progress(losses_accu) 99 | write_validation_loss(os.path.join(args.runs_root, 'validation_run.csv'), losses_accu, noise, 100 | checkpoint['epoch'], 101 | write_header=write_csv_header) 102 | write_csv_header = False 103 | 104 | if __name__ == '__main__': 105 | main() -------------------------------------------------------------------------------- /validate-trained-models_diff_jpeg2000_levels.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import pprint 4 | import argparse 5 | import torch 6 | import numpy as np 7 | import pickle 8 | import utils 9 | import csv 10 | import socket 11 | 12 | from model.hidden import Hidden 13 | from noise_layers.noiser import Noiser 14 | from average_meter import AverageMeter 15 | 16 | def write_validation_loss(file_name, losses_accu, experiment_name, epoch, write_header=False): 17 | with open(file_name, 'a', newline='') as csvfile: 18 | writer = csv.writer(csvfile) 19 | if write_header: 20 | row_to_write = ['experiment_name', 'epoch'] + [loss_name.strip() for loss_name in losses_accu.keys()] 21 | writer.writerow(row_to_write) 22 | row_to_write = [experiment_name, epoch] + ['{:.4f}'.format(loss_avg.avg) for loss_avg in losses_accu.values()] 23 | writer.writerow(row_to_write) 24 | 25 | def main(): 26 | # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 27 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 28 | 29 | parser = argparse.ArgumentParser(description='Training of HiDDeN nets') 30 | parser.add_argument('--hostname', default=socket.gethostname(), help='the host name of the running server') 31 | # parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).') 32 | parser.add_argument('--data-dir', '-d', required=True, type=str, help='The directory where the data is stored.') 33 | parser.add_argument('--runs_root', '-r', default=os.path.join('.', 'experiments'), type=str, 34 | help='The root folder where data about experiments are stored.') 35 | parser.add_argument('--batch-size', '-b', default=1, type=int, help='Validation batch size.') 36 | 37 | args = parser.parse_args() 38 | 39 | if args.hostname == 'ee898-System-Product-Name': 40 | args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet' 41 | args.hostname = 'ee898' 42 | elif args.hostname == 'DL178': 43 | args.data_dir = '/media/user/SSD1TB-2/ImageNet' 44 | else: 45 | args.data_dir = '/workspace/data_local/imagenet_pytorch' 46 | assert args.data_dir 47 | 48 | print_each = 25 49 | 50 | completed_runs = [o for o in os.listdir(args.runs_root) 51 | if os.path.isdir(os.path.join(args.runs_root, o)) and o != 'no-noise-defaults'] 52 | 53 | print(completed_runs) 54 | 55 | write_csv_header = True 56 | current_run = args.runs_root 57 | print(f'Run folder: {current_run}') 58 | options_file = os.path.join(current_run, 'options-and-config.pickle') 59 | train_options, hidden_config, noise_config = utils.load_options(options_file) 60 | train_options.train_folder = os.path.join(args.data_dir, 'val') 61 | train_options.validation_folder = os.path.join(args.data_dir, 'val') 62 | train_options.batch_size = args.batch_size 63 | checkpoint, chpt_file_name = utils.load_last_checkpoint(os.path.join(current_run, 'checkpoints')) 64 | print(f'Loaded checkpoint from file {chpt_file_name}') 65 | 66 | noiser = Noiser(noise_config, device, 'jpeg') 67 | model = Hidden(hidden_config, device, noiser, tb_logger=None) 68 | utils.model_from_checkpoint(model, checkpoint) 69 | 70 | print('Model loaded successfully. Starting validation run...') 71 | _, val_data = utils.get_data_loaders(hidden_config, train_options) 72 | file_count = len(val_data.dataset) 73 | if file_count % train_options.batch_size == 0: 74 | steps_in_epoch = file_count // train_options.batch_size 75 | else: 76 | steps_in_epoch = file_count // train_options.batch_size + 1 77 | 78 | with torch.no_grad(): 79 | noises = ['jpeg2000_100', 'jpeg2000_250', 'jpeg2000_500', 'jpeg2000_750', 'jpeg2000_900'] 80 | for noise in noises: 81 | losses_accu = {} 82 | step = 0 83 | for image, _ in val_data: 84 | step += 1 85 | image = image.to(device) 86 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 87 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch_specific_noise([image, message], noise=noise) 88 | if not losses_accu: # dict is empty, initialize 89 | for name in losses: 90 | losses_accu[name] = AverageMeter() 91 | for name, loss in losses.items(): 92 | losses_accu[name].update(loss) 93 | if step % print_each == 0 or step == steps_in_epoch: 94 | print(f'Step {step}/{steps_in_epoch}') 95 | utils.print_progress(losses_accu) 96 | print('-' * 40) 97 | 98 | # utils.print_progress(losses_accu) 99 | write_validation_loss(os.path.join(args.runs_root, 'validation_run.csv'), losses_accu, noise, 100 | checkpoint['epoch'], 101 | write_header=write_csv_header) 102 | write_csv_header = False 103 | 104 | if __name__ == '__main__': 105 | main() -------------------------------------------------------------------------------- /validate-trained-models_diff_corruptions.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import pprint 4 | import argparse 5 | import torch 6 | import numpy as np 7 | import pickle 8 | import utils 9 | import csv 10 | import socket 11 | 12 | from model.hidden import Hidden 13 | from noise_layers.noiser import Noiser 14 | from average_meter import AverageMeter 15 | 16 | def write_validation_loss(file_name, losses_accu, experiment_name, epoch, write_header=False): 17 | with open(file_name, 'a', newline='') as csvfile: 18 | writer = csv.writer(csvfile) 19 | if write_header: 20 | row_to_write = ['experiment_name', 'epoch'] + [loss_name.strip() for loss_name in losses_accu.keys()] 21 | writer.writerow(row_to_write) 22 | row_to_write = [experiment_name, epoch] + ['{:.4f}'.format(loss_avg.avg) for loss_avg in losses_accu.values()] 23 | writer.writerow(row_to_write) 24 | 25 | def main(): 26 | # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 27 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 28 | 29 | parser = argparse.ArgumentParser(description='Training of HiDDeN nets') 30 | parser.add_argument('--hostname', default=socket.gethostname(), help='the host name of the running server') 31 | # parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).') 32 | parser.add_argument('--data-dir', '-d', required=True, type=str, help='The directory where the data is stored.') 33 | parser.add_argument('--runs_root', '-r', default=os.path.join('.', 'experiments'), type=str, 34 | help='The root folder where data about experiments are stored.') 35 | parser.add_argument('--batch-size', '-b', default=1, type=int, help='Validation batch size.') 36 | 37 | args = parser.parse_args() 38 | 39 | if args.hostname == 'ee898-System-Product-Name': 40 | args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet' 41 | args.hostname = 'ee898' 42 | elif args.hostname == 'DL178': 43 | args.data_dir = '/media/user/SSD1TB-2/ImageNet' 44 | else: 45 | args.data_dir = '/workspace/data_local/imagenet_pytorch' 46 | assert args.data_dir 47 | 48 | print_each = 25 49 | 50 | completed_runs = [o for o in os.listdir(args.runs_root) 51 | if os.path.isdir(os.path.join(args.runs_root, o)) and o != 'no-noise-defaults'] 52 | 53 | print(completed_runs) 54 | 55 | write_csv_header = True 56 | current_run = args.runs_root 57 | print(f'Run folder: {current_run}') 58 | options_file = os.path.join(current_run, 'options-and-config.pickle') 59 | train_options, hidden_config, noise_config = utils.load_options(options_file) 60 | train_options.train_folder = os.path.join(args.data_dir, 'val') 61 | train_options.validation_folder = os.path.join(args.data_dir, 'val') 62 | train_options.batch_size = args.batch_size 63 | checkpoint, chpt_file_name = utils.load_last_checkpoint(os.path.join(current_run, 'checkpoints')) 64 | print(f'Loaded checkpoint from file {chpt_file_name}') 65 | 66 | noiser = Noiser(noise_config, device, 'jpeg') 67 | model = Hidden(hidden_config, device, noiser, tb_logger=None) 68 | utils.model_from_checkpoint(model, checkpoint) 69 | 70 | print('Model loaded successfully. Starting validation run...') 71 | _, val_data = utils.get_data_loaders(hidden_config, train_options) 72 | file_count = len(val_data.dataset) 73 | if file_count % train_options.batch_size == 0: 74 | steps_in_epoch = file_count // train_options.batch_size 75 | else: 76 | steps_in_epoch = file_count // train_options.batch_size + 1 77 | 78 | with torch.no_grad(): 79 | noises = ['jpeg_10', 'jpeg_25', 'jpeg_50', 'jpeg_75', 'jpeg_90', 80 | 'jpeg2000_100', 'jpeg2000_250', 'jpeg2000_500', 'jpeg2000_750', 'jpeg2000_900', 81 | 'webp_10', 'webp_25', 'webp_50', 'webp_75', 'webp_90'] 82 | for noise in noises: 83 | losses_accu = {} 84 | step = 0 85 | for image, _ in val_data: 86 | step += 1 87 | image = image.to(device) 88 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 89 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch_specific_noise([image, message], noise=noise) 90 | if not losses_accu: # dict is empty, initialize 91 | for name in losses: 92 | losses_accu[name] = AverageMeter() 93 | for name, loss in losses.items(): 94 | losses_accu[name].update(loss) 95 | if step % print_each == 0 or step == steps_in_epoch: 96 | print(f'Step {step}/{steps_in_epoch}') 97 | utils.print_progress(losses_accu) 98 | print('-' * 40) 99 | 100 | # utils.print_progress(losses_accu) 101 | write_validation_loss(os.path.join(args.runs_root, 'validation_run.csv'), losses_accu, noise, 102 | checkpoint['epoch'], 103 | write_header=write_csv_header) 104 | write_csv_header = False 105 | 106 | if __name__ == '__main__': 107 | main() -------------------------------------------------------------------------------- /noise_argparser.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import re 3 | from noise_layers.cropout import Cropout 4 | from noise_layers.crop import Crop 5 | from noise_layers.identity import Identity 6 | from noise_layers.dropout import Dropout 7 | from noise_layers.resize import Resize 8 | from noise_layers.quantization import Quantization 9 | from noise_layers.jpeg_compression import JpegCompression 10 | from noise_layers.jpeg_compression2 import JpegCompression2 11 | from noise_layers.jpeg_compression2000 import JpegCompression2000 12 | from noise_layers.gaussian import Gaussian 13 | from noise_layers.diff_jpeg import DiffJPEG 14 | from noise_layers.webp import WebP 15 | from noise_layers.mpeg4_compression import MPEG4 16 | from noise_layers.h264_compression import H264 17 | from noise_layers.xvid_compression import XVID 18 | from noise_layers.diff_quality_jpeg_compression2 import DiffQFJpegCompression2 19 | from noise_layers.diff_corruptions import DiffCorruptions 20 | 21 | 22 | def parse_pair(match_groups): 23 | heights = match_groups[0].split(',') 24 | hmin = float(heights[0]) 25 | hmax = float(heights[1]) 26 | widths = match_groups[1].split(',') 27 | wmin = float(widths[0]) 28 | wmax = float(widths[1]) 29 | return (hmin, hmax), (wmin, wmax) 30 | 31 | def parse_crop(crop_command): 32 | matches = re.match(r'crop\(\((\d+\.*\d*,\d+\.*\d*)\),\((\d+\.*\d*,\d+\.*\d*)\)\)', crop_command) 33 | (hmin, hmax), (wmin, wmax) = parse_pair(matches.groups()) 34 | return Crop((hmin, hmax), (wmin, wmax)) 35 | 36 | def parse_cropout(cropout_command): 37 | matches = re.match(r'cropout\(\((\d+\.*\d*,\d+\.*\d*)\),\((\d+\.*\d*,\d+\.*\d*)\)\)', cropout_command) 38 | (hmin, hmax), (wmin, wmax) = parse_pair(matches.groups()) 39 | return Cropout((hmin, hmax), (wmin, wmax)) 40 | 41 | 42 | def parse_dropout(dropout_command): 43 | matches = re.match(r'dropout\((\d+\.*\d*,\d+\.*\d*)\)', dropout_command) 44 | ratios = matches.groups()[0].split(',') 45 | keep_min = float(ratios[0]) 46 | keep_max = float(ratios[1]) 47 | return Dropout((keep_min, keep_max)) 48 | 49 | def parse_resize(resize_command): 50 | matches = re.match(r'resize\((\d+\.*\d*,\d+\.*\d*)\)', resize_command) 51 | ratios = matches.groups()[0].split(',') 52 | min_ratio = float(ratios[0]) 53 | max_ratio = float(ratios[1]) 54 | return Resize((min_ratio, max_ratio)) 55 | 56 | 57 | class NoiseArgParser(argparse.Action): 58 | def __init__(self, 59 | option_strings, 60 | dest, 61 | nargs=None, 62 | const=None, 63 | default=None, 64 | type=None, 65 | choices=None, 66 | required=False, 67 | help=None, 68 | metavar=None): 69 | argparse.Action.__init__(self, 70 | option_strings=option_strings, 71 | dest=dest, 72 | nargs=nargs, 73 | const=const, 74 | default=default, 75 | type=type, 76 | choices=choices, 77 | required=required, 78 | help=help, 79 | metavar=metavar, 80 | ) 81 | 82 | @staticmethod 83 | def parse_cropout_args(cropout_args): 84 | pass 85 | 86 | @staticmethod 87 | def parse_dropout_args(dropout_args): 88 | pass 89 | 90 | def __call__(self, parser, namespace, values, 91 | option_string=None): 92 | 93 | layers = [] 94 | split_commands = values[0].split('+') 95 | 96 | for command in split_commands: 97 | # remove all whitespace 98 | command = command.replace(' ', '') 99 | if command[:len('cropout')] == 'cropout': 100 | layers.append(parse_cropout(command)) 101 | elif command[:len('diff_qf_jpeg2')] == 'diff_qf_jpeg2': 102 | layers.append(DiffQFJpegCompression2()) 103 | elif command[:len('diff_corruptions')] == 'diff_corruptions': 104 | layers.append(DiffCorruptions()) 105 | elif command[:len('crop')] == 'crop': 106 | layers.append(parse_crop(command)) 107 | elif command[:len('dropout')] == 'dropout': 108 | layers.append(parse_dropout(command)) 109 | elif command[:len('resize')] == 'resize': 110 | layers.append(parse_resize(command)) 111 | elif command[:len('jpeg2000')] == 'jpeg2000': 112 | layers.append(JpegCompression2000()) 113 | elif command[:len('jpeg2')] == 'jpeg2': 114 | layers.append(JpegCompression2()) 115 | elif command[:len('jpeg')] == 'jpeg': 116 | layers.append('JpegPlaceholder') 117 | elif command[:len('quant')] == 'quant': 118 | layers.append('QuantizationPlaceholder') 119 | elif command[:len('combined2')] == 'combined2': 120 | layers.append('Combined2Placeholder') 121 | elif command[:len('identity')] == 'identity': 122 | layers.append(Identity()) 123 | elif command[:len('gaussian4')] == 'gaussian4': 124 | layers.append(Gaussian(3,4,3)) 125 | elif command[:len('gaussian')] == 'gaussian': 126 | layers.append(Gaussian()) 127 | elif command[:len('diff_jpeg')] == 'diff_jpeg': 128 | layers.append(DiffJPEG(128,128)) 129 | elif command[:len('webp')] == 'webp': 130 | layers.append(WebP()) 131 | elif command[:len('mpeg4')] == 'mpeg4': 132 | layers.append(MPEG4()) 133 | elif command[:len('h264')] == 'h264': 134 | layers.append(H264()) 135 | elif command[:len('xvid')] == 'xvid': 136 | layers.append(XVID()) 137 | else: 138 | raise ValueError('Command not recognized: \n{}'.format(command)) 139 | setattr(namespace, self.dest, layers) 140 | -------------------------------------------------------------------------------- /modules/decompression.py: -------------------------------------------------------------------------------- 1 | # Standard libraries 2 | import itertools 3 | import numpy as np 4 | # PyTorch 5 | import torch 6 | import torch.nn as nn 7 | # Local 8 | import utils 9 | 10 | 11 | class y_dequantize(nn.Module): 12 | """ Dequantize Y channel 13 | Inputs: 14 | image(tensor): batch x height x width 15 | factor(float): compression factor 16 | Outputs: 17 | image(tensor): batch x height x width 18 | """ 19 | def __init__(self, factor=1): 20 | super(y_dequantize, self).__init__() 21 | self.y_table = utils.y_table 22 | self.factor = factor 23 | 24 | def forward(self, image): 25 | return image * (self.y_table * self.factor) 26 | 27 | 28 | class c_dequantize(nn.Module): 29 | """ Dequantize CbCr channel 30 | Inputs: 31 | image(tensor): batch x height x width 32 | factor(float): compression factor 33 | Outputs: 34 | image(tensor): batch x height x width 35 | """ 36 | def __init__(self, factor=1): 37 | super(c_dequantize, self).__init__() 38 | self.factor = factor 39 | self.c_table = utils.c_table 40 | 41 | def forward(self, image): 42 | return image * (self.c_table * self.factor) 43 | 44 | 45 | class idct_8x8(nn.Module): 46 | """ Inverse discrete Cosine Transformation 47 | Input: 48 | dcp(tensor): batch x height x width 49 | Output: 50 | image(tensor): batch x height x width 51 | """ 52 | def __init__(self): 53 | super(idct_8x8, self).__init__() 54 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 55 | self.alpha = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha)).float()).cuda() 56 | tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) 57 | for x, y, u, v in itertools.product(range(8), repeat=4): 58 | tensor[x, y, u, v] = np.cos((2 * u + 1) * x * np.pi / 16) * np.cos( 59 | (2 * v + 1) * y * np.pi / 16) 60 | self.tensor = nn.Parameter(torch.from_numpy(tensor).float()).cuda() 61 | 62 | def forward(self, image): 63 | 64 | image = image * self.alpha 65 | result = 0.25 * torch.tensordot(image, self.tensor, dims=2) + 128 66 | result.view(image.shape) 67 | return result 68 | 69 | 70 | class block_merging(nn.Module): 71 | """ Merge pathces into image 72 | Inputs: 73 | patches(tensor) batch x height*width/64, height x width 74 | height(int) 75 | width(int) 76 | Output: 77 | image(tensor): batch x height x width 78 | """ 79 | def __init__(self): 80 | super(block_merging, self).__init__() 81 | 82 | def forward(self, patches, height, width): 83 | k = 8 84 | batch_size = patches.shape[0] 85 | image_reshaped = patches.view(batch_size, height//k, width//k, k, k) 86 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) 87 | return image_transposed.contiguous().view(batch_size, height, width) 88 | 89 | 90 | class chroma_upsampling(nn.Module): 91 | """ Upsample chroma layers 92 | Input: 93 | y(tensor): y channel image 94 | cb(tensor): cb channel 95 | cr(tensor): cr channel 96 | Ouput: 97 | image(tensor): batch x height x width x 3 98 | """ 99 | def __init__(self): 100 | super(chroma_upsampling, self).__init__() 101 | 102 | def forward(self, y, cb, cr): 103 | def repeat(x, k=2): 104 | height, width = x.shape[1:3] 105 | x = x.unsqueeze(-1) 106 | x = x.repeat(1, 1, k, k) 107 | x = x.view(-1, height * k, width * k) 108 | return x 109 | 110 | cb = repeat(cb) 111 | cr = repeat(cr) 112 | 113 | return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3) 114 | 115 | 116 | class ycbcr_to_rgb_jpeg(nn.Module): 117 | """ Converts YCbCr image to RGB JPEG 118 | Input: 119 | image(tensor): batch x height x width x 3 120 | Outpput: 121 | result(tensor): batch x 3 x height x width 122 | """ 123 | def __init__(self): 124 | super(ycbcr_to_rgb_jpeg, self).__init__() 125 | 126 | matrix = np.array( 127 | [[1., 0., 1.402], [1, -0.344136, -0.714136], [1, 1.772, 0]], 128 | dtype=np.float32).T 129 | self.shift = nn.Parameter(torch.tensor([0, -128., -128.])).cuda() 130 | self.matrix = nn.Parameter(torch.from_numpy(matrix)).cuda() 131 | 132 | def forward(self, image): 133 | result = torch.tensordot(image + self.shift, self.matrix, dims=1) 134 | #result = torch.from_numpy(result) 135 | result.view(image.shape) 136 | return result.permute(0, 3, 1, 2) 137 | 138 | 139 | class decompress_jpeg(nn.Module): 140 | """ Full JPEG decompression algortihm 141 | Input: 142 | compressed(dict(tensor)): batch x h*w/64 x 8 x 8 143 | rounding(function): rounding function to use 144 | factor(float): Compression factor 145 | Ouput: 146 | image(tensor): batch x 3 x height x width 147 | """ 148 | def __init__(self, height, width, rounding=torch.round, factor=1): 149 | super(decompress_jpeg, self).__init__() 150 | self.c_dequantize = c_dequantize(factor=factor) 151 | self.y_dequantize = y_dequantize(factor=factor) 152 | self.idct = idct_8x8() 153 | self.merging = block_merging() 154 | self.chroma = chroma_upsampling() 155 | self.colors = ycbcr_to_rgb_jpeg() 156 | 157 | self.height, self.width = height, width 158 | 159 | def forward(self, y, cb, cr): 160 | components = {'y': y, 'cb': cb, 'cr': cr} 161 | for k in components.keys(): 162 | if k in ('cb', 'cr'): 163 | comp = self.c_dequantize(components[k]) 164 | height, width = int(self.height/2), int(self.width/2) 165 | else: 166 | comp = self.y_dequantize(components[k]) 167 | height, width = self.height, self.width 168 | comp = self.idct(comp) 169 | components[k] = self.merging(comp, height, width) 170 | # 171 | image = self.chroma(components['y'], components['cb'], components['cr']) 172 | image = self.colors(image) 173 | 174 | image = torch.min(255*torch.ones_like(image), 175 | torch.max(torch.zeros_like(image), image)) 176 | return image/255 177 | -------------------------------------------------------------------------------- /modules/compression.py: -------------------------------------------------------------------------------- 1 | # Standard libraries 2 | import itertools 3 | import numpy as np 4 | # PyTorch 5 | import torch 6 | import torch.nn as nn 7 | # Local 8 | import utils 9 | 10 | 11 | class rgb_to_ycbcr_jpeg(nn.Module): 12 | """ Converts RGB image to YCbCr 13 | Input: 14 | image(tensor): batch x 3 x height x width 15 | Outpput: 16 | result(tensor): batch x height x width x 3 17 | """ 18 | def __init__(self): 19 | super(rgb_to_ycbcr_jpeg, self).__init__() 20 | matrix = np.array( 21 | [[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], 22 | [0.5, -0.418688, -0.081312]], dtype=np.float32).T 23 | self.shift = nn.Parameter(torch.tensor([0., 128., 128.])).cuda() 24 | # 25 | self.matrix = nn.Parameter(torch.from_numpy(matrix)).cuda() 26 | 27 | def forward(self, image): 28 | image = image.permute(0, 2, 3, 1) 29 | result = torch.tensordot(image, self.matrix, dims=1) + self.shift 30 | # result = torch.from_numpy(result) 31 | result.view(image.shape) 32 | return result 33 | 34 | 35 | 36 | class chroma_subsampling(nn.Module): 37 | """ Chroma subsampling on CbCv channels 38 | Input: 39 | image(tensor): batch x height x width x 3 40 | Output: 41 | y(tensor): batch x height x width 42 | cb(tensor): batch x height/2 x width/2 43 | cr(tensor): batch x height/2 x width/2 44 | """ 45 | def __init__(self): 46 | super(chroma_subsampling, self).__init__() 47 | 48 | def forward(self, image): 49 | image_2 = image.permute(0, 3, 1, 2).clone() 50 | avg_pool = nn.AvgPool2d(kernel_size=2, stride=(2, 2), 51 | count_include_pad=False) 52 | cb = avg_pool(image_2[:, 1, :, :].unsqueeze(1)) 53 | cr = avg_pool(image_2[:, 2, :, :].unsqueeze(1)) 54 | cb = cb.permute(0, 2, 3, 1) 55 | cr = cr.permute(0, 2, 3, 1) 56 | return image[:, :, :, 0], cb.squeeze(3), cr.squeeze(3) 57 | 58 | 59 | class block_splitting(nn.Module): 60 | """ Splitting image into patches 61 | Input: 62 | image(tensor): batch x height x width 63 | Output: 64 | patch(tensor): batch x h*w/64 x h x w 65 | """ 66 | def __init__(self): 67 | super(block_splitting, self).__init__() 68 | self.k = 8 69 | 70 | def forward(self, image): 71 | height, width = image.shape[1:3] 72 | batch_size = image.shape[0] 73 | image_reshaped = image.view(batch_size, height // self.k, self.k, -1, self.k) 74 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) 75 | return image_transposed.contiguous().view(batch_size, -1, self.k, self.k) 76 | 77 | 78 | class dct_8x8(nn.Module): 79 | """ Discrete Cosine Transformation 80 | Input: 81 | image(tensor): batch x height x width 82 | Output: 83 | dcp(tensor): batch x height x width 84 | """ 85 | def __init__(self): 86 | super(dct_8x8, self).__init__() 87 | tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) 88 | for x, y, u, v in itertools.product(range(8), repeat=4): 89 | tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos( 90 | (2 * y + 1) * v * np.pi / 16) 91 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 92 | # 93 | self.tensor = nn.Parameter(torch.from_numpy(tensor).float()).cuda() 94 | self.scale = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha) * 0.25).float() ).cuda() 95 | 96 | def forward(self, image): 97 | image = image - 128 98 | result = self.scale * torch.tensordot(image, self.tensor, dims=2) 99 | result.view(image.shape) 100 | return result 101 | 102 | 103 | class y_quantize(nn.Module): 104 | """ JPEG Quantization for Y channel 105 | Input: 106 | image(tensor): batch x height x width 107 | rounding(function): rounding function to use 108 | factor(float): Degree of compression 109 | Output: 110 | image(tensor): batch x height x width 111 | """ 112 | def __init__(self, rounding, factor=1): 113 | super(y_quantize, self).__init__() 114 | self.rounding = rounding 115 | self.factor = factor 116 | self.y_table = utils.y_table 117 | 118 | def forward(self, image): 119 | image = image.float() / (self.y_table * self.factor) 120 | image = self.rounding(image) 121 | return image 122 | 123 | 124 | class c_quantize(nn.Module): 125 | """ JPEG Quantization for CrCb channels 126 | Input: 127 | image(tensor): batch x height x width 128 | rounding(function): rounding function to use 129 | factor(float): Degree of compression 130 | Output: 131 | image(tensor): batch x height x width 132 | """ 133 | def __init__(self, rounding, factor=1): 134 | super(c_quantize, self).__init__() 135 | self.rounding = rounding 136 | self.factor = factor 137 | self.c_table = utils.c_table 138 | 139 | def forward(self, image): 140 | image = image.float() / (self.c_table * self.factor) 141 | image = self.rounding(image) 142 | return image 143 | 144 | 145 | class compress_jpeg(nn.Module): 146 | """ Full JPEG compression algortihm 147 | Input: 148 | imgs(tensor): batch x 3 x height x width 149 | rounding(function): rounding function to use 150 | factor(float): Compression factor 151 | Ouput: 152 | compressed(dict(tensor)): batch x h*w/64 x 8 x 8 153 | """ 154 | def __init__(self, rounding=torch.round, factor=1): 155 | super(compress_jpeg, self).__init__() 156 | self.l1 = nn.Sequential( 157 | rgb_to_ycbcr_jpeg(), 158 | chroma_subsampling() 159 | ) 160 | self.l2 = nn.Sequential( 161 | block_splitting(), 162 | dct_8x8() 163 | ) 164 | self.c_quantize = c_quantize(rounding=rounding, factor=factor) 165 | self.y_quantize = y_quantize(rounding=rounding, factor=factor) 166 | 167 | def forward(self, image): 168 | y, cb, cr = self.l1(image*255) 169 | components = {'y': y, 'cb': cb, 'cr': cr} 170 | for k in components.keys(): 171 | comp = self.l2(components[k]) 172 | if k in ('cb', 'cr'): 173 | comp = self.c_quantize(comp) 174 | else: 175 | comp = self.y_quantize(comp) 176 | 177 | components[k] = comp 178 | 179 | return components['y'], components['cb'], components['cr'] 180 | -------------------------------------------------------------------------------- /model/encoder_decoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from model.encoder import Encoder 4 | from model.decoder import Decoder 5 | from options import HiDDenConfiguration 6 | 7 | from noise_layers.noiser import Noiser 8 | from noise_layers.identity import Identity 9 | from noise_layers.jpeg_compression2000 import JpegCompression2000 10 | from noise_layers.jpeg_compression2 import JpegCompression2 11 | from noise_layers.jpeg_compression import JpegCompression 12 | from noise_layers.quantization import Quantization 13 | from noise_layers.dropout import Dropout 14 | from noise_layers.crop import Crop 15 | from noise_layers.cropout import Cropout 16 | from noise_layers.gaussian import Gaussian 17 | from noise_layers.webp import WebP 18 | from noise_layers.mpeg4_compression import MPEG4 19 | from noise_layers.h264_compression import H264 20 | from noise_layers.xvid_compression import XVID 21 | from noise_layers.diff_quality_jpeg_compression2 import DiffQFJpegCompression2 22 | from noise_layers.diff_corruptions import DiffCorruptions 23 | 24 | 25 | class EncoderDecoder(nn.Module): 26 | """ 27 | Combines Encoder->Noiser->Decoder into single pipeline. 28 | The input is the cover image and the watermark message. The module inserts the watermark into the image 29 | (obtaining encoded_image), then applies Noise layers (obtaining noised_image), then passes the noised_image 30 | to the Decoder which tries to recover the watermark (called decoded_message). The module outputs 31 | a three-tuple: (encoded_image, noised_image, decoded_message) 32 | """ 33 | def __init__(self, config: HiDDenConfiguration, noiser: Noiser): 34 | 35 | super(EncoderDecoder, self).__init__() 36 | self.encoder = Encoder(config) 37 | self.noiser = noiser 38 | 39 | self.decoder = Decoder(config) 40 | 41 | self.identity = Identity() 42 | self.dropout = Dropout([0.3, 0.3]) 43 | self.cropout = Cropout([0.547, 0.547], [0.547, 0.547]) 44 | self.crop = Crop([0.187, 0.187], [0.187, 0.187]) 45 | self.gaussian = Gaussian() 46 | self.jpeg = JpegCompression2() 47 | 48 | self.jpeg_10 = JpegCompression2(quality=10) 49 | self.jpeg_25 = JpegCompression2(quality=25) 50 | self.jpeg_50 = JpegCompression2(quality=50) 51 | self.jpeg_75 = JpegCompression2(quality=75) 52 | self.jpeg_90 = JpegCompression2(quality=90) 53 | 54 | self.jpeg2000_100 = JpegCompression2000(quality=100) 55 | self.jpeg2000_250 = JpegCompression2000(quality=250) 56 | self.jpeg2000_500 = JpegCompression2000(quality=500) 57 | self.jpeg2000_750 = JpegCompression2000(quality=750) 58 | self.jpeg2000_900 = JpegCompression2000(quality=900) 59 | 60 | self.webp_10 = WebP(quality=10) 61 | self.webp_25 = WebP(quality=25) 62 | self.webp_50 = WebP(quality=50) 63 | self.webp_75 = WebP(quality=75) 64 | self.webp_90 = WebP(quality=90) 65 | 66 | self.jpeg2000 = JpegCompression2000() 67 | self.webp = WebP() 68 | self.mpeg4 = MPEG4() 69 | self.h264 = H264() 70 | self.xvid = XVID() 71 | self.diff_qf_jpeg2 = DiffQFJpegCompression2() 72 | self.diff_corruptions = DiffCorruptions() 73 | 74 | def forward(self, image, message): 75 | encoded_image = self.encoder(image, message) 76 | # noised_and_cover = self.noiser([torch.clamp(encoded_image, 0., 1.), image]) # noised_and_cover = self.noiser([encoded_image, image]) # changed by chaoning 77 | noised_and_cover = self.noiser([encoded_image, image]) # changed by chaoning, to the original (quant does not work anyway) 78 | noised_image = noised_and_cover[0] 79 | decoded_message = self.decoder(noised_image) 80 | return encoded_image, noised_image, decoded_message 81 | 82 | def forward_specific_noiser(self, image, message, noiser='identity'): 83 | encoded_image = self.encoder(image, message) 84 | # encoded_image = torch.clamp(encoded_image, 0., 1.) # changed by chaoning 85 | if noiser == 'identity': 86 | noised_and_cover = self.identity([encoded_image, image]) 87 | elif noiser == 'dropout': 88 | noised_and_cover = self.dropout([encoded_image, image]) 89 | elif noiser == 'cropout': 90 | noised_and_cover = self.cropout([encoded_image, image]) 91 | elif noiser == 'crop': 92 | noised_and_cover = self.crop([encoded_image, image]) 93 | elif noiser == 'gaussian': 94 | noised_and_cover = self.gaussian([encoded_image, image]) 95 | elif noiser == 'jpeg': 96 | noised_and_cover = self.jpeg([encoded_image, image]) 97 | elif noiser == 'jpeg_10': 98 | noised_and_cover = self.jpeg_10([encoded_image, image]) 99 | elif noiser == 'jpeg_25': 100 | noised_and_cover = self.jpeg_25([encoded_image, image]) 101 | elif noiser == 'jpeg_50': 102 | noised_and_cover = self.jpeg_50([encoded_image, image]) 103 | elif noiser == 'jpeg_75': 104 | noised_and_cover = self.jpeg_75([encoded_image, image]) 105 | elif noiser == 'jpeg_90': 106 | noised_and_cover = self.jpeg_90([encoded_image, image]) 107 | elif noiser == 'jpeg2000': 108 | noised_and_cover = self.jpeg2000([encoded_image, image]) 109 | elif noiser == 'jpeg2000_100': 110 | noised_and_cover = self.jpeg2000_100([encoded_image, image]) 111 | elif noiser == 'jpeg2000_250': 112 | noised_and_cover = self.jpeg2000_250([encoded_image, image]) 113 | elif noiser == 'jpeg2000_500': 114 | noised_and_cover = self.jpeg2000_500([encoded_image, image]) 115 | elif noiser == 'jpeg2000_750': 116 | noised_and_cover = self.jpeg2000_750([encoded_image, image]) 117 | elif noiser == 'jpeg2000_900': 118 | noised_and_cover = self.jpeg2000_900([encoded_image, image]) 119 | elif noiser == 'webp': 120 | noised_and_cover = self.webp([encoded_image, image]) 121 | elif noiser == 'webp_10': 122 | noised_and_cover = self.webp_10([encoded_image, image]) 123 | elif noiser == 'webp_25': 124 | noised_and_cover = self.webp_25([encoded_image, image]) 125 | elif noiser == 'webp_50': 126 | noised_and_cover = self.webp_50([encoded_image, image]) 127 | elif noiser == 'webp_75': 128 | noised_and_cover = self.webp_75([encoded_image, image]) 129 | elif noiser == 'webp_90': 130 | noised_and_cover = self.webp_90([encoded_image, image]) 131 | elif noiser == 'mpeg4': 132 | noised_and_cover = self.mpeg4([encoded_image, image]) 133 | elif noiser == 'h264': 134 | noised_and_cover = self.h264([encoded_image, image]) 135 | elif noiser == 'xvid': 136 | noised_and_cover = self.xvid([encoded_image, image]) 137 | elif noiser == 'diff_qf_jpeg2': 138 | noised_and_cover = self.diff_qf_jpeg2([encoded_image, image]) 139 | elif noiser == 'diff_corruptions': 140 | noised_and_cover = self.diff_corruptions([encoded_image, image]) 141 | 142 | noised_image = noised_and_cover[0] 143 | decoded_message = self.decoder(noised_image) 144 | return encoded_image, noised_image, decoded_message -------------------------------------------------------------------------------- /noise_layers/jpeg_compression.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | def gen_filters(size_x: int, size_y: int, dct_or_idct_fun: callable) -> np.ndarray: 7 | tile_size_x = 8 8 | filters = np.zeros((size_x * size_y, size_x, size_y)) 9 | for k_y in range(size_y): 10 | for k_x in range(size_x): 11 | for n_y in range(size_y): 12 | for n_x in range(size_x): 13 | filters[k_y * tile_size_x + k_x, n_y, n_x] = dct_or_idct_fun(n_y, k_y, size_y) * dct_or_idct_fun(n_x, 14 | k_x, 15 | size_x) 16 | return filters 17 | 18 | # def zigzag_indices(shape: (int, int), count): 19 | # x_range, y_range = shape 20 | # index_order = sorted(((x, y) for x in range(x_range) for y in range(y_range)), 21 | # key=lambda p: (p[0] + p[1], -p[1] if (p[0] + p[1]) % 2 else p[1])) 22 | # 23 | # mask = np.zeros(shape) 24 | # for r, c in index_order[:count]: 25 | # mask[r,c] = 1 26 | # 27 | # return mask 28 | 29 | def get_jpeg_yuv_filter_mask(image_shape: tuple, window_size: int, keep_count: int): 30 | mask = np.zeros((window_size, window_size), dtype=np.uint8) 31 | 32 | index_order = sorted(((x, y) for x in range(window_size) for y in range(window_size)), 33 | key=lambda p: (p[0] + p[1], -p[1] if (p[0] + p[1]) % 2 else p[1])) 34 | 35 | for i, j in index_order[0:keep_count]: 36 | mask[i, j] = 1 37 | 38 | return np.tile(mask, (int(np.ceil(image_shape[0] / window_size)), 39 | int(np.ceil(image_shape[1] / window_size))))[0: image_shape[0], 0: image_shape[1]] 40 | 41 | 42 | def dct_coeff(n, k, N): 43 | return np.cos(np.pi / N * (n + 1. / 2.) * k) 44 | 45 | 46 | def idct_coeff(n, k, N): 47 | return (int(0 == n) * (- 1 / 2) + np.cos( 48 | np.pi / N * (k + 1. / 2.) * n)) * np.sqrt(1 / (2. * N)) 49 | 50 | 51 | def rgb2yuv(image_rgb, image_yuv_out): 52 | """ Transform the image from rgb to yuv """ 53 | image_yuv_out[:, 0, :, :] = 0.299 * image_rgb[:, 0, :, :].clone() + 0.587 * image_rgb[:, 1, :, :].clone() + 0.114 * image_rgb[:, 2, :, :].clone() 54 | image_yuv_out[:, 1, :, :] = -0.14713 * image_rgb[:, 0, :, :].clone() + -0.28886 * image_rgb[:, 1, :, :].clone() + 0.436 * image_rgb[:, 2, :, :].clone() 55 | image_yuv_out[:, 2, :, :] = 0.615 * image_rgb[:, 0, :, :].clone() + -0.51499 * image_rgb[:, 1, :, :].clone() + -0.10001 * image_rgb[:, 2, :, :].clone() 56 | 57 | 58 | def yuv2rgb(image_yuv, image_rgb_out): 59 | """ Transform the image from yuv to rgb """ 60 | image_rgb_out[:, 0, :, :] = image_yuv[:, 0, :, :].clone() + 1.13983 * image_yuv[:, 2, :, :].clone() 61 | image_rgb_out[:, 1, :, :] = image_yuv[:, 0, :, :].clone() + -0.39465 * image_yuv[:, 1, :, :].clone() + -0.58060 * image_yuv[:, 2, :, :].clone() 62 | image_rgb_out[:, 2, :, :] = image_yuv[:, 0, :, :].clone() + 2.03211 * image_yuv[:, 1, :, :].clone() 63 | 64 | 65 | class JpegCompression(nn.Module): 66 | def __init__(self, device, yuv_keep_weights = (25, 9, 9)): 67 | super(JpegCompression, self).__init__() 68 | self.device = device 69 | 70 | self.dct_conv_weights = torch.tensor(gen_filters(8, 8, dct_coeff), dtype=torch.float32).to(self.device) 71 | self.dct_conv_weights.unsqueeze_(1) 72 | self.idct_conv_weights = torch.tensor(gen_filters(8, 8, idct_coeff), dtype=torch.float32).to(self.device) 73 | self.idct_conv_weights.unsqueeze_(1) 74 | 75 | self.yuv_keep_weighs = yuv_keep_weights 76 | self.keep_coeff_masks = [] 77 | 78 | self.jpeg_mask = None 79 | 80 | # create a new large mask which we can use by slicing for images which are smaller 81 | self.create_mask((1000, 1000)) 82 | 83 | 84 | def create_mask(self, requested_shape): 85 | if self.jpeg_mask is None or requested_shape > self.jpeg_mask.shape[1:]: 86 | self.jpeg_mask = torch.empty((3,) + requested_shape, device=self.device) 87 | for channel, weights_to_keep in enumerate(self.yuv_keep_weighs): 88 | mask = torch.from_numpy(get_jpeg_yuv_filter_mask(requested_shape, 8, weights_to_keep)) 89 | self.jpeg_mask[channel] = mask 90 | 91 | def get_mask(self, image_shape): 92 | if self.jpeg_mask.shape < image_shape: 93 | self.create_mask(image_shape) 94 | # return the correct slice of it 95 | return self.jpeg_mask[:, :image_shape[1], :image_shape[2]].clone() 96 | 97 | 98 | def apply_conv(self, image, filter_type: str): 99 | 100 | if filter_type == 'dct': 101 | filters = self.dct_conv_weights 102 | elif filter_type == 'idct': 103 | filters = self.idct_conv_weights 104 | else: 105 | raise('Unknown filter_type value.') 106 | 107 | image_conv_channels = [] 108 | for channel in range(image.shape[1]): 109 | image_yuv_ch = image[:, channel, :, :].unsqueeze_(1) 110 | image_conv = F.conv2d(image_yuv_ch, filters, stride=8) 111 | image_conv = image_conv.permute(0, 2, 3, 1) 112 | image_conv = image_conv.view(image_conv.shape[0], image_conv.shape[1], image_conv.shape[2], 8, 8) 113 | image_conv = image_conv.permute(0, 1, 3, 2, 4) 114 | image_conv = image_conv.contiguous().view(image_conv.shape[0], 115 | image_conv.shape[1]*image_conv.shape[2], 116 | image_conv.shape[3]*image_conv.shape[4]) 117 | 118 | image_conv.unsqueeze_(1) 119 | 120 | # image_conv = F.conv2d() 121 | image_conv_channels.append(image_conv) 122 | 123 | image_conv_stacked = torch.cat(image_conv_channels, dim=1) 124 | 125 | return image_conv_stacked 126 | 127 | 128 | def forward(self, noised_and_cover): 129 | 130 | noised_image = noised_and_cover[0] 131 | # pad the image so that we can do dct on 8x8 blocks 132 | pad_height = (8 - noised_image.shape[2] % 8) % 8 133 | pad_width = (8 - noised_image.shape[3] % 8) % 8 134 | 135 | noised_image = nn.ZeroPad2d((0, pad_width, 0, pad_height))(noised_image) 136 | 137 | # convert to yuv 138 | image_yuv = torch.empty_like(noised_image) 139 | rgb2yuv(noised_image, image_yuv) 140 | 141 | assert image_yuv.shape[2] % 8 == 0 142 | assert image_yuv.shape[3] % 8 == 0 143 | 144 | # apply dct 145 | image_dct = self.apply_conv(image_yuv, 'dct') 146 | # get the jpeg-compression mask 147 | mask = self.get_mask(image_dct.shape[1:]) 148 | # multiply the dct-ed image with the mask. 149 | image_dct_mask = torch.mul(image_dct, mask) 150 | 151 | # apply inverse dct (idct) 152 | image_idct = self.apply_conv(image_dct_mask, 'idct') 153 | # transform from yuv to to rgb 154 | image_ret_padded = torch.empty_like(image_dct) 155 | yuv2rgb(image_idct, image_ret_padded) 156 | 157 | # un-pad 158 | noised_and_cover[0] = image_ret_padded[:, :, :image_ret_padded.shape[2]-pad_height, :image_ret_padded.shape[3]-pad_width].clone() 159 | 160 | return noised_and_cover -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pprint 3 | import argparse 4 | import torch 5 | import pickle 6 | import utils 7 | import logging 8 | import sys 9 | import socket 10 | import numpy as np 11 | 12 | from options import * 13 | from model.hidden import Hidden 14 | from noise_layers.noiser import Noiser 15 | from noise_argparser import NoiseArgParser 16 | 17 | from train import train, train_other_noises, train_own_noise 18 | 19 | 20 | def main(): 21 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 22 | 23 | parent_parser = argparse.ArgumentParser(description='Training of HiDDeN nets') 24 | subparsers = parent_parser.add_subparsers(dest='command', help='Sub-parser for commands') 25 | new_run_parser = subparsers.add_parser('new', help='starts a new run') 26 | new_run_parser.add_argument('--data-dir', '-d', required=True, type=str, 27 | help='The directory where the data is stored.') 28 | new_run_parser.add_argument('--batch-size', '-b', required=True, type=int, help='The batch size.') 29 | new_run_parser.add_argument('--epochs', '-e', default=300, type=int, help='Number of epochs to run the simulation.') 30 | new_run_parser.add_argument('--name', required=True, type=str, help='The name of the experiment.') 31 | new_run_parser.add_argument('--adv_loss', default=0, required=False, type=float, help='Coefficient of the adversarial loss.') 32 | new_run_parser.add_argument('--residual', default=0, required=False, type=int, help='If to use residual or not.') 33 | new_run_parser.add_argument('--video_dataset', default=0, required=False, type=int, help='If to use video dataset or not.') 34 | new_run_parser.add_argument('--save-dir', '-sd', default='runs', required=True, type=str, 35 | help='The save directory where the result is stored.') 36 | 37 | new_run_parser.add_argument('--size', '-s', default=128, type=int, 38 | help='The size of the images (images are square so this is height and width).') 39 | new_run_parser.add_argument('--message', '-m', default=30, type=int, help='The length in bits of the watermark.') 40 | new_run_parser.add_argument('--continue-from-folder', '-c', default='', type=str, 41 | help='The folder from where to continue a previous run. Leave blank if you are starting a new experiment.') 42 | # parser.add_argument('--tensorboard', dest='tensorboard', action='store_true', 43 | # help='If specified, use adds a Tensorboard log. On by default') 44 | new_run_parser.add_argument('--tensorboard', action='store_true', 45 | help='Use to switch on Tensorboard logging.') 46 | new_run_parser.add_argument('--enable-fp16', dest='enable_fp16', action='store_true', 47 | help='Enable mixed-precision training.') 48 | 49 | new_run_parser.add_argument('--noise', nargs='*', action=NoiseArgParser, 50 | help="Noise layers configuration. Use quotes when specifying configuration, e.g. 'cropout((0.55, 0.6), (0.55, 0.6))'") 51 | new_run_parser.add_argument('--hostname', default=socket.gethostname(), help='the host name of the running server') 52 | new_run_parser.add_argument('--cover-dependent', default=1, required=False, type=int, help='If to use cover dependent architecture or not.') 53 | new_run_parser.add_argument('--jpeg_type', '-j', required=False, type=str, default='jpeg', 54 | help='Jpeg type used in the combined2 noise.') 55 | 56 | new_run_parser.set_defaults(tensorboard=False) 57 | new_run_parser.set_defaults(enable_fp16=False) 58 | 59 | continue_parser = subparsers.add_parser('continue', help='Continue a previous run') 60 | continue_parser.add_argument('--folder', '-f', required=True, type=str, 61 | help='Continue from the last checkpoint in this folder.') 62 | continue_parser.add_argument('--data-dir', '-d', required=False, type=str, 63 | help='The directory where the data is stored. Specify a value only if you want to override the previous value.') 64 | continue_parser.add_argument('--epochs', '-e', required=False, type=int, 65 | help='Number of epochs to run the simulation. Specify a value only if you want to override the previous value.') 66 | 67 | # continue_parser.add_argument('--tensorboard', action='store_true', 68 | # help='Override the previous setting regarding tensorboard logging.') 69 | 70 | # Setting up a seed for debug 71 | seed = 123 72 | torch.manual_seed(seed) 73 | np.random.seed(seed) 74 | 75 | args = parent_parser.parse_args() 76 | checkpoint = None 77 | loaded_checkpoint_file_name = None 78 | print(args.cover_dependent) 79 | 80 | if not args.video_dataset: 81 | if args.hostname == 'ee898-System-Product-Name': 82 | args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet' 83 | args.hostname = 'ee898' 84 | elif args.hostname == 'DL178': 85 | args.data_dir = '/media/user/SSD1TB-2/ImageNet' 86 | else: 87 | args.data_dir = '/workspace/data_local/imagenet_pytorch' 88 | else: 89 | if args.hostname == 'ee898-System-Product-Name': 90 | args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet' 91 | args.hostname = 'ee898' 92 | elif args.hostname == 'DL178': 93 | args.data_dir = '/media/user/SSD1TB-2/ImageNet' 94 | else: 95 | args.data_dir = './oops_dataset/oops_video' 96 | assert args.data_dir 97 | 98 | if args.command == 'continue': 99 | this_run_folder = args.folder 100 | options_file = os.path.join(this_run_folder, 'options-and-config.pickle') 101 | train_options, hidden_config, noise_config = utils.load_options(options_file) 102 | checkpoint, loaded_checkpoint_file_name = utils.load_last_checkpoint(os.path.join(this_run_folder, 'checkpoints')) 103 | train_options.start_epoch = checkpoint['epoch'] + 1 104 | if args.data_dir is not None: 105 | train_options.train_folder = os.path.join(args.data_dir, 'train') 106 | train_options.validation_folder = os.path.join(args.data_dir, 'val') 107 | if args.epochs is not None: 108 | if train_options.start_epoch < args.epochs: 109 | train_options.number_of_epochs = args.epochs 110 | else: 111 | print(f'Command-line specifies of number of epochs = {args.epochs}, but folder={args.folder} ' 112 | f'already contains checkpoint for epoch = {train_options.start_epoch}.') 113 | exit(1) 114 | 115 | else: 116 | assert args.command == 'new' 117 | start_epoch = 1 118 | train_options = TrainingOptions( 119 | batch_size=args.batch_size, 120 | number_of_epochs=args.epochs, 121 | train_folder=os.path.join(args.data_dir, 'train'), 122 | validation_folder=os.path.join(args.data_dir, 'val'), 123 | runs_folder=os.path.join('.', args.save_dir), 124 | start_epoch=start_epoch, 125 | experiment_name=args.name, 126 | video_dataset=args.video_dataset) 127 | 128 | 129 | noise_config = args.noise if args.noise is not None else [] 130 | hidden_config = HiDDenConfiguration(H=args.size, W=args.size, 131 | message_length=args.message, 132 | encoder_blocks=4, encoder_channels=64, 133 | decoder_blocks=7, decoder_channels=64, 134 | use_discriminator=True, 135 | use_vgg=False, 136 | discriminator_blocks=3, discriminator_channels=64, 137 | decoder_loss=1, 138 | encoder_loss=0.7, 139 | adversarial_loss=args.adv_loss, 140 | cover_dependent=args.cover_dependent, 141 | residual=args.residual, 142 | enable_fp16=args.enable_fp16 143 | ) 144 | 145 | this_run_folder = utils.create_folder_for_run(train_options.runs_folder, args.name) 146 | with open(os.path.join(this_run_folder, 'options-and-config.pickle'), 'wb+') as f: 147 | pickle.dump(train_options, f) 148 | pickle.dump(noise_config, f) 149 | pickle.dump(hidden_config, f) 150 | 151 | 152 | logging.basicConfig(level=logging.INFO, 153 | format='%(message)s', 154 | handlers=[ 155 | logging.FileHandler(os.path.join(this_run_folder, f'{train_options.experiment_name}.log')), 156 | logging.StreamHandler(sys.stdout) 157 | ]) 158 | if (args.command == 'new' and args.tensorboard) or \ 159 | (args.command == 'continue' and os.path.isdir(os.path.join(this_run_folder, 'tb-logs'))): 160 | logging.info('Tensorboard is enabled. Creating logger.') 161 | from tensorboard_logger import TensorBoardLogger 162 | tb_logger = TensorBoardLogger(os.path.join(this_run_folder, 'tb-logs')) 163 | else: 164 | tb_logger = None 165 | 166 | 167 | noiser = Noiser(noise_config, device, args.jpeg_type) 168 | model = Hidden(hidden_config, device, noiser, tb_logger) 169 | 170 | if args.command == 'continue': 171 | # if we are continuing, we have to load the model params 172 | assert checkpoint is not None 173 | logging.info(f'Loading checkpoint from file {loaded_checkpoint_file_name}') 174 | utils.model_from_checkpoint(model, checkpoint) 175 | 176 | logging.info('HiDDeN model: {}\n'.format(model.to_stirng())) 177 | logging.info('Model Configuration:\n') 178 | logging.info(pprint.pformat(vars(hidden_config))) 179 | logging.info('\nNoise configuration:\n') 180 | logging.info(pprint.pformat(str(noise_config))) 181 | logging.info('\nTraining train_options:\n') 182 | logging.info(pprint.pformat(vars(train_options))) 183 | 184 | # train(model, device, hidden_config, train_options, this_run_folder, tb_logger) 185 | # train_other_noises(model, device, hidden_config, train_options, this_run_folder, tb_logger) 186 | if str(args.noise[0]) == "WebP()": 187 | noise = 'webp' 188 | elif str(args.noise[0]) == "JpegCompression2000()": 189 | noise = 'jpeg2000' 190 | elif str(args.noise[0]) == "MPEG4()": 191 | noise = 'mpeg4' 192 | elif str(args.noise[0]) == "H264()": 193 | noise = 'h264' 194 | elif str(args.noise[0]) == "XVID()": 195 | noise = 'xvid' 196 | elif str(args.noise[0]) == "DiffQFJpegCompression2()": 197 | noise = 'diff_qf_jpeg2' 198 | elif str(args.noise[0]) == "DiffCorruptions()": 199 | noise = 'diff_corruptions' 200 | else: 201 | noise = 'jpeg' 202 | train_own_noise(model, device, hidden_config, train_options, this_run_folder, tb_logger, noise) 203 | 204 | if __name__ == '__main__': 205 | main() -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import re 4 | import csv 5 | import time 6 | import pickle 7 | import logging 8 | 9 | import torch 10 | from torchvision import datasets, transforms 11 | import torchvision.utils 12 | from torch.utils import data 13 | import torch.nn.functional as F 14 | 15 | from options import HiDDenConfiguration, TrainingOptions 16 | from model.hidden import Hidden 17 | import torch.nn as nn 18 | 19 | from oops_dataset import OOPSDataset 20 | 21 | 22 | def image_to_tensor(image): 23 | """ 24 | Transforms a numpy-image into torch tensor 25 | :param image: (batch_size x height x width x channels) uint8 array 26 | :return: (batch_size x channels x height x width) torch tensor in range [-1.0, 1.0] 27 | """ 28 | image_tensor = torch.Tensor(image) 29 | image_tensor.unsqueeze_(0) 30 | image_tensor = image_tensor.permute(0, 3, 1, 2) 31 | image_tensor = image_tensor / 127.5 - 1 32 | return image_tensor 33 | 34 | 35 | def tensor_to_image(tensor): 36 | """ 37 | Transforms a torch tensor into numpy uint8 array (image) 38 | :param tensor: (batch_size x channels x height x width) torch tensor in range [-1.0, 1.0] 39 | :return: (batch_size x height x width x channels) uint8 array 40 | """ 41 | image = tensor.permute(0, 2, 3, 1).cpu().numpy() 42 | image = (image + 1) * 127.5 43 | return np.clip(image, 0, 255).astype(np.uint8) 44 | 45 | 46 | def save_images(original_images, watermarked_images, epoch, folder, resize_to=None): 47 | images = original_images[:original_images.shape[0], :, :, :].cpu() 48 | watermarked_images = watermarked_images[:watermarked_images.shape[0], :, :, :].cpu() 49 | 50 | # # scale values to range [0, 1] from original range of [-1, 1] 51 | # images = (images + 1) / 2 52 | # watermarked_images = (watermarked_images + 1) / 2 53 | 54 | if resize_to is not None: 55 | images = F.interpolate(images, size=resize_to) 56 | watermarked_images = F.interpolate(watermarked_images, size=resize_to) 57 | 58 | stacked_images = torch.cat([images, watermarked_images], dim=0) 59 | filename = os.path.join(folder, 'epoch-{}.png'.format(epoch)) 60 | torchvision.utils.save_image(stacked_images, filename, original_images.shape[0], normalize=False) 61 | 62 | 63 | def sorted_nicely(l): 64 | """ Sort the given iterable in the way that humans expect.""" 65 | convert = lambda text: int(text) if text.isdigit() else text 66 | alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] 67 | return sorted(l, key=alphanum_key) 68 | 69 | 70 | def last_checkpoint_from_folder(folder: str): 71 | last_file = sorted_nicely(os.listdir(folder))[-1] 72 | last_file = os.path.join(folder, last_file) 73 | return last_file 74 | 75 | 76 | def save_checkpoint(model: Hidden, experiment_name: str, epoch: int, checkpoint_folder: str): 77 | """ Saves a checkpoint at the end of an epoch. """ 78 | if not os.path.exists(checkpoint_folder): 79 | os.makedirs(checkpoint_folder) 80 | 81 | checkpoint_filename = f'{experiment_name}--epoch-{epoch}.pyt' 82 | checkpoint_filename = os.path.join(checkpoint_folder, checkpoint_filename) 83 | logging.info('Saving checkpoint to {}'.format(checkpoint_filename)) 84 | checkpoint = { 85 | 'enc-dec-model': model.encoder_decoder.state_dict(), 86 | 'enc-dec-optim': model.optimizer_enc_dec.state_dict(), 87 | 'discrim-model': model.discriminator.state_dict(), 88 | 'discrim-optim': model.optimizer_discrim.state_dict(), 89 | 'epoch': epoch 90 | } 91 | torch.save(checkpoint, checkpoint_filename) 92 | logging.info('Saving checkpoint done.') 93 | 94 | 95 | # def load_checkpoint(hidden_net: Hidden, options: Options, this_run_folder: str): 96 | def load_last_checkpoint(checkpoint_folder): 97 | """ Load the last checkpoint from the given folder """ 98 | last_checkpoint_file = last_checkpoint_from_folder(checkpoint_folder) 99 | checkpoint = torch.load(last_checkpoint_file) 100 | 101 | return checkpoint, last_checkpoint_file 102 | 103 | 104 | def model_from_checkpoint(hidden_net, checkpoint): 105 | """ Restores the hidden_net object from a checkpoint object """ 106 | hidden_net.encoder_decoder.load_state_dict(checkpoint['enc-dec-model']) 107 | hidden_net.optimizer_enc_dec.load_state_dict(checkpoint['enc-dec-optim']) 108 | hidden_net.discriminator.load_state_dict(checkpoint['discrim-model']) 109 | hidden_net.optimizer_discrim.load_state_dict(checkpoint['discrim-optim']) 110 | 111 | 112 | def load_options(options_file_name) -> (TrainingOptions, HiDDenConfiguration, dict): 113 | """ Loads the training, model, and noise configurations from the given folder """ 114 | with open(os.path.join(options_file_name), 'rb') as f: 115 | train_options = pickle.load(f) 116 | noise_config = pickle.load(f) 117 | hidden_config = pickle.load(f) 118 | # for backward-capability. Some models were trained and saved before .enable_fp16 was added 119 | if not hasattr(hidden_config, 'enable_fp16'): 120 | setattr(hidden_config, 'enable_fp16', False) 121 | 122 | return train_options, hidden_config, noise_config 123 | 124 | 125 | def get_data_loaders(hidden_config: HiDDenConfiguration, train_options: TrainingOptions): 126 | """ Get torch data loaders for training and validation. The data loaders take a crop of the image, 127 | transform it into tensor, and normalize it.""" 128 | # data_transforms = { 129 | # 'train': transforms.Compose([ 130 | # transforms.RandomCrop((hidden_config.H, hidden_config.W), pad_if_needed=True), 131 | # transforms.ToTensor(), 132 | # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) 133 | # ]), 134 | # 'test': transforms.Compose([ 135 | # transforms.CenterCrop((hidden_config.H, hidden_config.W)), 136 | # transforms.ToTensor(), 137 | # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) 138 | # ]) 139 | # } 140 | 141 | # data_transforms = { 142 | # 'train': transforms.Compose([ 143 | # transforms.Resize([hidden_config.H, hidden_config.W]), 144 | # transforms.ToTensor(), 145 | # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) 146 | # ]), 147 | # 'test': transforms.Compose([ 148 | # transforms.Resize([hidden_config.H, hidden_config.W]), 149 | # transforms.ToTensor(), 150 | # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) 151 | # ]) 152 | # } 153 | 154 | data_transforms = { 155 | 'train': transforms.Compose([ 156 | transforms.Resize([hidden_config.H, hidden_config.W]), 157 | transforms.ToTensor() 158 | ]), 159 | 'test': transforms.Compose([ 160 | transforms.Resize([hidden_config.H, hidden_config.W]), 161 | transforms.ToTensor() 162 | ]) 163 | } 164 | 165 | try: 166 | if not train_options.video_dataset: 167 | train_images = datasets.ImageFolder(train_options.train_folder, data_transforms['train']) 168 | train_loader = torch.utils.data.DataLoader(train_images, batch_size=train_options.batch_size, shuffle=True, 169 | num_workers=4) 170 | 171 | validation_images = datasets.ImageFolder(train_options.validation_folder, data_transforms['test']) 172 | validation_loader = torch.utils.data.DataLoader(validation_images, batch_size=train_options.batch_size, 173 | shuffle=False, num_workers=4) 174 | else: 175 | train_images = OOPSDataset(train_options.train_folder, data_transforms['train']) 176 | train_loader = train_images 177 | # train_loader = torch.utils.data.DataLoader(train_images, batch_size=train_options.batch_size, shuffle=True, 178 | # num_workers=4) 179 | 180 | validation_images = OOPSDataset(train_options.validation_folder, data_transforms['test']) 181 | validation_loader = validation_images 182 | # validation_loader = torch.utils.data.DataLoader(validation_images, batch_size=train_options.batch_size, 183 | # shuffle=False, num_workers=4) 184 | except AttributeError: 185 | train_images = datasets.ImageFolder(train_options.train_folder, data_transforms['train']) 186 | train_loader = torch.utils.data.DataLoader(train_images, batch_size=train_options.batch_size, shuffle=True, 187 | num_workers=4) 188 | 189 | validation_images = datasets.ImageFolder(train_options.validation_folder, data_transforms['test']) 190 | validation_loader = torch.utils.data.DataLoader(validation_images, batch_size=train_options.batch_size, 191 | shuffle=False, num_workers=4) 192 | 193 | return train_loader, validation_loader 194 | 195 | 196 | def log_progress(losses_accu): 197 | log_print_helper(losses_accu, logging.info) 198 | 199 | 200 | def print_progress(losses_accu): 201 | log_print_helper(losses_accu, print) 202 | 203 | 204 | def log_print_helper(losses_accu, log_or_print_func): 205 | max_len = max([len(loss_name) for loss_name in losses_accu]) 206 | for loss_name, loss_value in losses_accu.items(): 207 | log_or_print_func(loss_name.ljust(max_len + 4) + '{:.4f}'.format(loss_value.avg)) 208 | 209 | 210 | def create_folder_for_run(runs_folder, experiment_name): 211 | if not os.path.exists(runs_folder): 212 | os.makedirs(runs_folder) 213 | 214 | this_run_folder = os.path.join(runs_folder, f'{experiment_name} {time.strftime("%Y.%m.%d--%H-%M-%S")}') 215 | 216 | os.makedirs(this_run_folder) 217 | os.makedirs(os.path.join(this_run_folder, 'checkpoints')) 218 | os.makedirs(os.path.join(this_run_folder, 'images')) 219 | 220 | return this_run_folder 221 | 222 | 223 | def write_losses(file_name, losses_accu, epoch, duration): 224 | with open(file_name, 'a', newline='') as csvfile: 225 | writer = csv.writer(csvfile) 226 | if epoch == 1: 227 | row_to_write = ['epoch'] + [loss_name.strip() for loss_name in losses_accu.keys()] + ['duration'] 228 | writer.writerow(row_to_write) 229 | row_to_write = [epoch] + ['{:.4f}'.format(loss_avg.avg) for loss_avg in losses_accu.values()] + [ 230 | '{:.0f}'.format(duration)] 231 | writer.writerow(row_to_write) 232 | 233 | y_table = np.array( 234 | [[16, 11, 10, 16, 24, 40, 51, 61], [12, 12, 14, 19, 26, 58, 60, 235 | 55], [14, 13, 16, 24, 40, 57, 69, 56], 236 | [14, 17, 22, 29, 51, 87, 80, 62], [18, 22, 37, 56, 68, 109, 103, 237 | 77], [24, 35, 55, 64, 81, 104, 113, 92], 238 | [49, 64, 78, 87, 103, 121, 120, 101], [72, 92, 95, 98, 112, 100, 103, 99]], 239 | dtype=np.float32).T 240 | 241 | y_table = nn.Parameter(torch.from_numpy(y_table)).cuda() 242 | # 243 | c_table = np.empty((8, 8), dtype=np.float32) 244 | c_table.fill(99) 245 | c_table[:4, :4] = np.array([[17, 18, 24, 47], [18, 21, 26, 66], 246 | [24, 26, 56, 99], [47, 66, 99, 99]]).T 247 | c_table = nn.Parameter(torch.from_numpy(c_table)).cuda() -------------------------------------------------------------------------------- /model/hidden.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | 5 | from options import HiDDenConfiguration 6 | from model.discriminator import Discriminator 7 | from model.encoder_decoder import EncoderDecoder 8 | from vgg_loss import VGGLoss 9 | from noise_layers.noiser import Noiser 10 | 11 | 12 | class Hidden: 13 | def __init__(self, configuration: HiDDenConfiguration, device: torch.device, noiser: Noiser, tb_logger): 14 | """ 15 | :param configuration: Configuration for the net, such as the size of the input image, number of channels in the intermediate layers, etc. 16 | :param device: torch.device object, CPU or GPU 17 | :param noiser: Object representing stacked noise layers. 18 | :param tb_logger: Optional TensorboardX logger object, if specified -- enables Tensorboard logging 19 | """ 20 | super(Hidden, self).__init__() 21 | 22 | self.encoder_decoder = EncoderDecoder(configuration, noiser).to(device) 23 | self.discriminator = Discriminator(configuration).to(device) 24 | self.optimizer_enc_dec = torch.optim.Adam(self.encoder_decoder.parameters()) 25 | self.optimizer_discrim = torch.optim.Adam(self.discriminator.parameters()) 26 | 27 | if configuration.use_vgg: 28 | self.vgg_loss = VGGLoss(3, 1, False) 29 | self.vgg_loss.to(device) 30 | else: 31 | self.vgg_loss = None 32 | 33 | self.config = configuration 34 | self.device = device 35 | 36 | self.bce_with_logits_loss = nn.BCEWithLogitsLoss().to(device) 37 | self.mse_loss = nn.MSELoss().to(device) 38 | 39 | # Defined the labels used for training the discriminator/adversarial loss 40 | self.cover_label = 1 41 | self.encoded_label = 0 42 | 43 | self.tb_logger = tb_logger 44 | if tb_logger is not None: 45 | from tensorboard_logger import TensorBoardLogger 46 | encoder_final = self.encoder_decoder.encoder._modules['final_layer'] 47 | encoder_final.weight.register_hook(tb_logger.grad_hook_by_name('grads/encoder_out')) 48 | decoder_final = self.encoder_decoder.decoder._modules['linear'] 49 | decoder_final.weight.register_hook(tb_logger.grad_hook_by_name('grads/decoder_out')) 50 | discrim_final = self.discriminator._modules['linear'] 51 | discrim_final.weight.register_hook(tb_logger.grad_hook_by_name('grads/discrim_out')) 52 | 53 | 54 | def train_on_batch(self, batch: list): 55 | """ 56 | Trains the network on a single batch consisting of images and messages 57 | :param batch: batch of training data, in the form [images, messages] 58 | :return: dictionary of error metrics from Encoder, Decoder, and Discriminator on the current batch 59 | """ 60 | images, messages = batch 61 | 62 | batch_size = images.shape[0] 63 | self.encoder_decoder.train() 64 | self.discriminator.train() 65 | with torch.enable_grad(): 66 | # ---------------- Train the discriminator ----------------------------- 67 | self.optimizer_discrim.zero_grad() 68 | # train on cover 69 | d_target_label_cover = torch.full((batch_size, 1), self.cover_label, device=self.device) 70 | d_target_label_encoded = torch.full((batch_size, 1), self.encoded_label, device=self.device) 71 | g_target_label_encoded = torch.full((batch_size, 1), self.cover_label, device=self.device) 72 | 73 | d_on_cover = self.discriminator(images) 74 | d_loss_on_cover = self.bce_with_logits_loss(d_on_cover, d_target_label_cover) 75 | d_loss_on_cover.backward() 76 | 77 | # train on fake 78 | encoded_images, noised_images, decoded_messages = self.encoder_decoder(images, messages) 79 | d_on_encoded = self.discriminator(encoded_images.detach()) 80 | d_loss_on_encoded = self.bce_with_logits_loss(d_on_encoded, d_target_label_encoded) 81 | 82 | d_loss_on_encoded.backward() 83 | self.optimizer_discrim.step() 84 | 85 | # --------------Train the generator (encoder-decoder) --------------------- 86 | self.optimizer_enc_dec.zero_grad() 87 | # target label for encoded images should be 'cover', because we want to fool the discriminator 88 | d_on_encoded_for_enc = self.discriminator(encoded_images) 89 | g_loss_adv = self.bce_with_logits_loss(d_on_encoded_for_enc, g_target_label_encoded) 90 | 91 | if self.vgg_loss == None: 92 | g_loss_enc = self.mse_loss(encoded_images, images) 93 | else: 94 | vgg_on_cov = self.vgg_loss(images) 95 | vgg_on_enc = self.vgg_loss(encoded_images) 96 | g_loss_enc = self.mse_loss(vgg_on_cov, vgg_on_enc) 97 | 98 | g_loss_dec = self.mse_loss(decoded_messages, messages) 99 | g_loss = self.config.adversarial_loss * g_loss_adv + self.config.encoder_loss * g_loss_enc \ 100 | + self.config.decoder_loss * g_loss_dec 101 | 102 | g_loss.backward() 103 | self.optimizer_enc_dec.step() 104 | 105 | decoded_rounded = decoded_messages.detach().cpu().numpy().round().clip(0, 1) 106 | bitwise_avg_err = np.sum(np.abs(decoded_rounded - messages.detach().cpu().numpy())) / ( 107 | batch_size * messages.shape[1]) 108 | 109 | losses = { 110 | 'loss ': g_loss.item(), 111 | 'encoder_mse ': g_loss_enc.item(), 112 | 'dec_mse ': g_loss_dec.item(), 113 | 'bitwise-error ': bitwise_avg_err, 114 | 'adversarial_bce': g_loss_adv.item(), 115 | 'discr_cover_bce': d_loss_on_cover.item(), 116 | 'discr_encod_bce': d_loss_on_encoded.item() 117 | } 118 | return losses, (encoded_images, noised_images, decoded_messages) 119 | 120 | def validate_on_batch(self, batch: list): 121 | """ 122 | Runs validation on a single batch of data consisting of images and messages 123 | :param batch: batch of validation data, in form [images, messages] 124 | :return: dictionary of error metrics from Encoder, Decoder, and Discriminator on the current batch 125 | """ 126 | # if TensorboardX logging is enabled, save some of the tensors. 127 | if self.tb_logger is not None: 128 | encoder_final = self.encoder_decoder.encoder._modules['final_layer'] 129 | self.tb_logger.add_tensor('weights/encoder_out', encoder_final.weight) 130 | decoder_final = self.encoder_decoder.decoder._modules['linear'] 131 | self.tb_logger.add_tensor('weights/decoder_out', decoder_final.weight) 132 | discrim_final = self.discriminator._modules['linear'] 133 | self.tb_logger.add_tensor('weights/discrim_out', discrim_final.weight) 134 | 135 | images, messages = batch 136 | 137 | batch_size = images.shape[0] 138 | 139 | self.encoder_decoder.eval() 140 | self.discriminator.eval() 141 | with torch.no_grad(): 142 | d_target_label_cover = torch.full((batch_size, 1), self.cover_label, device=self.device) 143 | d_target_label_encoded = torch.full((batch_size, 1), self.encoded_label, device=self.device) 144 | g_target_label_encoded = torch.full((batch_size, 1), self.cover_label, device=self.device) 145 | 146 | d_on_cover = self.discriminator(images) 147 | d_loss_on_cover = self.bce_with_logits_loss(d_on_cover, d_target_label_cover) 148 | 149 | encoded_images, noised_images, decoded_messages = self.encoder_decoder(images, messages) 150 | 151 | d_on_encoded = self.discriminator(encoded_images) 152 | d_loss_on_encoded = self.bce_with_logits_loss(d_on_encoded, d_target_label_encoded) 153 | 154 | d_on_encoded_for_enc = self.discriminator(encoded_images) 155 | g_loss_adv = self.bce_with_logits_loss(d_on_encoded_for_enc, g_target_label_encoded) 156 | 157 | if self.vgg_loss is None: 158 | g_loss_enc = self.mse_loss(encoded_images, images) 159 | else: 160 | vgg_on_cov = self.vgg_loss(images) 161 | vgg_on_enc = self.vgg_loss(encoded_images) 162 | g_loss_enc = self.mse_loss(vgg_on_cov, vgg_on_enc) 163 | 164 | g_loss_dec = self.mse_loss(decoded_messages, messages) 165 | g_loss = self.config.adversarial_loss * g_loss_adv + self.config.encoder_loss * g_loss_enc \ 166 | + self.config.decoder_loss * g_loss_dec 167 | 168 | decoded_rounded = decoded_messages.detach().cpu().numpy().round().clip(0, 1) 169 | bitwise_avg_err = np.sum(np.abs(decoded_rounded - messages.detach().cpu().numpy())) / ( 170 | batch_size * messages.shape[1]) 171 | 172 | losses = { 173 | 'loss ': g_loss.item(), 174 | 'encoder_mse ': g_loss_enc.item(), 175 | 'dec_mse ': g_loss_dec.item(), 176 | 'bitwise-error ': bitwise_avg_err, 177 | 'adversarial_bce': g_loss_adv.item(), 178 | 'discr_cover_bce': d_loss_on_cover.item(), 179 | 'discr_encod_bce': d_loss_on_encoded.item() 180 | } 181 | return losses, (encoded_images, noised_images, decoded_messages) 182 | 183 | def validate_on_batch_specific_noise(self, batch: list, noise='identity'): 184 | """ 185 | Runs validation on a single batch of data consisting of images and messages 186 | :param batch: batch of validation data, in form [images, messages] 187 | :return: dictionary of error metrics from Encoder, Decoder, and Discriminator on the current batch 188 | """ 189 | # if TensorboardX logging is enabled, save some of the tensors. 190 | if self.tb_logger is not None: 191 | encoder_final = self.encoder_decoder.encoder._modules['final_layer'] 192 | self.tb_logger.add_tensor('weights/encoder_out', encoder_final.weight) 193 | decoder_final = self.encoder_decoder.decoder._modules['linear'] 194 | self.tb_logger.add_tensor('weights/decoder_out', decoder_final.weight) 195 | discrim_final = self.discriminator._modules['linear'] 196 | self.tb_logger.add_tensor('weights/discrim_out', discrim_final.weight) 197 | 198 | images, messages = batch 199 | 200 | batch_size = images.shape[0] 201 | 202 | self.encoder_decoder.eval() 203 | self.discriminator.eval() 204 | with torch.no_grad(): 205 | d_target_label_cover = torch.full((batch_size, 1), self.cover_label, device=self.device) 206 | d_target_label_encoded = torch.full((batch_size, 1), self.encoded_label, device=self.device) 207 | g_target_label_encoded = torch.full((batch_size, 1), self.cover_label, device=self.device) 208 | 209 | d_on_cover = self.discriminator(images) 210 | d_loss_on_cover = self.bce_with_logits_loss(d_on_cover, d_target_label_cover) 211 | 212 | encoded_images, noised_images, decoded_messages = self.encoder_decoder.forward_specific_noiser(images, messages, noiser=noise) 213 | 214 | d_on_encoded = self.discriminator(encoded_images) 215 | d_loss_on_encoded = self.bce_with_logits_loss(d_on_encoded, d_target_label_encoded) 216 | 217 | d_on_encoded_for_enc = self.discriminator(encoded_images) 218 | g_loss_adv = self.bce_with_logits_loss(d_on_encoded_for_enc, g_target_label_encoded) 219 | 220 | if self.vgg_loss is None: 221 | g_loss_enc = self.mse_loss(encoded_images, images) 222 | else: 223 | vgg_on_cov = self.vgg_loss(images) 224 | vgg_on_enc = self.vgg_loss(encoded_images) 225 | g_loss_enc = self.mse_loss(vgg_on_cov, vgg_on_enc) 226 | 227 | g_loss_dec = self.mse_loss(decoded_messages, messages) 228 | g_loss = self.config.adversarial_loss * g_loss_adv + self.config.encoder_loss * g_loss_enc \ 229 | + self.config.decoder_loss * g_loss_dec 230 | 231 | decoded_rounded = decoded_messages.detach().cpu().numpy().round().clip(0, 1) 232 | bitwise_avg_err = np.sum(np.abs(decoded_rounded - messages.detach().cpu().numpy())) / ( 233 | batch_size * messages.shape[1]) 234 | 235 | losses = { 236 | 'loss ': g_loss.item(), 237 | 'encoder_mse ': g_loss_enc.item(), 238 | 'dec_mse ': g_loss_dec.item(), 239 | 'bitwise-error ': bitwise_avg_err, 240 | 'adversarial_bce': g_loss_adv.item(), 241 | 'discr_cover_bce': d_loss_on_cover.item(), 242 | 'discr_encod_bce': d_loss_on_encoded.item() 243 | } 244 | return losses, (encoded_images, noised_images, decoded_messages) 245 | 246 | def to_stirng(self): 247 | return '{}\n{}'.format(str(self.encoder_decoder), str(self.discriminator)) 248 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import torch 4 | import random 5 | import numpy as np 6 | import utils 7 | import logging 8 | from collections import defaultdict 9 | 10 | from options import * 11 | from model.hidden import Hidden 12 | from average_meter import AverageMeter 13 | 14 | 15 | def train(model: Hidden, 16 | device: torch.device, 17 | hidden_config: HiDDenConfiguration, 18 | train_options: TrainingOptions, 19 | this_run_folder: str, 20 | tb_logger): 21 | """ 22 | Trains the HiDDeN model 23 | :param model: The model 24 | :param device: torch.device object, usually this is GPU (if avaliable), otherwise CPU. 25 | :param hidden_config: The network configuration 26 | :param train_options: The training settings 27 | :param this_run_folder: The parent folder for the current training run to store training artifacts/results/logs. 28 | :param tb_logger: TensorBoardLogger object which is a thin wrapper for TensorboardX logger. 29 | Pass None to disable TensorboardX logging 30 | :return: 31 | """ 32 | 33 | train_data, val_data = utils.get_data_loaders(hidden_config, train_options) 34 | file_count = len(train_data.dataset) 35 | if file_count % train_options.batch_size == 0: 36 | steps_in_epoch = file_count // train_options.batch_size 37 | else: 38 | steps_in_epoch = file_count // train_options.batch_size + 1 39 | steps_in_epoch = 313 40 | 41 | print_each = 10 42 | images_to_save = 8 43 | saved_images_size = (512, 512) 44 | 45 | for epoch in range(train_options.start_epoch, train_options.number_of_epochs + 1): 46 | logging.info('\nStarting epoch {}/{}'.format(epoch, train_options.number_of_epochs)) 47 | logging.info('Batch size = {}\nSteps in epoch = {}'.format(train_options.batch_size, steps_in_epoch)) 48 | training_losses = defaultdict(AverageMeter) 49 | epoch_start = time.time() 50 | step = 1 51 | for image, _ in train_data: 52 | image = image.to(device) 53 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 54 | losses, _ = model.train_on_batch([image, message]) 55 | 56 | for name, loss in losses.items(): 57 | training_losses[name].update(loss) 58 | if step % print_each == 0 or step == steps_in_epoch: 59 | #import pdb; pdb.set_trace() 60 | logging.info( 61 | 'Epoch: {}/{} Step: {}/{}'.format(epoch, train_options.number_of_epochs, step, steps_in_epoch)) 62 | utils.log_progress(training_losses) 63 | logging.info('-' * 40) 64 | step += 1 65 | if step == steps_in_epoch: 66 | break 67 | 68 | train_duration = time.time() - epoch_start 69 | logging.info('Epoch {} training duration {:.2f} sec'.format(epoch, train_duration)) 70 | logging.info('-' * 40) 71 | utils.write_losses(os.path.join(this_run_folder, 'train.csv'), training_losses, epoch, train_duration) 72 | if tb_logger is not None: 73 | tb_logger.save_losses(training_losses, epoch) 74 | tb_logger.save_grads(epoch) 75 | tb_logger.save_tensors(epoch) 76 | 77 | first_iteration = True 78 | validation_losses = defaultdict(AverageMeter) 79 | logging.info('Running validation for epoch {}/{}'.format(epoch, train_options.number_of_epochs)) 80 | step = 1 81 | for image, _ in val_data: 82 | image = image.to(device) 83 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 84 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch([image, message]) 85 | for name, loss in losses.items(): 86 | validation_losses[name].update(loss) 87 | if first_iteration: 88 | if hidden_config.enable_fp16: 89 | image = image.float() 90 | encoded_images = encoded_images.float() 91 | utils.save_images(image.cpu()[:images_to_save, :, :, :], 92 | encoded_images[:images_to_save, :, :, :].cpu(), 93 | epoch, 94 | os.path.join(this_run_folder, 'images'), resize_to=saved_images_size) 95 | first_iteration = False 96 | step += 1 97 | if step == steps_in_epoch // 10: 98 | break 99 | 100 | utils.log_progress(validation_losses) 101 | logging.info('-' * 40) 102 | utils.save_checkpoint(model, train_options.experiment_name, epoch, os.path.join(this_run_folder, 'checkpoints')) 103 | utils.write_losses(os.path.join(this_run_folder, 'validation.csv'), validation_losses, epoch, 104 | time.time() - epoch_start) 105 | 106 | def train_other_noises(model: Hidden, 107 | device: torch.device, 108 | hidden_config: HiDDenConfiguration, 109 | train_options: TrainingOptions, 110 | this_run_folder: str, 111 | tb_logger): 112 | """ 113 | Trains the HiDDeN model 114 | :param model: The model 115 | :param device: torch.device object, usually this is GPU (if avaliable), otherwise CPU. 116 | :param hidden_config: The network configuration 117 | :param train_options: The training settings 118 | :param this_run_folder: The parent folder for the current training run to store training artifacts/results/logs. 119 | :param tb_logger: TensorBoardLogger object which is a thin wrapper for TensorboardX logger. 120 | Pass None to disable TensorboardX logging 121 | :return: 122 | """ 123 | 124 | train_data, val_data = utils.get_data_loaders(hidden_config, train_options) 125 | file_count = len(train_data.dataset) 126 | if file_count % train_options.batch_size == 0: 127 | steps_in_epoch = file_count // train_options.batch_size 128 | else: 129 | steps_in_epoch = file_count // train_options.batch_size + 1 130 | steps_in_epoch = 313 131 | 132 | print_each = 10 133 | images_to_save = 8 134 | saved_images_size = (512, 512) 135 | 136 | for epoch in range(train_options.start_epoch, train_options.number_of_epochs + 1): 137 | logging.info('\nStarting epoch {}/{}'.format(epoch, train_options.number_of_epochs)) 138 | logging.info('Batch size = {}\nSteps in epoch = {}'.format(train_options.batch_size, steps_in_epoch)) 139 | training_losses = defaultdict(AverageMeter) 140 | epoch_start = time.time() 141 | step = 1 142 | for image, _ in train_data: 143 | image = image.to(device) 144 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 145 | losses, _ = model.train_on_batch([image, message]) 146 | 147 | for name, loss in losses.items(): 148 | training_losses[name].update(loss) 149 | if step % print_each == 0 or step == steps_in_epoch: 150 | #import pdb; pdb.set_trace() 151 | logging.info( 152 | 'Epoch: {}/{} Step: {}/{}'.format(epoch, train_options.number_of_epochs, step, steps_in_epoch)) 153 | utils.log_progress(training_losses) 154 | logging.info('-' * 40) 155 | step += 1 156 | if step == steps_in_epoch: 157 | break 158 | 159 | train_duration = time.time() - epoch_start 160 | logging.info('Epoch {} training duration {:.2f} sec'.format(epoch, train_duration)) 161 | logging.info('-' * 40) 162 | utils.write_losses(os.path.join(this_run_folder, 'train.csv'), training_losses, epoch, train_duration) 163 | if tb_logger is not None: 164 | tb_logger.save_losses(training_losses, epoch) 165 | tb_logger.save_grads(epoch) 166 | tb_logger.save_tensors(epoch) 167 | 168 | noises = ['identity', 'dropout', 'cropout', 'crop', 'gaussian', 'jpeg'] 169 | for noise in noises: 170 | first_iteration = True 171 | validation_losses = defaultdict(AverageMeter) 172 | logging.info('Running validation for epoch {}/{} for noise {}'.format(epoch, train_options.number_of_epochs, noise)) 173 | step = 1 174 | for image, _ in val_data: 175 | image = image.to(device) 176 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 177 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch_specific_noise([image, message], noise=noise) 178 | for name, loss in losses.items(): 179 | validation_losses[name].update(loss) 180 | if first_iteration: 181 | if hidden_config.enable_fp16: 182 | image = image.float() 183 | encoded_images = encoded_images.float() 184 | utils.save_images(image.cpu()[:images_to_save, :, :, :], 185 | encoded_images[:images_to_save, :, :, :].cpu(), 186 | epoch, 187 | os.path.join(this_run_folder, 'images'), resize_to=saved_images_size) 188 | first_iteration = False 189 | step += 1 190 | if step == steps_in_epoch // 10: 191 | break 192 | 193 | utils.log_progress(validation_losses) 194 | logging.info('-' * 40) 195 | utils.save_checkpoint(model, train_options.experiment_name, epoch, os.path.join(this_run_folder, 'checkpoints')) 196 | utils.write_losses(os.path.join(this_run_folder, 'validation_' + noise + '.csv'), validation_losses, epoch, 197 | time.time() - epoch_start) 198 | 199 | def train_own_noise(model: Hidden, 200 | device: torch.device, 201 | hidden_config: HiDDenConfiguration, 202 | train_options: TrainingOptions, 203 | this_run_folder: str, 204 | tb_logger, 205 | noise): 206 | """ 207 | Trains the HiDDeN model 208 | :param model: The model 209 | :param device: torch.device object, usually this is GPU (if avaliable), otherwise CPU. 210 | :param hidden_config: The network configuration 211 | :param train_options: The training settings 212 | :param this_run_folder: The parent folder for the current training run to store training artifacts/results/logs. 213 | :param tb_logger: TensorBoardLogger object which is a thin wrapper for TensorboardX logger. 214 | Pass None to disable TensorboardX logging 215 | :return: 216 | """ 217 | 218 | train_data, val_data = utils.get_data_loaders(hidden_config, train_options) 219 | file_count = len(train_data.dataset) 220 | if file_count % train_options.batch_size == 0: 221 | steps_in_epoch = file_count // train_options.batch_size 222 | else: 223 | steps_in_epoch = file_count // train_options.batch_size + 1 224 | steps_in_epoch = 313 225 | 226 | print_each = 10 227 | images_to_save = 8 228 | saved_images_size = (512, 512) # for qualitative check purpose to use a larger size 229 | 230 | for epoch in range(train_options.start_epoch, train_options.number_of_epochs + 1): 231 | logging.info('\nStarting epoch {}/{}'.format(epoch, train_options.number_of_epochs)) 232 | logging.info('Batch size = {}\nSteps in epoch = {}'.format(train_options.batch_size, steps_in_epoch)) 233 | training_losses = defaultdict(AverageMeter) 234 | 235 | if train_options.video_dataset: 236 | random.shuffle(train_data.dataset) 237 | 238 | epoch_start = time.time() 239 | step = 1 240 | for image, _ in train_data: 241 | image = image.to(device) 242 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 243 | losses, _ = model.train_on_batch([image, message]) 244 | 245 | for name, loss in losses.items(): 246 | training_losses[name].update(loss) 247 | if step % print_each == 0 or step == steps_in_epoch: 248 | #import pdb; pdb.set_trace() 249 | logging.info( 250 | 'Epoch: {}/{} Step: {}/{}'.format(epoch, train_options.number_of_epochs, step, steps_in_epoch)) 251 | utils.log_progress(training_losses) 252 | logging.info('-' * 40) 253 | step += 1 254 | if step == steps_in_epoch: 255 | break 256 | 257 | train_duration = time.time() - epoch_start 258 | logging.info('Epoch {} training duration {:.2f} sec'.format(epoch, train_duration)) 259 | logging.info('-' * 40) 260 | utils.write_losses(os.path.join(this_run_folder, 'train.csv'), training_losses, epoch, train_duration) 261 | if tb_logger is not None: 262 | tb_logger.save_losses(training_losses, epoch) 263 | tb_logger.save_grads(epoch) 264 | tb_logger.save_tensors(epoch) 265 | 266 | first_iteration = True 267 | validation_losses = defaultdict(AverageMeter) 268 | logging.info('Running validation for epoch {}/{} for noise {}'.format(epoch, train_options.number_of_epochs, noise)) 269 | step = 1 270 | for image, _ in val_data: 271 | image = image.to(device) 272 | message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device) 273 | losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch_specific_noise([image, message], noise=noise) 274 | for name, loss in losses.items(): 275 | validation_losses[name].update(loss) 276 | if first_iteration: 277 | if hidden_config.enable_fp16: 278 | image = image.float() 279 | encoded_images = encoded_images.float() 280 | utils.save_images(image.cpu()[:images_to_save, :, :, :], 281 | encoded_images[:images_to_save, :, :, :].cpu(), 282 | epoch, 283 | os.path.join(this_run_folder, 'images'), resize_to=saved_images_size) 284 | first_iteration = False 285 | step += 1 286 | if step == steps_in_epoch // 10: 287 | break 288 | 289 | utils.log_progress(validation_losses) 290 | logging.info('-' * 40) 291 | utils.save_checkpoint(model, train_options.experiment_name, epoch, os.path.join(this_run_folder, 'checkpoints')) 292 | utils.write_losses(os.path.join(this_run_folder, 'validation_' + noise + '.csv'), validation_losses, epoch, 293 | time.time() - epoch_start) --------------------------------------------------------------------------------