utils.py 文件源码

def X2img(X, image_name, mod='rgb'):
    if mod=='bgr':
        (b,g,r) = torch.chunk(X, 3)
        X = torch.cat((r,g,b))
    img = X.clone().cpu().clamp(0,255).numpy()
    img = img.transpose(1,2,0).astype('uint8')
    img = Image.fromarray(img)
    img.save(image_name)

# load image