Added face aging model and interface with Gradio

.gitignore

@@ -0,0 +1,5 @@
+# Virtual Environment
+venv
+
+__pycache__
+

app.py

@@ -1,7 +1,39 @@
+import os
+import torch
+from models import UNet
+from test_functions import process_image
+from PIL import Image
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+from huggingface_hub import hf_hub_download
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+MODEL_PATH = hf_hub_download(repo_id="Robys01/face-aging", filename="best_unet_model.pth")
+print(f"Model downloaded to {MODEL_PATH}")
+
+model = UNet()  
+model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu"), weights_only=False))
+model.eval()
+
+def age_image(image: Image.Image, source_age: int, target_age: int) -> Image.Image:
+    # Ensure the image is in RGB or grayscale; if not, convert it.
+    if image.mode not in ["RGB", "L"]:
+        print(f"Converting image from {image.mode} to RGB")
+        image = image.convert("RGB")
+
+    processed_image = process_image(model, image, source_age, target_age)
+    return processed_image
+
+iface = gr.Interface(
+    fn=age_image,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Slider(10, 90, value=20, step=1, label="Current age", info="Choose the current age"),
+        gr.Slider(10, 90, value=70, step=1, label="Target age", info="Choose the desired age")
+    ],
+    outputs=gr.Image(type="pil", label="Aged Image"),
+    title="Face Aging Demo",
+    description="Upload an image along with a source age approximation and a target age to generate an aged version of the face."
+)
+
+if __name__ == "__main__":
+    iface.launch()

models.py

@@ -0,0 +1,98 @@
+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)

requirements.txt

@@ -0,0 +1,66 @@
+-f https://download.pytorch.org/whl/cpu/torch_stable.html
+torch==2.3.1+cpu
+torchvision==0.18.1+cpu
+aiofiles==23.2.1
+annotated-types==0.7.0
+antialiased-cnns==0.3
+anyio==4.8.0
+certifi==2025.1.31
+charset-normalizer==3.4.1
+click==8.1.8
+dlib==19.24.6
+face-recognition==1.3.0
+face_recognition_models==0.3.0
+fastapi==0.115.8
+ffmpy==0.5.0
+filelock==3.17.0
+fsspec==2025.2.0
+gradio==5.15.0
+gradio_client==1.7.0
+h11==0.14.0
+httpcore==1.0.7
+httptools==0.6.4
+httpx==0.28.1
+huggingface-hub==0.28.1
+idna==3.10
+Jinja2==3.1.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.4.2
+numpy==2.2.2
+orjson==3.10.15
+packaging==24.2
+pandas==2.2.3
+pillow==11.1.0
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydub==0.25.1
+Pygments==2.19.1
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.20
+pytz==2025.1
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.4
+ruff==0.9.6
+safehttpx==0.1.6
+semantic-version==2.10.0
+setuptools==75.8.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+starlette==0.45.3
+sympy==1.13.1
+tomlkit==0.13.2
+tqdm==4.67.1
+typer==0.15.1
+typing_extensions==4.12.2
+tzdata==2025.1
+urllib3==2.3.0
+uvicorn==0.34.0
+uvloop==0.21.0
+watchfiles==1.0.4
+websockets==14.2

test_functions.py

@@ -0,0 +1,100 @@
+import face_recognition
+import numpy as np
+import torch
+from torch.autograd import Variable
+from torchvision import transforms
+from PIL import Image
+
+mask_file = torch.from_numpy(np.array(Image.open('assets/mask1024.jpg').convert('L'))) / 255
+small_mask_file = torch.from_numpy(np.array(Image.open('assets/mask512.jpg').convert('L'))) / 255
+
+def sliding_window_tensor(input_tensor, window_size, stride, your_model, mask=mask_file, small_mask=small_mask_file):
+    """
+    Apply aging operation on input tensor using a sliding-window method. This operation is done on the GPU, if available.
+    """
+
+    input_tensor = input_tensor.to(next(your_model.parameters()).device)
+    mask = mask.to(next(your_model.parameters()).device)
+    small_mask = small_mask.to(next(your_model.parameters()).device)
+
+    n, c, h, w = input_tensor.size()
+    output_tensor = torch.zeros((n, 3, h, w), dtype=input_tensor.dtype, device=input_tensor.device)
+
+    count_tensor = torch.zeros((n, 3, h, w), dtype=torch.float32, device=input_tensor.device)
+
+    add = 2 if window_size % stride != 0 else 1
+
+    for y in range(0, h - window_size + add, stride):
+        for x in range(0, w - window_size + add, stride):
+            window = input_tensor[:, :, y:y + window_size, x:x + window_size]
+
+            # Apply the same preprocessing as during training
+            input_variable = Variable(window, requires_grad=False)  # Assuming GPU is available
+
+            # Forward pass
+            with torch.no_grad():
+                output = your_model(input_variable)
+
+            output_tensor[:, :, y:y + window_size, x:x + window_size] += output * small_mask
+            count_tensor[:, :, y:y + window_size, x:x + window_size] += small_mask
+
+    count_tensor = torch.clamp(count_tensor, min=1.0)
+
+    # Average the overlapping regions
+    output_tensor /= count_tensor
+
+    # Apply mask
+    output_tensor *= mask
+
+    return output_tensor.cpu()
+
+
+def process_image(your_model, image, source_age, target_age=0,
+                  window_size=512, stride=256, steps=18):
+
+    input_size = (1024, 1024)
+
+    # image = face_recognition.load_image_file(filename)
+    image = np.array(image)
+
+    fl = face_recognition.face_locations(image)[0]
+
+    # calculate margins
+    margin_y_t = int((fl[2] - fl[0]) * .63 * .85)  # larger as the forehead is often cut off
+    margin_y_b = int((fl[2] - fl[0]) * .37 * .85)
+    margin_x = int((fl[1] - fl[3]) // (2 / .85))
+    margin_y_t += 2 * margin_x - margin_y_t - margin_y_b  # make sure square is preserved
+
+    l_y = max([fl[0] - margin_y_t, 0])
+    r_y = min([fl[2] + margin_y_b, image.shape[0]])
+    l_x = max([fl[3] - margin_x, 0])
+    r_x = min([fl[1] + margin_x, image.shape[1]])
+
+    # crop image
+    cropped_image = image[l_y:r_y, l_x:r_x, :]
+
+    # Resizing
+    orig_size = cropped_image.shape[:2]
+
+    cropped_image = transforms.ToTensor()(cropped_image)
+
+    cropped_image_resized = transforms.Resize(input_size, interpolation=Image.BILINEAR, antialias=True)(cropped_image)
+
+    source_age_channel = torch.full_like(cropped_image_resized[:1, :, :], source_age / 100)
+    target_age_channel = torch.full_like(cropped_image_resized[:1, :, :], target_age / 100)
+    input_tensor = torch.cat([cropped_image_resized, source_age_channel, target_age_channel], dim=0).unsqueeze(0)
+
+    image = transforms.ToTensor()(image)
+
+    # performing actions on image
+    aged_cropped_image = sliding_window_tensor(input_tensor, window_size, stride, your_model)
+
+    # resize back to original size
+    aged_cropped_image_resized = transforms.Resize(orig_size, interpolation=Image.BILINEAR, antialias=True)(
+        aged_cropped_image)
+
+    # re-apply
+    image[:, l_y:r_y, l_x:r_x] += aged_cropped_image_resized.squeeze(0)
+    image = torch.clamp(image, 0, 1)
+
+    return transforms.functional.to_pil_image(image)