methods/alphaface/eval.py

import argparse
import pdb

import matplotlib.pyplot as plt
import torch
from Models.Swapper_AlphaFace import build_AlphaFace
from utils.utils_config import get_config
from utils.utils_distributed_sampler import setup_seed
import os
from dataset.get_dataloader import get_dataloader
import torchvision.transforms as transforms
from PIL import Image
from Objectives.Loss import *
from torchvision.utils import make_grid
from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from tqdm import tqdm
from pathlib import Path
from typing import Iterable, List, Union


def list_images(
    directory: Union[str, Path],
    extensions: Union[Iterable[str], None] = None,
    recursive: bool = False
) -> List[Path]:
    """
    Return a sorted list of image files found in *directory*.

    Parameters
    ----------
    directory : str | pathlib.Path
        Folder to search.
    extensions : Iterable[str] | None, default None
        File-name extensions to keep (case-insensitive, with or without the dot).
        If None, a sensible default set is used.
    recursive : bool, default False
        If True, search all nested sub-directories; otherwise, look only at the top
        level of *directory*.

    Returns
    -------
    list[pathlib.Path]
        Absolute, sorted paths to the matching image files.
    """
    # Default, widely-used image extensions
    default_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff",
                    ".gif", ".webp", ".ppm", ".pgm"}
    exts = default_exts if extensions is None else {
        (e.lower() if e.startswith(".") else f".{e.lower()}") for e in extensions
    }

    directory = Path(directory).expanduser().resolve()

    if recursive:
        paths = (p for p in directory.rglob("*") if p.is_file())
    else:
        paths = (p for p in directory.iterdir() if p.is_file())

    return sorted(p for p in paths if p.suffix.lower() in exts)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def normalize_by_127_5(img):
    """
    Normalize a tensor image by dividing by its maximum value.

    Args:
        img (torch.Tensor): Input image tensor.

    Returns:
        torch.Tensor: Normalized image tensor with values in [0, 1].
    """
    img = (img * 255.0).int()
    return (img / 127.5) - 1.0  # Return unchanged if max is 0 (e.g., all-zero image)


# Return unchanged if max is 0 (e.g., all-zero image)


def tensor2img(tensor):
    """
    Normalize a tensor image by dividing by its maximum value.

    Args:
        img (torch.Tensor): Input image tensor.

    Returns:
        torch.Tensor: Normalized image tensor with values in [0, 1].
    """
    return (tensor * 255.0).int()  # Return unchanged if max is 0 (e.g., all-zero image)


def eval_alphaface(config, model):
    writer = SummaryWriter(config.tb_dir)

    if not os.path.exists(config.output):
        os.makedirs(config.output)
        print(f"Directory '{config.output}' created.")
    else:
        print(f"Directory '{config.output}' already exists.")
   
    print('Resuming from checkpoint...from %s'%(config.model_path))
    dict_checkpoint = torch.load(config.model_path)
    print('from '+config.model_path)
    model.Swapper.load_state_dict(dict_checkpoint["swapper"])

    # Get source identity code for specific face identity swapping
    model = model.cuda()
    print(f"Number of parameters in Swapper: {count_parameters(model.Swapper)}")
    model.Swapper.eval()
    model.Id_encoder.eval()


    t_transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor()
    ])

    s_transform =transforms.Compose([
        transforms.Resize((112, 112)),
        transforms.ToTensor(),
        transforms.Lambda(normalize_by_127_5)
    ])

    src_img_list = list_images(config.src_path)
    tar_img_list = list_images(config.tar_path)
    #pdb.set_trace()
    #pdb.set_trace()
    for src_img_file in src_img_list:
        img1_s = s_transform(Image.open(src_img_file)).unsqueeze(0)
        img1_s = img1_s.cuda()
        
        for tar_img_file in tar_img_list:
        
            img2_t = t_transform(Image.open(tar_img_file)).unsqueeze(0)
            img2_t = img2_t.cuda()
            
            swapped_2_1 = model(img2_t, img1_s)
                

            # If the tensor is batched, take the first item
            if swapped_2_1.dim() == 4:
                swapped_2_1 = swapped_2_1[0]
        
            # If the tensor is in [C, H, W] format, convert to [H, W, C]
            if swapped_2_1.shape[0] in {1, 3, 4}:
                swapped_2_1 = swapped_2_1.permute(1, 2, 0)
        
            # Clamp to valid range and convert to uint8
            if swapped_2_1.max() <= 1.0:
                swapped_2_1 = swapped_2_1 * 255.0
        
            swapped_2_1 = swapped_2_1.clamp(0, 255).byte()
        
            # Convert to numpy and PIL
            np_image = swapped_2_1.cpu().numpy()
            img = Image.fromarray(np_image)
        
            # If single channel, convert to 'L' mode explicitly
            if np_image.shape[2] == 1:
                img = img.convert("L")
            #pdb.set_trace()
            save_img_path = str(src_img_file).split('/')[-1].split('.')[0]+'_'+str(tar_img_file).split('/')[-1]
        
            img.save(os.path.join(config.output,save_img_path))
            print(f"Saved image to {os.path.join(config.output,save_img_path)}")
            
            #print('source: %s | target: %s > %s'%(src_img_file,tar_img_list,os.path.join(config.output,save_img_path)))
    
     


def main(args):
    cfg = get_config(args.config)
    setup_seed(seed=cfg.seed, cuda_deterministic=False)

    if cfg.tensorboard == True:
        writer = SummaryWriter(cfg.tb_path)
    else:
        writer = None
    os.makedirs(cfg.log_dir, exist_ok=True)

    # Get dataloader for train

    
    # Model initialisation
    print('Preparing the student model')
    alphaface = build_AlphaFace(config=cfg).to('cuda')


    eval_alphaface(cfg, alphaface)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Distributed Arcface Training in Pytorch")
    parser.add_argument("config", type=str, default='./configs/test_config', help="py config file")
    main(parser.parse_args())