├── .DS_Store ├── README.md ├── __pycache__ ├── main.cpython-310.pyc └── schemas.cpython-310.pyc ├── assets ├── 25_0_combined_A.png ├── 392_1_combined_B.png ├── 58_2_combined_A.png ├── 62_1_combined_A.png ├── 62_1_combined_B.png ├── 91_1_combined_B.png ├── 96_1_combined_A.png ├── blue1.png ├── ddf_train_loss_images.png ├── ddf_val_loss_acc.png ├── ezgif.com-animated-gif-maker.gif ├── loss.png ├── lr.png ├── output_e18000.png ├── pipeline.png ├── swapping.png └── upload.png ├── backend ├── .gitignore ├── __init__.py ├── ddf │ ├── DiffJPEG_pytorch │ │ ├── DiffJPEG.py │ │ ├── README.md │ │ ├── compression.py │ │ ├── decompression.py │ │ ├── modules │ │ │ ├── __init__.py │ │ │ ├── compression.py │ │ │ └── decompression.py │ │ └── utils.py │ ├── DualDefense_gan_fs.py │ ├── __init__.py │ ├── adv.py │ ├── assets │ │ ├── 25_0_combined_A.png │ │ ├── 58_2_combined_A.png │ │ ├── byeon_total5.jpg │ │ ├── cha_total2.jpg │ │ ├── output_A2B.png │ │ ├── output_B2A.png │ │ └── output_e18000.png │ ├── config.yaml │ ├── config2.yaml │ ├── data │ │ └── haarcascade_frontalface_alt.xml │ ├── dataset.py │ ├── decoder_fs.py │ ├── encoder_fs.py │ ├── engine │ │ ├── dwt.py │ │ └── utils.py │ ├── facenet_pytorch │ │ └── models │ │ │ ├── inception_resnet_v1.py │ │ │ ├── mtcnn.py │ │ │ └── utils │ │ │ ├── detect_face.py │ │ │ ├── download.py │ │ │ ├── tensorflow2pytorch.py │ │ │ └── training.py │ ├── faceswap_pytorch │ │ ├── __init__.py │ │ ├── image_augmentation.py │ │ ├── models.py │ │ ├── padding_same_conv.py │ │ ├── train.py │ │ ├── training_data.py │ │ ├── umeyama.py │ │ └── util.py │ ├── loss.py │ ├── main.py │ ├── scripts │ │ ├── run_test.sh │ │ └── run_train.sh │ ├── test.py │ ├── train.py │ └── upload.js ├── main.py ├── posts │ └── posts.json ├── requirements.txt └── schemas.py ├── frontend ├── .gitignore ├── README.md ├── components.json ├── eslint.config.mjs ├── lib │ └── utils.ts ├── next.config.ts ├── package-lock.json ├── package.json ├── postcss.config.mjs ├── public │ ├── byeon │ │ ├── byeon_1.png │ │ ├── byeon_2.png │ │ ├── byeon_3.png │ │ ├── byeon_4.png │ │ └── byeon_5.png │ ├── cha │ │ ├── chaeunwoo2.jpg │ │ ├── chaeunwoo3.png │ │ ├── chaeunwoo4.png │ │ └── chaeunwoo5.png │ ├── chu │ │ ├── chuu1.png │ │ ├── chuu2.png │ │ ├── chuu3.png │ │ ├── chuu4.png │ │ ├── chuu5.png │ │ ├── chuu6.png │ │ ├── chuu7.png │ │ ├── chuu8.png │ │ └── chuu9.png │ ├── file.svg │ ├── globe.svg │ ├── next.svg │ ├── panel.png │ ├── vercel.svg │ ├── win │ │ ├── winter1.png │ │ ├── winter2.png │ │ └── winter3.png │ └── window.svg ├── src │ ├── components │ │ ├── Info.js │ │ └── Navbar.js │ ├── pages │ │ ├── _app.js │ │ ├── _document.js │ │ ├── api │ │ │ └── images.js │ │ ├── attack │ │ │ └── index.js │ │ ├── index.js │ │ ├── panel.js │ │ └── upload.js │ └── styles │ │ └── globals.css ├── tailwind.config.js ├── tailwind.config.ts └── tsconfig.json ├── package-lock.json ├── package.json └── requirements.txt /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/.DS_Store -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # deepsafe 2 | ## Introduction 3 | The misuse of deepfake technology poses serious security threats, such as identity theft. DeepSafe is a defense system that employs an Adversarial Robust Watermarking technique to disrupt face-swapping models while enabling image source tracking. Based on the [Dual Defense framework](https://ieeexplore.ieee.org/document/10486948), this system utilizes the Original-domain Feature Emulation Attack (OFEA) method to embed an invisible watermark, preventing face swapping while ensuring identity traceability. 4 | 5 |
6 | gif 7 |

Original Image / Encoded / Protected 🔒/ Attacked 🔓

8 |
9 | 10 | *DeepSafe offers a two-layered defense mechanism against deepfake attacks:* 11 | ### 1️⃣ Preemptive Deepfake Prevention ✅ 12 | - When users upload their images, DeepSafe embeds an invisible adversarial watermark that disrupts face-swapping models. 13 | - The watermarked image is uploaded to the platform, ensuring protection. 14 | - If an attacker attempts face-swapping, the image degrades or fails to swap properly due to the embedded watermark. 15 | 16 | ### 2️⃣ Post-attack Detection & Identity Tracking ⚠️ 17 | - Even if the face-swapping model processes the image, DeepSafe preserves watermark information. 18 | - The watermark can be extracted from manipulated images, allowing for identity tracking and alerting affected users 19 | 20 | 21 | #### Our Pipeline looks following: 22 |

23 | 24 |

25 | 26 | ## ⚒️ How to Play 27 | Install all the requirements in `requirements.txt` with created conda env. 28 | ``` 29 | cd deepsafe/backend 30 | uvicorn main:app --reload 31 | 32 | cd deepsafe/frontend 33 | npm run dev 34 | ``` 35 | ### 1. Move to home page 🏠 36 | ![blue](./assets/blue1.png) 37 | ### 2. Upload your post ✍️ 38 | Watermark will be encoded into the raw image. 39 | 40 | ![upload](./assets/upload.png) 41 | ### 3. Try Face Swapping 🚨 42 | If the generated image is damaged, it means the image is protected well. 43 | 44 | ![swap](./assets/swapping.png) 45 | 46 | 47 | ## 📂 Dataset & Preprocessing 48 | - 4 identities: Winter, Chuu, Cha Eun-woo, and Byeon Woo-seok (800 images per person, total: 4,000 images)- Extracted frames from YouTube videos (Winter: 35, Chuu: 48, Cha Eun-woo: 31, Byeon Woo-seok: 48 videos) 49 | - Preprocessing: 50 | - Extracted frames every 20 frames using OpenCV 51 | - Face cropping: resized to 256x256 52 | - Final resizing: 160x160 for model training 53 | 54 | ## 💫 Model Training & Results 55 | ### 1) Faceswap 56 | - Training Epochs: 47,000 57 | - Loss Functions: 58 | - L1 Loss (weight = 1.0) 59 | - VGG Perceptual Loss (weight = 0.08) 60 | - Optimizer: Adam (lr=5e-5) with MultiStepLR scheduler 61 | - Baseline Code: [Faceswap Deepfake Pytorch](https://github.com/Oldpan/Faceswap-Deepfake-Pytorch) 62 | 63 | 64 |

65 | 66 | 67 |

68 | 69 |

70 | 71 |

72 | 73 | ### 2) Dual Defense 74 | - loss 75 | - alpha * `image_loss` + beta * `message_loss` 76 | - `image_loss` = lambda_en * image_encoded_loss + lambda_s * ssim_loss + lambda_adv * image_adv_logits_loss 77 | - `message_loss` = L_wm_en + L_wm_adv 78 | - lambda(beta)=1, alpha=2 79 | - Train:Valid:Test=0.8:0.1:0.1 80 | - message size = 4 (4bits, watermark information) 81 | - batch size = 16 82 | 83 | 84 |

85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 |

Original Image → Encoded ✅ → Protected 😀🔒 → Attacked 🥵🔓

105 |

