Function code2img

vqvae/api.py:32–44

Convert a batch of code to imgs Args: model: ... code: [b, h, w] or [b, h*w] LongTensor

(model, code)

    return id_t1.view(img.shape[0], -1) 

def code2img(model, code):
    '''Convert a batch of code to imgs
    Args:
        model: ...
        code: [b, h, w] or [b, h*w] LongTensor
    '''
    if len(code.shape) == 2:
        s = int(math.sqrt(len(code.view(-1))) + 1e-5)
        code = code.view(code.shape[0], s, s)
    with torch.no_grad():
        out = model.decode_code(code)
        out = out * torch.tensor([0.30379, 0.32279, 0.32800], device=out.device).view(1, -1, 1, 1) + torch.tensor([0.79093, 0.76271, 0.75340], device=out.device).view(1, -1, 1, 1)
    return out