Add files using upload-large-folder tool

112139c verified 3 months ago

5.5 kB

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
import torch
import numpy as np
from scipy import linalg
from tqdm import tqdm

import sys
sys.path.append(".")
sys.path.append("..")

from configs.paths_config import val_paths
from models.Inception.inception import InceptionV3
from models.triencoder import TriEncoder
from datasets.dataset import ImageFolderDataset
from torch.utils.data import DataLoader
from configs.paths_config import dataset_paths, val_paths
# valuate each pretrain weight...
# self used script ...

device = 'cuda'

def gen_rand_pose(pitch_range=26, yaw_range=36, cx_range=0.2, cy_range=0.2, fov_range=4.8, mode="yaw"):
    # set range
    # pitch:  +-26, yaw: +-36/+-49
    if mode == "yaw":
        return (torch.rand(1, device=device) - 0.5) * (pitch_range / 180 * torch.pi) + torch.pi/2
    elif mode == "pitch":
        return (torch.rand(1, device=device) - 0.5) * (yaw_range / 180 * torch.pi) + torch.pi/2
    else:
        raise NotImplementedError

def calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6):
    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)

    if not np.isfinite(cov_sqrt).all():
        print("product of cov matrices is singular")
        offset = np.eye(sample_cov.shape[0]) * eps
        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))

    if np.iscomplexobj(cov_sqrt):
        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
            m = np.max(np.abs(cov_sqrt.imag))

            raise ValueError(f"Imaginary component {m}")

        cov_sqrt = cov_sqrt.real

    mean_diff = sample_mean - real_mean
    mean_norm = mean_diff @ mean_diff

    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)

    fid = mean_norm + trace

    return fid

@torch.no_grad()
def compute_fid(net, inception, model_path, dataloader, batch_size, n_sample=10000):
    basename = os.path.basename(model_path)

    n_batch = n_sample // batch_size
    resid = n_sample - (n_batch * batch_size)

    batch_sizes = [batch_size] * n_batch
    batch_sizes += [resid] if resid > 0 else []

    sample_mean = np.load(val_paths['ffhq_mean'])
    sample_cov = np.load(val_paths['ffhq_cov'])

    mean_savepath = f"assets/temp/{basename}_fid_pred_sample_mean{n_sample}.npy"
    cov_savepath = f"assets/temp/{basename}_fid_pred_sample_cov{n_sample}.npy"
    if os.path.isfile(mean_savepath) and os.path.isfile(cov_savepath):
        pred_sample_mean = np.load(mean_savepath)
        pred_sample_cov = np.load(cov_savepath)
    else:
        print("pred.....")
        pred_features = []
        # pred_features = np.zeros((n_sample, 2048))
        # for idx, batch in enumerate(tqdm(batch_sizes)):
        for batch in tqdm(batch_sizes):
            # _, encoder_input, camera_param, ws = net.sample_triplane(batch, gen_rand_pose(pitch_range=26, mode='pitch'), gen_rand_pose(yaw_range=49, mode='yaw'))
            # gen_triplanes = net.encoder(encoder_input['image'])
            # render_res = net.triplane_renderer(gen_triplanes, camera_param)
            # img = render_res['image']
            image, label = next(iter(dataloader))
            image = image.to(device) # [B, 3, 512, 512] range: (0, 255) -> (-1, 1) should divide 127.5 and sub 1
            label = label.to(device) # [B, 25]
            gen_triplanes = net.encoder((image.float()/127.5 - 1))
            render_res = net.triplane_renderer(gen_triplanes, label)
            img = render_res['image']
            feat = inception(img)[0].view(img.shape[0], -1)
            pred_features.append(feat.to("cpu"))
            # pred_features[idx * batch_size: idx * batch_size + batch] = feat.to("cpu").numpy()
        # import pdb; pdb.set_trace()
        pred_features = torch.cat(pred_features, 0).numpy()
        pred_sample_mean = np.mean(pred_features, 0)
        pred_sample_cov = np.cov(pred_features, rowvar=False)

        np.save(f"assets/temp/{basename}_fid_pred_sample_mean{n_sample}.npy", pred_sample_mean)
        np.save(f"assets/temp/{basename}_fid_pred_sample_cov{n_sample}.npy", pred_sample_cov)
        print("pred done")
    print("calc fid....")
    # ---------------------------------------------------------------------------------------------------------
    fid = calc_fid(pred_sample_mean, pred_sample_cov, sample_mean, sample_cov)
    print("fid:", fid)
    pass


def testfid50k(model_path):
    # load model
    batch_size=16
    model_dict = torch.load(model_path)
    net = TriEncoder(mode="Normal").to(device)
    net.encoder.load_state_dict(model_dict['encoder_state_dict'])
    net.triplane_renderer.load_state_dict(model_dict['renderer_state_dict'])

    path = dataset_paths["ffhq512"]
    image_datasets = ImageFolderDataset(path=path, resolution=None, use_labels=True)
    dataloader = DataLoader(dataset=image_datasets, batch_size=batch_size, shuffle=True, num_workers=8)

    inception = InceptionV3([3], normalize_input=False).to(device)
    inception.eval()
    compute_fid(net, inception, model_path, dataloader, batch_size=batch_size, n_sample=50000)
    del net, inception

def main():
    ckpt_path = "exp10/checkpoints"
    ckpt_list = os.listdir(ckpt_path)
    ckpt_list = ['iteration_140000.pt']
    for ckpt in ckpt_list:
        if ckpt.endswith(".pt"):
            model_path = os.path.join(ckpt_path, ckpt)
            print(model_path)
            print("test fid 50k...")
            testfid50k(model_path)
    pass

