|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
import pdb |
|
import shutil |
|
import sys |
|
|
|
sys.path.append("./") |
|
from collections import defaultdict |
|
|
|
import numpy as np |
|
import torch |
|
import torch.nn.functional as F |
|
from PIL import Image |
|
from prettytable import PrettyTable |
|
from torch.utils.data import Dataset |
|
from torchmetrics.image import StructuralSimilarityIndexMeasure |
|
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity |
|
from torchvision import transforms |
|
from tqdm import tqdm |
|
from tqlt import utils as tu |
|
|
|
from openlrm.models.arcface_utils import ResNetArcFace |
|
from openlrm.utils.face_detector import FaceDetector |
|
|
|
device = "cuda" |
|
model_path = "./pretrained_models/gagatracker/vgghead/vgg_heads_l.trcd" |
|
face_detector = FaceDetector(model_path=model_path, device=device) |
|
|
|
id_face_net = ResNetArcFace() |
|
id_face_net.cuda() |
|
id_face_net.eval() |
|
|
|
|
|
def get_image_paths_current_dir(folder_path): |
|
image_extensions = { |
|
".jpg", |
|
".jpeg", |
|
".png", |
|
".gif", |
|
".bmp", |
|
".tiff", |
|
".webp", |
|
".jfif", |
|
} |
|
|
|
return sorted( |
|
[ |
|
os.path.join(folder_path, f) |
|
for f in os.listdir(folder_path) |
|
if os.path.splitext(f)[1].lower() in image_extensions |
|
] |
|
) |
|
|
|
|
|
def write_json(path, x): |
|
"""write a json file. |
|
|
|
Args: |
|
path (str): path to write json file. |
|
x (dict): dict to write. |
|
""" |
|
import json |
|
|
|
with open(path, "w") as f: |
|
json.dump(x, f, indent=2) |
|
|
|
|
|
def crop_face_image(image_path): |
|
rgb = np.array(Image.open(image_path)) |
|
rgb = torch.from_numpy(rgb).permute(2, 0, 1) |
|
bbox = face_detector(rgb) |
|
head_rgb = rgb[:, int(bbox[1]) : int(bbox[3]), int(bbox[0]) : int(bbox[2])] |
|
head_rgb = head_rgb.permute(1, 2, 0) |
|
head_rgb = head_rgb.cpu().numpy() |
|
return head_rgb |
|
|
|
|
|
def gray_resize_for_identity(out, size=128): |
|
out_gray = ( |
|
0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :] |
|
) |
|
out_gray = out_gray.unsqueeze(1) |
|
out_gray = F.interpolate( |
|
out_gray, (size, size), mode="bilinear", align_corners=False |
|
) |
|
return out_gray |
|
|
|
|
|
@torch.no_grad() |
|
def eval(input_folder, target_folder, device="cuda"): |
|
|
|
gt_imgs = get_image_paths_current_dir(target_folder) |
|
result_imgs = get_image_paths_current_dir(input_folder) |
|
|
|
gt_imgs = gt_imgs[::4] |
|
result_imgs = result_imgs[::4] |
|
|
|
if "visualization" in result_imgs[-1]: |
|
result_imgs = result_imgs[:-1] |
|
|
|
if len(gt_imgs) != len(result_imgs): |
|
return -1 |
|
|
|
to_tensor = transforms.ToTensor() |
|
|
|
face_id_loss_list = [] |
|
|
|
for input_img, gt_img in zip(result_imgs, gt_imgs): |
|
|
|
try: |
|
input_img = crop_face_image(input_img) |
|
input_head_tensor = gray_resize_for_identity( |
|
to_tensor(input_img).unsqueeze(0).to(device) |
|
) |
|
input_head_feature = id_face_net(input_head_tensor).detach() |
|
|
|
head_img = crop_face_image(gt_img) |
|
head_img = to_tensor(head_img).unsqueeze(0).to(device) |
|
src_head_tensor = gray_resize_for_identity(head_img) |
|
src_head_feature = id_face_net(src_head_tensor).detach() |
|
|
|
face_id_loss = F.l1_loss(input_head_feature, src_head_feature) |
|
face_id_loss_list.append(face_id_loss.item()) |
|
except: |
|
continue |
|
|
|
if len(face_id_loss_list) > 0: |
|
|
|
return np.mean(face_id_loss_list) |
|
else: |
|
return -1 |
|
|
|
|
|
def get_parse(): |
|
import argparse |
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("-f1", "--folder1", type=str, required=True) |
|
parser.add_argument("-f2", "--folder2", type=str, required=True) |
|
parser.add_argument("--pad", action="store_true") |
|
parser.add_argument("--debug", action="store_true") |
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
opt = get_parse() |
|
|
|
input_folder = opt.folder1 |
|
target_folder = opt.folder2 |
|
|
|
save_folder = os.path.join("./exps/metricsanigs", "psnr_results") |
|
os.makedirs(save_folder, exist_ok=True) |
|
|
|
input_folders = tu.next_folders(input_folder) |
|
|
|
results_dict = defaultdict(dict) |
|
|
|
results_dict = defaultdict(dict) |
|
face_similarity_list = [] |
|
|
|
for input_folder in input_folders: |
|
|
|
item_basename = tu.basename(input_folder) |
|
|
|
mask_item_folder = None |
|
input_item_folder = os.path.join(input_folder, "rgb") |
|
target_item_folder = os.path.join(target_folder, item_basename) |
|
|
|
if os.path.exists(input_item_folder) and os.path.exists(target_item_folder): |
|
|
|
fs_ = eval(input_item_folder, target_item_folder) |
|
|
|
if fs_ == -1: |
|
continue |
|
|
|
face_similarity_list.append(fs_) |
|
|
|
results_dict[item_basename]["face_similarity"] = fs_ |
|
if opt.debug: |
|
break |
|
print(results_dict) |
|
|
|
results_dict["all_mean"]["face_similarity"] = np.mean(face_similarity_list) |
|
|
|
write_json(os.path.join(save_folder, "face_similarity.json"), results_dict) |
|
|
