add necessary files

app.py

@@ -3,12 +3,8 @@ import torch
 import argparse
 import git
 
-git.Repo.clone_from("https://github.com/timroelofs123/face_reaging.git", "./face_reaging")
 git.Repo.clone_from("https://huggingface.co/timroelofs123/face_re-aging", "./hf")
 
-import sys
-sys.path.append("./face_reaging")
-
 from model.models import UNet
 from scripts.test_functions import process_image
 

model/losses.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+import lpips  # LPIPS library for perceptual loss
+
+class GeneratorLoss(nn.Module):
+    def __init__(self, discriminator_model, l1_weight=1.0, perceptual_weight=1.0, adversarial_weight=0.05,
+                 device="cpu"):
+        super(GeneratorLoss, self).__init__()
+        self.discriminator_model = discriminator_model
+        self.l1_weight = l1_weight
+        self.perceptual_weight = perceptual_weight
+        self.adversarial_weight = adversarial_weight
+        self.criterion_l1 = nn.L1Loss()
+        self.criterion_adversarial = nn.BCEWithLogitsLoss()
+        self.criterion_perceptual = lpips.LPIPS(net='vgg').to(device)
+
+    def forward(self, output, target, source):
+        # L1 loss
+
+        l1_loss = self.criterion_l1(output, target)
+
+        # Perceptual loss
+        perceptual_loss = torch.mean(self.criterion_perceptual(output, target))
+
+        # Adversarial loss
+        fake_input = torch.cat([output, source[:, 4:5, :, :]], dim=1)
+        fake_prediction = self.discriminator_model(fake_input)
+
+        adversarial_loss = self.criterion_adversarial(fake_prediction, torch.ones_like(fake_prediction))
+
+        # Combine losses
+        generator_loss = self.l1_weight * l1_loss + self.perceptual_weight * perceptual_loss + \
+                         self.adversarial_weight * adversarial_loss
+
+        return generator_loss, l1_loss, perceptual_loss, adversarial_loss
+
+class DiscriminatorLoss(nn.Module):
+    def __init__(self, discriminator_model, fake_weight=1.0, real_weight=2.0, mock_weight=.5):
+        super(DiscriminatorLoss, self).__init__()
+        self.discriminator_model = discriminator_model
+        self.criterion_adversarial = nn.BCEWithLogitsLoss()
+        self.fake_weight = fake_weight
+        self.real_weight = real_weight
+        self.mock_weight = mock_weight
+
+    def forward(self, output, target, source):
+        # Adversarial loss
+        fake_input = torch.cat([output, source[:, 4:5, :, :]], dim=1)  # prediction img with target age
+        real_input = torch.cat([target, source[:, 4:5, :, :]], dim=1)  # target img with target age
+
+        mock_input1 = torch.cat([source[:, :3, :, :], source[:, 4:5, :, :]], dim=1)  # source img with target age
+        mock_input2 = torch.cat([target, source[:, 3:4, :, :]], dim=1)  # target img with source age
+        mock_input3 = torch.cat([output, source[:, 3:4, :, :]], dim=1)  # prediction img with source age
+        mock_input4 = torch.cat([target, source[:, 3:4, :, :]], dim=1)  # target img with target age
+
+        fake_pred, real_pred = self.discriminator_model(fake_input), self.discriminator_model(real_input)
+        mock_pred1, mock_pred2, mock_pred3, mock_pred4 = (self.discriminator_model(mock_input1),
+                                                          self.discriminator_model(mock_input2),
+                                                          self.discriminator_model(mock_input3),
+                                                          self.discriminator_model(mock_input4))
+
+        discriminator_loss = (self.fake_weight * self.criterion_adversarial(fake_pred, torch.zeros_like(fake_pred)) +
+                              self.real_weight * self.criterion_adversarial(real_pred, torch.ones_like(real_pred)) +
+                              self.mock_weight * self.criterion_adversarial(mock_pred1, torch.zeros_like(mock_pred1)) +
+                              self.mock_weight * self.criterion_adversarial(mock_pred2, torch.zeros_like(mock_pred2)) +
+                              self.mock_weight * self.criterion_adversarial(mock_pred3, torch.zeros_like(mock_pred3)) +
+                              self.mock_weight * self.criterion_adversarial(mock_pred4, torch.zeros_like(mock_pred4))
+                              )
+
+        return discriminator_loss

