test_functions.py

import os

import cv2
import numpy as np
import torch
from PIL import Image
from torch.autograd import Variable
from torchvision import transforms


def _mask_path(filename: str) -> str:
    assets_path = os.path.join("assets", filename)
    if os.path.exists(assets_path):
        return assets_path
    return filename


mask_file = torch.from_numpy(
    np.array(Image.open(_mask_path("mask1024.jpg")).convert("L"))
).float() / 255.0

small_mask_file = torch.from_numpy(
    np.array(Image.open(_mask_path("mask512.jpg")).convert("L"))
).float() / 255.0


def _detect_face_box(image_rgb: np.ndarray):
    """Detect face using OpenCV Haar Cascade. Returns (top, right, bottom, left)."""
    gray = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY)
    cascade_path = os.path.join(
        cv2.data.haarcascades,
        "haarcascade_frontalface_default.xml",
    )
    face_cascade = cv2.CascadeClassifier(cascade_path)

    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(60, 60),
    )

    if len(faces) == 0:
        return None

    # Choose the largest detected face.
    x, y, w, h = max(faces, key=lambda f: f[2] * f[3])
    return (y, x + w, y + h, x)


def _fallback_center_face_box(image_rgb: np.ndarray):
    """Fallback crop when no face is detected, so API still returns an image."""
    h, w = image_rgb.shape[:2]
    box_size = int(min(w, h) * 0.72)
    cx, cy = w // 2, h // 2
    left = max(cx - box_size // 2, 0)
    right = min(cx + box_size // 2, w)
    top = max(cy - box_size // 2, 0)
    bottom = min(cy + box_size // 2, h)
    return (top, right, bottom, left)


def sliding_window_tensor(
    input_tensor,
    window_size,
    stride,
    your_model,
    mask=mask_file,
    small_mask=small_mask_file,
):
    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]
            input_variable = Variable(window, requires_grad=False)

            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)
    output_tensor /= count_tensor
    output_tensor *= mask

    return output_tensor.cpu()


def full_image_tensor(
    input_tensor,
    your_model,
    mask=mask_file,
):
    """Run the model once on the full 1024x1024 tensor.

    This replaces the old sliding-window path that ran the model 9 times
    and blended overlapping patches. It is much faster, but output can be
    slightly different from the previous blended result.
    """
    device = next(your_model.parameters()).device
    input_tensor = input_tensor.to(device)
    mask = mask.to(device)

    with torch.no_grad():
        output_tensor = your_model(input_tensor)

    output_tensor = output_tensor * mask
    return output_tensor.cpu()


def process_image(your_model, image, source_age, target_age=0, window_size=512, stride=256):
    input_size = (1024, 1024)

    image_np = np.array(image.convert("RGB"))

    fl = _detect_face_box(image_np)
    if fl is None:
        fl = _fallback_center_face_box(image_np)

    margin_y_t = int((fl[2] - fl[0]) * 0.63 * 0.85)
    margin_y_b = int((fl[2] - fl[0]) * 0.37 * 0.85)
    margin_x = int((fl[1] - fl[3]) // (2 / 0.85))
    margin_y_t += 2 * margin_x - margin_y_t - margin_y_b

    l_y = max(fl[0] - margin_y_t, 0)
    r_y = min(fl[2] + margin_y_b, image_np.shape[0])
    l_x = max(fl[3] - margin_x, 0)
    r_x = min(fl[1] + margin_x, image_np.shape[1])

    cropped_image = image_np[l_y:r_y, l_x:r_x, :]
    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)

    original_image_tensor = transforms.ToTensor()(image_np)

    # Fast path: run the model once on the full 1024x1024 crop.
    # Old rollback path:
    # aged_cropped_image = sliding_window_tensor(input_tensor, window_size, stride, your_model)
    aged_cropped_image = full_image_tensor(
        input_tensor,
        your_model,
    )

    aged_cropped_image_resized = transforms.Resize(
        orig_size,
        interpolation=Image.BILINEAR,
        antialias=True,
    )(aged_cropped_image)

    original_image_tensor[:, l_y:r_y, l_x:r_x] += aged_cropped_image_resized.squeeze(0)
    original_image_tensor = torch.clamp(original_image_tensor, 0, 1)

    return transforms.functional.to_pil_image(original_image_tensor)