import json
import os
import pathlib
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image
from pytorch_fid.inception import InceptionV3
from scipy import linalg
from torch.nn.functional import adaptive_avg_pool2d
from tools.metrics.utils import tracker
try:
from tqdm import tqdm
except ImportError:
# If tqdm is not available, provide a mock version of it
def tqdm(x):
return x
IMAGE_EXTENSIONS = {"bmp", "jpg", "jpeg", "pgm", "png", "ppm", "tif", "tiff", "webp"}
class ImagePathDataset(torch.utils.data.Dataset):
def __init__(self, files, transforms=None):
self.files = files
self.transforms = transforms
def __len__(self):
return len(self.files)
def __getitem__(self, i):
path = self.files[i]
try:
img = Image.open(path)
assert img.mode == "RGB"
if self.transforms is not None:
img = self.transforms(img)
except Exception as e:
raise FileNotFoundError(path, "\n", e)
return img
def get_activations(files, model, batch_size=50, dims=2048, device="cpu", num_workers=1):
model.eval()
if batch_size > len(files):
print("Warning: batch size is bigger than the data size. " "Setting batch size to data size")
batch_size = len(files)
transform = T.Compose(
[
T.Resize(args.img_size), # Image.BICUBIC
T.CenterCrop(args.img_size),
T.ToTensor(),
]
)
dataset = ImagePathDataset(files, transforms=transform)
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=False, drop_last=False, num_workers=num_workers
)
pred_arr = np.empty((len(files), dims))
start_idx = 0
for batch in tqdm(dataloader, desc=f"FID: {args.exp_name}", position=args.gpu_id, leave=True):
batch = batch.to(device)
with torch.no_grad():
pred = model(batch)[0]
# If model output is not scalar, apply global spatial average pooling.
# This happens if you choose a dimensionality not equal 2048.
if pred.size(2) != 1 or pred.size(3) != 1:
pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
pred = pred.squeeze(3).squeeze(2).cpu().numpy()
pred_arr[start_idx : start_idx + pred.shape[0]] = pred
start_idx = start_idx + pred.shape[0]
return pred_arr
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
mu1 = np.atleast_1d(mu1)
mu2 = np.atleast_1d(mu2)
sigma1 = np.atleast_2d(sigma1)
sigma2 = np.atleast_2d(sigma2)
assert mu1.shape == mu2.shape, "Training and test mean vectors have different lengths"
assert sigma1.shape == sigma2.shape, "Training and test covariances have different dimensions"
diff = mu1 - mu2
# Product might be almost singular
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
if not np.isfinite(covmean).all():
msg = ("fid calculation produces singular product; " "adding %s to diagonal of cov estimates") % eps
print(msg)
offset = np.eye(sigma1.shape[0]) * eps
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
# Numerical error might give slight imaginary component
if np.iscomplexobj(covmean):
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
m = np.max(np.abs(covmean.imag))
raise ValueError(f"Imaginary component {m}")
covmean = covmean.real
tr_covmean = np.trace(covmean)
return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
def calculate_activation_statistics(files, model, batch_size=50, dims=2048, device="cpu", num_workers=1):
act = get_activations(files, model, batch_size, dims, device, num_workers)
mu = np.mean(act, axis=0)
sigma = np.cov(act, rowvar=False)
return mu, sigma
def compute_statistics_of_path(path, model, batch_size, dims, device, num_workers=1, flag="ref"):
if path.endswith(".npz"):
print("loaded from npz files")
with np.load(path) as f:
m, s = f["mu"][:], f["sigma"][:]
elif path.endswith(".json"):
with open(path) as file:
data_dict = json.load(file)
all_lines = list(data_dict.keys())[:sample_nums]
files = []
if isinstance(all_lines, list):
for k in all_lines:
v = data_dict[k]
if "PG-eval-data" in args.img_path:
img_path = os.path.join(args.img_path, v["category"], f"{k}.jpg")
else:
img_path = os.path.join(args.img_path, args.exp_name, f"{k}.jpg")
files.append(img_path)
elif isinstance(all_lines, dict):
assert sample_nums >= 30_000, ValueError(f"{sample_nums} is not supported for json files")
for k, v in all_lines.items():
if "PG-eval-data" in args.img_path:
img_path = os.path.join(args.img_path, v["category"], f"{k}.jpg")
else:
img_path = os.path.join(args.img_path, args.exp_name, f"{k}.jpg")
files.append(img_path)
files = sorted(files)
m, s = calculate_activation_statistics(files, model, batch_size, dims, device, num_workers)
else:
path = pathlib.Path(path)
files = sorted([file for ext in IMAGE_EXTENSIONS for file in path.glob(f"*.{ext}")])
m, s = calculate_activation_statistics(files, model, batch_size, dims, device, num_workers)
return m, s
def calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1):
"""Calculates the FID of two paths"""
for p in paths:
if not os.path.exists(p):
raise RuntimeError("Invalid path: %s" % p)
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
model = InceptionV3([block_idx]).to(device)
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size, dims, device, num_workers, flag="ref")
m2, s2 = compute_statistics_of_path(paths[1], model, batch_size, dims, device, num_workers, flag="gen")
fid_value = calculate_frechet_distance(m1, s1, m2, s2)
return fid_value
def save_fid_stats(paths, batch_size, device, dims, num_workers=1):
"""Calculates the FID of two paths"""
if not os.path.exists(paths[0]):
raise RuntimeError("Invalid path: %s" % paths[0])
if os.path.exists(paths[1]):
raise RuntimeError("Existing output file: %s" % paths[1])
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
model = InceptionV3([block_idx]).to(device)
print(f"Saving statistics for {paths[0]}")
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size, dims, device, num_workers, flag="ref")
np.savez_compressed(paths[1], mu=m1, sigma=s1)
def main():
txt_path = args.txt_path if args.txt_path is not None else args.img_path
save_txt_path = os.path.join(txt_path, f"{args.exp_name}_sample{sample_nums}.txt")
if os.path.exists(save_txt_path):
with open(save_txt_path) as f:
fid_value = f.readlines()[0].strip()
print(f"FID {fid_value}: {args.exp_name}")
return {args.exp_name: float(fid_value)}
if args.device is None:
device = torch.device("cuda" if (torch.cuda.is_available()) else "cpu")
else:
device = torch.device(args.device)
if args.num_workers is None:
try:
num_cpus = len(os.sched_getaffinity(0))
except AttributeError:
num_cpus = os.cpu_count()
num_workers = min(num_cpus, 8) if num_cpus is not None else 0
else:
num_workers = args.num_workers
if args.save_stats:
save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers)
return
fid_value = calculate_fid_given_paths(args.path, args.batch_size, device, args.dims, num_workers)
print(f"FID {fid_value}: {args.exp_name}")
with open(save_txt_path, "w") as file:
file.write(str(fid_value))
return {args.exp_name: fid_value}
def parse_args():
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument("--batch-size", type=int, default=50, help="Batch size to use")
parser.add_argument(
"--num-workers", type=int, help="Number of processes to use for data loading. Defaults to `min(8, num_cpus)`"
)
parser.add_argument("--img_size", type=int, default=512)
parser.add_argument("--device", type=str, default="cuda", help="Device to use. Like cuda, cuda:0 or cpu")
parser.add_argument("--img_path", type=str, default=None)
parser.add_argument("--exp_name", type=str, default="Sana")
parser.add_argument("--txt_path", type=str, default=None)
parser.add_argument("--sample_nums", type=int, default=30_000)
parser.add_argument(
"--dims",
type=int,
default=2048,
choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
help="Dimensionality of Inception features to use. By default, uses pool3 features",
)
parser.add_argument(
"--save-stats",
action="store_true",
help="Generate an npz archive from a directory of samples. The first path is used as input and the second as output.",
)
parser.add_argument("--stat", action="store_true")
parser.add_argument(
"--path", type=str, nargs=2, default=["", ""], help="Paths to the generated images or to .npz statistic files"
)
# online logging setting
parser.add_argument("--log_metric", type=str, default="metric")
parser.add_argument("--gpu_id", type=int, default=0)
parser.add_argument("--log_fid", action="store_true")
parser.add_argument("--suffix_label", type=str, default="", help="used for fid online log")
parser.add_argument("--tracker_pattern", type=str, default="epoch_step", help="used for fid online log")
parser.add_argument(
"--report_to",
type=str,
default=None,
help=(
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
),
)
parser.add_argument(
"--tracker_project_name",
type=str,
default="t2i-evit-baseline",
help=(
"The `project_name` argument passed to Accelerator.init_trackers for"
" more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"
),
)
parser.add_argument(
"--name",
type=str,
default="baseline",
help=("Wandb Project Name"),
)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
sample_nums = args.sample_nums
if args.stat:
if args.device is None:
device = torch.device("cuda" if (torch.cuda.is_available()) else "cpu")
else:
device = torch.device(args.device)
if args.num_workers is None:
try:
num_cpus = len(os.sched_getaffinity(0))
except AttributeError:
num_cpus = os.cpu_count()
num_workers = min(num_cpus, 8) if num_cpus is not None else 0
else:
num_workers = args.num_workers
save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers)
else:
print(args.path, args.exp_name)
args.exp_name = os.path.basename(args.exp_name) or os.path.dirname(args.exp_name)
fid_result = main()
if args.log_fid:
tracker(args, fid_result, args.suffix_label, pattern=args.tracker_pattern, metric="FID")