how to produce a single pic, not a video

how to produce the final aligned pic？ about gangealing HOT 2 CLOSED

wpeebles commented on May 20, 2024

how to produce the final aligned pic？

from gangealing.

Comments (2)

wpeebles commented on May 20, 2024

Here's a minimal code snippet for applying a pre-trained Spatial Transformer (non-clustering) to one image:

from models import get_stn
from utils.download import download_model
from utils.vis_tools.helpers import load_pil
from torchvision.utils import save_image

resolution = 512  # resolution the input image will be resized to (this can be any power of 2)
input_img = load_pil('my_image.png', resolution)  # load the input image and resize to (1, C, resolution, resolution)
ckpt = download_model('cat')  # download model weights
stn = get_stn(['similarity', 'flow'], flow_size=128, supersize=resolution).to('cuda')  # instantiate STN
stn.load_state_dict(ckpt['t_ema'])  # load weights
aligned_img = stn(input_img, iters=3, output_resolution=resolution)  # forward pass through the STN
save_image(aligned_img, 'output.png', normalize=True, range=(-1, 1))  # save to disk

If you're using the celeba or cub models, use iters=1 instead. If your input image isn't square you may want to pad or crop it beforehand. Also, stn supports batch mode, so input_img can be an (N, C, H, W) tensor containing multiple images, in which case aligned_image will also be (N, C, H, W).

from gangealing.