def main():
"""Main function."""
args = parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
assert os.path.isfile(args.image_path)
output_dir = args.output_dir or f'results/inversion/test'
if not os.path.exists(output_dir):
os.makedirs(output_dir)
inverter = StyleGANInverter(
args.model_name,
mode=args.mode,
learning_rate=args.learning_rate,
iteration=args.num_iterations,
reconstruction_loss_weight=1.0,
perceptual_loss_weight=args.loss_weight_feat,
regularization_loss_weight=args.loss_weight_enc,
clip_loss_weight=args.loss_weight_clip,
description=args.description,
logger=None)
image_size = inverter.G.resolution
# Invert the given image.
image = resize_image(load_image(args.image_path), (image_size, image_size))
_, viz_results = inverter.easy_invert(image, num_viz=args.num_results)
if args.mode == 'man':
image_name = os.path.splitext(os.path.basename(args.image_path))[0]
else:
image_name = 'gen'
save_image(f'{output_dir}/{image_name}_enc.png', viz_results[1])
save_image(f'{output_dir}/{image_name}_inv.png', viz_results[-1])
print(f'save {image_name} in {output_dir}')
def main():
"""Main function."""
inverter = StyleGANInverter(model_name,
mode=mode,
learning_rate=ini_lr,
iteration=step,
reconstruction_loss_weight=lambda_l2,
perceptual_loss_weight=lambda_feat,
regularization_loss_weight=lambda_enc,
clip_loss_weight=lambda_clip,
description=description)
image_size = inverter.G.resolution
text_inputs = torch.cat([clip.tokenize(description)]).cuda()
# Invert images.
# uploaded_file = uploaded_file.read()
if uploaded_file is not None:
image = Image.open(uploaded_file)
# st.image(image, caption='Uploaded Image.', use_column_width=True)
# st.write("")
st.write("Just a second...")
image = resize_image(np.array(image), (image_size, image_size))
_, viz_results = inverter.easy_invert(image, 1)
if mode == 'man':
final_result = np.hstack([image, viz_results[-1]])
else:
final_result = np.hstack([viz_results[1], viz_results[-1]])
# return final_result
with st.beta_container():
st.image(final_result, use_column_width=True)
def predict(self, image, description):
self.args.description = description
self.args.image_path = image
inverter = StyleGANInverter(
self.args.model_name,
mode=self.args.mode,
learning_rate=self.args.learning_rate,
iteration=self.args.num_iterations,
reconstruction_loss_weight=1.0,
perceptual_loss_weight=self.args.loss_weight_feat,
regularization_loss_weight=self.args.loss_weight_enc,
clip_loss_weight=self.args.loss_weight_clip,
description=self.args.description,
logger=None,
)
image_size = inverter.G.resolution
# Invert the given image.
image = resize_image(
load_image(str(self.args.image_path)), (image_size, image_size)
)
_, viz_results = inverter.easy_invert(image, num_viz=self.args.num_results)
out_path = Path(tempfile.mkdtemp()) / "out.png"
save_image(str(out_path), viz_results[-1])
return out_path
def main():
"""Main function."""
args = parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
assert os.path.exists(args.image_list)
image_list_name = os.path.splitext(os.path.basename(args.image_list))[0]
output_dir = args.output_dir or f'results/inversion/{image_list_name}'
logger = setup_logger(output_dir, 'inversion.log', 'inversion_logger')
logger.info(f'Loading model.')
inverter = StyleGANInverter(
args.model_name,
learning_rate=args.learning_rate,
iteration=args.num_iterations,
reconstruction_loss_weight=1.0,
perceptual_loss_weight=args.loss_weight_feat,
regularization_loss_weight=args.loss_weight_enc,
logger=logger)
image_size = inverter.G.resolution
# Load image list.
logger.info(f'Loading image list.')
image_list = []
with open(args.image_list, 'r') as f:
for line in f:
image_list.append(line.strip())
# Initialize visualizer.
save_interval = args.num_iterations // args.num_results
headers = ['Name', 'Original Image', 'Encoder Output']
for step in range(1, args.num_iterations + 1):
if step == args.num_iterations or step % save_interval == 0:
headers.append(f'Step {step:06d}')
viz_size = None if args.viz_size == 0 else args.viz_size
visualizer = HtmlPageVisualizer(
num_rows=len(image_list), num_cols=len(headers), viz_size=viz_size)
visualizer.set_headers(headers)
# Invert images.
logger.info(f'Start inversion.')
latent_codes = []
for img_idx in tqdm(range(len(image_list)), leave=False):
image_path = image_list[img_idx]
image_name = os.path.splitext(os.path.basename(image_path))[0]
image = resize_image(load_image(image_path), (image_size, image_size))
code, viz_results = inverter.easy_invert(image, num_viz=args.num_results)
latent_codes.append(code)
save_image(f'{output_dir}/{image_name}_ori.png', image)
save_image(f'{output_dir}/{image_name}_enc.png', viz_results[1])
save_image(f'{output_dir}/{image_name}_inv.png', viz_results[-1])
visualizer.set_cell(img_idx, 0, text=image_name)
visualizer.set_cell(img_idx, 1, image=image)
for viz_idx, viz_img in enumerate(viz_results[1:]):
visualizer.set_cell(img_idx, viz_idx + 2, image=viz_img)
# Save results.
os.system(f'cp {args.image_list} {output_dir}/image_list.txt')
np.save(f'{output_dir}/inverted_codes.npy',
np.concatenate(latent_codes, axis=0))
visualizer.save(f'{output_dir}/inversion.html')
def setup_inverter(model_name,
num_iterations=100,
regularization_loss_weight=2):
global inverter
setup_model(model_name)
inverter = StyleGANInverter(
models[model_name]['name'],
learning_rate=0.01,
iteration=num_iterations,
reconstruction_loss_weight=1.0,
perceptual_loss_weight=5e-5,
regularization_loss_weight=regularization_loss_weight)
def main():
"""Main function."""
args = parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
assert os.path.exists(args.target_list)
target_list_name = os.path.splitext(os.path.basename(args.target_list))[0]
assert os.path.exists(args.context_list)
context_list_name = os.path.splitext(os.path.basename(args.context_list))[0]
output_dir = args.output_dir or f'results/diffusion'
job_name = f'{target_list_name}_TO_{context_list_name}'
logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger')
logger.info(f'Loading model.')
inverter = StyleGANInverter(
args.model_name,
learning_rate=args.learning_rate,
iteration=args.num_iterations,
reconstruction_loss_weight=1.0,
perceptual_loss_weight=args.loss_weight_feat,
regularization_loss_weight=0.0,
logger=logger)
image_size = inverter.G.resolution
# Load image list.
logger.info(f'Loading target images and context images.')
target_list = []
with open(args.target_list, 'r') as f:
for line in f:
target_list.append(line.strip())
num_targets = len(target_list)
context_list = []
with open(args.context_list, 'r') as f:
for line in f:
context_list.append(line.strip())
num_contexts = len(context_list)
num_pairs = num_targets * num_contexts
# Initialize visualizer.
save_interval = args.num_iterations // args.num_results
headers = ['Target Image', 'Context Image', 'Stitched Image',
'Encoder Output']
for step in range(1, args.num_iterations + 1):
if step == args.num_iterations or step % save_interval == 0:
headers.append(f'Step {step:06d}')
viz_size = None if args.viz_size == 0 else args.viz_size
visualizer = HtmlPageVisualizer(
num_rows=num_pairs, num_cols=len(headers), viz_size=viz_size)
visualizer.set_headers(headers)
# Diffuse images.
logger.info(f'Start diffusion.')
latent_codes = []
for target_idx in tqdm(range(num_targets), desc='Target ID', leave=False):
# Load target.
target_image = resize_image(load_image(target_list[target_idx]),
(image_size, image_size))
visualizer.set_cell(target_idx * num_contexts, 0, image=target_image)
for context_idx in tqdm(range(num_contexts), desc='Context ID',
leave=False):
row_idx = target_idx * num_contexts + context_idx
context_image = resize_image(load_image(context_list[context_idx]),
(image_size, image_size))
visualizer.set_cell(row_idx, 1, image=context_image)
code, viz_results = inverter.easy_diffuse(target=target_image,
context=context_image,
center_x=args.center_x,
center_y=args.center_y,
crop_x=args.crop_size,
crop_y=args.crop_size,
num_viz=args.num_results)
for viz_idx, viz_img in enumerate(viz_results):
visualizer.set_cell(row_idx, viz_idx + 2, image=viz_img)
latent_codes.append(code)
# Save results.
os.system(f'cp {args.target_list} {output_dir}/target_list.txt')
os.system(f'cp {args.context_list} {output_dir}/context_list.txt')
np.save(f'{output_dir}/{job_name}_inverted_codes.npy',
np.concatenate(latent_codes, axis=0))
visualizer.save(f'{output_dir}/{job_name}.html')