model/models.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+import antialiased_cnns
+
+
+class DownLayer(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DownLayer, self).__init__()
+        self.layer = nn.Sequential(
+            nn.MaxPool2d(kernel_size=2, stride=1),
+            antialiased_cnns.BlurPool(in_channels, stride=2),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.LeakyReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.LeakyReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class UpLayer(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UpLayer, self).__init__()
+        # Conv transpose upsampling
+
+        self.blur_upsample = nn.Sequential(
+            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2, padding=0),
+            antialiased_cnns.BlurPool(out_channels, stride=1)
+        )
+
+        self.layer = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.LeakyReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.LeakyReLU(inplace=True)
+        )
+
+    def forward(self, x, skip):
+        x = self.blur_upsample(x)
+        x = torch.cat([x, skip], dim=1)  # Concatenate with skip connection
+        return self.layer(x)
+
+
+class UNet(nn.Module):
+    def __init__(self):
+        super(UNet, self).__init__()
+        self.init_conv = nn.Sequential(
+            nn.Conv2d(5, 64, kernel_size=3, padding=1),  # output: 512 x 512 x 64
+            nn.LeakyReLU(inplace=True),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),  # output: 512 x 512 x 64
+            nn.LeakyReLU(inplace=True)
+        )
+
+        self.down1 = DownLayer(64, 128)  # output: 256 x 256 x 128
+        self.down2 = DownLayer(128, 256)  # output: 128 x 128 x 256
+        self.down3 = DownLayer(256, 512)  # output: 64 x 64 x 512
+        self.down4 = DownLayer(512, 1024)  # output: 32 x 32 x 1024
+        self.up1 = UpLayer(1024, 512)  # output: 64 x 64 x 512
+        self.up2 = UpLayer(512, 256)  # output: 128 x 128 x 256
+        self.up3 = UpLayer(256, 128)  # output: 256 x 256 x 128
+        self.up4 = UpLayer(128, 64)  # output: 512 x 512 x 64
+        self.final_conv = nn.Conv2d(64, 3, kernel_size=1)  # output: 512 x 512 x 3
+
+    def forward(self, x):
+        x0 = self.init_conv(x)
+        x1 = self.down1(x0)
+        x2 = self.down2(x1)
+        x3 = self.down3(x2)
+        x4 = self.down4(x3)
+        x = self.up1(x4, x3)
+        x = self.up2(x, x2)
+        x = self.up3(x, x1)
+        x = self.up4(x, x0)
+        x = self.final_conv(x)
+        return x
+
+
+class PatchGANDiscriminator(nn.Module):
+    def __init__(self, input_channels=3):
+        super(PatchGANDiscriminator, self).__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(input_channels, 64, kernel_size=4, stride=2, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Conv2d(256, 1, kernel_size=4, stride=1, padding=1)
+            # Output layer with 1 channel for binary classification
+        )
+
+    def forward(self, x):
+        return self.model(x)

utils/dataloader.py

@@ -0,0 +1,63 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from PIL import Image
+import os
+import random
+from pathlib import Path
+
+
+# Define the transformations
+transform = transforms.Compose([
+    transforms.RandomRotation(degrees=10),
+    transforms.RandomCrop(512),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+    transforms.ToTensor(),
+])
+
+class CustomDataset(Dataset):
+    def __init__(self, root_dir, transform=None):
+        self.root_dir = root_dir
+        self.transform = transform
+        self.image_folders = [folder for folder in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, folder))]
+
+    def __len__(self):
+        return len(self.image_folders)
+
+    def __getitem__(self, idx):
+        folder_name = self.image_folders[idx]
+        folder_path = os.path.join(self.root_dir, folder_name)
+
+        # # Get the list of image filenames in the folder
+        # image_filenames = [f"{i}.jpg" for i in range(0, 101, 10)]
+        image_filenames = os.listdir(folder_path)
+
+        # Pick two random assets from the folder
+        source_image_name, target_image_name = random.sample(image_filenames, 2)
+        # source_image_name, target_image_name = '20.jpg', '80.jpg'
+
+        source_age = int(Path(source_image_name).stem) / 100
+        target_age = int(Path(target_image_name).stem) / 100
+
+        # Randomly select two assets from the folder
+        source_image_path = os.path.join(folder_path, source_image_name)
+        target_image_path = os.path.join(folder_path, target_image_name)
+
+        source_image = Image.open(source_image_path).convert('RGB')
+        target_image = Image.open(target_image_path).convert('RGB')
+
+        # Apply the same random crop and augmentations to both assets
+        if self.transform:
+            seed = torch.randint(0, 2 ** 32 - 1, (1,)).item()
+            torch.manual_seed(seed)
+            source_image = self.transform(source_image)
+            torch.manual_seed(seed)
+            target_image = self.transform(target_image)
+
+        source_age_channel = torch.full_like(source_image[:1, :, :], source_age)
+        target_age_channel = torch.full_like(source_image[:1, :, :], target_age)
+
+        # Concatenate the age channels with the source_image
+        source_image = torch.cat([source_image, source_age_channel, target_age_channel], dim=0)
+
+        return source_image, target_image