if __name__ == '__main__':
    main()
    pass

	import os
	os.environ['CUDA_VISIBLE_DEVICES'] = '3'
	import torch
	import numpy as np
	from scipy import linalg
	from tqdm import tqdm

	import sys
	sys.path.append(".")
	sys.path.append("..")

	from configs.paths_config import val_paths
	from models.Inception.inception import InceptionV3
	from models.triencoder import TriEncoder
	from datasets.dataset import ImageFolderDataset
	from torch.utils.data import DataLoader
	from configs.paths_config import dataset_paths, val_paths
	# valuate each pretrain weight...
	# self used script ...

	device = 'cuda'

	def gen_rand_pose(pitch_range=26, yaw_range=36, cx_range=0.2, cy_range=0.2, fov_range=4.8, mode="yaw"):
	# set range
	# pitch: +-26, yaw: +-36/+-49
	if mode == "yaw":
	return (torch.rand(1, device=device) - 0.5) * (pitch_range / 180 * torch.pi) + torch.pi/2
	elif mode == "pitch":
	return (torch.rand(1, device=device) - 0.5) * (yaw_range / 180 * torch.pi) + torch.pi/2
	else:
	raise NotImplementedError

	def calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6):
	cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)

	if not np.isfinite(cov_sqrt).all():
	print("product of cov matrices is singular")
	offset = np.eye(sample_cov.shape[0]) * eps
	cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))

	if np.iscomplexobj(cov_sqrt):
	if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
	m = np.max(np.abs(cov_sqrt.imag))

	raise ValueError(f"Imaginary component {m}")

	cov_sqrt = cov_sqrt.real

	mean_diff = sample_mean - real_mean
	mean_norm = mean_diff @ mean_diff

	trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)

	fid = mean_norm + trace

	return fid

	@torch.no_grad()
	def compute_fid(net, inception, model_path, dataloader, batch_size, n_sample=10000):
	basename = os.path.basename(model_path)

	n_batch = n_sample // batch_size
	resid = n_sample - (n_batch * batch_size)

	batch_sizes = [batch_size] * n_batch
	batch_sizes += [resid] if resid > 0 else []

	sample_mean = np.load(val_paths['ffhq_mean'])
	sample_cov = np.load(val_paths['ffhq_cov'])

	mean_savepath = f"assets/temp/{basename}_fid_pred_sample_mean{n_sample}.npy"
	cov_savepath = f"assets/temp/{basename}_fid_pred_sample_cov{n_sample}.npy"
	if os.path.isfile(mean_savepath) and os.path.isfile(cov_savepath):
	pred_sample_mean = np.load(mean_savepath)
	pred_sample_cov = np.load(cov_savepath)
	else:
	print("pred.....")
	pred_features = []
	# pred_features = np.zeros((n_sample, 2048))
	# for idx, batch in enumerate(tqdm(batch_sizes)):
	for batch in tqdm(batch_sizes):
	# _, encoder_input, camera_param, ws = net.sample_triplane(batch, gen_rand_pose(pitch_range=26, mode='pitch'), gen_rand_pose(yaw_range=49, mode='yaw'))
	# gen_triplanes = net.encoder(encoder_input['image'])
	# render_res = net.triplane_renderer(gen_triplanes, camera_param)
	# img = render_res['image']
	image, label = next(iter(dataloader))
	image = image.to(device) # [B, 3, 512, 512] range: (0, 255) -> (-1, 1) should divide 127.5 and sub 1
	label = label.to(device) # [B, 25]
	gen_triplanes = net.encoder((image.float()/127.5 - 1))
	render_res = net.triplane_renderer(gen_triplanes, label)
	img = render_res['image']
	feat = inception(img)[0].view(img.shape[0], -1)
	pred_features.append(feat.to("cpu"))
	# pred_features[idx * batch_size: idx * batch_size + batch] = feat.to("cpu").numpy()
	# import pdb; pdb.set_trace()
	pred_features = torch.cat(pred_features, 0).numpy()
	pred_sample_mean = np.mean(pred_features, 0)
	pred_sample_cov = np.cov(pred_features, rowvar=False)

	np.save(f"assets/temp/{basename}_fid_pred_sample_mean{n_sample}.npy", pred_sample_mean)
	np.save(f"assets/temp/{basename}_fid_pred_sample_cov{n_sample}.npy", pred_sample_cov)
	print("pred done")
	print("calc fid....")
	# ---------------------------------------------------------------------------------------------------------
	fid = calc_fid(pred_sample_mean, pred_sample_cov, sample_mean, sample_cov)
	print("fid:", fid)
	pass


	def testfid50k(model_path):
	# load model
	batch_size=16
	model_dict = torch.load(model_path)
	net = TriEncoder(mode="Normal").to(device)
	net.encoder.load_state_dict(model_dict['encoder_state_dict'])
	net.triplane_renderer.load_state_dict(model_dict['renderer_state_dict'])

	path = dataset_paths["ffhq512"]
	image_datasets = ImageFolderDataset(path=path, resolution=None, use_labels=True)
	dataloader = DataLoader(dataset=image_datasets, batch_size=batch_size, shuffle=True, num_workers=8)

	inception = InceptionV3([3], normalize_input=False).to(device)
	inception.eval()
	compute_fid(net, inception, model_path, dataloader, batch_size=batch_size, n_sample=50000)
	del net, inception

	def main():
	ckpt_path = "exp10/checkpoints"
	ckpt_list = os.listdir(ckpt_path)
	ckpt_list = ['iteration_140000.pt']
	for ckpt in ckpt_list:
	if ckpt.endswith(".pt"):
	model_path = os.path.join(ckpt_path, ckpt)
	print(model_path)
	print("test fid 50k...")
	testfid50k(model_path)
	pass

	if __name__ == '__main__':
	main()
	pass