ic_gan / predict.py

ArantxaCasanova

First model version

a00ee36 over 2 years ago

No virus

11.6 kB

	# Copyright (c) Facebook, Inc. and its affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the license found in the
	# LICENSE file in the root directory of this source tree.

	import sys
	import os
	import tempfile
	import warnings
	from pathlib import Path
	import nltk
	import torch
	from torch import nn
	import torchvision.transforms as transforms
	import numpy as np
	import imageio
	from PIL import Image as Image_PIL
	from scipy.stats import truncnorm
	from nltk.corpus import wordnet as wn
	import cma
	import sklearn.metrics
	import cog

	sys.path.insert(0, "stylegan2_ada_pytorch")
	from pytorch_pretrained_biggan import convert_to_images, utils
	import inference.utils as inference_utils
	import data_utils.utils as data_utils

	NORM_MEAN = torch.Tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
	NORM_STD = torch.Tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

	nltk.download("wordnet")
	IND2NAME = {
	index: wn.of2ss("%08dn" % offset).lemma_names()[0]
	for offset, index in utils.IMAGENET.items()
	}
	NAME2IND = dict([(value, key) for key, value in IND2NAME.items()])

	CLASS_NAMES = sorted(list(IND2NAME.values()))


	class Predictor(cog.Predictor):
	def setup(self):
	torch.manual_seed(np.random.randint(sys.maxsize))
	warnings.simplefilter("ignore", cma.evolution_strategy.InjectionWarning)
	self.last_gen_model = None
	self.last_feature_extractor = None
	self.model = None
	self.feature_extractor = None
	self.noise_size = 128
	self.batch_size = 4
	self.size = 256

	@cog.input("image", type=Path, help="Input image Instance")
	@cog.input("gen_model", type=str, options=["icgan", "cc_icgan"], default="icgan",
	help='Select type of IC-GAN model. "icgan" is conditioned on the input image; '
	'"cc_icgan" is conditioned on both input image and a conditional_class')
	@cog.input("conditional_class", type=str, default=None, options=CLASS_NAMES,
	help="Choose conditional class. Only valid for gen_model=cc_icgan")
	@cog.input("num_samples", type=int, default=1, options=[1, 4, 9, 16],
	help="number of samples generated")
	@cog.input("seed", type=int, default=0, help="seed=0 means no seed")
	def predict(self, image, gen_model="icgan", conditional_class=None, num_samples=1, seed=0):
	assert isinstance(seed, int), "seed should be an integer"
	if gen_model == 'cc_icgan':
	assert conditional_class is not None, 'please set conditional_class for cc_icgan'
	num_samples_ranked = num_samples
	experiment_name = (
	"icgan_biggan_imagenet_res256"
	if gen_model == "icgan"
	else "cc_icgan_biggan_imagenet_res256"
	)
	num_samples_total = num_samples * 10
	truncation = 0.7
	if conditional_class is not None:
	class_index = NAME2IND[conditional_class]

	input_image_instance = str(image)

	if gen_model == "icgan":
	class_index = None

	if seed == 0:
	seed = None

	state = None if not seed else np.random.RandomState(seed)
	np.random.seed(seed)

	feature_extractor_name = ("classification" if gen_model == "cc_icgan" else "selfsupervised")

	# Load feature extractor (outlier filtering and optionally input image feature extraction)
	self.feature_extractor, self.last_feature_extractor = load_feature_extractor(
	gen_model, self.last_feature_extractor, self.feature_extractor)
	# Load features
	if input_image_instance not in ["None", "", None]:
	print("Obtaining instance features from input image!")
	input_feature_index = None
	input_image_tensor = preprocess_input_image(input_image_instance, self.size)
	with torch.no_grad():
	input_features, _ = self.feature_extractor(input_image_tensor.cuda())
	input_features /= torch.linalg.norm(input_features, dim=-1, keepdims=True)
	elif input_feature_index is not None:
	print("Selecting an instance from pre-extracted vectors!")
	input_features = np.load(
	"stored_instances/imagenet_res"
	+ str(self.size)
	+ "_rn50_"
	+ feature_extractor_name
	+ "_kmeans_k1000_instance_features.npy",
	allow_pickle=True,
	).item()["instance_features"][input_feature_index: input_feature_index + 1]
	else:
	input_features = None

	# Load generative model
	self.model, self.last_gen_model = load_generative_model(
	gen_model, self.last_gen_model, experiment_name, self.model)
	# Prepare other variables

	replace_to_inplace_relu(self.model)

	# Create noise, instance and class vector
	noise_vector = truncnorm.rvs(
	-2 * truncation,
	2 * truncation,
	size=(num_samples_total, self.noise_size),
	random_state=state,
	).astype(np.float32)
	noise_vector = torch.tensor(noise_vector, requires_grad=False, device="cuda")
	if input_features is not None:
	instance_vector = torch.tensor(
	input_features, requires_grad=False, device="cuda"
	).repeat(num_samples_total, 1)
	else:
	instance_vector = None
	if class_index is not None:
	input_label = torch.LongTensor([class_index] * num_samples_total)
	else:
	input_label = None
	if input_feature_index is not None:
	print("Conditioning on instance with index: ", input_feature_index)

	all_outs, all_dists = [], []
	for i_bs in range(num_samples_total // self.batch_size + 1):
	start = i_bs * self.batch_size
	end = min(start + self.batch_size, num_samples_total)
	if start == end:
	break
	out = get_output(
	noise_vector[start:end],
	input_label[start:end] if input_label is not None else None,
	instance_vector[start:end] if instance_vector is not None else None,
	self.model,
	truncation,
	channels=3,
	)

	if instance_vector is not None:
	# Get features from generated images + feature extractor
	out_ = preprocess_generated_image(out)
	with torch.no_grad():
	out_features, _ = self.feature_extractor(out_.cuda())
	out_features /= torch.linalg.norm(out_features, dim=-1, keepdims=True)
	dists = sklearn.metrics.pairwise_distances(
	out_features.cpu(),
	instance_vector[start:end].cpu(),
	metric="euclidean",
	n_jobs=-1,
	)
	all_dists.append(np.diagonal(dists))
	all_outs.append(out.detach().cpu())
	del out
	all_outs = torch.cat(all_outs)
	all_dists = np.concatenate(all_dists)

	# Order samples by distance to conditioning feature vector and select only num_samples_ranked images
	selected_idxs = np.argsort(all_dists)[:num_samples_ranked]
	# Create figure
	row_i, col_i, i_im = 0, 0, 0
	all_images_mosaic = np.zeros(
	(
	3,
	self.size * (int(np.sqrt(num_samples_ranked))),
	self.size * (int(np.sqrt(num_samples_ranked))),
	)
	)
	for j in selected_idxs:
	all_images_mosaic[
	:,
	row_i * self.size: row_i * self.size + self.size,
	col_i * self.size: col_i * self.size + self.size,
	] = all_outs[j]
	if row_i == int(np.sqrt(num_samples_ranked)) - 1:
	row_i = 0
	if col_i == int(np.sqrt(num_samples_ranked)) - 1:
	col_i = 0
	else:
	col_i += 1
	else:
	row_i += 1
	i_im += 1

	out_path = Path(tempfile.mkdtemp()) / "out.png"
	save(all_images_mosaic[np.newaxis, ...], str(out_path), torch_format=False)
	return out_path


	def replace_to_inplace_relu(model):
	for child_name, child in model.named_children():
	if isinstance(child, nn.ReLU):
	setattr(model, child_name, nn.ReLU(inplace=False))
	else:
	replace_to_inplace_relu(child)


	def save(out, name=None, torch_format=True):
	if torch_format:
	with torch.no_grad():
	out = out.cpu().numpy()
	img = convert_to_images(out)[0]
	if name:
	imageio.imwrite(name, np.asarray(img))
	return img


	def load_icgan(experiment_name, root_=""):
	root = os.path.join(root_, experiment_name)
	config = torch.load("%s/%s.pth" % (root, "state_dict_best0"))["config"]

	config["weights_root"] = root_
	config["model_backbone"] = "biggan"
	config["experiment_name"] = experiment_name
	G, config = inference_utils.load_model_inference(config)
	G.cuda()
	G.eval()
	return G


	def get_output(noise_vector, input_label, input_features, model, truncation, channels):
	# stochastic_truncation = False as how it is set in colab
	noise_vector = noise_vector.clamp(-2 * truncation, 2 * truncation)
	if input_label is not None:
	input_label = torch.LongTensor(input_label)
	else:
	input_label = None

	out = model(
	noise_vector,
	input_label.cuda() if input_label is not None else None,
	input_features.cuda() if input_features is not None else None,
	)

	if channels == 1:
	out = out.mean(dim=1, keepdim=True)
	out = out.repeat(1, 3, 1, 1)
	return out


	def load_generative_model(gen_model, last_gen_model, experiment_name, model):
	# Load generative model
	if gen_model != last_gen_model:
	model = load_icgan(experiment_name, root_="./")
	last_gen_model = gen_model
	return model, last_gen_model


	def load_feature_extractor(gen_model, last_feature_extractor, feature_extractor):
	# Load feature extractor to obtain instance features
	feat_ext_name = "classification" if gen_model == "cc_icgan" else "selfsupervised"
	if last_feature_extractor != feat_ext_name:
	if feat_ext_name == "classification":
	feat_ext_path = ""
	else:
	feat_ext_path = "swav_pretrained.pth.tar"
	last_feature_extractor = feat_ext_name
	feature_extractor = data_utils.load_pretrained_feature_extractor(
	feat_ext_path, feature_extractor=feat_ext_name
	)
	feature_extractor.eval()
	return feature_extractor, last_feature_extractor


	def preprocess_input_image(input_image_path, size):
	pil_image = Image_PIL.open(input_image_path).convert("RGB")
	transform_list = transforms.Compose(
	[
	data_utils.CenterCropLongEdge(),
	transforms.Resize((size, size)),
	transforms.ToTensor(),
	transforms.Normalize(NORM_MEAN, NORM_STD),
	]
	)
	tensor_image = transform_list(pil_image)
	tensor_image = torch.nn.functional.interpolate(
	tensor_image.unsqueeze(0), 224, mode="bicubic", align_corners=True
	)
	return tensor_image


	def preprocess_generated_image(image):
	transform_list = transforms.Normalize(NORM_MEAN, NORM_STD)
	image = transform_list(image * 0.5 + 0.5)
	image = torch.nn.functional.interpolate(
	image, 224, mode="bicubic", align_corners=True
	)
	return image