def generate_batch_translation_dataset(self, batch_size, truncation=0.4):
# generate noise -> BxN(2)x128
noise_A = truncated_noise_sample(batch_size=batch_size,
truncation=truncation)
noise_B = truncated_noise_sample(batch_size=batch_size,
truncation=truncation)
n_A = torch.from_numpy(noise_A).unsqueeze(0)
n_B = torch.from_numpy(noise_B).unsqueeze(0)
n_AB = torch.cat([n_A, n_B], dim=0).to(self.device)
n_AB = n_AB.view((N // 2) * batch_size, -1) # (B*2)x128
# generate class -> BxN(2)x1000
idx_list = range(len(self.class_list))
idx_ab = [random.sample(idx_list, 2)
for i in range(batch_size)] # 2 pair idx
class_ab = [[
self.class_list[idx_ab[i][0]], self.class_list[idx_ab[i][1]]
] for i in range(len(idx_ab))]
class_ab = np.array(class_ab).transpose()
class_A = one_hot_from_int(class_ab[0],
batch_size=batch_size) # Bx1000 numpy array
class_B = one_hot_from_int(class_ab[1], batch_size=batch_size)
c_A = torch.from_numpy(class_A).unsqueeze(0)
c_B = torch.from_numpy(class_B).unsqueeze(0)
c_AB = torch.cat([c_A, c_B], dim=0).to(self.device)
c_AB = c_AB.view((N // 2) * batch_size, -1) # (B*2)x1000
with torch.no_grad():
output = self.big_gan(n_AB[A + B + A + B], c_AB[A + B + B + A],
truncation)
# generate class idx
idx_ab = torch.from_numpy(np.array(idx_ab).transpose()) # 2xB
return output.view(N, batch_size, 3, output.size(2),
output.size(3)), idx_ab.long()
def generate_batch_morphing_dataset(self,
batch_size,
interp_noise=True,
interp_class=True,
truncation=0.4):
# generate noise -> BxN(3)x128
noise_A = truncated_noise_sample(batch_size=batch_size,
truncation=truncation)
noise_B = truncated_noise_sample(
batch_size=batch_size,
truncation=truncation) if interp_noise else 0
n_A = torch.from_numpy(noise_A).expand(N, batch_size, -1)
n_B = torch.from_numpy(noise_B).expand(
N, batch_size, -1) if interp_noise else n_A.clone()
# generate class -> BxN(3)x1000
idx_list = range(len(self.class_list))
idx_ab = [random.sample(idx_list, 2)
for i in range(batch_size)] # 2 pair idx
class_ab = [[
self.class_list[idx_ab[i][0]], self.class_list[idx_ab[i][1]]
] for i in range(len(idx_ab))]
class_ab = np.array(class_ab).transpose()
class_A = one_hot_from_int(class_ab[0],
batch_size=batch_size) # Bx1000 numpy array
class_B = one_hot_from_int(class_ab[1], batch_size=batch_size)
c_A = torch.from_numpy(class_A).expand(N, batch_size, -1)
c_B = torch.from_numpy(class_B).expand(
N, batch_size, -1) if interp_class else c_A.clone()
# # sample index for intermediate image
# idx_0 = random.randint(0, n_frames-i_frames)
# idx_1 = idx_0 + i_frames - 1
# idx_i = random.randint(idx_0, idx_1) # sample index for intermediate image, including end points
# interpolation -> BxN(3)xZ
# t = torch.linspace(start=0, end=1.0, steps=n_frames)
# t = t[[idx_0, idx_i, idx_1]].unsqueeze(1).unsqueeze(1).expand(3, batch_size, 1)
a = torch.cat([
torch.zeros(1, batch_size, 1),
torch.rand(1, batch_size, 1),
torch.ones(1, batch_size, 1)
])
n = self.interp_func(n_A, n_B, a).to(self.device)
c = self.interp_func(c_A, c_B, a).to(self.device)
n = n.view(N * batch_size, -1) # (B*N(3))xZ
c = c.view(N * batch_size, -1)
with torch.no_grad():
output = self.big_gan(n, c, truncation)
# generate class idx
idx_ab = torch.from_numpy(np.array(idx_ab).transpose()) # 2xB
return output.view(N, batch_size, 3, output.size(2),
output.size(3)), a[1].squeeze(dim=1), idx_ab.long()
def generate_class_AB(self, class_A, class_B, batch_size=0):
if batch_size == 0:
self.class_A = one_hot_from_int(class_A)
self.class_B = one_hot_from_int(class_B)
else: # for translation dataset
class_A = one_hot_from_int(class_A, batch_size=batch_size)
class_B = one_hot_from_int(class_B, batch_size=batch_size)
# generate class -> BxN(1)x1000
c_A = torch.from_numpy(class_A).unsqueeze(0)
c_B = torch.from_numpy(class_B).unsqueeze(0)
c_AB = torch.cat([c_A, c_B], dim=0).to(self.device)
self.c_AB = c_AB.view(2 * batch_size, -1) # (B*2)x1000
def generate_images(num_images, class_index, model, truncation: float = 1.):
print("[info] initializing input variables...")
batch_size = min(10, num_images)
class_vec = one_hot_from_int(class_index, batch_size=num_images) * .97
noise_vec = truncated_noise_sample(truncation=truncation,
batch_size=num_images)
noise_tensor = torch.tensor(noise_vec, dtype=torch.float)
class_tensor = torch.tensor(class_vec, dtype=torch.float)
print('[info] preparing model inference...')
noise_tensor = noise_tensor.to(device)
class_tensor = class_tensor.to(device)
print(
f'[info] generating {num_images} images of class index {class_index}...'
)
images = []
for i in range(0, num_images, batch_size):
torch.cuda.empty_cache()
noise_feed = noise_tensor[i:i + batch_size]
class_feed = class_tensor[i:i + batch_size]
with torch.no_grad():
output = model(noise_feed, class_feed, truncation)
output_cpu = output.cpu().data.numpy()
for out in output_cpu:
image = np.array(toimage(out))
images.append(image)
print('[info] done.\n')
torch.cuda.empty_cache()
return images
def translate_sentiment(self, audio_sentiment, noise=0.1):
# Step 1: Find the closest word match to the audio sentiment
dists, idx = self.database.kneighbors(audio_sentiment)
song_words = [self.words[cur_index[0]] for cur_index in idx]
song_words = [x.split(',')[0].lower() for x in song_words]
word_ids = [self.class_ids[cur_index] for cur_index in idx]
class_vectors = []
noise_vectors = []
# Calculate the weighted nearest class neighbors
for i in range(len(word_ids)):
cur_class_vector = one_hot_from_int(word_ids[i],
batch_size=len(word_ids[i]))
# If k = 1
if len(word_ids[i]) == 1:
cur_class_vector = cur_class_vector[0]
else:
similarities = np.exp(-dists[i] / self.temperature)
similarities = similarities / np.sum(similarities)
similarities = similarities.reshape((-1, 1))
cur_class_vector = np.sum(cur_class_vector * similarities,
axis=0)
class_vectors.append(cur_class_vector)
noise_vectors.append(
truncated_noise_sample(truncation=noise, batch_size=1)[0])
class_vectors = np.float32(class_vectors)
noise_vectors = np.float32(noise_vectors)
return class_vectors, noise_vectors, song_words
def forward(self, z, y):
class_vector = one_hot_from_int(y, y.size(0))
class_vector = torch.from_numpy(class_vector).cuda()
to_return = self.model(z, class_vector, 1.)
return torch.nn.functional.interpolate(
to_return,
size=int(self.imagenet_preprocessing_str[-3:]),
mode='bilinear') #, antialias = True)
def generate_t(self, class_A=0, class_B=1, class_t=0.5, truncation=0.4):
self.class_vector_A = torch.from_numpy(one_hot_from_int(class_A))
self.class_vector_B = torch.from_numpy(one_hot_from_int(class_B))
# Interpolation
class_vector_t = lerp(self.class_vector_A, self.class_vector_B,
torch.tensor(class_t, dtype=torch.float32))
# GPU.
class_vector_t = class_vector_t.to(self.device)
# Generate images
with torch.no_grad():
output = self.big_gan(self.noise_vector,
class_vector_t,
truncation=0.4)
# save_as_images(output)
return output
def set_output_class(self, class_id):
if isinstance(class_id, int):
self.v_class = torch.from_numpy(one_hot_from_int([class_id])).to(
self.device)
self.outclass = f"class{class_id}"
elif isinstance(class_id, str):
self.outclass = class_id.replace(" ", "_")
self.v_class = torch.from_numpy(one_hot_from_names([class_id])).to(
self.device)
else:
raise RuntimeError(f"Unknown class identifier {class_id}")
def get_sample_image(self):
n = truncated_noise_sample(batch_size=self.batch_size,
truncation=self.truncation)
c_idx = np.random.choice(self.class_list, size=1)
c = one_hot_from_int(c_idx, batch_size=self.batch_size)
# np to torch
n = torch.from_numpy(n).to(self.device)
c = torch.from_numpy(c).to(self.device)
with torch.no_grad():
output = self.big_gan(n, c, self.truncation)
return output
def evaluate_inversion(args, inverted_net_path):
# Load saved inverted net
device = 'cuda:{}'.format(
args.gpu_ids[0]) if len(args.gpu_ids) > 0 else 'cpu'
ckpt_dict = torch.load(inverted_net_path, map_location=device)
# Build model, load parameters
model_args = ckpt_dict['model_args']
inverted_net = models.ResNet18(**model_args)
inverted_net = nn.DataParallel(inverted_net, args.gpu_ids)
inverted_net.load_state_dict(ckpt_dict['model_state'])
import pdb
pdb.set_trace()
# Get test images (CelebA)
initial_generated_image_dir = '/deep/group/sharonz/generator/z_test_images/'
initial_generated_image_name = '058004_crop.jpg'
initial_generated_image = util.get_image(initial_generated_image_dir,
initial_generated_image_name)
initial_generated_image = initial_generated_image / 255.
intiial_generated_image = initial_generated_image.cuda()
inverted_noise = inverted_net(initial_generated_image)
if 'BigGAN' in args.model:
class_vector = one_hot_from_int(207, batch_size=batch_size)
class_vector = torch.from_numpy(class_vector)
num_params = int(''.join(filter(str.isdigit, args.model)))
generator = BigGAN.from_pretrained(f'biggan-deep-{num_params}')
generator = generator.to(args.device)
generated_image = generator.forward(inverted_noise, class_vector,
args.truncation)
# Get difference btw initial and subsequent generated image
# Save both
return
def __getitem__(self, idx):
if idx > self.n_iters:
raise StopIteration
# Sample the space
idx = np.random.randint(0, 1000, self.in_batch)
class_vectors = one_hot_from_int(idx, batch_size=self.in_batch)
noise_vectors = truncated_noise_sample(truncation=self.noise,
batch_size=self.in_batch)
class_vectors = torch.tensor(class_vectors).to(config['device'])
noise_vectors = torch.tensor(noise_vectors).to(config['device'])
with torch.no_grad():
output = self.gan_model(noise_vectors, class_vectors, self.noise)
# Convert to PIL Image
output = output.detach().cpu().numpy()
output = np.uint8(np.clip(((output + 1) / 2.0) * 256, 0, 255))
output = output.transpose((0, 2, 3, 1))
images = []
# Pre-process each image to feed them into image sentiment analysis model
for i in range(output.shape[0]):
cur_img = Image.fromarray(output[i])
cur_img = self.transform_image(cur_img)
images.append(cur_img)
images = torch.stack(images).to(config['device'])
# Feed-forward image sentiment analysis
sentiment = self.sentiment_model(images)
class_vectors = class_vectors.cpu().detach().numpy()
noise_vectors = noise_vectors.cpu().detach().numpy()
sentiment = sentiment.cpu().detach().numpy()
return class_vectors, noise_vectors, sentiment
def test(args):
# Get loader for z-test
loader = get_loader(args, phase='test')
batch_size = args.batch_size
class_vector = None
# TODO: make into function that takes in args.model and returns the pretrained model
# and also consider whether it's class conditional and what kind of class conditional (how many classes) -> probably just imagenet now, actually maybe cifar-10 too
# and also consider add truncation sampling as option too - this should return model, z_test noise vec, and class_vec (optionally)
if args.ckpt_path and not args.use_pretrained:
model, ckpt_info = ModelSaver.load_model(args.ckpt_path, args.gpu_ids)
else:
if 'BigGAN' in args.model:
num_params = int(''.join(filter(str.isdigit, args.model)))
if 'perturbation' in args.loss_fn:
# Use custom BigGAN with Perturbation Net wrapper
model = models.BigGANPerturbationNet.from_pretrained(
f'biggan-deep-{num_params}')
else:
# Use pretrained BigGAN from package
model = BigGAN.from_pretrained(f'biggan-deep-{num_params}')
z_test = truncated_noise_sample(truncation=args.truncation,
batch_size=batch_size)
z_test = torch.from_numpy(z_test)
# Get class conditional label
# 981 is baseball player
# 207 is golden retriever
# TODO: Conditional generation only
class_vector = one_hot_from_int(207, batch_size=batch_size)
class_vector = torch.from_numpy(class_vector)
elif 'WGAN-GP' in args.model:
generator_path = "/deep/group/gen-eval/model-training/src/GAN_models/improved-wgan-pytorch/experiments/exp4_wgan_gp/generator.pt"
model = torch.load(generator_path)
z_test = torch.randn(batch_size, 128)
elif 'BEGAN' in args.model:
generator_path = "/deep/group/gen-eval/model-training/src/GAN_models/BEGAN-pytorch/trained_models/64/models/gen_97000.pth"
model = models.BEGANGenerator()
model.load_state_dict(torch.load(generator_path))
z_test = np.random.uniform(-1, 1, size=(batch_size, 64))
z_test = torch.FloatTensor(z_test)
# Freeze model instead of using .eval()
for param in model.parameters():
param.requires_grad = False
# If using perturbation net, learn perturbation layers
if 'perturbation' in args.loss_fn:
trainable_params = []
for name, param in model.named_parameters():
if 'perturb' in name:
param.requires_grad = True
trainable_params.append(param)
print(f'Number of trainable params: {len(trainable_params)}')
model = nn.DataParallel(model, args.gpu_ids)
model = model.to(args.device)
# Loss functions
if 'mse' in args.loss_fn:
pixel_criterion = torch.nn.MSELoss().to(args.device)
else:
pixel_criterion = torch.nn.L1Loss().to(args.device)
if 'perceptual' in args.loss_fn:
# Combination pixel-perceptual loss - Sec 3.2. By default, uses pixel L1.
perceptual_criterion = torch.nn.L1Loss().to(args.device)
perceptual_loss_weight = args.perceptual_loss_weight
vgg_feature_extractor = models.VGGFeatureExtractor().to(args.device)
vgg_feature_extractor.eval()
elif 'perturbation' in args.loss_fn:
# Perturbation network R. By default, uses pixel L1.
# Sec 3.3: http://ganpaint.io/Bau_et_al_Semantic_Photo_Manipulation_preprint.pdf
reg_loss_weight = args.reg_loss_weight
# z_loss_fn = util.get_loss_fn(args.loss_fn, args)
max_z_test_loss = 100. # TODO: actually put max value possible here
# Get logger, saver
logger = TestLogger(args)
# saver = ModelSaver(args) TODO: saver for perturbation network R
print(f'Logs: {logger.log_dir}')
print(f'Ckpts: {args.save_dir}')
# Run z-test in batches
logger.log_hparams(args)
while not logger.is_finished_training():
logger.start_epoch()
for _, z_test_target, mask in loader:
logger.start_iter()
if torch.cuda.is_available():
mask = mask.cuda()
z_test = z_test.cuda()
z_test_target = z_test_target.cuda()
#class_vector = class_vector.cuda()
masked_z_test_target = z_test_target * mask
obscured_z_test_target = z_test_target * (1.0 - mask)
if 'perturbation' in args.loss_fn:
# With backprop on only trainable parameters in perturbation net
params = trainable_params + [z_test.requires_grad_()]
z_optimizer = util.get_optimizer(params, args)
else:
# With backprop on only the input z, run one step of z-test and get z-loss
z_optimizer = util.get_optimizer([z_test.requires_grad_()],
args)
with torch.set_grad_enabled(True):
if class_vector is not None:
z_probs = model.forward(z_test, class_vector,
args.truncation).float()
z_probs = (z_probs + 1) / 2.
else:
z_probs = model.forward(z_test).float()
# Calculate the masked loss using z-test vector
masked_z_probs = z_probs * mask
z_loss = torch.zeros(1, requires_grad=True).to(args.device)
pixel_loss = torch.zeros(1, requires_grad=True).to(args.device)
pixel_loss = pixel_criterion(masked_z_probs,
masked_z_test_target)
if 'perceptual' in args.loss_fn:
z_probs_features = vgg_feature_extractor(masked_z_probs)
z_test_features = vgg_feature_extractor(
masked_z_test_target).detach()
perceptual_loss = torch.zeros(1, requires_grad=True).to(
args.device)
perceptual_loss = perceptual_criterion(
z_probs_features, z_test_features)
z_loss = pixel_loss + perceptual_loss_weight * perceptual_loss
elif 'perturbation' in args.loss_fn:
reg_loss = torch.zeros(1,
requires_grad=True).to(args.device)
for name, param in model.named_parameters():
if 'perturb' in name:
delta = param - 1
reg_loss += torch.pow(delta, 2).mean() #sum()
z_loss = pixel_loss + reg_loss_weight * reg_loss
else:
z_loss = pixel_loss
# Backprop on z-test vector
z_loss.backward()
z_optimizer.step()
z_optimizer.zero_grad()
# Compute the full loss (without mask) and obscured loss (loss only on masked region)
# For logging and final evaluation (obscured loss is final MSE), so not in backprop loop
full_z_loss = torch.zeros(1)
full_pixel_loss = torch.zeros(1)
full_pixel_loss = pixel_criterion(z_probs, z_test_target) #.mean()
obscured_z_probs = z_probs * (1.0 - mask)
obscured_z_loss = torch.zeros(1)
obscured_pixel_loss = torch.zeros(1)
obscured_pixel_loss = pixel_criterion(
obscured_z_probs, obscured_z_test_target) #.mean()
if 'perceptual' in args.loss_fn:
# Full loss
z_probs_full_features = vgg_feature_extractor(z_probs).detach()
z_test_full_features = vgg_feature_extractor(
z_test_target).detach()
full_perceptual_loss = torch.zeros(1)
full_perceptual_loss = perceptual_criterion(
z_probs_full_features, z_test_full_features)
full_z_loss = full_pixel_loss + perceptual_loss_weight * full_perceptual_loss
# Obscured loss
z_probs_obscured_features = vgg_feature_extractor(
z_probs).detach()
z_test_obscured_features = vgg_feature_extractor(
z_test_target).detach()
obscured_perceptual_loss = torch.zeros(1)
obscured_perceptual_loss = perceptual_criterion(
z_probs_obscured_features, z_test_obscured_features)
obscured_z_loss = obscured_pixel_loss + perceptual_loss_weight * obscured_perceptual_loss
elif 'perturbation' in args.loss_fn:
full_z_loss = full_pixel_loss + reg_loss_weight * reg_loss
obscured_z_loss = obscured_pixel_loss + reg_loss_weight * reg_loss
else:
full_z_loss = full_pixel_loss
obscured_z_loss = obscured_pixel_loss
"""# TODO: z_loss is not always MSE anymore - figure out desired metric
if z_loss < max_z_test_loss:
# Save MSE on obscured region # TODO: z_loss is not always MSE anymore - figure out desired metric
final_metrics = {'z-loss': z_loss.item(), 'obscured-z-loss': obscured_z_loss.item()}
logger._log_scalars(final_metrics)
print('Recall (z loss - non obscured loss - if MSE)', z_loss)
print('Precision (MSE value on masked region)', obscured_z_loss)
"""
# Log both train and eval model settings, and visualize their outputs
logger.log_status(
masked_probs=masked_z_probs,
masked_loss=z_loss,
masked_test_target=masked_z_test_target,
full_probs=z_probs,
full_loss=full_z_loss,
full_test_target=z_test_target,
obscured_probs=obscured_z_probs,
obscured_loss=obscured_z_loss,
obscured_test_target=obscured_z_test_target,
save_preds=args.save_preds,
)
logger.end_iter()
logger.end_epoch()
# Last log after everything completes
logger.log_status(
masked_probs=masked_z_probs,
masked_loss=z_loss,
masked_test_target=masked_z_test_target,
full_probs=z_probs,
full_loss=full_z_loss,
full_test_target=z_test_target,
obscured_probs=obscured_z_probs,
obscured_loss=obscured_z_loss,
obscured_test_target=obscured_z_test_target,
save_preds=args.save_preds,
force_visualize=True,
)
(classname, muembd, *xlim, *ylim, truncation,
np.random.randint(10000))))
plt.show()
#%%
plt.figure(figsize=[4, 4])
plt.imshow(imgs[0])
plt.axis("image")
plt.axis("off")
plt.tight_layout()
plt.show()
#%%
savedir = r"C:\Users\zhanq\OneDrive - Washington University in St. Louis\Generator_Testing\BigGAN256"
truncation = 0.7
muembd = 0.1
classname = [1, 2]
class_vector = one_hot_from_int([classname], batch_size=len(classname))
ebd_class = model.embeddings(torch.from_numpy(class_vector).cuda())
ebd_vecs = muembd * torch.randn(
(10, ebd_class.shape[1]
)).cuda() + ebd_class # add Gaussian perturbation around
with torch.no_grad():
output = model.generator(
torch.cat((torch.zeros_like(ebd_vecs), ebd_vecs), dim=1),
truncation)
imgs = convert_to_images(output.cpu())
figh = plt.figure(figsize=[25, 3])
gs = figh.add_gridspec(1, len(imgs)) # 1d interpolation
for i, img in enumerate(imgs):
plt.subplot(gs[i])
plt.imshow(img)
def train_inverted_net(args):
# Start by training an external model on samples of G(z) -> z inversion
model = util.get_invert_model(args)
model = nn.DataParallel(model, args.gpu_ids)
model = model.to(args.device)
print(f'{args.invert_model} num params {count_parameters(model)}')
generator = util.get_model(args)
if generator is not None:
generator = nn.DataParallel(generator, args.gpu_ids)
generator = generator.to(args.device)
print(f'{args.model} num params {count_parameters(generator)}')
else:
# Load saved pairings (ProGAN/StyleGAN)
pairing_dir = '/deep/group/gen-eval/model-training/src/GAN_models/stylegan'
pairing_path = f'{pairing_dir}/otavio_sampled_output/pairing.csv'
pairings = pd.read_csv(pairing_path)
num_pairings = len(pairings)
noise_targets = pairings['noise']
image_inputs = pairings['image']
if 'BigGAN' in args.model:
class_vector = one_hot_from_int(207, batch_size=args.batch_size)
class_vector = torch.from_numpy(class_vector)
class_vector = class_vector.cuda()
# TODO: remove bc cant use gpu in laoder i don't think
#loader = get_loader(args, phase='invert')
#logger = TestLogger(args)
#logger.log_hparams(args)
criterion = torch.nn.MSELoss().to(args.device)
optimizer = util.get_optimizer(model.parameters(), args)
for i in range(args.num_invert_epochs):
if generator is not None:
noise_target = util.get_noise(args)
image_input = generator.forward(noise_target).float()
image_input = (image_input + 1.) / 2.
else:
# TODO: make into loader
idx = i % num_pairings
noise_target = np.load(f'{pairing_dir}/{noise_targets[idx]}')
noise_target = torch.from_numpy(noise_target).float()
print(f'noise target shape {noise_target.shape}')
image_input = np.array(
Image.open(f'{pairing_dir}/{image_inputs[idx]}'))
image_input = torch.from_numpy(image_input / 255.)
image_input = image_input.float().unsqueeze(0)
image_input = image_input.permute(0, 3, 1, 2)
noise_target = noise_target.cuda()
image_input = image_input.cuda()
with torch.set_grad_enabled(True):
probs = model.forward(image_input)
loss = torch.zeros(1, requires_grad=True).to(args.device)
loss = criterion(probs, noise_target)
print(f'iter {i}: loss = {loss}')
loss.backward()
optimizer.step()
optimizer.zero_grad()
if i % 1 == 0:
corres_image_input = image_input.detach().cpu()
corres_np = util.convert_image_from_tensor(corres_image_input)
# Run check - saving image
if 'BigGAN' in args.model:
predicted_image = generator.forward(probs, class_vector,
truncation).float()
else:
if generator is not None:
predicted_image = generator.forward(probs).float()
predicted_image = predicted_image.detach().cpu()
predicted_image = (predicted_image + 1) / 2.
predicted_np = util.convert_image_from_tensor(
predicted_image)
if len(predicted_np.shape) == 4:
predicted_np = predicted_np[0]
corres_np = corres_np[0]
visuals = util.concat_images([predicted_np, corres_np])
visuals_pil = Image.fromarray(visuals)
timestamp = datetime.now().strftime('%b%d_%H%M%S%f')
visuals_image_dir = f'predicted_inversion_images/{args.model}'
os.makedirs(visuals_image_dir, exist_ok=True)
visuals_image_path = f'{visuals_image_dir}/{timestamp}_{i}.png'
visuals_pil.save(visuals_image_path)
print(f'Saved {visuals_image_path}')
else:
# Save noise vector - do forward separately in tf env
probs = probs.detach().cpu().numpy()
pred_noise_dir = f'predicted_inversion_noise/{args.model}'
os.makedirs(pred_noise_dir, exist_ok=True)
pred_noise_path = f'{pred_noise_dir}/{args.model}_noise_{i}.npy'
np.save(pred_noise_path, probs)
print(f'Saved {pred_noise_path}')
if i % 1 == 0:
corres_image_input = image_input.detach().cpu()
corres_np = util.convert_image_from_tensor(corres_image_input)
if len(corres_np.shape) == 4:
corres_np = corres_np[0]
corres_pil = Image.fromarray(corres_np)
timestamp = datetime.now().strftime('%b%d_%H%M%S%f')
corres_image_dir = f'generated_images/{args.model}'
os.makedirs(corres_image_dir, exist_ok=True)
corres_image_path = f'{corres_image_dir}/{timestamp}_{i}.png'
corres_pil.save(corres_image_path)
# saver = ModelSaver(args)
global_step = args.num_invert_epochs
ckpt_dict = {
'ckpt_info': {
'global_step': global_step
},
'model_name': model.module.__class__.__name__,
'model_args': model.module.args_dict(),
'model_state': model.to('cpu').state_dict(),
'optimizer': optimizer.state_dict(),
}
ckpt_dir = os.path.join(args.save_dir, f'{args.model}')
os.makedirs(ckpt_dir, exist_ok=True)
ckpt_path = os.path.join(
ckpt_dir, f'{args.invert_model}_step_{global_step}.pth.tar')
torch.save(ckpt_dict, ckpt_path)
print(f'Saved model to {ckpt_path}')
import pdb
pdb.set_trace()
return model
import os
from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
save_as_images, display_in_terminal, one_hot_from_int)
# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
import logging
logging.basicConfig(level=logging.INFO)
# Load pre-trained model tokenizer (vocabulary)
model = BigGAN.from_pretrained('biggan-deep-512')
# Prepare a input
truncation = 0.4
# class_vector = one_hot_from_names(['soap bubble', 'coffee', 'mushroom'], batch_size=3) # pip install nltk
class_vector = one_hot_from_int([445, 445, 445], batch_size=3)
noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3)
# All in tensors
noise_vector = torch.from_numpy(noise_vector)
class_vector = torch.from_numpy(class_vector)
# If you have a GPU, put everything on cuda
noise_vector = noise_vector.to('cuda')
class_vector = class_vector.to('cuda')
model.to('cuda')
# Generate an image
with torch.no_grad():
output = model(noise_vector, class_vector, truncation)
def create_class_vector(x, device, batch_size: int = 1):
class_vector = one_hot_from_int(x, batch_size)
return torch.from_numpy(class_vector).to(device)
plt.hist(dist_mat[dist_mat != 0], bins=50)
plt.title("Distribution of Distance Matrix of BigGAN-Deep256")
plt.xlabel("L2 Distance")
plt.savefig(join(homedir, "EmbedVectDistanceDist.png"))
plt.show()
#%%
#%%
from nltk.corpus import wordnet as wn
import pickle, urllib
ImageNet_Classname = pickle.load(
urllib.request.urlopen(
'https://gist.githubusercontent.com/yrevar/6135f1bd8dcf2e0cc683/raw/d133d61a09d7e5a3b36b8c111a8dd5c4b5d560ee/imagenet1000_clsid_to_human.pkl'
))
#%% Visualize Images for each class
truncation = 0.7
class_vector = one_hot_from_int(list(range(1000)), batch_size=1000)
sample_n = class_vector.shape[0]
# noise_vector = np.zeros((sample_n, 128))
numbering = list(range(1000))
#%%
noise_vec = truncated_noise_sample(batch_size=1,
dim_z=128,
truncation=0.7,
seed=None)
#%%
batch = 10
csr = 0
csr_end = 0
while csr_end < sample_n:
csr_end = min(csr + batch, sample_n)
# imgs = BigGAN_render(class_vector[csr:csr_end, :], noise_vector[csr:csr_end, :], truncation)