106 | 107 | ## Directory Structure 108 | - `frontend` : frontend codes (stack; next.js) 109 | - `backend/ddf` : Dual defense code aligning with backend 110 | - `backend/` : backend codes (stack; FastAPI) 111 | - Note: You can see full training codes of faceswap in `feature/faceswap` branch. 112 | 113 | ## References 114 | - Dual Defense official codes: https://github.com/Xming-zzz/DualDefense 115 | 116 | 117 | -------------------------------------------------------------------------------- /__pycache__/main.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/__pycache__/main.cpython-310.pyc -------------------------------------------------------------------------------- /__pycache__/schemas.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/__pycache__/schemas.cpython-310.pyc -------------------------------------------------------------------------------- /assets/25_0_combined_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/25_0_combined_A.png -------------------------------------------------------------------------------- /assets/392_1_combined_B.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/392_1_combined_B.png -------------------------------------------------------------------------------- /assets/58_2_combined_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/58_2_combined_A.png -------------------------------------------------------------------------------- /assets/62_1_combined_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/62_1_combined_A.png -------------------------------------------------------------------------------- /assets/62_1_combined_B.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/62_1_combined_B.png -------------------------------------------------------------------------------- /assets/91_1_combined_B.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/91_1_combined_B.png -------------------------------------------------------------------------------- /assets/96_1_combined_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/96_1_combined_A.png -------------------------------------------------------------------------------- /assets/blue1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/blue1.png -------------------------------------------------------------------------------- /assets/ddf_train_loss_images.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/ddf_train_loss_images.png -------------------------------------------------------------------------------- /assets/ddf_val_loss_acc.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/ddf_val_loss_acc.png -------------------------------------------------------------------------------- /assets/ezgif.com-animated-gif-maker.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/ezgif.com-animated-gif-maker.gif -------------------------------------------------------------------------------- /assets/loss.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/loss.png -------------------------------------------------------------------------------- /assets/lr.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/lr.png -------------------------------------------------------------------------------- /assets/output_e18000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/output_e18000.png -------------------------------------------------------------------------------- /assets/pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/pipeline.png -------------------------------------------------------------------------------- /assets/swapping.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/swapping.png -------------------------------------------------------------------------------- /assets/upload.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/assets/upload.png -------------------------------------------------------------------------------- /backend/.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | .ipynb_checkpoints 3 | .mypy_cache 4 | .vscode 5 | __pycache__ 6 | .pytest_cache 7 | htmlcov 8 | dist 9 | site 10 | .coverage* 11 | coverage.xml 12 | .netlify 13 | test.db 14 | log.txt 15 | Pipfile.lock 16 | env3.* 17 | env 18 | docs_build 19 | site_build 20 | venv 21 | docs.zip 22 | archive.zip 23 | 24 | # vim temporary files 25 | *~ 26 | .*.sw? 27 | .cache 28 | 29 | # macOS 30 | .DS_Store -------------------------------------------------------------------------------- /backend/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/__init__.py -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/DiffJPEG.py: -------------------------------------------------------------------------------- 1 | # Pytorch 2 | import torch 3 | import torch.nn as nn 4 | # Local 5 | from DiffJPEG_pytorch.compression import compress_jpeg 6 | from DiffJPEG_pytorch.decompression import decompress_jpeg 7 | from DiffJPEG_pytorch.utils import diff_round, quality_to_factor 8 | 9 | 10 | class DiffJPEG(nn.Module): 11 | def __init__(self, height, width, differentiable=True, quality=80): 12 | ''' Initialize the DiffJPEG layer 13 | Inputs: 14 | height(int): Original image hieght 15 | width(int): Original image width 16 | differentiable(bool): If true uses custom differentiable 17 | rounding function, if false uses standrard torch.round 18 | quality(float): Quality factor for jpeg compression scheme. 19 | ''' 20 | super(DiffJPEG, self).__init__() 21 | if differentiable: 22 | rounding = diff_round 23 | else: 24 | rounding = torch.round 25 | self.rounding = rounding 26 | self.height = height 27 | self.width = width 28 | self.factor = quality_to_factor(quality) 29 | # self.compress = compress_jpeg(rounding=rounding, factor=factor) 30 | # self.decompress = decompress_jpeg(height, width, rounding=rounding,factor=factor) 31 | 32 | def forward(self, x): 33 | y, cb, cr = compress_jpeg(x, self.rounding, self.factor) 34 | recovered = decompress_jpeg(y, cb, cr, self.height, self.width, self.factor) 35 | # y, cb, cr = self.compress(x) 36 | # recovered = self.decompress(y, cb, cr) 37 | return recovered -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/README.md: -------------------------------------------------------------------------------- 1 | # DiffJPEG: A PyTorch implementation 2 | 3 | This is a pytorch implementation of differentiable jpeg compression algorithm. This work is based on the discussion in this [paper](https://machine-learning-and-security.github.io/papers/mlsec17_paper_54.pdf). The work relies heavily on the tensorflow implementation in this [repository](https://github.com/rshin/differentiable-jpeg) 4 | 5 | ## Requirements 6 | - Pytorch 1.0.0 7 | - numpy 1.15.4 8 | 9 | ## Use 10 | 11 | DiffJPEG functions as a standard pytorch module/layer. To use, first import the layer and then initialize with the desired parameters: 12 | - differentaible(bool): If true uses custom differentiable rounding function, if false uses standrard torch.round 13 | - quality(float): Quality factor for jpeg compression scheme. 14 | 15 | ``` python 16 | from DiffJPEG import DiffJPEG 17 | jpeg = DiffJPEG(hieght=224, width=224, differentiable=True, quality=80) 18 | ``` 19 | -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/compression.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Note to user: This file, while functional, is not fully differentiable in 4 | PyTorch and is not easily moved to and from the gpu. For updated version use 5 | the source copde in modules and updated DiffJPEG module. 6 | 7 | """ 8 | # Standard libraries 9 | import itertools 10 | import numpy as np 11 | # PyTorch 12 | import torch 13 | import torch.nn as nn 14 | # Local 15 | import DiffJPEG_pytorch.utils as utils 16 | 17 | 18 | def rgb_to_ycbcr(image): 19 | """ Converts RGB image to YCbCr 20 | Input: 21 | image(tensor): batch x 3 x height x width 22 | Outpput: 23 | result(tensor): batch x height x width x 3 24 | """ 25 | matrix = np.array( 26 | [[65.481, 128.553, 24.966], [-37.797, -74.203, 112.], 27 | [112., -93.786, -18.214]], 28 | dtype=np.float32).T / 255 29 | shift = [16., 128., 128.] 30 | image = image 31 | image = image.permute(0, 2, 3, 1) 32 | result = torch.tensordot(image, torch.from_numpy(matrix), dims=1) + shift 33 | # result = torch.from_numpy(result) 34 | result.view(image.shape) 35 | return result 36 | 37 | 38 | def rgb_to_ycbcr_jpeg(image): 39 | """ Converts RGB image to YCbCr 40 | Input: 41 | image(tensor): batch x 3 x height x width 42 | Outpput: 43 | result(tensor): batch x height x width x 3 44 | """ 45 | matrix = np.array( 46 | [[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], 47 | [0.5, -0.418688, -0.081312]], 48 | dtype=np.float32).T 49 | shift =torch.tensor( [0., 128., 128.]) 50 | image = image.permute(0, 2, 3, 1) 51 | device = image.device 52 | result = torch.tensordot(image, torch.from_numpy(matrix).to(device), dims=1) + shift.to(device) 53 | # result = torch.from_numpy(result) 54 | result.view(image.shape) 55 | return result 56 | 57 | 58 | def chroma_subsampling(image): 59 | """ Chroma subsampling on CbCv channels 60 | Input: 61 | image(tensor): batch x height x width x 3 62 | Output: 63 | y(tensor): batch x height x width 64 | cb(tensor): batch x height/2 x width/2 65 | cr(tensor): batch x height/2 x width/2 66 | """ 67 | image_2 = image.permute(0, 3, 1, 2).clone() 68 | avg_pool = nn.AvgPool2d(kernel_size=2, stride=(2, 2), 69 | count_include_pad=False) 70 | cb = avg_pool(image_2[:, 1, :, :].unsqueeze(1)) 71 | cr = avg_pool(image_2[:, 2, :, :].unsqueeze(1)) 72 | cb = cb.permute(0, 2, 3, 1) 73 | cr = cr.permute(0, 2, 3, 1) 74 | return image[:, :, :, 0], cb.squeeze(3), cr.squeeze(3) 75 | 76 | 77 | def block_splitting(image): 78 | """ Splitting image into patches 79 | Input: 80 | image(tensor): batch x height x width 81 | Output: 82 | patch(tensor): batch x h*w/64 x h x w 83 | """ 84 | k = 8 85 | height, width = image.shape[1:3] 86 | batch_size = image.shape[0] 87 | image_reshaped = image.view(batch_size, height // k, k, -1, k) 88 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) 89 | return image_transposed.contiguous().view(batch_size, -1, k, k) 90 | 91 | 92 | def dct_8x8_ref(image): 93 | """ Reference Discrete Cosine Transformation 94 | Input: 95 | image(tensor): batch x height x width 96 | Output: 97 | dcp(tensor): batch x height x width 98 | """ 99 | image = image - 128 100 | result = np.zeros((8, 8), dtype=np.float32) 101 | for u, v in itertools.product(range(8), range(8)): 102 | value = 0 103 | for x, y in itertools.product(range(8), range(8)): 104 | value += image[x, y] * np.cos((2 * x + 1) * u * 105 | np.pi / 16) * np.cos((2 * y + 1) * v * np.pi / 16) 106 | result[u, v] = value 107 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 108 | scale = np.outer(alpha, alpha) * 0.25 109 | return result * scale 110 | 111 | def dct_8x8(image): 112 | """ Discrete Cosine Transformation 113 | Input: 114 | image(tensor): batch x height x width (device에 위치) 115 | Output: 116 | dcp(tensor): batch x height x width 117 | """ 118 | image = image - 128 119 | tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) 120 | for x, y, u, v in itertools.product(range(8), repeat=4): 121 | tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos( 122 | (2 * y + 1) * v * np.pi / 16) 123 | tensor = torch.tensor(tensor, dtype=torch.float32, device=image.device) # GPU로 이동 124 | alpha = torch.tensor([1. / np.sqrt(2)] + [1] * 7, dtype=torch.float32, device=image.device) 125 | scale = torch.outer(alpha, alpha) * 0.25 # NumPy 대신 PyTorch의 `torch.outer` 사용 126 | result = scale * torch.tensordot(image, tensor, dims=2) 127 | return result.view(image.shape) 128 | 129 | # def dct_8x8(image): 130 | # """ Discrete Cosine Transformation 131 | # Input: 132 | # image(tensor): batch x height x width 133 | # Output: 134 | # dcp(tensor): batch x height x width 135 | # """ 136 | 137 | # image = image - 128 138 | # tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) 139 | # print(f">>DEBUG) {image.device}") 140 | # for x, y, u, v in itertools.product(range(8), repeat=4): 141 | # tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos( 142 | # (2 * y + 1) * v * np.pi / 16) 143 | # alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 144 | # scale = np.outer(alpha, alpha) * 0.25 145 | 146 | # result = scale * torch.tensordot(image, tensor, dims=2) 147 | # #result = torch.from_numpy(result) 148 | # result.view(image.shape) 149 | # return result 150 | 151 | 152 | def y_quantize(image, rounding, factor=1): 153 | """ JPEG Quantization for Y channel 154 | Input: 155 | image(tensor): batch x height x width 156 | rounding(function): rounding function to use 157 | factor(float): Degree of compression 158 | Output: 159 | image(tensor): batch x height x width 160 | """ 161 | image = image.float() / (utils.y_table.to(image.device) * factor) 162 | image = rounding(image) 163 | return image 164 | 165 | 166 | def c_quantize(image, rounding, factor=1): 167 | """ JPEG Quantization for CrCb channels 168 | Input: 169 | image(tensor): batch x height x width 170 | rounding(function): rounding function to use 171 | factor(float): Degree of compression 172 | Output: 173 | image(tensor): batch x height x width 174 | """ 175 | image = image.float() / (utils.c_table.to(image.device) * factor) 176 | image = rounding(image) 177 | return image 178 | 179 | 180 | def compress_jpeg(imgs, rounding=torch.round, factor=1, device=None): 181 | """ Full JPEG compression algortihm 182 | Input: 183 | imgs(tensor): batch x 3 x height x width 184 | rounding(function): rounding function to use 185 | factor(float): Compression factor 186 | Ouput: 187 | compressed(dict(tensor)): batch x h*w/64 x 8 x 8 188 | """ 189 | temp = rgb_to_ycbcr_jpeg(imgs*255) 190 | y, cb, cr = chroma_subsampling(temp) 191 | components = {'y': y, 'cb': cb, 'cr': cr} 192 | for k in components.keys(): 193 | comp = block_splitting(components[k]) 194 | comp = dct_8x8(comp) 195 | comp = c_quantize(comp, rounding, factor=factor) if k in ( 196 | 'cb', 'cr') else y_quantize(comp, rounding, factor=factor) 197 | 198 | components[k] = comp 199 | return components['y'], components['cb'], components['cr'] 200 | -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/decompression.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Note to user: This file, while functional, is not fully differentiable in 4 | PyTorch and is not easily moved to and from the gpu. For updated version use 5 | the source copde in modules and updated DiffJPEG module. 6 | 7 | """ 8 | # Standard libraries 9 | import itertools 10 | import numpy as np 11 | # PyTorch 12 | import torch 13 | # Local 14 | import DiffJPEG_pytorch.utils as utils 15 | 16 | 17 | def y_dequantize(image, factor=1): 18 | """ Dequantize Y channel 19 | Inputs: 20 | image(tensor): batch x height x width 21 | factor(float): compression factor 22 | Outputs: 23 | image(tensor): batch x height x width 24 | 25 | """ 26 | return image * (utils.y_table.to(image.device) * factor) 27 | 28 | 29 | def c_dequantize(image, factor=1): 30 | """ Dequantize CbCr channel 31 | Inputs: 32 | image(tensor): batch x height x width 33 | factor(float): compression factor 34 | Outputs: 35 | image(tensor): batch x height x width 36 | 37 | """ 38 | return image * (utils.c_table.to(image.device) * factor) 39 | 40 | 41 | def idct_8x8_ref(image): 42 | """ Reference Inverse Discrete Cosine Transformation 43 | Input: 44 | dcp(tensor): batch x height x width 45 | Output: 46 | image(tensor): batch x height x width 47 | """ 48 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 49 | alpha = np.outer(alpha, alpha) 50 | image = image * alpha 51 | 52 | result = np.zeros((8, 8), dtype=np.float32) 53 | for u, v in itertools.product(range(8), range(8)): 54 | value = 0 55 | for x, y in itertools.product(range(8), range(8)): 56 | value += image[x, y] * np.cos((2 * u + 1) * x * np.pi / 16) * np.cos( 57 | (2 * v + 1) * y * np.pi / 16) 58 | result[u, v] = value 59 | return result * 0.25 + 128 60 | 61 | 62 | def idct_8x8(image): 63 | """ Inverse discrete Cosine Transformation 64 | Input: 65 | dcp(tensor): batch x height x width 66 | Output: 67 | image(tensor): batch x height x width 68 | """ 69 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 70 | alpha = torch.tensor(alpha, dtype=torch.float32, device=image.device) 71 | 72 | alpha = torch.outer(alpha, alpha) 73 | image = image * alpha 74 | 75 | tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) 76 | 77 | for x, y, u, v in itertools.product(range(8), repeat=4): 78 | tensor[x, y, u, v] = np.cos((2 * u + 1) * x * np.pi / 16) * np.cos( 79 | (2 * v + 1) * y * np.pi / 16) 80 | tensor = torch.tensor(tensor, dtype=torch.float32, device=image.device) 81 | result = 0.25 * torch.tensordot(image, tensor, dims=2) + 128 82 | # result = torch.from_numpy(result) 83 | result.view(image.shape) 84 | return result 85 | 86 | 87 | def block_merging(patches, height, width): 88 | """ Merge pathces into image 89 | Inputs: 90 | patches(tensor) batch x height*width/64, height x width 91 | height(int) 92 | width(int) 93 | Output: 94 | image(tensor): batch x height x width 95 | """ 96 | k = 8 97 | batch_size = patches.shape[0] 98 | image_reshaped = patches.view(batch_size, height//k, width//k, k, k) 99 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) 100 | return image_transposed.contiguous().view(batch_size, height, width) 101 | 102 | 103 | def chroma_upsampling(y, cb, cr): 104 | """ Upsample chroma layers 105 | Input: 106 | y(tensor): y channel image 107 | cb(tensor): cb channel 108 | cr(tensor): cr channel 109 | Ouput: 110 | image(tensor): batch x height x width x 3 111 | """ 112 | def repeat(x, k=2): 113 | height, width = x.shape[1:3] 114 | x = x.unsqueeze(-1) 115 | x = x.repeat(1, 1, k, k) 116 | x = x.view(-1, height * k, width * k) 117 | return x 118 | 119 | cb = repeat(cb) 120 | cr = repeat(cr) 121 | 122 | return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3) 123 | 124 | 125 | def ycbcr_to_rgb(image): 126 | """ Converts YCbCr image to RGB 127 | Input: 128 | image(tensor): batch x height x width x 3 129 | Outpput: 130 | result(tensor): batch x 3 x height x width 131 | """ 132 | matrix = np.array( 133 | [[298.082, 0, 408.583], [298.082, -100.291, -208.120], 134 | [298.082, 516.412, 0]], 135 | dtype=np.float32).T / 256 136 | shift = [-222.921, 135.576, -276.836] 137 | 138 | result = torch.tensordot(image, matrix, dims=1) + shift 139 | #result = torch.from_numpy(result) 140 | result.view(image.shape) 141 | return result.permute(0, 3, 1, 2) 142 | 143 | 144 | def ycbcr_to_rgb_jpeg(image): 145 | """ Converts YCbCr image to RGB JPEG 146 | Input: 147 | image(tensor): batch x height x width x 3 148 | Outpput: 149 | result(tensor): batch x 3 x height x width 150 | """ 151 | matrix = np.array( 152 | [[1., 0., 1.402], [1, -0.344136, -0.714136], [1, 1.772, 0]], 153 | dtype=np.float32).T 154 | shift =torch.tensor( [0, -128, -128]).to(image.device) 155 | 156 | result = torch.tensordot(image + shift, torch.from_numpy(matrix).to(image.device), dims=1) 157 | #result = torch.from_numpy(result) 158 | result.view(image.shape) 159 | return result.permute(0, 3, 1, 2) 160 | 161 | 162 | def decompress_jpeg(y, cb, cr, height, width, factor=1): 163 | """ Full JPEG decompression algortihm 164 | Input: 165 | compressed(dict(tensor)): batch x h*w/64 x 8 x 8 166 | rounding(function): rounding function to use 167 | factor(float): Compression factor 168 | Ouput: 169 | image(tensor): batch x 3 x height x width 170 | """ 171 | upresults = {} 172 | components = {'y': y, 'cb': cb, 'cr': cr} 173 | for k in components.keys(): 174 | comp = c_dequantize(components[k], factor) if k in ( 175 | 'cb', 'cr') else y_dequantize(components[k], factor) 176 | comp = idct_8x8(comp) 177 | comp = block_merging(comp, int(height/2), int(width/2) 178 | ) if k in ('cb', 'cr') else block_merging(comp, height, width) 179 | upresults[k] = comp 180 | image = chroma_upsampling(upresults['y'], upresults['cb'], upresults['cr']) 181 | image = ycbcr_to_rgb_jpeg(image) 182 | image = torch.min(255*torch.ones_like(image), 183 | torch.max(torch.zeros_like(image), image)) 184 | return image/255 185 | -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # python3 2 | from DiffJPEG.compression import compress_jpeg 3 | from DiffJPEG.decompression import decompress_jpeg 4 | -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/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.])) 24 | # 25 | self.matrix = nn.Parameter(torch.from_numpy(matrix)) 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()) 94 | self.scale = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha) * 0.25).float() ) 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 | -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/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 | """ 20 | def __init__(self, factor=1): 21 | super(y_dequantize, self).__init__() 22 | self.y_table = utils.y_table 23 | self.factor = factor 24 | 25 | def forward(self, image): 26 | return image * (self.y_table * self.factor) 27 | 28 | 29 | class c_dequantize(nn.Module): 30 | """ Dequantize CbCr channel 31 | Inputs: 32 | image(tensor): batch x height x width 33 | factor(float): compression factor 34 | Outputs: 35 | image(tensor): batch x height x width 36 | 37 | """ 38 | def __init__(self, factor=1): 39 | super(c_dequantize, self).__init__() 40 | self.factor = factor 41 | self.c_table = utils.c_table 42 | 43 | def forward(self, image): 44 | return image * (self.c_table * self.factor) 45 | 46 | 47 | class idct_8x8(nn.Module): 48 | """ Inverse discrete Cosine Transformation 49 | Input: 50 | dcp(tensor): batch x height x width 51 | Output: 52 | image(tensor): batch x height x width 53 | """ 54 | def __init__(self): 55 | super(idct_8x8, self).__init__() 56 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) 57 | self.alpha = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha)).float()) 58 | tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) 59 | for x, y, u, v in itertools.product(range(8), repeat=4): 60 | tensor[x, y, u, v] = np.cos((2 * u + 1) * x * np.pi / 16) * np.cos( 61 | (2 * v + 1) * y * np.pi / 16) 62 | self.tensor = nn.Parameter(torch.from_numpy(tensor).float()) 63 | 64 | def forward(self, image): 65 | 66 | image = image * self.alpha 67 | result = 0.25 * torch.tensordot(image, self.tensor, dims=2) + 128 68 | result.view(image.shape) 69 | return result 70 | 71 | 72 | class block_merging(nn.Module): 73 | """ Merge pathces into image 74 | Inputs: 75 | patches(tensor) batch x height*width/64, height x width 76 | height(int) 77 | width(int) 78 | Output: 79 | image(tensor): batch x height x width 80 | """ 81 | def __init__(self): 82 | super(block_merging, self).__init__() 83 | 84 | def forward(self, patches, height, width): 85 | k = 8 86 | batch_size = patches.shape[0] 87 | image_reshaped = patches.view(batch_size, height//k, width//k, k, k) 88 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) 89 | return image_transposed.contiguous().view(batch_size, height, width) 90 | 91 | 92 | class chroma_upsampling(nn.Module): 93 | """ Upsample chroma layers 94 | Input: 95 | y(tensor): y channel image 96 | cb(tensor): cb channel 97 | cr(tensor): cr channel 98 | Ouput: 99 | image(tensor): batch x height x width x 3 100 | """ 101 | def __init__(self): 102 | super(chroma_upsampling, self).__init__() 103 | 104 | def forward(self, y, cb, cr): 105 | def repeat(x, k=2): 106 | height, width = x.shape[1:3] 107 | x = x.unsqueeze(-1) 108 | x = x.repeat(1, 1, k, k) 109 | x = x.view(-1, height * k, width * k) 110 | return x 111 | 112 | cb = repeat(cb) 113 | cr = repeat(cr) 114 | 115 | return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3) 116 | 117 | 118 | class ycbcr_to_rgb_jpeg(nn.Module): 119 | """ Converts YCbCr image to RGB JPEG 120 | Input: 121 | image(tensor): batch x height x width x 3 122 | Outpput: 123 | result(tensor): batch x 3 x height x width 124 | """ 125 | def __init__(self): 126 | super(ycbcr_to_rgb_jpeg, self).__init__() 127 | 128 | matrix = np.array( 129 | [[1., 0., 1.402], [1, -0.344136, -0.714136], [1, 1.772, 0]], 130 | dtype=np.float32).T 131 | self.shift = nn.Parameter(torch.tensor([0, -128., -128.])) 132 | self.matrix = nn.Parameter(torch.from_numpy(matrix)) 133 | 134 | def forward(self, image): 135 | result = torch.tensordot(image + self.shift, self.matrix, dims=1) 136 | #result = torch.from_numpy(result) 137 | result.view(image.shape) 138 | return result.permute(0, 3, 1, 2) 139 | 140 | 141 | class decompress_jpeg(nn.Module): 142 | """ Full JPEG decompression algortihm 143 | Input: 144 | compressed(dict(tensor)): batch x h*w/64 x 8 x 8 145 | rounding(function): rounding function to use 146 | factor(float): Compression factor 147 | Ouput: 148 | image(tensor): batch x 3 x height x width 149 | """ 150 | def __init__(self, height, width, rounding=torch.round, factor=1): 151 | super(decompress_jpeg, self).__init__() 152 | self.c_dequantize = c_dequantize(factor=factor) 153 | self.y_dequantize = y_dequantize(factor=factor) 154 | self.idct = idct_8x8() 155 | self.merging = block_merging() 156 | self.chroma = chroma_upsampling() 157 | self.colors = ycbcr_to_rgb_jpeg() 158 | 159 | self.height, self.width = height, width 160 | 161 | def forward(self, y, cb, cr): 162 | components = {'y': y, 'cb': cb, 'cr': cr} 163 | for k in components.keys(): 164 | if k in ('cb', 'cr'): 165 | comp = self.c_dequantize(components[k]) 166 | height, width = int(self.height/2), int(self.width/2) 167 | else: 168 | comp = self.y_dequantize(components[k]) 169 | height, width = self.height, self.width 170 | comp = self.idct(comp) 171 | components[k] = self.merging(comp, height, width) 172 | # 173 | image = self.chroma(components['y'], components['cb'], components['cr']) 174 | image = self.colors(image) 175 | 176 | image = torch.min(255*torch.ones_like(image), 177 | torch.max(torch.zeros_like(image), image)) 178 | return image/255 179 | -------------------------------------------------------------------------------- /backend/ddf/DiffJPEG_pytorch/utils.py: -------------------------------------------------------------------------------- 1 | # Standard libraries 2 | import numpy as np 3 | # PyTorch 4 | import torch 5 | import torch.nn as nn 6 | 7 | y_table = np.array( 8 | [[16, 11, 10, 16, 24, 40, 51, 61], [12, 12, 14, 19, 26, 58, 60, 9 | 55], [14, 13, 16, 24, 40, 57, 69, 56], 10 | [14, 17, 22, 29, 51, 87, 80, 62], [18, 22, 37, 56, 68, 109, 103, 11 | 77], [24, 35, 55, 64, 81, 104, 113, 92], 12 | [49, 64, 78, 87, 103, 121, 120, 101], [72, 92, 95, 98, 112, 100, 103, 99]], 13 | dtype=np.float32).T 14 | 15 | y_table = nn.Parameter(torch.from_numpy(y_table)) 16 | # 17 | c_table = np.empty((8, 8), dtype=np.float32) 18 | c_table.fill(99) 19 | c_table[:4, :4] = np.array([[17, 18, 24, 47], [18, 21, 26, 66], 20 | [24, 26, 56, 99], [47, 66, 99, 99]]).T 21 | c_table = nn.Parameter(torch.from_numpy(c_table)) 22 | 23 | 24 | def diff_round(x): 25 | """ Differentiable rounding function 26 | Input: 27 | x(tensor) 28 | Output: 29 | x(tensor) 30 | """ 31 | return torch.round(x) + (x - torch.round(x))**3 32 | 33 | 34 | def quality_to_factor(quality): 35 | """ Calculate factor corresponding to quality 36 | Input: 37 | quality(float): Quality for jpeg compression 38 | Output: 39 | factor(float): Compression factor 40 | """ 41 | if quality < 50: 42 | quality = 5000. / quality 43 | else: 44 | quality = 200. - quality*2 45 | return quality / 100. 46 | -------------------------------------------------------------------------------- /backend/ddf/DualDefense_gan_fs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from ddf.adv import Adversary_Init 4 | from ddf.encoder_fs import ResNetUNet 5 | from ddf.decoder_fs import Decoder 6 | from ddf.faceswap_pytorch.models import Autoencoder_df 7 | 8 | CHECKPOINTS = { 9 | 'trump_cage' : '/data1/yoojinoh/def/241116_faceswap_160/checkpoint/autoencoder_160.t7', 10 | 'winter_karina': '/data1/yoojinoh/def/250101_faceswap_160/checkpoint/ae_win_kar_160.t7', 11 | 'byeon_cha': '/data1/yoojinoh/def/train/0128_fc_bc_v2/checkpoint/bc_ckpt_best.t7', 12 | 'win_chuu': '/home/yoojinoh/Others/ckpts/faceswap/wc_ckpt_best_82000.t7' 13 | } 14 | 15 | 16 | class DualDefense(nn.Module): 17 | def __init__(self, message_size, in_channels,device=False,model_type=None): 18 | super().__init__() 19 | self.encoder = ResNetUNet(message_size) 20 | self.df_model = Autoencoder_df() 21 | self.decoder = Decoder(message_size) 22 | self.adv_model = Adversary_Init(in_channels) 23 | checkpoint = torch.load(CHECKPOINTS[model_type]) 24 | self.df_model.load_state_dict(checkpoint['state']) 25 | 26 | if device: 27 | self.encoder = self.encoder.to(device) 28 | self.df_model = self.df_model.to(device) 29 | self.decoder = self.decoder.to(device) 30 | self.adv_model = self.adv_model.to(device) 31 | 32 | def encode(self, x, message): 33 | return self.encoder(x, message) 34 | 35 | def deepfake1(self, x, type): 36 | return self.df_model(x, type) 37 | 38 | def adv(self, x): 39 | return self.adv_model(x) 40 | 41 | def decode(self, x): 42 | return self.decoder(x) 43 | -------------------------------------------------------------------------------- /backend/ddf/__init__.py: -------------------------------------------------------------------------------- 1 | """Codebase for Dual Defense""" 2 | 3 | from .main import crop_and_encode_image, USER_WATERMARK_IDS, apply_faceswap 4 | 5 | __all__ = ["crop_and_encode_image", "USER_WATERMARK_IDS", "apply_faceswap"] 6 | -------------------------------------------------------------------------------- /backend/ddf/adv.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class Adversary_Init(nn.Module): 4 | """Discriminator, used in L_D of eq(8)""" 5 | def __init__(self, in_channels): 6 | super(Adversary_Init, self).__init__() 7 | self.model = nn.Sequential( 8 | # First Conv-BN-ReLU block 9 | nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=3, stride=1, padding=1), 10 | nn.BatchNorm2d(64), 11 | nn.ReLU(), 12 | 13 | # Second Conv-BN-ReLU block 14 | nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), 15 | nn.BatchNorm2d(128), 16 | nn.ReLU(), 17 | 18 | # Third Conv-BN-ReLU block 19 | nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1), 20 | nn.BatchNorm2d(256), 21 | nn.ReLU(), 22 | 23 | # Average pooling layer 24 | nn.AdaptiveAvgPool2d((1, 1)), 25 | 26 | # Fully connected layer with Sigmoid 27 | nn.Flatten(), 28 | nn.Linear(256, 1), 29 | nn.Sigmoid() 30 | ) 31 | 32 | def forward(self, x): 33 | return self.model(x) -------------------------------------------------------------------------------- /backend/ddf/assets/25_0_combined_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/25_0_combined_A.png -------------------------------------------------------------------------------- /backend/ddf/assets/58_2_combined_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/58_2_combined_A.png -------------------------------------------------------------------------------- /backend/ddf/assets/byeon_total5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/byeon_total5.jpg -------------------------------------------------------------------------------- /backend/ddf/assets/cha_total2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/cha_total2.jpg -------------------------------------------------------------------------------- /backend/ddf/assets/output_A2B.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/output_A2B.png -------------------------------------------------------------------------------- /backend/ddf/assets/output_B2A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/output_B2A.png -------------------------------------------------------------------------------- /backend/ddf/assets/output_e18000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/assets/output_e18000.png -------------------------------------------------------------------------------- /backend/ddf/config.yaml: -------------------------------------------------------------------------------- 1 | # config.yaml 2 | train: 3 | batch_size: 16 # 32 4 | lr: 0.0005 5 | epoch: 2500 6 | start_decode: 30 # d_iter 7 | clip: 15 8 | message_size: 15 # 4 9 | T_max: 50 10 | name: bc-ckpt-ddf 11 | lambda_val: 2 # beta 12 | alpha_val: 0.5 # alpha 13 | 14 | trump_path: /data1/yoojinoh/def/new_data/faceDataset2/byeonFace 15 | cage_path: /data1/yoojinoh/def/new_data/faceDataset2/chaFace 16 | 17 | test: 18 | batch_size: 2 # 16 19 | message_size: 4 20 | quality: 50 # JPEG quality 21 | 22 | model_path: /data1/yoojinoh/def/train/0127_bc_ddf/bc-ckpt-ddf/ckpt_best_lam2_al0.5.pt 23 | trump_path: /data1/yoojinoh/def/new_data/faceDataset2/byeonFace 24 | cage_path: /data1/yoojinoh/def/new_data/faceDataset2/chaFace 25 | save_path: /data1/yoojinoh/def/eval/0128_bc_ddf -------------------------------------------------------------------------------- /backend/ddf/config2.yaml: -------------------------------------------------------------------------------- 1 | # config.yaml 2 | train: 3 | batch_size: 16 # 32 4 | lr: 0.0005 5 | epoch: 2500 6 | start_decode: 30 # d_iter 7 | clip: 15 8 | message_size: 4 # 4 9 | T_max: 50 10 | name: bc-ckpt-ddf 11 | lambda_val: 2 # beta 12 | alpha_val: 0.5 # alpha 13 | 14 | trump_path: /data1/yoojinoh/def/new_data/faceDataset2/byeonFace 15 | cage_path: /data1/yoojinoh/def/new_data/faceDataset2/chaFace 16 | 17 | test: 18 | batch_size: 2 # 16 19 | message_size: 15 20 | quality: 50 # JPEG quality 21 | 22 | model_path: /home/yoojinoh/Others/ckpts/winchuu_ckpt_best_img_lam1_al2.pt 23 | trump_path: /data1/yoojinoh/def/new_data/faceDataset2/winterFace 24 | cage_path: /data1/yoojinoh/def/new_data/faceDataset2/chuuFace 25 | save_path: /data1/yoojinoh/def/eval/winchuu_0208 26 | 27 | model_type: win_chuu -------------------------------------------------------------------------------- /backend/ddf/dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from torch.utils.data import Dataset, DataLoader 4 | from PIL import Image 5 | import cv2 6 | from torchvision.transforms import transforms 7 | from faceswap_pytorch.models import Autoencoder, toTensor, var_to_np 8 | from faceswap_pytorch.training_data import get_training_data 9 | from faceswap_pytorch.umeyama import * 10 | import torch 11 | 12 | transform_train = transforms.Compose([ 13 | transforms.ToTensor(), 14 | transforms.Resize((160, 160)), 15 | transforms.RandomHorizontalFlip(0.5), 16 | transforms.RandomRotation((-10, 10)), 17 | transforms.ColorJitter(brightness=(0.8, 1.2), contrast=(0.8, 1.2), saturation=(0.8, 1.2), hue=(-0.1, 0.1)) 18 | ]) 19 | 20 | transform_val = transforms.Compose([ 21 | transforms.ToTensor(), 22 | transforms.Resize((160, 160)), 23 | ]) 24 | 25 | transform_test = transforms.Compose([ 26 | transforms.ToTensor(), 27 | transforms.Resize((160, 160)), ### ADDED(YOOJIN) 28 | ]) 29 | 30 | class CustomDataset(Dataset): 31 | def __init__(self, images, transform=None, type='train'): 32 | self.images = images 33 | self.transform = transform 34 | self.type = type 35 | 36 | def __len__(self): 37 | return len(self.images) 38 | 39 | def __getitem__(self, idx): 40 | # image = Image.open(self.ids[idx]) 41 | image = Image.fromarray(cv2.imread(self.images[idx])) 42 | 43 | # image = cv2.imread(self.ids[idx]) 44 | 45 | if self.type == "train": 46 | if self.transform: 47 | image = self.transform(image) 48 | if self.type=="test": 49 | if self.transform: 50 | image = self.transform(image) 51 | return image 52 | 53 | 54 | def split_dataset(path, train_transform=None, valid_transform=None, test_transform=None, val_ratio=0.2, test_ratio=0.2): 55 | images = [x.path for x in os.scandir(path) if x.name.endswith(".jpg") or x.name.endswith(".png")] 56 | images = [os.path.join(root, file) for root,_,files in os.walk(path) for file in files if file.endswith(".jpg") or file.endswith(".png")] 57 | total_len = len(images) 58 | train_images = images[: int(total_len * (1 - val_ratio - test_ratio))] 59 | val_images = images[int(total_len * (1 - val_ratio - test_ratio)): int(total_len * (1 - test_ratio))] 60 | test_images = images[int(total_len * (1 - test_ratio)):] 61 | print(f"# of Train:{len(train_images)} | Valid: {len(val_images)} | Test: {len(test_images)}") 62 | return CustomDataset(train_images, train_transform), CustomDataset(val_images, valid_transform), CustomDataset(test_images, test_transform, "test") 63 | 64 | 65 | ''' 66 | trans_test = transforms.Compose([ 67 | transforms.ToTensor() 68 | ]) 69 | 70 | background = cv2.imread('') 71 | cv2.imwrite('tmp1.png', background) 72 | yuv_background = cv2.cvtColor(background, cv2.COLOR_BGR2YUV) # 73 | Y, U, V = yuv_background[..., 0], yuv_background[..., 1], yuv_background[..., 2] 74 | YY = trans_test(Y) 75 | YYY = YY.permute(1, 2, 0).squeeze().numpy() * 255 76 | print(yuv_background[..., 0] == YYY) 77 | x = yuv_background[..., 0] == YYY 78 | 79 | yuv_background[..., 0] = YYY 80 | mm = cv2.cvtColor(yuv_background, cv2.COLOR_YUV2BGR) 81 | cv2.imwrite('tmp2.png', mm) 82 | 83 | ''' 84 | 85 | # get pair of random warped images from aligened face image 86 | def random_warp(image): 87 | if isinstance(image, torch.Tensor): 88 | image = image.detach().cpu().permute(1, 2, 0).cpu().numpy() 89 | assert image.shape == (256, 256, 3) 90 | range_ = np.linspace(128 - 80, 128 + 80, 5) 91 | mapx = np.broadcast_to(range_, (5, 5)) 92 | mapy = mapx.T 93 | 94 | mapx = mapx + np.random.normal(size=(5, 5), scale=5) 95 | mapy = mapy + np.random.normal(size=(5, 5), scale=5) 96 | 97 | interp_mapx = cv2.resize(mapx, (80, 80))[8:72, 8:72].astype('float32') 98 | interp_mapy = cv2.resize(mapy, (80, 80))[8:72, 8:72].astype('float32') 99 | 100 | warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR) 101 | 102 | src_points = np.stack([mapx.ravel(), mapy.ravel()], axis=-1) 103 | dst_points = np.mgrid[0:65:16, 0:65:16].T.reshape(-1, 2) 104 | mat = umeyama(src_points, dst_points, True)[0:2] 105 | 106 | target_image = cv2.warpAffine(image, mat, (64, 64)) 107 | 108 | return warped_image, target_image 109 | -------------------------------------------------------------------------------- /backend/ddf/decoder_fs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Conv2dBatchNorm2dReLU(nn.Module): 6 | def __init__(self, dim_in, dim_out): 7 | super().__init__() 8 | self.layer = nn.Sequential( 9 | nn.Conv2d(dim_in, dim_out, kernel_size=3, padding=1, bias=False), 10 | nn.BatchNorm2d(dim_out), 11 | nn.ReLU() 12 | ) 13 | 14 | def forward(self, x): 15 | return self.layer(x) 16 | 17 | #TODO(Yoojin): Check Watermark Decoder 18 | class Decoder(nn.Module): 19 | def __init__(self, message_size): 20 | super().__init__() 21 | # self.layer1 = Conv2dBatchNorm2dReLU(1, 32) 22 | self.layer1 = Conv2dBatchNorm2dReLU(3, 32) 23 | self.layer2 = Conv2dBatchNorm2dReLU(32, 64) 24 | self.layer3 = Conv2dBatchNorm2dReLU(64, 128) 25 | self.layer4 = Conv2dBatchNorm2dReLU(128, 256) 26 | self.layer5 = Conv2dBatchNorm2dReLU(256, 512) 27 | self.layer6 = Conv2dBatchNorm2dReLU(512, 1024) 28 | self.layer7 = Conv2dBatchNorm2dReLU(1024, 1024) 29 | self.max_pool_layer = nn.MaxPool2d(2) 30 | self.linear = nn.Linear(1024 * 20 * 20, message_size) 31 | self.sigmoid = nn.Sigmoid() 32 | self.dropout = nn.Dropout(p=0.5) # d 33 | 34 | for layer in self.modules(): 35 | if isinstance(layer, (nn.Conv2d)): 36 | nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu') 37 | elif isinstance(layer, (nn.Linear)): 38 | nn.init.xavier_normal_(layer.weight) 39 | nn.init.zeros_(layer.bias) 40 | 41 | def forward(self, x): # input: (16, 3, 64, 64) 42 | x = self.layer1(x) # (16, 32, 64, 64) 43 | x = self.layer2(x) # (16, 64, 64, 64) 44 | x = self.max_pool_layer(x) # (b, 64, 32, 32) 45 | x = self.layer3(x) # (16, 128, 32, 32) 46 | x = self.layer4(x) # (16, 256, 32, 32) 47 | x = self.max_pool_layer(x) # (b, 256, 16, 16) 48 | x = self.layer5(x) # (16, 512, 16, 16) 49 | x = self.layer6(x) # (b, 1024, 32, 32) ---- (16, 1024, 16, 16) 50 | x = self.max_pool_layer(x) # (b, 1024, 8, 8) --- (16, 1024, 8, 8) 51 | x = self.layer7(x) # (b, 1024, 8, 8) --- (16, 1024, 8, 8) 52 | x = x.flatten(start_dim=1) # (16, 1024*8*8) 53 | x = self.dropout(x) # (16, 1024*8*8) 54 | x = self.linear(x) # (b, message_size) 55 | return self.sigmoid(x) 56 | -------------------------------------------------------------------------------- /backend/ddf/encoder_fs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torchvision import models 4 | 5 | 6 | def convrelu(in_channels, out_channels, kernel, padding): 7 | return nn.Sequential( 8 | nn.Conv2d(in_channels, out_channels, kernel, padding=padding), 9 | nn.ReLU(inplace=True), 10 | ) 11 | 12 | 13 | class ResNetUNet(nn.Module): 14 | def __init__(self, message_size): 15 | super().__init__() 16 | 17 | self.message_size = message_size 18 | self.base_model = models.resnet18(pretrained=False) 19 | self.base_layers = list(self.base_model.children()) 20 | 21 | self.layer0 = nn.Sequential(*self.base_layers[:3]) # size=(N, 64, x.H/2, x.W/2) 22 | self.layer0_1x1 = convrelu(64 + message_size, 64, 1, 0) 23 | self.layer1 = nn.Sequential(*self.base_layers[3:5]) # size=(N, 64, x.H/4, x.W/4) 24 | self.layer1_1x1 = convrelu(64 + message_size, 64, 1, 0) 25 | self.layer2 = self.base_layers[5] # size=(N, 128, x.H/8, x.W/8) 26 | self.layer2_1x1 = convrelu(128 + message_size, 128, 1, 0) 27 | self.layer3 = self.base_layers[6] # size=(N, 256, x.H/16, x.W/16) 28 | self.layer3_1x1 = convrelu(256 + message_size, 256, 1, 0) 29 | self.layer4 = self.base_layers[7] # size=(N, 512, x.H/32, x.W/32) 30 | self.layer4_1x1 = convrelu(512 + message_size, 512, 1, 0) 31 | 32 | self.cbam2 = CBAM(channel=128) 33 | 34 | self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) 35 | 36 | self.conv_up3 = convrelu(256 + 512, 512, 3, 1) 37 | self.conv_up2 = convrelu(128 + 512, 256, 3, 1) 38 | self.conv_up1 = convrelu(64 + 256, 256, 3, 1) 39 | self.conv_up0 = convrelu(64 + 256, 128, 3, 1) 40 | 41 | self.conv_original_size0 = convrelu(3, 64, 3, 1) 42 | self.conv_original_size1 = convrelu(64, 64, 3, 1) 43 | self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1) 44 | 45 | self.conv_last = nn.Conv2d(64, 3, 1) 46 | 47 | def forward(self, input, message): 48 | batch_size, _, h, w = input.shape 49 | x_original = self.conv_original_size0(input) 50 | x_original = self.conv_original_size1(x_original) #size=(N, 64, 256, 256) 51 | 52 | layer0 = self.layer0(input) # size=(N, 64, 128, 128) 53 | layer1 = self.layer1(layer0) # size=(N, 64, 64, 64) 54 | layer2 = self.layer2(layer1) # size=(N, 128, 32, 32) 55 | layer2 = self.cbam2(layer2) # size=(N, 128, 32, 32) 56 | layer3 = self.layer3(layer2) # size=(N, 256, 16, 16) 57 | layer4 = self.layer4(layer3) # size=(N, 512, 8, 8) 58 | layer4 = self.layer4_1x1(torch.cat( 59 | (layer4, message.expand(h // 32, w // 32, batch_size, self.message_size).permute(2, 3, 0, 1).contiguous()), 60 | dim=1)) # size=(N, 512, 8, 8) 61 | x = self.upsample(layer4) # size=(N, 512, 16, 16) 62 | 63 | layer3 = self.layer3_1x1(torch.cat( 64 | (layer3, message.expand(h // 16, w // 16, batch_size, self.message_size).permute(2, 3, 0, 1).contiguous()), 65 | dim=1)) # size=(N, 256, 16, 16) 66 | x = torch.cat([x, layer3], dim=1) # size=(N, 512+256, 16, 16) 67 | x = self.conv_up3(x) # size=(N, 512, 16, 16) 68 | x = self.upsample(x) # size=(N, 512, 32, 32) 69 | 70 | layer2 = self.layer2_1x1(torch.cat( 71 | (layer2, message.expand(h // 8, w // 8, batch_size, self.message_size).permute(2, 3, 0, 1).contiguous()), 72 | dim=1)) 73 | x = torch.cat([x, layer2], dim=1) 74 | x = self.conv_up2(x) 75 | x = self.upsample(x) # size=(N, 256, 64, 64) 76 | 77 | layer1 = self.layer1_1x1(torch.cat( 78 | (layer1, message.expand(h // 4, w // 4, batch_size, self.message_size).permute(2, 3, 0, 1).contiguous()), 79 | dim=1)) 80 | x = torch.cat([x, layer1], dim=1) 81 | x = self.conv_up1(x) 82 | x = self.upsample(x) # size=(N, 256, 128, 128) 83 | 84 | layer0 = self.layer0_1x1(torch.cat( 85 | (layer0, message.expand(h // 2, w // 2, batch_size, self.message_size).permute(2, 3, 0, 1).contiguous()), 86 | dim=1)) 87 | x = torch.cat([x, layer0], dim=1) 88 | x = self.conv_up0(x) 89 | x = self.upsample(x) 90 | 91 | x = torch.cat([x, x_original], dim=1) # size=(N, 256, 128, 128) 92 | x = self.conv_original_size2(x) 93 | out = self.conv_last(x) 94 | 95 | return out # size=(N, 3, 256, 256) 96 | 97 | 98 | class ChannelAttentionModule(nn.Module): 99 | def __init__(self, channel, ratio=16): 100 | super(ChannelAttentionModule, self).__init__() 101 | self.avg_pool = nn.AdaptiveAvgPool2d(1) 102 | self.max_pool = nn.AdaptiveMaxPool2d(1) 103 | self.shared_MLP = nn.Sequential( 104 | nn.Conv2d(channel, channel // ratio, 1, bias=False), 105 | nn.ReLU(), 106 | nn.Conv2d(channel // ratio, channel, 1, bias=False) 107 | ) 108 | self.sigmoid = nn.Sigmoid() 109 | 110 | def forward(self, x): 111 | avgout = self.shared_MLP(self.avg_pool(x)) 112 | maxout = self.shared_MLP(self.max_pool(x)) 113 | return self.sigmoid(avgout + maxout) 114 | 115 | 116 | class SpatialAttentionModule(nn.Module): 117 | def __init__(self): 118 | super(SpatialAttentionModule, self).__init__() 119 | self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3) 120 | self.sigmoid = nn.Sigmoid() 121 | # self.dropout = nn.Dropout(p=0.9) # d 122 | 123 | def forward(self, x): 124 | avgout = torch.mean(x, dim=1, keepdim=True) 125 | maxout, _ = torch.max(x, dim=1, keepdim=True) 126 | out = torch.cat([avgout, maxout], dim=1) 127 | out = self.sigmoid(self.conv2d(out)) 128 | # out = self.dropout(out) 129 | return out 130 | 131 | 132 | class CBAM(nn.Module): 133 | def __init__(self, channel): 134 | super(CBAM, self).__init__() 135 | self.channel_attention = ChannelAttentionModule(channel) 136 | self.spatial_attention = SpatialAttentionModule() 137 | 138 | def forward(self, x): 139 | out = self.channel_attention(x) * x 140 | out = self.spatial_attention(out) * out 141 | return out -------------------------------------------------------------------------------- /backend/ddf/engine/dwt.py: -------------------------------------------------------------------------------- 1 | """ 2 | Reproduced 3 | """ 4 | import torch 5 | import pywt 6 | import torch.nn as nn 7 | from pytorch_wavelets import DWTForward # Forward DWT 8 | import numpy as np 9 | import pywt 10 | import pytorch_wavelets.dwt.lowlevel as lowlevel 11 | import torch 12 | import cv2 13 | 14 | # def get_y_channels(data): 15 | # outputs = [] 16 | # for x in data: 17 | # ycrcb_image = cv2.cvtColor(x, cv2.COLOR_BGR2YCrCb) # BGR -> YCrCb 18 | # y_channel = ycrcb_image[:, :, 0] 19 | # y_tensor = torch.tensor(y_channel, dtype=torch.float32).unsqueeze(0).unsqueeze(0) 20 | # outputs.append(y_tensor) 21 | # return outputs 22 | 23 | def get_y_channels(data): 24 | outputs = [] 25 | for x in data: 26 | # Assuming x is in the shape (C, H, W) and in RGB format 27 | if x.shape[0] != 3: 28 | raise ValueError("Input tensor must have 3 channels (RGB).") 29 | 30 | # Extract RGB channels 31 | R, G, B = x[0, :, :], x[1, :, :], x[2, :, :] 32 | 33 | # Compute Y channel using the formula for Y in YCrCb 34 | # Y = 0.299 * R + 0.587 * G + 0.114 * B 35 | Y = 0.299 * R + 0.587 * G + 0.114 * B 36 | 37 | # Convert Y to a tensor with shape (1, 1, H, W) 38 | y_tensor = Y.unsqueeze(0).unsqueeze(0) # Add channel and batch dimensions 39 | 40 | outputs.append(y_tensor) 41 | return outputs 42 | 43 | def get_dwt(dwt, data): 44 | LL_lst = [] 45 | for y in data: 46 | yl, _ = dwt(y) 47 | LL_lst.append(yl) 48 | return LL_lst 49 | 50 | class DWTForward(nn.Module): 51 | """ Performs a 2d DWT Forward decomposition of an image 52 | 53 | Args: 54 | J (int): Number of levels of decomposition 55 | wave (str or pywt.Wavelet or tuple(ndarray)): Which wavelet to use. 56 | Can be: 57 | 1) a string to pass to pywt.Wavelet constructor 58 | 2) a pywt.Wavelet class 59 | 3) a tuple of numpy arrays, either (h0, h1) or (h0_col, h1_col, h0_row, h1_row) 60 | mode (str): 'zero', 'symmetric', 'reflect' or 'periodization'. The 61 | padding scheme 62 | """ 63 | def __init__(self, J=1, wave='haar', mode='zero'): 64 | super().__init__() 65 | if isinstance(wave, str): 66 | wave = pywt.Wavelet(wave) 67 | if isinstance(wave, pywt.Wavelet): 68 | h0_col, h1_col = wave.dec_lo, wave.dec_hi 69 | h0_row, h1_row = h0_col, h1_col 70 | else: 71 | if len(wave) == 2: 72 | h0_col, h1_col = wave[0], wave[1] 73 | h0_row, h1_row = h0_col, h1_col 74 | elif len(wave) == 4: 75 | h0_col, h1_col = wave[0], wave[1] 76 | h0_row, h1_row = wave[2], wave[3] 77 | 78 | # Prepare the filters 79 | filts = lowlevel.prep_filt_afb2d(h0_col, h1_col, h0_row, h1_row) 80 | self.register_buffer('h0_col', filts[0]) 81 | self.register_buffer('h1_col', filts[1]) 82 | self.register_buffer('h0_row', filts[2]) 83 | self.register_buffer('h1_row', filts[3]) 84 | self.J = J 85 | self.mode = mode 86 | 87 | def forward(self, x): 88 | """ Forward pass of the DWT. 89 | 90 | Args: 91 | x (tensor): Input of shape :math:`(N, C_{in}, H_{in}, W_{in})` 92 | 93 | Returns: 94 | (yl, yh) 95 | tuple of lowpass (yl) and bandpass (yh) coefficients. 96 | yh is a list of length J with the first entry 97 | being the finest scale coefficients. yl has shape 98 | :math:`(N, C_{in}, H_{in}', W_{in}')` and yh has shape 99 | :math:`list(N, C_{in}, 3, H_{in}'', W_{in}'')`. The new 100 | dimension in yh iterates over the LH, HL and HH coefficients. 101 | 102 | Note: 103 | :math:`H_{in}', W_{in}', H_{in}'', W_{in}''` denote the correctly 104 | downsampled shapes of the DWT pyramid. 105 | """ 106 | yh = [] 107 | ll = x 108 | mode = lowlevel.mode_to_int(self.mode) 109 | 110 | # Do a multilevel transform 111 | for j in range(self.J): 112 | # Do 1 level of the transform 113 | ll, high = lowlevel.AFB2D.apply( 114 | ll, self.h0_col, self.h1_col, self.h0_row, self.h1_row, mode) 115 | yh.append(high) 116 | 117 | return ll, yh 118 | 119 | 120 | # class DWTForward_Init(nn.Module): 121 | # def __init__(self, J=1, mode='zero', wave='haar'): 122 | # super(DWTForward_Init, self).__init__() 123 | # self.dwt = DWTForward(J=J, wave=wave) 124 | 125 | # def forward(self, x): 126 | # Yl, Yh = self.dwt(x) # Yl: 저주파(LL), Yh: 고주파(LH, HL, HH) 127 | # return Yl, Yh 128 | 129 | # class DWTForward_Init(nn.Module): 130 | # def __init__(self, J=1, mode='zero', wave='haar'): 131 | # super(DWTForward_Init, self).__init__() 132 | # self.J = J 133 | # self.mode = mode 134 | # self.wave = wave 135 | 136 | # def forward(self, x): 137 | # coeffs = pywt.wavedec2(x.squeeze().cpu().numpy(), 138 | # self.wave, 139 | # mode=self.mode, 140 | # level=self.J) 141 | # # 저주파 성분(LL)과 고주파 성분(LH, HL, HH)을 분리 142 | # Yl = torch.tensor(coeffs[0]).unsqueeze(0).to(x.device) # LL 성분 143 | # Yh = [torch.tensor(detail).unsqueeze(0).to(x.device) for detail in coeffs[1:]] # LH, HL, HH 성분 144 | # return Yl, Yh 145 | 146 | class DWTInverse_Init(nn.Module): 147 | def __init__(self, mode='zero', wave='haar'): 148 | super(DWTInverse_Init, self).__init__() 149 | self.mode = mode 150 | self.wave = wave 151 | 152 | def forward(self, Yl, Yh): 153 | # 역변환을 위해 저주파와 고주파 성분을 결합 154 | coeffs = [Yl.squeeze().cpu().numpy()] + [detail.squeeze().cpu().numpy() for detail in Yh] 155 | rec_image = pywt.waverec2(coeffs, self.wave, mode=self.mode) 156 | return torch.tensor(rec_image).unsqueeze(0).to(Yl.device) 157 | -------------------------------------------------------------------------------- /backend/ddf/facenet_pytorch/models/inception_resnet_v1.py: -------------------------------------------------------------------------------- 1 | import os 2 | import requests 3 | from requests.adapters import HTTPAdapter 4 | 5 | import torch 6 | from torch import nn 7 | from torch.nn import functional as F 8 | 9 | from .utils.download import download_url_to_file 10 | 11 | 12 | class BasicConv2d(nn.Module): 13 | 14 | def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0): 15 | super().__init__() 16 | self.conv = nn.Conv2d( 17 | in_planes, out_planes, 18 | kernel_size=kernel_size, stride=stride, 19 | padding=padding, bias=False 20 | ) # verify bias false 21 | self.bn = nn.BatchNorm2d( 22 | out_planes, 23 | eps=0.001, # value found in tensorflow 24 | momentum=0.1, # default pytorch value 25 | affine=True 26 | ) 27 | self.relu = nn.ReLU(inplace=False) 28 | 29 | def forward(self, x): 30 | x = self.conv(x) 31 | x = self.bn(x) 32 | x = self.relu(x) 33 | return x 34 | 35 | 36 | class Block35(nn.Module): 37 | 38 | def __init__(self, scale=1.0): 39 | super().__init__() 40 | 41 | self.scale = scale 42 | 43 | self.branch0 = BasicConv2d(256, 32, kernel_size=1, stride=1) 44 | 45 | self.branch1 = nn.Sequential( 46 | BasicConv2d(256, 32, kernel_size=1, stride=1), 47 | BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1) 48 | ) 49 | 50 | self.branch2 = nn.Sequential( 51 | BasicConv2d(256, 32, kernel_size=1, stride=1), 52 | BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1), 53 | BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1) 54 | ) 55 | 56 | self.conv2d = nn.Conv2d(96, 256, kernel_size=1, stride=1) 57 | self.relu = nn.ReLU(inplace=False) 58 | 59 | def forward(self, x): 60 | x0 = self.branch0(x) 61 | x1 = self.branch1(x) 62 | x2 = self.branch2(x) 63 | out = torch.cat((x0, x1, x2), 1) 64 | out = self.conv2d(out) 65 | out = out * self.scale + x 66 | out = self.relu(out) 67 | return out 68 | 69 | 70 | class Block17(nn.Module): 71 | 72 | def __init__(self, scale=1.0): 73 | super().__init__() 74 | 75 | self.scale = scale 76 | 77 | self.branch0 = BasicConv2d(896, 128, kernel_size=1, stride=1) 78 | 79 | self.branch1 = nn.Sequential( 80 | BasicConv2d(896, 128, kernel_size=1, stride=1), 81 | BasicConv2d(128, 128, kernel_size=(1,7), stride=1, padding=(0,3)), 82 | BasicConv2d(128, 128, kernel_size=(7,1), stride=1, padding=(3,0)) 83 | ) 84 | 85 | self.conv2d = nn.Conv2d(256, 896, kernel_size=1, stride=1) 86 | self.relu = nn.ReLU(inplace=False) 87 | 88 | def forward(self, x): 89 | x0 = self.branch0(x) 90 | x1 = self.branch1(x) 91 | out = torch.cat((x0, x1), 1) 92 | out = self.conv2d(out) 93 | out = out * self.scale + x 94 | out = self.relu(out) 95 | return out 96 | 97 | 98 | class Block8(nn.Module): 99 | 100 | def __init__(self, scale=1.0, noReLU=False): 101 | super().__init__() 102 | 103 | self.scale = scale 104 | self.noReLU = noReLU 105 | 106 | self.branch0 = BasicConv2d(1792, 192, kernel_size=1, stride=1) 107 | 108 | self.branch1 = nn.Sequential( 109 | BasicConv2d(1792, 192, kernel_size=1, stride=1), 110 | BasicConv2d(192, 192, kernel_size=(1,3), stride=1, padding=(0,1)), 111 | BasicConv2d(192, 192, kernel_size=(3,1), stride=1, padding=(1,0)) 112 | ) 113 | 114 | self.conv2d = nn.Conv2d(384, 1792, kernel_size=1, stride=1) 115 | if not self.noReLU: 116 | self.relu = nn.ReLU(inplace=False) 117 | 118 | def forward(self, x): 119 | x0 = self.branch0(x) 120 | x1 = self.branch1(x) 121 | out = torch.cat((x0, x1), 1) 122 | out = self.conv2d(out) 123 | out = out * self.scale + x 124 | if not self.noReLU: 125 | out = self.relu(out) 126 | return out 127 | 128 | 129 | class Mixed_6a(nn.Module): 130 | 131 | def __init__(self): 132 | super().__init__() 133 | 134 | self.branch0 = BasicConv2d(256, 384, kernel_size=3, stride=2) 135 | 136 | self.branch1 = nn.Sequential( 137 | BasicConv2d(256, 192, kernel_size=1, stride=1), 138 | BasicConv2d(192, 192, kernel_size=3, stride=1, padding=1), 139 | BasicConv2d(192, 256, kernel_size=3, stride=2) 140 | ) 141 | 142 | self.branch2 = nn.MaxPool2d(3, stride=2) 143 | 144 | def forward(self, x): 145 | x0 = self.branch0(x) 146 | x1 = self.branch1(x) 147 | x2 = self.branch2(x) 148 | out = torch.cat((x0, x1, x2), 1) 149 | return out 150 | 151 | 152 | class Mixed_7a(nn.Module): 153 | 154 | def __init__(self): 155 | super().__init__() 156 | 157 | self.branch0 = nn.Sequential( 158 | BasicConv2d(896, 256, kernel_size=1, stride=1), 159 | BasicConv2d(256, 384, kernel_size=3, stride=2) 160 | ) 161 | 162 | self.branch1 = nn.Sequential( 163 | BasicConv2d(896, 256, kernel_size=1, stride=1), 164 | BasicConv2d(256, 256, kernel_size=3, stride=2) 165 | ) 166 | 167 | self.branch2 = nn.Sequential( 168 | BasicConv2d(896, 256, kernel_size=1, stride=1), 169 | BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1), 170 | BasicConv2d(256, 256, kernel_size=3, stride=2) 171 | ) 172 | 173 | self.branch3 = nn.MaxPool2d(3, stride=2) 174 | 175 | def forward(self, x): 176 | x0 = self.branch0(x) 177 | x1 = self.branch1(x) 178 | x2 = self.branch2(x) 179 | x3 = self.branch3(x) 180 | out = torch.cat((x0, x1, x2, x3), 1) 181 | return out 182 | 183 | 184 | class InceptionResnetV1(nn.Module): 185 | """Inception Resnet V1 model with optional loading of pretrained weights. 186 | 187 | Model parameters can be loaded based on pretraining on the VGGFace2 or CASIA-Webface 188 | datasets. Pretrained state_dicts are automatically downloaded on model instantiation if 189 | requested and cached in the torch cache. Subsequent instantiations use the cache rather than 190 | redownloading. 191 | 192 | Keyword Arguments: 193 | pretrained {str} -- Optional pretraining dataset. Either 'vggface2' or 'casia-webface'. 194 | (default: {None}) 195 | classify {bool} -- Whether the model should output classification probabilities or feature 196 | embeddings. (default: {False}) 197 | num_classes {int} -- Number of output classes. If 'pretrained' is set and num_classes not 198 | equal to that used for the pretrained model, the final linear layer will be randomly 199 | initialized. (default: {None}) 200 | dropout_prob {float} -- Dropout probability. (default: {0.6}) 201 | """ 202 | def __init__(self, pretrained=None, classify=False, num_classes=None, dropout_prob=0.6, device=None): 203 | super().__init__() 204 | 205 | # Set simple attributes 206 | self.pretrained = pretrained 207 | self.classify = classify 208 | self.num_classes = num_classes 209 | 210 | if pretrained == 'vggface2': 211 | tmp_classes = 8631 212 | elif pretrained == 'casia-webface': 213 | tmp_classes = 10575 214 | elif pretrained is None and self.classify and self.num_classes is None: 215 | raise Exception('If "pretrained" is not specified and "classify" is True, "num_classes" must be specified') 216 | 217 | 218 | # Define layers 219 | self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2) 220 | self.conv2d_2a = BasicConv2d(32, 32, kernel_size=3, stride=1) 221 | self.conv2d_2b = BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1) 222 | self.maxpool_3a = nn.MaxPool2d(3, stride=2) 223 | self.conv2d_3b = BasicConv2d(64, 80, kernel_size=1, stride=1) 224 | self.conv2d_4a = BasicConv2d(80, 192, kernel_size=3, stride=1) 225 | self.conv2d_4b = BasicConv2d(192, 256, kernel_size=3, stride=2) 226 | self.repeat_1 = nn.Sequential( 227 | Block35(scale=0.17), 228 | Block35(scale=0.17), 229 | Block35(scale=0.17), 230 | Block35(scale=0.17), 231 | Block35(scale=0.17), 232 | ) 233 | self.mixed_6a = Mixed_6a() 234 | self.repeat_2 = nn.Sequential( 235 | Block17(scale=0.10), 236 | Block17(scale=0.10), 237 | Block17(scale=0.10), 238 | Block17(scale=0.10), 239 | Block17(scale=0.10), 240 | Block17(scale=0.10), 241 | Block17(scale=0.10), 242 | Block17(scale=0.10), 243 | Block17(scale=0.10), 244 | Block17(scale=0.10), 245 | ) 246 | self.mixed_7a = Mixed_7a() 247 | self.repeat_3 = nn.Sequential( 248 | Block8(scale=0.20), 249 | Block8(scale=0.20), 250 | Block8(scale=0.20), 251 | Block8(scale=0.20), 252 | Block8(scale=0.20), 253 | ) 254 | self.block8 = Block8(noReLU=True) 255 | self.avgpool_1a = nn.AdaptiveAvgPool2d(1) 256 | self.dropout = nn.Dropout(dropout_prob) 257 | self.last_linear = nn.Linear(1792, 512, bias=False) 258 | self.last_bn = nn.BatchNorm1d(512, eps=0.001, momentum=0.1, affine=True) 259 | 260 | if pretrained is not None: 261 | self.logits = nn.Linear(512, tmp_classes) 262 | load_weights(self, pretrained) 263 | 264 | if self.classify and self.num_classes is not None: 265 | self.logits = nn.Linear(512, self.num_classes) 266 | 267 | self.device = torch.device('cpu') 268 | if device is not None: 269 | self.device = device 270 | self.to(device) 271 | 272 | def forward(self, x): 273 | """Calculate embeddings or logits given a batch of input image tensors. 274 | 275 | Arguments: 276 | x {torch.tensor} -- Batch of image tensors representing faces. 277 | 278 | Returns: 279 | torch.tensor -- Batch of embedding vectors or multinomial logits. 280 | """ 281 | x = self.conv2d_1a(x) 282 | x = self.conv2d_2a(x) 283 | x = self.conv2d_2b(x) 284 | x = self.maxpool_3a(x) 285 | x = self.conv2d_3b(x) 286 | x = self.conv2d_4a(x) 287 | x = self.conv2d_4b(x) 288 | x = self.repeat_1(x) 289 | x = self.mixed_6a(x) 290 | x = self.repeat_2(x) 291 | x = self.mixed_7a(x) 292 | x = self.repeat_3(x) 293 | x = self.block8(x) 294 | x = self.avgpool_1a(x) 295 | x = self.dropout(x) 296 | x = self.last_linear(x.view(x.shape[0], -1)) 297 | x = self.last_bn(x) 298 | if self.classify: 299 | x = self.logits(x) 300 | else: 301 | x = F.normalize(x, p=2, dim=1) 302 | return x 303 | 304 | 305 | def load_weights(mdl, name): 306 | """Download pretrained state_dict and load into model. 307 | 308 | Arguments: 309 | mdl {torch.nn.Module} -- Pytorch model. 310 | name {str} -- Name of dataset that was used to generate pretrained state_dict. 311 | 312 | Raises: 313 | ValueError: If 'pretrained' not equal to 'vggface2' or 'casia-webface'. 314 | """ 315 | if name == 'vggface2': 316 | path = 'https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180402-114759-vggface2.pt' 317 | elif name == 'casia-webface': 318 | path = 'https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180408-102900-casia-webface.pt' 319 | else: 320 | raise ValueError('Pretrained models only exist for "vggface2" and "casia-webface"') 321 | 322 | model_dir = os.path.join(get_torch_home(), 'checkpoints') 323 | os.makedirs(model_dir, exist_ok=True) 324 | 325 | cached_file = os.path.join(model_dir, os.path.basename(path)) 326 | if not os.path.exists(cached_file): 327 | download_url_to_file(path, cached_file) 328 | 329 | state_dict = torch.load(cached_file) 330 | mdl.load_state_dict(state_dict) 331 | 332 | 333 | def get_torch_home(): 334 | torch_home = os.path.expanduser( 335 | os.getenv( 336 | 'TORCH_HOME', 337 | os.path.join(os.getenv('XDG_CACHE_HOME', '~/.cache'), 'torch') 338 | ) 339 | ) 340 | return torch_home 341 | -------------------------------------------------------------------------------- /backend/ddf/facenet_pytorch/models/utils/detect_face.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn.functional import interpolate 3 | from torchvision.transforms import functional as F 4 | from torchvision.ops.boxes import batched_nms 5 | from PIL import Image 6 | import numpy as np 7 | import os 8 | import math 9 | 10 | # OpenCV is optional, but required if using numpy arrays instead of PIL 11 | try: 12 | import cv2 13 | except: 14 | pass 15 | 16 | def fixed_batch_process(im_data, model): 17 | batch_size = 512 18 | out = [] 19 | for i in range(0, len(im_data), batch_size): 20 | batch = im_data[i:(i+batch_size)] 21 | out.append(model(batch)) 22 | 23 | return tuple(torch.cat(v, dim=0) for v in zip(*out)) 24 | 25 | def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): 26 | if isinstance(imgs, (np.ndarray, torch.Tensor)): 27 | if isinstance(imgs,np.ndarray): 28 | imgs = torch.as_tensor(imgs.copy(), device=device) 29 | 30 | if isinstance(imgs,torch.Tensor): 31 | imgs = torch.as_tensor(imgs, device=device) 32 | 33 | if len(imgs.shape) == 3: 34 | imgs = imgs.unsqueeze(0) 35 | else: 36 | if not isinstance(imgs, (list, tuple)): 37 | imgs = [imgs] 38 | if any(img.size != imgs[0].size for img in imgs): 39 | raise Exception("MTCNN batch processing only compatible with equal-dimension images.") 40 | imgs = np.stack([np.uint8(img) for img in imgs]) 41 | imgs = torch.as_tensor(imgs.copy(), device=device) 42 | 43 | 44 | 45 | model_dtype = next(pnet.parameters()).dtype 46 | imgs = imgs.permute(0, 3, 1, 2).type(model_dtype) 47 | 48 | batch_size = len(imgs) 49 | h, w = imgs.shape[2:4] 50 | m = 12.0 / minsize 51 | minl = min(h, w) 52 | minl = minl * m 53 | 54 | # Create scale pyramid 55 | scale_i = m 56 | scales = [] 57 | while minl >= 12: 58 | scales.append(scale_i) 59 | scale_i = scale_i * factor 60 | minl = minl * factor 61 | 62 | # First stage 63 | boxes = [] 64 | image_inds = [] 65 | 66 | scale_picks = [] 67 | 68 | all_i = 0 69 | offset = 0 70 | for scale in scales: 71 | im_data = imresample(imgs, (int(h * scale + 1), int(w * scale + 1))) 72 | im_data = (im_data - 127.5) * 0.0078125 73 | reg, probs = pnet(im_data) 74 | 75 | boxes_scale, image_inds_scale = generateBoundingBox(reg, probs[:, 1], scale, threshold[0]) 76 | boxes.append(boxes_scale) 77 | image_inds.append(image_inds_scale) 78 | 79 | pick = batched_nms(boxes_scale[:, :4], boxes_scale[:, 4], image_inds_scale, 0.5) 80 | scale_picks.append(pick + offset) 81 | offset += boxes_scale.shape[0] 82 | 83 | boxes = torch.cat(boxes, dim=0) 84 | image_inds = torch.cat(image_inds, dim=0) 85 | 86 | scale_picks = torch.cat(scale_picks, dim=0) 87 | 88 | # NMS within each scale + image 89 | boxes, image_inds = boxes[scale_picks], image_inds[scale_picks] 90 | 91 | 92 | # NMS within each image 93 | pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7) 94 | boxes, image_inds = boxes[pick], image_inds[pick] 95 | 96 | regw = boxes[:, 2] - boxes[:, 0] 97 | regh = boxes[:, 3] - boxes[:, 1] 98 | qq1 = boxes[:, 0] + boxes[:, 5] * regw 99 | qq2 = boxes[:, 1] + boxes[:, 6] * regh 100 | qq3 = boxes[:, 2] + boxes[:, 7] * regw 101 | qq4 = boxes[:, 3] + boxes[:, 8] * regh 102 | boxes = torch.stack([qq1, qq2, qq3, qq4, boxes[:, 4]]).permute(1, 0) 103 | boxes = rerec(boxes) 104 | y, ey, x, ex = pad(boxes, w, h) 105 | 106 | # Second stage 107 | if len(boxes) > 0: 108 | im_data = [] 109 | for k in range(len(y)): 110 | if ey[k] > (y[k] - 1) and ex[k] > (x[k] - 1): 111 | img_k = imgs[image_inds[k], :, (y[k] - 1):ey[k], (x[k] - 1):ex[k]].unsqueeze(0) 112 | im_data.append(imresample(img_k, (24, 24))) 113 | im_data = torch.cat(im_data, dim=0) 114 | im_data = (im_data - 127.5) * 0.0078125 115 | 116 | # This is equivalent to out = rnet(im_data) to avoid GPU out of memory. 117 | out = fixed_batch_process(im_data, rnet) 118 | 119 | out0 = out[0].permute(1, 0) 120 | out1 = out[1].permute(1, 0) 121 | score = out1[1, :] 122 | ipass = score > threshold[1] 123 | boxes = torch.cat((boxes[ipass, :4], score[ipass].unsqueeze(1)), dim=1) 124 | image_inds = image_inds[ipass] 125 | mv = out0[:, ipass].permute(1, 0) 126 | 127 | # NMS within each image 128 | pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7) 129 | boxes, image_inds, mv = boxes[pick], image_inds[pick], mv[pick] 130 | boxes = bbreg(boxes, mv) 131 | boxes = rerec(boxes) 132 | 133 | # Third stage 134 | points = torch.zeros(0, 5, 2, device=device) 135 | if len(boxes) > 0: 136 | y, ey, x, ex = pad(boxes, w, h) 137 | im_data = [] 138 | for k in range(len(y)): 139 | if ey[k] > (y[k] - 1) and ex[k] > (x[k] - 1): 140 | img_k = imgs[image_inds[k], :, (y[k] - 1):ey[k], (x[k] - 1):ex[k]].unsqueeze(0) 141 | im_data.append(imresample(img_k, (48, 48))) 142 | im_data = torch.cat(im_data, dim=0) 143 | im_data = (im_data - 127.5) * 0.0078125 144 | 145 | # This is equivalent to out = onet(im_data) to avoid GPU out of memory. 146 | out = fixed_batch_process(im_data, onet) 147 | 148 | out0 = out[0].permute(1, 0) 149 | out1 = out[1].permute(1, 0) 150 | out2 = out[2].permute(1, 0) 151 | score = out2[1, :] 152 | points = out1 153 | ipass = score > threshold[2] 154 | points = points[:, ipass] 155 | boxes = torch.cat((boxes[ipass, :4], score[ipass].unsqueeze(1)), dim=1) 156 | image_inds = image_inds[ipass] 157 | mv = out0[:, ipass].permute(1, 0) 158 | 159 | w_i = boxes[:, 2] - boxes[:, 0] + 1 160 | h_i = boxes[:, 3] - boxes[:, 1] + 1 161 | points_x = w_i.repeat(5, 1) * points[:5, :] + boxes[:, 0].repeat(5, 1) - 1 162 | points_y = h_i.repeat(5, 1) * points[5:10, :] + boxes[:, 1].repeat(5, 1) - 1 163 | points = torch.stack((points_x, points_y)).permute(2, 1, 0) 164 | boxes = bbreg(boxes, mv) 165 | 166 | # NMS within each image using "Min" strategy 167 | # pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7) 168 | pick = batched_nms_numpy(boxes[:, :4], boxes[:, 4], image_inds, 0.7, 'Min') 169 | boxes, image_inds, points = boxes[pick], image_inds[pick], points[pick] 170 | 171 | boxes = boxes.cpu().numpy() 172 | points = points.cpu().numpy() 173 | 174 | image_inds = image_inds.cpu() 175 | 176 | batch_boxes = [] 177 | batch_points = [] 178 | for b_i in range(batch_size): 179 | b_i_inds = np.where(image_inds == b_i) 180 | batch_boxes.append(boxes[b_i_inds].copy()) 181 | batch_points.append(points[b_i_inds].copy()) 182 | 183 | batch_boxes, batch_points = np.array(batch_boxes, dtype=object), np.array(batch_points, dtype=object) 184 | 185 | return batch_boxes, batch_points 186 | 187 | 188 | def bbreg(boundingbox, reg): 189 | if reg.shape[1] == 1: 190 | reg = torch.reshape(reg, (reg.shape[2], reg.shape[3])) 191 | 192 | w = boundingbox[:, 2] - boundingbox[:, 0] + 1 193 | h = boundingbox[:, 3] - boundingbox[:, 1] + 1 194 | b1 = boundingbox[:, 0] + reg[:, 0] * w 195 | b2 = boundingbox[:, 1] + reg[:, 1] * h 196 | b3 = boundingbox[:, 2] + reg[:, 2] * w 197 | b4 = boundingbox[:, 3] + reg[:, 3] * h 198 | boundingbox[:, :4] = torch.stack([b1, b2, b3, b4]).permute(1, 0) 199 | 200 | return boundingbox 201 | 202 | 203 | def generateBoundingBox(reg, probs, scale, thresh): 204 | stride = 2 205 | cellsize = 12 206 | 207 | reg = reg.permute(1, 0, 2, 3) 208 | 209 | mask = probs >= thresh 210 | mask_inds = mask.nonzero() 211 | image_inds = mask_inds[:, 0] 212 | score = probs[mask] 213 | reg = reg[:, mask].permute(1, 0) 214 | bb = mask_inds[:, 1:].type(reg.dtype).flip(1) 215 | q1 = ((stride * bb + 1) / scale).floor() 216 | q2 = ((stride * bb + cellsize - 1 + 1) / scale).floor() 217 | boundingbox = torch.cat([q1, q2, score.unsqueeze(1), reg], dim=1) 218 | return boundingbox, image_inds 219 | 220 | 221 | def nms_numpy(boxes, scores, threshold, method): 222 | if boxes.size == 0: 223 | return np.empty((0, 3)) 224 | 225 | x1 = boxes[:, 0].copy() 226 | y1 = boxes[:, 1].copy() 227 | x2 = boxes[:, 2].copy() 228 | y2 = boxes[:, 3].copy() 229 | s = scores 230 | area = (x2 - x1 + 1) * (y2 - y1 + 1) 231 | 232 | I = np.argsort(s) 233 | pick = np.zeros_like(s, dtype=np.int16) 234 | counter = 0 235 | while I.size > 0: 236 | i = I[-1] 237 | pick[counter] = i 238 | counter += 1 239 | idx = I[0:-1] 240 | 241 | xx1 = np.maximum(x1[i], x1[idx]).copy() 242 | yy1 = np.maximum(y1[i], y1[idx]).copy() 243 | xx2 = np.minimum(x2[i], x2[idx]).copy() 244 | yy2 = np.minimum(y2[i], y2[idx]).copy() 245 | 246 | w = np.maximum(0.0, xx2 - xx1 + 1).copy() 247 | h = np.maximum(0.0, yy2 - yy1 + 1).copy() 248 | 249 | inter = w * h 250 | if method == 'Min': 251 | o = inter / np.minimum(area[i], area[idx]) 252 | else: 253 | o = inter / (area[i] + area[idx] - inter) 254 | I = I[np.where(o <= threshold)] 255 | 256 | pick = pick[:counter].copy() 257 | return pick 258 | 259 | 260 | def batched_nms_numpy(boxes, scores, idxs, threshold, method): 261 | device = boxes.device 262 | if boxes.numel() == 0: 263 | return torch.empty((0,), dtype=torch.int64, device=device) 264 | # strategy: in order to perform NMS independently per class. 265 | # we add an offset to all the boxes. The offset is dependent 266 | # only on the class idx, and is large enough so that boxes 267 | # from different classes do not overlap 268 | max_coordinate = boxes.max() 269 | offsets = idxs.to(boxes) * (max_coordinate + 1) 270 | boxes_for_nms = boxes + offsets[:, None] 271 | boxes_for_nms = boxes_for_nms.cpu().numpy() 272 | scores = scores.cpu().numpy() 273 | keep = nms_numpy(boxes_for_nms, scores, threshold, method) 274 | return torch.as_tensor(keep, dtype=torch.long, device=device) 275 | 276 | 277 | def pad(boxes, w, h): 278 | boxes = boxes.trunc().int().cpu().numpy() 279 | x = boxes[:, 0] 280 | y = boxes[:, 1] 281 | ex = boxes[:, 2] 282 | ey = boxes[:, 3] 283 | 284 | x[x < 1] = 1 285 | y[y < 1] = 1 286 | ex[ex > w] = w 287 | ey[ey > h] = h 288 | 289 | return y, ey, x, ex 290 | 291 | 292 | def rerec(bboxA): 293 | h = bboxA[:, 3] - bboxA[:, 1] 294 | w = bboxA[:, 2] - bboxA[:, 0] 295 | 296 | l = torch.max(w, h) 297 | bboxA[:, 0] = bboxA[:, 0] + w * 0.5 - l * 0.5 298 | bboxA[:, 1] = bboxA[:, 1] + h * 0.5 - l * 0.5 299 | bboxA[:, 2:4] = bboxA[:, :2] + l.repeat(2, 1).permute(1, 0) 300 | 301 | return bboxA 302 | 303 | 304 | def imresample(img, sz): 305 | im_data = interpolate(img, size=sz, mode="area") 306 | return im_data 307 | 308 | 309 | def crop_resize(img, box, image_size): 310 | if isinstance(img, np.ndarray): 311 | img = img[box[1]:box[3], box[0]:box[2]] 312 | out = cv2.resize( 313 | img, 314 | (image_size, image_size), 315 | interpolation=cv2.INTER_AREA 316 | ).copy() 317 | elif isinstance(img, torch.Tensor): 318 | img = img[box[1]:box[3], box[0]:box[2]] 319 | out = imresample( 320 | img.permute(2, 0, 1).unsqueeze(0).float(), 321 | (image_size, image_size) 322 | ).byte().squeeze(0).permute(1, 2, 0) 323 | else: 324 | out = img.crop(box).copy().resize((image_size, image_size), Image.BILINEAR) 325 | return out 326 | 327 | 328 | def save_img(img, path): 329 | if isinstance(img, np.ndarray): 330 | cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) 331 | else: 332 | img.save(path) 333 | 334 | 335 | def get_size(img): 336 | if isinstance(img, (np.ndarray, torch.Tensor)): 337 | return img.shape[1::-1] 338 | else: 339 | return img.size 340 | 341 | 342 | def extract_face(img, box, image_size=160, margin=0, save_path=None): 343 | """Extract face + margin from PIL Image given bounding box. 344 | 345 | Arguments: 346 | img {PIL.Image} -- A PIL Image. 347 | box {numpy.ndarray} -- Four-element bounding box. 348 | image_size {int} -- Output image size in pixels. The image will be square. 349 | margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. 350 | Note that the application of the margin differs slightly from the davidsandberg/facenet 351 | repo, which applies the margin to the original image before resizing, making the margin 352 | dependent on the original image size. 353 | save_path {str} -- Save path for extracted face image. (default: {None}) 354 | 355 | Returns: 356 | torch.tensor -- tensor representing the extracted face. 357 | """ 358 | margin = [ 359 | margin * (box[2] - box[0]) / (image_size - margin), 360 | margin * (box[3] - box[1]) / (image_size - margin), 361 | ] 362 | raw_image_size = get_size(img) 363 | box = [ 364 | int(max(box[0] - margin[0] / 2, 0)), 365 | int(max(box[1] - margin[1] / 2, 0)), 366 | int(min(box[2] + margin[0] / 2, raw_image_size[0])), 367 | int(min(box[3] + margin[1] / 2, raw_image_size[1])), 368 | ] 369 | 370 | face = crop_resize(img, box, image_size) 371 | 372 | if save_path is not None: 373 | os.makedirs(os.path.dirname(save_path) + "/", exist_ok=True) 374 | save_img(face, save_path) 375 | 376 | if isinstance(face, np.ndarray) or isinstance(face, Image.Image): 377 | face = F.to_tensor(np.float32(face)) 378 | elif isinstance(face, torch.Tensor): 379 | face = face.float() 380 | else: 381 | raise NotImplementedError 382 | 383 | return face 384 | -------------------------------------------------------------------------------- /backend/ddf/facenet_pytorch/models/utils/download.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import os 3 | import shutil 4 | import sys 5 | import tempfile 6 | 7 | from urllib.request import urlopen, Request 8 | 9 | try: 10 | from tqdm.auto import tqdm # automatically select proper tqdm submodule if available 11 | except ImportError: 12 | from tqdm import tqdm 13 | 14 | 15 | def download_url_to_file(url, dst, hash_prefix=None, progress=True): 16 | r"""Download object at the given URL to a local path. 17 | Args: 18 | url (string): URL of the object to download 19 | dst (string): Full path where object will be saved, e.g. `/tmp/temporary_file` 20 | hash_prefix (string, optional): If not None, the SHA256 downloaded file should start with `hash_prefix`. 21 | Default: None 22 | progress (bool, optional): whether or not to display a progress bar to stderr 23 | Default: True 24 | Example: 25 | >>> torch.hub.download_url_to_file('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth', '/tmp/temporary_file') 26 | """ 27 | file_size = None 28 | # We use a different API for python2 since urllib(2) doesn't recognize the CA 29 | # certificates in older Python 30 | req = Request(url, headers={"User-Agent": "torch.hub"}) 31 | u = urlopen(req) 32 | meta = u.info() 33 | if hasattr(meta, 'getheaders'): 34 | content_length = meta.getheaders("Content-Length") 35 | else: 36 | content_length = meta.get_all("Content-Length") 37 | if content_length is not None and len(content_length) > 0: 38 | file_size = int(content_length[0]) 39 | 40 | # We deliberately save it in a temp file and move it after 41 | # download is complete. This prevents a local working checkpoint 42 | # being overridden by a broken download. 43 | dst = os.path.expanduser(dst) 44 | dst_dir = os.path.dirname(dst) 45 | f = tempfile.NamedTemporaryFile(delete=False, dir=dst_dir) 46 | 47 | try: 48 | if hash_prefix is not None: 49 | sha256 = hashlib.sha256() 50 | with tqdm(total=file_size, disable=not progress, 51 | unit='B', unit_scale=True, unit_divisor=1024) as pbar: 52 | while True: 53 | buffer = u.read(8192) 54 | if len(buffer) == 0: 55 | break 56 | f.write(buffer) 57 | if hash_prefix is not None: 58 | sha256.update(buffer) 59 | pbar.update(len(buffer)) 60 | 61 | f.close() 62 | if hash_prefix is not None: 63 | digest = sha256.hexdigest() 64 | if digest[:len(hash_prefix)] != hash_prefix: 65 | raise RuntimeError('invalid hash value (expected "{}", got "{}")' 66 | .format(hash_prefix, digest)) 67 | shutil.move(f.name, dst) 68 | finally: 69 | f.close() 70 | if os.path.exists(f.name): 71 | os.remove(f.name) 72 | -------------------------------------------------------------------------------- /backend/ddf/facenet_pytorch/models/utils/training.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import time 4 | 5 | 6 | class Logger(object): 7 | 8 | def __init__(self, mode, length, calculate_mean=False): 9 | self.mode = mode 10 | self.length = length 11 | self.calculate_mean = calculate_mean 12 | if self.calculate_mean: 13 | self.fn = lambda x, i: x / (i + 1) 14 | else: 15 | self.fn = lambda x, i: x 16 | 17 | def __call__(self, loss, metrics, i): 18 | track_str = '\r{} | {:5d}/{:<5d}| '.format(self.mode, i + 1, self.length) 19 | loss_str = 'loss: {:9.4f} | '.format(self.fn(loss, i)) 20 | metric_str = ' | '.join('{}: {:9.4f}'.format(k, self.fn(v, i)) for k, v in metrics.items()) 21 | print(track_str + loss_str + metric_str + ' ', end='') 22 | if i + 1 == self.length: 23 | print('') 24 | 25 | 26 | class BatchTimer(object): 27 | """Batch timing class. 28 | Use this class for tracking training and testing time/rate per batch or per sample. 29 | 30 | Keyword Arguments: 31 | rate {bool} -- Whether to report a rate (batches or samples per second) or a time (seconds 32 | per batch or sample). (default: {True}) 33 | per_sample {bool} -- Whether to report times or rates per sample or per batch. 34 | (default: {True}) 35 | """ 36 | 37 | def __init__(self, rate=True, per_sample=True): 38 | self.start = time.time() 39 | self.end = None 40 | self.rate = rate 41 | self.per_sample = per_sample 42 | 43 | def __call__(self, y_pred, y): 44 | self.end = time.time() 45 | elapsed = self.end - self.start 46 | self.start = self.end 47 | self.end = None 48 | 49 | if self.per_sample: 50 | elapsed /= len(y_pred) 51 | if self.rate: 52 | elapsed = 1 / elapsed 53 | 54 | return torch.tensor(elapsed) 55 | 56 | 57 | def accuracy(logits, y): 58 | _, preds = torch.max(logits, 1) 59 | return (preds == y).float().mean() 60 | 61 | 62 | def pass_epoch( 63 | model, loss_fn, loader, optimizer=None, scheduler=None, 64 | batch_metrics={'time': BatchTimer()}, show_running=True, 65 | device='cpu', writer=None 66 | ): 67 | """Train or evaluate over a data epoch. 68 | 69 | Arguments: 70 | model {torch.nn.Module} -- Pytorch model. 71 | loss_fn {callable} -- A function to compute (scalar) loss. 72 | loader {torch.utils.data.DataLoader} -- A pytorch data loader. 73 | 74 | Keyword Arguments: 75 | optimizer {torch.optim.Optimizer} -- A pytorch optimizer. 76 | scheduler {torch.optim.lr_scheduler._LRScheduler} -- LR scheduler (default: {None}) 77 | batch_metrics {dict} -- Dictionary of metric functions to call on each batch. The default 78 | is a simple timer. A progressive average of these metrics, along with the average 79 | loss, is printed every batch. (default: {{'time': iter_timer()}}) 80 | show_running {bool} -- Whether or not to print losses and metrics for the current batch 81 | or rolling averages. (default: {False}) 82 | device {str or torch.device} -- Device for pytorch to use. (default: {'cpu'}) 83 | writer {torch.utils.tensorboard.SummaryWriter} -- Tensorboard SummaryWriter. (default: {None}) 84 | 85 | Returns: 86 | tuple(torch.Tensor, dict) -- A tuple of the average loss and a dictionary of average 87 | metric values across the epoch. 88 | """ 89 | 90 | mode = 'Train' if model.training else 'Valid' 91 | logger = Logger(mode, length=len(loader), calculate_mean=show_running) 92 | loss = 0 93 | metrics = {} 94 | 95 | for i_batch, (x, y) in enumerate(loader): 96 | x = x.to(device) 97 | y = y.to(device) 98 | y_pred = model(x) 99 | loss_batch = loss_fn(y_pred, y) 100 | 101 | if model.training: 102 | loss_batch.backward() 103 | optimizer.step() 104 | optimizer.zero_grad() 105 | 106 | metrics_batch = {} 107 | for metric_name, metric_fn in batch_metrics.items(): 108 | metrics_batch[metric_name] = metric_fn(y_pred, y).detach().cpu() 109 | metrics[metric_name] = metrics.get(metric_name, 0) + metrics_batch[metric_name] 110 | 111 | if writer is not None and model.training: 112 | if writer.iteration % writer.interval == 0: 113 | writer.add_scalars('loss', {mode: loss_batch.detach().cpu()}, writer.iteration) 114 | for metric_name, metric_batch in metrics_batch.items(): 115 | writer.add_scalars(metric_name, {mode: metric_batch}, writer.iteration) 116 | writer.iteration += 1 117 | 118 | loss_batch = loss_batch.detach().cpu() 119 | loss += loss_batch 120 | if show_running: 121 | logger(loss, metrics, i_batch) 122 | else: 123 | logger(loss_batch, metrics_batch, i_batch) 124 | 125 | if model.training and scheduler is not None: 126 | scheduler.step() 127 | 128 | loss = loss / (i_batch + 1) 129 | metrics = {k: v / (i_batch + 1) for k, v in metrics.items()} 130 | 131 | if writer is not None and not model.training: 132 | writer.add_scalars('loss', {mode: loss.detach()}, writer.iteration) 133 | for metric_name, metric in metrics.items(): 134 | writer.add_scalars(metric_name, {mode: metric}) 135 | 136 | return loss, metrics 137 | 138 | 139 | def collate_pil(x): 140 | out_x, out_y = [], [] 141 | for xx, yy in x: 142 | out_x.append(xx) 143 | out_y.append(yy) 144 | return out_x, out_y 145 | -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/finallyupper/deepsafe/5379bed6e70214b3a9798f476d888d24e7341437/backend/ddf/faceswap_pytorch/__init__.py -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/image_augmentation.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy 3 | import torch 4 | from faceswap_pytorch.umeyama import umeyama 5 | 6 | random_transform_args = { 7 | 'rotation_range': 10, 8 | 'zoom_range': 0.05, 9 | 'shift_range': 0.05, 10 | 'random_flip': 0.4, 11 | } 12 | 13 | 14 | def random_transform(image, rotation_range, zoom_range, shift_range, random_flip): 15 | #Added(Yoojin) 16 | if isinstance(image, torch.Tensor): 17 | image = image.detach().cpu().permute(1, 2, 0).cpu().numpy() # (C, H, W) -> (H, W, C) 18 | 19 | h, w = image.shape[0:2] 20 | rotation = numpy.random.uniform(-rotation_range, rotation_range) 21 | scale = numpy.random.uniform(1 - zoom_range, 1 + zoom_range) 22 | tx = numpy.random.uniform(-shift_range, shift_range) * w 23 | ty = numpy.random.uniform(-shift_range, shift_range) * h 24 | mat = cv2.getRotationMatrix2D((w // 2, h // 2), rotation, scale) 25 | mat[:, 2] += (tx, ty) 26 | result = cv2.warpAffine(image, mat, (w, h), borderMode=cv2.BORDER_REPLICATE) 27 | if numpy.random.random() < random_flip: 28 | result = result[:, ::-1] 29 | return result 30 | 31 | # def random_transform(image, rotation_range, zoom_range, shift_range, random_flip): 32 | # if isinstance(image, torch.Tensor): 33 | # image = image.detach().cpu().permute(1, 2, 0).cpu().numpy() # Convert tensor to numpy 34 | 35 | # h, w = image.shape[0:2] 36 | # rotation = numpy.random.uniform(-rotation_range, rotation_range) 37 | # scale = numpy.random.uniform(1 - zoom_range, 1 + zoom_range) 38 | # tx = numpy.random.uniform(-shift_range, shift_range) * w 39 | # ty = numpy.random.uniform(-shift_range, shift_range) * h 40 | # mat = cv2.getRotationMatrix2D((w // 2, h // 2), rotation, scale) 41 | # mat[:, 2] += (tx, ty) 42 | # result = cv2.warpAffine(image, mat, (w, h), borderMode=cv2.BORDER_REPLICATE) 43 | # if numpy.random.random() < random_flip: 44 | # result = result[:, ::-1] 45 | # return result 46 | 47 | 48 | def random_warp(image): 49 | assert image.shape == (256, 256, 3) # Ensure input is 256x256 50 | 51 | # Generate random displacement maps 52 | range_ = numpy.linspace(128 - 80, 128 + 80, 5) # Grid points 53 | mapx = numpy.broadcast_to(range_, (5, 5)) 54 | mapy = mapx.T 55 | 56 | # Add random noise to displacement maps 57 | mapx = mapx + numpy.random.normal(size=(5, 5), scale=5) 58 | mapy = mapy + numpy.random.normal(size=(5, 5), scale=5) 59 | 60 | # Interpolate the displacement maps to the desired resolution (160x160) 61 | interp_mapx = cv2.resize(mapx, (160, 160)).astype('float32') # Changed size to 160x160 62 | interp_mapy = cv2.resize(mapy, (160, 160)).astype('float32') # Changed size to 160x160 63 | 64 | # Apply remapping to generate the warped image 65 | warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR) 66 | 67 | # Define source and destination points for affine transformation 68 | src_points = numpy.stack([mapx.ravel(), mapy.ravel()], axis=-1) 69 | dst_points = numpy.mgrid[0:129:32, 0:129:32].T.reshape(-1, 2) # Adjusted for 160x160 70 | 71 | # Compute affine transformation matrix 72 | mat = umeyama(src_points, dst_points, True)[0:2] 73 | 74 | # Apply affine transformation to produce the target image 75 | target_image = cv2.warpAffine(image, mat, (160, 160)) # Resize to 160x160 76 | 77 | return warped_image, target_image 78 | 79 | 80 | # # get pair of random warped images from aligened face image 81 | # def random_warp(image): 82 | # assert image.shape == (256, 256, 3) 83 | # range_ = numpy.linspace(128 - 80, 128 + 80, 5) 84 | # mapx = numpy.broadcast_to(range_, (5, 5)) 85 | # mapy = mapx.T 86 | 87 | # mapx = mapx + numpy.random.normal(size=(5, 5), scale=5) 88 | # mapy = mapy + numpy.random.normal(size=(5, 5), scale=5) 89 | 90 | # interp_mapx = cv2.resize(mapx, (80, 80))[8:72, 8:72].astype('float32') 91 | # interp_mapy = cv2.resize(mapy, (80, 80))[8:72, 8:72].astype('float32') 92 | 93 | # warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR) 94 | 95 | # src_points = numpy.stack([mapx.ravel(), mapy.ravel()], axis=-1) 96 | # dst_points = numpy.mgrid[0:65:16, 0:65:16].T.reshape(-1, 2) 97 | # mat = umeyama(src_points, dst_points, True)[0:2] 98 | 99 | # target_image = cv2.warpAffine(image, mat, (64, 64)) 100 | 101 | # return warped_image, target_image 102 | -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.utils.data 3 | from torch import nn, optim 4 | from ddf.faceswap_pytorch.padding_same_conv import Conv2d 5 | 6 | def toTensor(img): 7 | img = torch.from_numpy(img.transpose((0, 3, 1, 2))) 8 | return img 9 | 10 | 11 | def var_to_np(img_var): 12 | return img_var.data.cpu().numpy() 13 | 14 | 15 | class _ConvLayer(nn.Sequential): 16 | def __init__(self, input_features, output_features): 17 | super(_ConvLayer, self).__init__() 18 | self.add_module('conv2', Conv2d(input_features, 19 | output_features, 20 | kernel_size=5, 21 | stride=2)) 22 | self.add_module('leakyrelu', nn.LeakyReLU(0.1, inplace=True)) 23 | 24 | 25 | class _UpScale(nn.Sequential): 26 | def __init__(self, input_features, output_features): # 1024, 512 27 | super(_UpScale, self).__init__() 28 | self.add_module('conv2_', Conv2d(input_features, output_features * 4, 29 | kernel_size=3)) 30 | self.add_module('leakyrelu', nn.LeakyReLU(0.1, inplace=True)) 31 | self.add_module('pixelshuffler', _PixelShuffler()) 32 | 33 | 34 | class Flatten(nn.Module): 35 | def forward(self, input): # 16x1024x16x16 36 | output = input.reshape(input.size(0), -1) # Modified(Yoojin): output = input.view(input.size(0), -1) 37 | #output = input.view(input.size(0), -1) 38 | return output # 16x(1024*16*16) 39 | 40 | 41 | class Reshape(nn.Module): 42 | def forward(self, input): 43 | output = input.view(-1, 1024, 10, 10) # channel * 4 * 4 44 | return output 45 | 46 | 47 | class _PixelShuffler(nn.Module): 48 | def forward(self, input): 49 | batch_size, c, h, w = input.size() 50 | rh, rw = (2, 2) 51 | oh, ow = h * rh, w * rw 52 | oc = c // (rh * rw) 53 | out = input.view(batch_size, rh, rw, oc, h, w) 54 | out = out.permute(0, 3, 4, 1, 5, 2).contiguous() 55 | out = out.view(batch_size, oc, oh, ow) # channel first 56 | 57 | return out 58 | 59 | 60 | class Autoencoder(nn.Module): 61 | def __init__(self): 62 | super(Autoencoder, self).__init__() 63 | 64 | self.encoder = nn.Sequential( 65 | _ConvLayer(3, 128), 66 | _ConvLayer(128, 256), 67 | _ConvLayer(256, 512), 68 | _ConvLayer(512, 1024), 69 | Flatten(), 70 | nn.Linear(1024 * 10 * 10,1024), 71 | nn.Linear(1024,1024 * 10 * 10), 72 | Reshape(), 73 | _UpScale(1024, 512), 74 | ) 75 | 76 | self.decoder_A = nn.Sequential( 77 | _UpScale(512, 256), 78 | _UpScale(256, 128), 79 | _UpScale(128, 64), 80 | Conv2d(64, 3, kernel_size=5, padding=1), 81 | nn.Sigmoid(), 82 | ) 83 | 84 | self.decoder_B = nn.Sequential( 85 | _UpScale(512, 256), 86 | _UpScale(256, 128), 87 | _UpScale(128, 64), 88 | Conv2d(64, 3, kernel_size=5, padding=1), 89 | nn.Sigmoid(), 90 | ) 91 | 92 | def forward(self, x, select='A'): 93 | #print(f'Input shape: {x.shape}') 94 | for i, layer in enumerate(self.encoder): 95 | x = layer(x) 96 | #print(f"Shape after encoder layer {i + 1}: {x.shape}") 97 | 98 | # Decoder 99 | if select == 'A': 100 | for i, layer in enumerate(self.decoder_A): 101 | x = layer(x) 102 | # print(f"Shape after decoder_A layer {i + 1}: {x.shape}") 103 | else: 104 | for i, layer in enumerate(self.decoder_B): 105 | x = layer(x) 106 | # print(f"Shape after decoder_B layer {i + 1}: {x.shape}") 107 | return x 108 | 109 | class Autoencoder_df(nn.Module): 110 | def __init__(self): 111 | super(Autoencoder_df, self).__init__() 112 | 113 | self.encoder = nn.Sequential( 114 | _ConvLayer(3, 128), 115 | _ConvLayer(128, 256), 116 | _ConvLayer(256, 512), 117 | _ConvLayer(512, 1024), 118 | Flatten(), 119 | nn.Linear(1024 * 10 * 10, 1024), 120 | nn.Linear(1024, 1024 * 10 * 10), 121 | Reshape(), 122 | _UpScale(1024, 512), 123 | ) 124 | 125 | self.decoder_A = nn.Sequential( 126 | _UpScale(512, 256), 127 | _UpScale(256, 128), 128 | _UpScale(128, 64), 129 | Conv2d(64, 3, kernel_size=5, padding=1), 130 | nn.Sigmoid(), 131 | ) 132 | 133 | self.decoder_B = nn.Sequential( 134 | _UpScale(512, 256), 135 | _UpScale(256, 128), 136 | _UpScale(128, 64), 137 | Conv2d(64, 3, kernel_size=5, padding=1), 138 | nn.Sigmoid(), 139 | ) 140 | 141 | def forward(self, x, select='A'): 142 | # Encoder 143 | encoded = self.encoder(x) 144 | 145 | # Decoder 146 | if select == 'A': 147 | for i, layer in enumerate(self.decoder_A): 148 | encoded = layer(encoded) 149 | # Extract k-1 feature map (last before Sigmoid) 150 | if i == len(self.decoder_A) - 2: 151 | k_minus_1_feature_map = encoded 152 | else: 153 | for i, layer in enumerate(self.decoder_B): 154 | encoded = layer(encoded) 155 | # Extract k-1 feature map (last before Sigmoid) 156 | if i == len(self.decoder_B) - 2: 157 | k_minus_1_feature_map = encoded 158 | 159 | # Final output (after Sigmoid) 160 | swapped_image = encoded 161 | return x, k_minus_1_feature_map, swapped_image #CHECK(Yoojin) -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/padding_same_conv.py: -------------------------------------------------------------------------------- 1 | # modify con2d function to use same padding 2 | # code referd to @famssa in 'https://github.com/pytorch/pytorch/issues/3867' 3 | # and tensorflow source code 4 | 5 | import torch.utils.data 6 | from torch.nn import functional as F 7 | 8 | import math 9 | import torch 10 | from torch.nn.parameter import Parameter 11 | from torch.nn.functional import pad 12 | from torch.nn.modules import Module 13 | from torch.nn.modules.utils import _single, _pair, _triple 14 | 15 | 16 | class _ConvNd(Module): 17 | 18 | def __init__(self, in_channels, out_channels, kernel_size, stride, 19 | padding, dilation, transposed, output_padding, groups, bias): 20 | super(_ConvNd, self).__init__() 21 | if in_channels % groups != 0: 22 | raise ValueError('in_channels must be divisible by groups') 23 | if out_channels % groups != 0: 24 | raise ValueError('out_channels must be divisible by groups') 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.kernel_size = kernel_size 28 | self.stride = stride 29 | self.padding = padding 30 | self.dilation = dilation 31 | self.transposed = transposed 32 | self.output_padding = output_padding 33 | self.groups = groups 34 | if transposed: 35 | self.weight = Parameter(torch.Tensor( 36 | in_channels, out_channels // groups, *kernel_size)) 37 | else: 38 | self.weight = Parameter(torch.Tensor( 39 | out_channels, in_channels // groups, *kernel_size)) 40 | if bias: 41 | self.bias = Parameter(torch.Tensor(out_channels)) 42 | else: 43 | self.register_parameter('bias', None) 44 | self.reset_parameters() 45 | 46 | def reset_parameters(self): 47 | n = self.in_channels 48 | for k in self.kernel_size: 49 | n *= k 50 | stdv = 1. / math.sqrt(n) 51 | self.weight.data.uniform_(-stdv, stdv) 52 | if self.bias is not None: 53 | self.bias.data.uniform_(-stdv, stdv) 54 | 55 | def __repr__(self): 56 | s = ('{name}({in_channels}, {out_channels}, kernel_size={kernel_size}' 57 | ', stride={stride}') 58 | if self.padding != (0,) * len(self.padding): 59 | s += ', padding={padding}' 60 | if self.dilation != (1,) * len(self.dilation): 61 | s += ', dilation={dilation}' 62 | if self.output_padding != (0,) * len(self.output_padding): 63 | s += ', output_padding={output_padding}' 64 | if self.groups != 1: 65 | s += ', groups={groups}' 66 | if self.bias is None: 67 | s += ', bias=False' 68 | s += ')' 69 | return s.format(name=self.__class__.__name__, **self.__dict__) 70 | 71 | 72 | class Conv2d(_ConvNd): 73 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, 74 | padding=0, dilation=1, groups=1, bias=True): 75 | kernel_size = _pair(kernel_size) 76 | stride = _pair(stride) 77 | padding = _pair(padding) 78 | dilation = _pair(dilation) 79 | super(Conv2d, self).__init__( 80 | in_channels, out_channels, kernel_size, stride, padding, dilation, 81 | False, _pair(0), groups, bias) 82 | 83 | def forward(self, input): 84 | return conv2d_same_padding(input, self.weight, self.bias, self.stride, 85 | self.padding, self.dilation, self.groups) 86 | 87 | 88 | # custom con2d, because pytorch don't have "padding='same'" option. 89 | def conv2d_same_padding(input, weight, bias=None, stride=1, padding=1, dilation=1, groups=1): 90 | 91 | input_rows = input.size(2) 92 | filter_rows = weight.size(2) 93 | effective_filter_size_rows = (filter_rows - 1) * dilation[0] + 1 94 | out_rows = (input_rows + stride[0] - 1) // stride[0] 95 | padding_needed = max(0, (out_rows - 1) * stride[0] + effective_filter_size_rows - 96 | input_rows) 97 | padding_rows = max(0, (out_rows - 1) * stride[0] + 98 | (filter_rows - 1) * dilation[0] + 1 - input_rows) 99 | rows_odd = (padding_rows % 2 != 0) 100 | padding_cols = max(0, (out_rows - 1) * stride[0] + 101 | (filter_rows - 1) * dilation[0] + 1 - input_rows) 102 | cols_odd = (padding_rows % 2 != 0) 103 | 104 | if rows_odd or cols_odd: 105 | input = pad(input, [0, int(cols_odd), 0, int(rows_odd)]) 106 | 107 | return F.conv2d(input, weight, bias, stride, 108 | padding=(padding_rows // 2, padding_cols // 2), 109 | dilation=dilation, groups=groups) 110 | 111 | -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/train.py: -------------------------------------------------------------------------------- 1 | # author:oldpan 2 | # data:2018-4-16 3 | # Just for study and research 4 | 5 | from __future__ import print_function 6 | import argparse 7 | import os 8 | 9 | import cv2 10 | import numpy as np 11 | import torch 12 | 13 | import torch.utils.data 14 | from torch import nn, optim 15 | from torch.nn import functional as F 16 | import torch.backends.cudnn as cudnn 17 | 18 | from faceswap_pytorch.models import Autoencoder, toTensor, var_to_np 19 | from faceswap_pytorch.util import get_image_paths, load_images, stack_images 20 | from faceswap_pytorch.training_data import get_training_data 21 | 22 | parser = argparse.ArgumentParser(description='DeepFake-Pytorch') 23 | parser.add_argument('--batch-size', type=int, default=64, metavar='N', 24 | help='input batch size for training (default: 64)') 25 | parser.add_argument('--epochs', type=int, default=100000, metavar='N', 26 | help='number of epochs to train (default: 10000)') 27 | parser.add_argument('--no-cuda', action='store_true', default=False, 28 | help='enables CUDA training') 29 | parser.add_argument('--seed', type=int, default=1, metavar='S', 30 | help='random seed (default: 1)') 31 | parser.add_argument('--log-interval', type=int, default=100, metavar='N', 32 | help='how many batches to wait before logging training status') 33 | 34 | args = parser.parse_args() 35 | args.cuda = not args.no_cuda and torch.cuda.is_available() 36 | 37 | if args.cuda is True: 38 | print('===> Using GPU to train') 39 | device = torch.device('cuda:5') 40 | cudnn.benchmark = True 41 | else: 42 | print('===> Using CPU to train') 43 | 44 | torch.manual_seed(args.seed) 45 | if args.cuda: 46 | torch.cuda.manual_seed(args.seed) 47 | 48 | print('===> Loading datasets') 49 | images_A = get_image_paths("/data/yoojinoh/def/celeb/winterEnd") 50 | images_B = get_image_paths("/data/yoojinoh/def/celeb/karinaEnd") 51 | images_A = load_images(images_A) / 255.0 52 | images_B = load_images(images_B) / 255.0 53 | images_A += images_B.mean(axis=(0, 1, 2)) - images_A.mean(axis=(0, 1, 2)) 54 | 55 | model = Autoencoder().to(device) 56 | 57 | start_epoch = 0 58 | print('===> Start from scratch') 59 | criterion = nn.L1Loss() 60 | optimizer_1 = optim.Adam([{'params': model.encoder.parameters()}, 61 | {'params': model.decoder_A.parameters()}] 62 | , lr=5e-5, betas=(0.5, 0.999)) 63 | optimizer_2 = optim.Adam([{'params': model.encoder.parameters()}, 64 | {'params': model.decoder_B.parameters()}] 65 | , lr=5e-5, betas=(0.5, 0.999)) 66 | 67 | if __name__ == "__main__": 68 | 69 | print('Start training, press \'q\' to stop') 70 | 71 | for epoch in range(start_epoch, args.epochs): 72 | batch_size = args.batch_size 73 | 74 | warped_A, target_A = get_training_data(images_A, batch_size) 75 | warped_B, target_B = get_training_data(images_B, batch_size) 76 | #print(f'{warped_A.shape}, {warped_B.shape}') 77 | warped_A, target_A = toTensor(warped_A), toTensor(target_A) 78 | warped_B, target_B = toTensor(warped_B), toTensor(target_B) 79 | 80 | if args.cuda: 81 | warped_A = warped_A.to(device).float() 82 | target_A = target_A.to(device).float() 83 | warped_B = warped_B.to(device).float() 84 | target_B = target_B.to(device).float() 85 | 86 | optimizer_1.zero_grad() 87 | optimizer_2.zero_grad() 88 | 89 | warped_A = model(warped_A, 'A') 90 | warped_B = model(warped_B, 'B') 91 | 92 | loss1 = criterion(warped_A, target_A) 93 | loss2 = criterion(warped_B, target_B) 94 | loss = loss1.item() + loss2.item() 95 | loss1.backward() 96 | loss2.backward() 97 | optimizer_1.step() 98 | optimizer_2.step() 99 | print('epoch: {}, lossA:{}, lossB:{}'.format(epoch, loss1.item(), loss2.item())) 100 | 101 | if epoch % args.log_interval == 0: 102 | 103 | test_A_ = target_A[0:14] 104 | test_B_ = target_B[0:14] 105 | test_A = var_to_np(target_A[0:14]) 106 | test_B = var_to_np(target_B[0:14]) 107 | print('===> Saving models...') 108 | state = { 109 | 'state': model.state_dict(), 110 | 'epoch': epoch 111 | } 112 | if not os.path.isdir('/data/yoojinoh/def/250101_faceswap_160/checkpoint'): 113 | os.mkdir('/data/yoojinoh/def/250101_faceswap_160/checkpoint') 114 | torch.save(state, '/data/yoojinoh/def/250101_faceswap_160/checkpoint/ae_win_kar_160.t7') 115 | 116 | figure_A = np.stack([ 117 | test_A, 118 | var_to_np(model(test_A_, 'A')), 119 | var_to_np(model(test_A_, 'B')), 120 | ], axis=1) 121 | figure_B = np.stack([ 122 | test_B, 123 | var_to_np(model(test_B_, 'B')), 124 | var_to_np(model(test_B_, 'A')), 125 | ], axis=1) 126 | 127 | figure = np.concatenate([figure_A, figure_B], axis=0) 128 | figure = figure.transpose((0, 1, 3, 4, 2)) 129 | figure = figure.reshape((4, 7) + figure.shape[1:]) 130 | figure = stack_images(figure) 131 | 132 | figure = np.clip(figure * 255, 0, 255).astype('uint8') 133 | 134 | os.makedirs('/data/yoojinoh/def/250101_faceswap_160/train_outputs/', exist_ok=True) 135 | cv2.imwrite("/data/yoojinoh/def/250101_faceswap_160/train_outputs/output.png", figure) 136 | 137 | key = cv2.waitKey(1) 138 | if key == ord('q'): 139 | exit() 140 | -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/training_data.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | from faceswap_pytorch.image_augmentation import random_transform 3 | from faceswap_pytorch.image_augmentation import random_warp 4 | 5 | random_transform_args = { 6 | 'rotation_range': 10, 7 | 'zoom_range': 0.05, 8 | 'shift_range': 0.05, 9 | 'random_flip': 0.4, 10 | } 11 | 12 | 13 | def get_training_data(images, batch_size): 14 | indices = numpy.random.randint(len(images), size=batch_size) 15 | for i, index in enumerate(indices): 16 | image = images[index] 17 | image = random_transform(image, **random_transform_args) 18 | warped_img, target_img = random_warp(image) 19 | 20 | if i == 0: 21 | warped_images = numpy.empty((batch_size,) + warped_img.shape, warped_img.dtype) 22 | target_images = numpy.empty((batch_size,) + target_img.shape, warped_img.dtype) 23 | 24 | warped_images[i] = warped_img 25 | target_images[i] = target_img 26 | 27 | return warped_images, target_images # 16x64x54x3, 16x64x54x3, 28 | -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/umeyama.py: -------------------------------------------------------------------------------- 1 | # # License (Modified BSD) # Copyright (C) 2011, the scikit-image team All rights reserved. # # Redistribution and 2 | # use in source and binary forms, with or without modification, are permitted provided that the following conditions 3 | # are met: # # Redistributions of source code must retain the above copyright notice, this list of conditions and the 4 | # following disclaimer. # Redistributions in binary form must reproduce the above copyright notice, this list of 5 | # conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 6 | # Neither the name of skimage nor the names of its contributors may be used to endorse or promote products derived 7 | # from this software without specific prior written permission. # THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' 8 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND 9 | # FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, 10 | # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE 11 | # GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 12 | # LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY 13 | # OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 14 | 15 | # umeyama function from scikit-image/skimage/transform/_geometric.py 16 | 17 | import numpy as np 18 | 19 | 20 | def umeyama(src, dst, estimate_scale): 21 | """Estimate N-D similarity transformation with or without scaling. 22 | Parameters 23 | ---------- 24 | src : (M, N) array 25 | Source coordinates. 26 | dst : (M, N) array 27 | Destination coordinates. 28 | estimate_scale : bool 29 | Whether to estimate scaling factor. 30 | Returns 31 | ------- 32 | T : (N + 1, N + 1) 33 | The homogeneous similarity transformation matrix. The matrix contains 34 | NaN values only if the problem is not well-conditioned. 35 | References 36 | ---------- 37 | .. [1] "Least-squares estimation of transformation parameters between two 38 | point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573 39 | """ 40 | 41 | num = src.shape[0] 42 | dim = src.shape[1] 43 | 44 | # Compute mean of src and dst. 45 | src_mean = src.mean(axis=0) 46 | dst_mean = dst.mean(axis=0) 47 | 48 | # Subtract mean from src and dst. 49 | src_demean = src - src_mean 50 | dst_demean = dst - dst_mean 51 | 52 | # Eq. (38). 53 | A = np.dot(dst_demean.T, src_demean) / num 54 | 55 | # Eq. (39). 56 | d = np.ones((dim,), dtype=np.double) 57 | if np.linalg.det(A) < 0: 58 | d[dim - 1] = -1 59 | 60 | T = np.eye(dim + 1, dtype=np.double) 61 | 62 | U, S, V = np.linalg.svd(A) 63 | 64 | # Eq. (40) and (43). 65 | rank = np.linalg.matrix_rank(A) 66 | if rank == 0: 67 | return np.nan * T 68 | elif rank == dim - 1: 69 | if np.linalg.det(U) * np.linalg.det(V) > 0: 70 | T[:dim, :dim] = np.dot(U, V) 71 | else: 72 | s = d[dim - 1] 73 | d[dim - 1] = -1 74 | T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V)) 75 | d[dim - 1] = s 76 | else: 77 | T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T)) 78 | 79 | if estimate_scale: 80 | # Eq. (41) and (42). 81 | scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d) 82 | else: 83 | scale = 1.0 84 | 85 | T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) 86 | T[:dim, :dim] *= scale 87 | 88 | return T 89 | -------------------------------------------------------------------------------- /backend/ddf/faceswap_pytorch/util.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy 3 | import os 4 | 5 | 6 | def get_image_paths(directory): 7 | return [x.path for x in os.scandir(directory) if x.name.endswith(".jpg") or x.name.endswith(".png") or x.name.endswith(".JPG")] 8 | 9 | 10 | def load_images(image_paths, convert=None): 11 | iter_all_images = (cv2.resize(cv2.imread(fn), (256, 256)) for fn in image_paths) 12 | if convert: 13 | iter_all_images = (convert(img) for img in iter_all_images) 14 | for i, image in enumerate(iter_all_images): 15 | if i == 0: 16 | all_images = numpy.empty((len(image_paths),) + image.shape, dtype=image.dtype) 17 | all_images[i] = image 18 | return all_images 19 | 20 | 21 | def get_transpose_axes(n): 22 | if n % 2 == 0: 23 | y_axes = list(range(1, n - 1, 2)) 24 | x_axes = list(range(0, n - 1, 2)) 25 | else: 26 | y_axes = list(range(0, n - 1, 2)) 27 | x_axes = list(range(1, n - 1, 2)) 28 | return y_axes, x_axes, [n - 1] 29 | 30 | 31 | def stack_images(images): 32 | images_shape = numpy.array(images.shape) 33 | new_axes = get_transpose_axes(len(images_shape)) 34 | new_shape = [numpy.prod(images_shape[x]) for x in new_axes] 35 | return numpy.transpose( 36 | images, 37 | axes=numpy.concatenate(new_axes) 38 | ).reshape(new_shape) 39 | -------------------------------------------------------------------------------- /backend/ddf/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class L1_Charbonnier_loss(torch.nn.Module): 5 | """L1 Charbonnierloss.""" 6 | def __init__(self): 7 | super(L1_Charbonnier_loss, self).__init__() 8 | self.eps = 1e-6 9 | 10 | def forward(self, X, Y): 11 | diff = torch.add(X, -Y) 12 | error = torch.sqrt(diff * diff + self.eps) 13 | loss = torch.mean(error) 14 | return loss 15 | 16 | 17 | class GWLoss(nn.Module): 18 | """Gradient Weighted Loss""" 19 | def __init__(self, reduction='mean'): 20 | super(GWLoss, self).__init__() 21 | self.w = c.w 22 | self.reduction = reduction 23 | sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float) 24 | sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float) 25 | self.weight_x = nn.Parameter(data=sobel_x, requires_grad=False) 26 | self.weight_y = nn.Parameter(data=sobel_y, requires_grad=False) 27 | def forward(self, x1, x2): 28 | b, c, w, h = x1.shape 29 | weight_x = self.weight_x.expand(c, 1, 3, 3).type_as(x1) 30 | weight_y = self.weight_y.expand(c, 1, 3, 3).type_as(x1) 31 | Ix1 = F.conv2d(x1, weight_x, stride=1, padding=1, groups=c) 32 | Ix2 = F.conv2d(x2, weight_x, stride=1, padding=1, groups=c) 33 | Iy1 = F.conv2d(x1, weight_y, stride=1, padding=1, groups=c) 34 | Iy2 = F.conv2d(x2, weight_y, stride=1, padding=1, groups=c) 35 | dx = torch.abs(Ix1 - Ix2) 36 | dy = torch.abs(Iy1 - Iy2) 37 | # loss = torch.exp(2*(dx + dy)) * torch.abs(x1 - x2) 38 | loss = (1 + self.w * dx) * (1 + self.w * dy) * torch.abs(x1 - x2) 39 | if self.reduction == 'mean': 40 | return torch.mean(loss) 41 | else: 42 | return torch.sum(loss) -------------------------------------------------------------------------------- /backend/ddf/main.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | import cv2 4 | from ddf.DualDefense_gan_fs import DualDefense 5 | import warnings 6 | warnings.filterwarnings('ignore') 7 | from PIL import Image 8 | from torchvision.transforms import transforms 9 | import numpy as np 10 | from ddf.engine.utils import blend_image 11 | 12 | _DDF_MODEL_MAPPINGS = { 13 | "byeon_cha": "/home/yoojinoh/Others/ckpts/ckpt_best_img_lam1_al2.pt", 14 | "win_chuu": "/home/yoojinoh/Others/ckpts/winchuu_ckpt_best_img_lam1_al2.pt", 15 | } 16 | 17 | _TEST_CONFIGS = { 18 | "message_size": {'byeon_cha': 4, 'win_chuu': 15}, 19 | "height": 160, 20 | "width": 160 21 | } 22 | 23 | USER_WATERMARK_IDS = { 24 | "byeon": [0., 1., 0., 1.], 25 | "cha": [1., 1., 0., 0.], 26 | "win": [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], 27 | "chu": [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.] 28 | } 29 | 30 | def assign_messages(users, device='cpu'): 31 | num_users = len(users) 32 | max_users = 16 33 | assert num_users <= max_users, f"The number of maximum users is {max_users}." 34 | unique_messages = torch.eye(max_users, MESSAGE_SIZE, dtype=torch.float).to(device) 35 | user_messages = {user: unique_messages[i].tolist() for i, user in enumerate(users)} 36 | #TODO(Yoojin): Save info to DB 37 | return user_messages 38 | 39 | def find_message2user(message:list)->str: 40 | for key, value in USER_WATERMARK_IDS.items(): 41 | if value == message: 42 | return key 43 | else: 44 | if len(message) == 4: 45 | if message[0] == 0. and message[1] == 1.: return 'byeon' 46 | else: return 'cha' 47 | elif len(message) == 15: 48 | if message[0] == 0. and message[1] == 0.: return 'win' 49 | else: return 'chu' 50 | return None # ex. 'byeon' 51 | 52 | 53 | def encode_image(model, image, message): 54 | """insert watermark into original image""" 55 | if len(image.shape) == 3: 56 | image = image.unsqueeze(0) 57 | return model.encode(image, message) 58 | 59 | def save_image(image, save_path): 60 | cv2.imwrite(save_path, image) 61 | print(f'[INFO] Saved image to {save_path}') 62 | 63 | def decode_image(model, encoded_image): 64 | return model.decode(encoded_image) 65 | 66 | def get_transform(bgr_image, height, width): 67 | transform = transforms.Compose([ 68 | transforms.ToTensor(), 69 | transforms.Resize((height, width))] 70 | ) 71 | return transform(bgr_image) 72 | 73 | def load_image(image_path): 74 | image = Image.fromarray(cv2.imread(image_path)) #BGR image 75 | return image 76 | 77 | def load_model(model_type, mode = 'eval', device=None):# Modified 78 | model = DualDefense(_TEST_CONFIGS['message_size'][model_type], in_channels=3, model_type=model_type, device=device) 79 | model.encoder.load_state_dict(torch.load(_DDF_MODEL_MAPPINGS[model_type])['encoder'], strict=False) 80 | model.decoder.load_state_dict(torch.load(_DDF_MODEL_MAPPINGS[model_type])['decoder'], strict=False) 81 | 82 | if mode == 'eval': 83 | model.encoder.eval() 84 | model.decoder.eval() 85 | return model 86 | 87 | def apply_faceswap(model_type, swapped_image_path, src_path, tgt_path, src_user, encoded=True): 88 | """ 89 | Apply dual defense and save the faceswapped image (swap face a with face b) 90 | 91 | Args: 92 | model_type (str) : type of dual defense model to apply. (byeon_cha, win_chu) 93 | swapped_image_path (str) : path to save the face swapped image 94 | src_path (str): path of (encoded) image a 95 | tgt_path (str): path of (encoded) image b 96 | encoded (bool): if the given image is encoded or not 97 | """ 98 | device = torch.device(f'cuda' if torch.cuda.is_available() else 'cpu') 99 | assert model_type in _DDF_MODEL_MAPPINGS.keys() 100 | model = load_model(model_type, mode='eval', device=device) 101 | 102 | original_src_image, src_image, src_coord = crop_image(src_path) # load_image(src_path) 103 | original_tgt_image, tgt_image, tgt_coord = crop_image(tgt_path) # load_image(tgt_path) 104 | 105 | src_image = Image.fromarray(src_image) 106 | tgt_image = Image.fromarray(tgt_image) 107 | 108 | src_image = get_transform(src_image, _TEST_CONFIGS['height'], _TEST_CONFIGS['width']).to(device).unsqueeze(0) 109 | tgt_image = get_transform(tgt_image, _TEST_CONFIGS['height'], _TEST_CONFIGS['width']).to(device).unsqueeze(0) 110 | 111 | with torch.no_grad(): 112 | src_image = src_image.to(device) 113 | tgt_image = tgt_image.to(device) 114 | 115 | # Apply faceswap 116 | assert src_user in ['cha', 'byeon', 'win', 'chu'], f"There is no user named {src_user}" 117 | if src_user in ['cha', 'chu']: 118 | print(f'[DEBUG]Src user is given {src_user}') 119 | _, _, _image_a_deepfake = model.deepfake1(src_image, 'A') 120 | _, _, _image_b_deepfake = model.deepfake1(tgt_image, 'B') 121 | else: # ['byeon', 'win'] 122 | _, _, _image_a_deepfake = model.deepfake1(src_image, 'B') 123 | _, _, _image_b_deepfake = model.deepfake1(tgt_image, 'A') 124 | 125 | image_a_deepfake =(_image_a_deepfake[0] * 255).permute(1, 2, 0).detach().cpu().numpy() 126 | image_b_deepfake =(_image_b_deepfake[0] * 255).permute(1, 2, 0).detach().cpu().numpy() 127 | 128 | save_image(image_a_deepfake,swapped_image_path) # os.path.join(swapped_image_path, 'output_A2B.png')) 129 | 130 | if encoded: 131 | _pred_a_message = model.decode(_image_a_deepfake) 132 | 133 | pred_a_message = list(map(int, _pred_a_message[0])) 134 | 135 | usera = find_message2user(pred_a_message) 136 | 137 | print(f'>> Someone tried to make deepfake with user {usera} ') 138 | src_output_path = swapped_image_path # os.path.join(swapped_image_path, 'output_A2B.png') 139 | return {"source_image_url": src_output_path, 140 | "user_prediction": f"🚨Someone tried to make deepfake with user {usera} \n➡️ PROTECTED 👩‍✈️!"} 141 | def restore_original(original_image, encoded_face, coord, result_path): 142 | (x, y, w, h) = coord 143 | 144 | encoded_face = (encoded_face * 255).permute(1, 2, 0).detach().cpu().numpy() 145 | encoded_face = np.clip(encoded_face, 0, 255) 146 | 147 | resized_encoded_face = cv2.resize(encoded_face, (w, h)) 148 | 149 | result_image = original_image.astype(np.float32) 150 | result_image[y:y+h, x:x+w] = resized_encoded_face 151 | save_image(result_image, result_path) 152 | print("Restored to original image") 153 | 154 | def crop_and_encode_image(model_type, image_path, message, device, alpha=1.0): 155 | """ 156 | Inputs original image, and save image with encoded face 157 | args: 158 | model_type (str): ['byeon_cha', 'win_chuu'] 159 | image_path (str): path of the image to encode 160 | message (list) : [1., 0., 1] ->ex. torch.randint(0, 2, (1, 4), dtype=torch.float).to(device).detach() 161 | device : cpu or gpu 162 | alpha (float) : the amount to blend encoded image with original image 163 | """ 164 | original_image, cropped_face, (x, y, w, h) = crop_image(image_path) 165 | transformed_cropped_face = get_transform(cropped_face, 160, 160).to(device) 166 | 167 | model = load_model(model_type, 'eval', device) 168 | message = torch.FloatTensor(message).to(device).detach() 169 | encoded_face = encode_image(model, transformed_cropped_face, message)[0] 170 | 171 | #Modified(Yoojin) 172 | encoded_face = blend_image(encoded_face, transformed_cropped_face).to(device) 173 | 174 | print(f'Successfully encoded image with message {message}') 175 | result_path = os.path.join(os.path.dirname(image_path), 'encoded_' + os.path.basename(image_path)) 176 | restore_original(original_image, encoded_face, (x, y, w, h), result_path) 177 | 178 | from mtcnn import MTCNN 179 | import cv2 180 | 181 | def crop_image(image_path): 182 | detector = MTCNN() 183 | print(f"[DEBUG] Image path = {image_path}") 184 | image = cv2.imread(image_path) 185 | if image is None: 186 | raise ValueError(f"Error loading image from path: {image_path}") 187 | 188 | rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 189 | detections = detector.detect_faces(rgb_image) 190 | 191 | if not detections: 192 | raise ValueError("No face detected in the image.") 193 | x, y, w, h = detections[0]['box'] 194 | cropped_face = image[y:y+h, x:x+w] 195 | 196 | return image, cropped_face, (x, y, w, h) 197 | 198 | 199 | # if __name__ == "__main__": 200 | # device = torch.device(f'cuda' if torch.cuda.is_available() else 'cpu') 201 | # #TODO(Yoojin): 개별 user를 위한 watermark information으로 수정 202 | # crop_and_encode_image('byeon_cha', 203 | # '/home/yoojinoh/Others/byeon_total5.jpg', 204 | # USER_WATERMARK_IDS['byeon'], 205 | # device) 206 | # crop_and_encode_image('byeon_cha', 207 | # '/home/yoojinoh/Others/cha_total2.jpg', 208 | # USER_WATERMARK_IDS['cha'], 209 | # device) 210 | # apply_faceswap(model_type="byeon_cha", swapped_image_path="/home/yoojinoh/Others/", 211 | # src_path="/home/yoojinoh/Others/encoded_byeon_total5.jpg", 212 | # tgt_path="/home/yoojinoh/Others/encoded_cha_total2.jpg") 213 | -------------------------------------------------------------------------------- /backend/ddf/scripts/run_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | python test.py \ 4 | --config_path "/home/yoojinoh/Others/deepsafe/config.yaml" \ 5 | --gpus "2" -------------------------------------------------------------------------------- /backend/ddf/scripts/run_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | python train.py \ 4 | --config_path "/home/yoojinoh/Others/deepsafe/config.yaml" \ 5 | --save_path "/data1/yoojinoh/def/train/0129_bc_ddf" \ 6 | --gpus "0" -------------------------------------------------------------------------------- /backend/ddf/upload.js: -------------------------------------------------------------------------------- 1 | import { useState, useEffect } from "react"; 2 | import { useRouter } from "next/router"; 3 | import Navbar from "../components/Navbar"; 4 | 5 | export default function Upload() { 6 | const router = useRouter(); 7 | const [image, setImage] = useState(null); 8 | const [imagePreview, setImagePreview] = useState(""); 9 | const [title, setTitle] = useState(""); 10 | const [content, setContent] = useState(""); 11 | const [isLoading, setIsLoading] = useState(false); 12 | const [imageUrl, setImageUrl] = useState(""); // To store the image URL 13 | const [showModal, setShowModal] = useState(false); // To control the modal visibility 14 | const publicImages = { 15 | chu: ["/chu/chuu1.jpg", "/chu/chuu2.jpg", "/chu/chuu3.jpg", "/chu/chuu4.jpg"], 16 | cha: ["/cha/chaeunwoo1.jpg", "/cha/chaeunwoo2.jpg", "/cha/chaeunwoo3.jpg", "/cha/chaeunwoo4.jpg"], 17 | byeon: ["/byeon/byeon_1.png", "/byeon/byeon_2.png", "/byeon/byeon_3.png", "/byeon/byeon_4.png"], 18 | winter: ["/winter/winter1.jpg", "/winter/winter2.jpg", "winter/winter3.jpg", "winter/winter4.jpg"], 19 | }; 20 | 21 | const [activeTab, setActiveTab] = useState("chu"); 22 | 23 | 24 | const handlePublicImageSelect = (img) => { 25 | setImagePreview(img); 26 | // Fetch the image as a blob and set it as the image file 27 | fetch(img) 28 | .then((response) => response.blob()) 29 | .then((blob) => { 30 | const file = new File([blob], "publicImage.jpg", { type: blob.type }); 31 | setImage(file); 32 | }) 33 | .catch((error) => console.error("Error fetching image:", error)); 34 | setShowModal(false); 35 | }; 36 | 37 | 38 | // Handle image upload 39 | useEffect(() => { 40 | if(!image) return; 41 | uploadImage(); 42 | }, [image]); 43 | 44 | 45 | // Upload image and get the image URL 46 | const uploadImage = async () => { 47 | if (!image) { 48 | alert("Please upload an image first."); 49 | return; 50 | } 51 | 52 | setIsLoading(true); 53 | try { 54 | const formData = new FormData(); 55 | formData.append("image", image); 56 | 57 | const response = await fetch("http://localhost:8000/upload-image", { 58 | method: "POST", 59 | body: image, 60 | }); 61 | 62 | if (!response.ok) { 63 | throw new Error("Failed to upload image."); 64 | } 65 | 66 | const data = await response.json(); // Get the response data 67 | setImageUrl(data.image_url); // Set the image URL returned from the server 68 | } catch (error) { 69 | console.error("Error uploading image:", error); 70 | alert("Failed to upload image. Please try again."); 71 | } finally { 72 | setIsLoading(false); 73 | } 74 | }; 75 | 76 | // Submit post with title, content, and image URL 77 | const submitPost = async () => { 78 | if (!imageUrl || !title || !content) { 79 | alert("Please fill in all fields and upload an image."); 80 | return; 81 | } 82 | 83 | setIsLoading(true); 84 | try { 85 | const response = await fetch("http://localhost:8000/upload-post", { 86 | method: "POST", 87 | headers: { 88 | "Content-Type": "application/json", 89 | }, 90 | body: JSON.stringify({ 91 | user: activeTab, 92 | title, 93 | content, 94 | image_url: imageUrl, // Use the image URL obtained from image upload 95 | }), 96 | }); 97 | 98 | if (!response.ok) { 99 | throw new Error("Failed to upload post."); 100 | } 101 | 102 | // Redirect to home page after successful post upload 103 | alert("🔒 Image Encoded base on User info!") 104 | router.push("/"); 105 | } catch (error) { 106 | console.error("Error uploading post:", error); 107 | alert("Failed to upload post. Please try again."); 108 | } finally { 109 | setIsLoading(false); 110 | } 111 | }; 112 | 113 | return ( 114 |
115 | {/* Outer Container with Gradient Border */} 116 | 117 |
118 | {/* Inner Box */} 119 |
120 | {/* Main Content */} 121 | 122 |
123 |

Upload Post

124 |

125 | Upload an image and add a title and content for your post. 126 |

127 | 128 | 129 |
130 | setTitle(e.target.value)} 135 | className="w-full p-3 pl-5 rounded-full border border-gray-300 shadow-md" 136 | /> 137 |
138 | 139 |
140 | {!imagePreview && ( 141 |
142 | 148 |
149 | )} 150 | 151 | {showModal && ( 152 |
153 |
154 |

Select an image.

155 |
156 | {Object.keys(publicImages).map((tab) => ( 157 | 164 | ))} 165 |
166 |
167 | {publicImages[activeTab].map((img, index) => ( 168 | {`Public handlePublicImageSelect(img)} 174 | /> 175 | ))} 176 |
177 | 183 |
184 |
185 | )} 186 | 187 | {imagePreview && ( 188 |
189 | {/* X 버튼 */} 190 | 196 | Uploaded Preview 201 |
202 | )} 203 |
204 | 205 | {/* Content Input */} 206 |
207 |