Swin2SR 使用记录

背景

部分图片分辨率比较低可以使用这个模型来提升图片的分辨率

环境准备

pip install -q git+https://github.com/huggingface/transformers.git

使用

from transformers import Swin2SRForImageSuperResolution
model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64")

import requests
from PIL import Image

# url = "https://huggingface.co/spaces/jjourney1125/swin2sr/resolve/main/samples/butterfly.jpg"
# image = Image.open(requests.get(url, stream=True).raw)

image = Image.open("C.png").convert("RGB")
# image

from transformers import Swin2SRImageProcessor 

processor = Swin2SRImageProcessor()
pixel_values = processor(image, return_tensors="pt").pixel_values
# print(pixel_values.shape)

import torch

with torch.no_grad():
  outputs = model(pixel_values)

import numpy as np

output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
output = np.moveaxis(output, source=0, destination=-1)
output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
Image.fromarray(output)

参考