Transformers 文件

影像到影像任務指南

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

影像到影像任務指南

影像到影像任務是指應用程式接收影像並輸出另一影像的任務。這包括各種子任務,如影像增強(超解析度、低光照增強、去雨等)、影像修復等。

本指南將向您展示如何:

  • 使用影像到影像管道進行超解析度任務,
  • 無需管道即可執行相同任務的影像到影像模型。

請注意,截至本指南釋出時,`image-to-image` 管道僅支援超解析度任務。

我們先安裝必要的庫。

pip install transformers

我們現在可以使用 Swin2SR 模型 初始化管道。然後,我們可以透過呼叫管道並傳入影像進行推理。截至目前,此管道僅支援 Swin2SR 模型

from transformers import pipeline
import torch
from accelerate.test_utils.testing import get_backend
# automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
device, _, _ = get_backend()
pipe = pipeline(task="image-to-image", model="caidas/swin2SR-lightweight-x2-64", device=device)

現在,讓我們載入一張影像。

from PIL import Image
import requests

url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/cat.jpg"
image = Image.open(requests.get(url, stream=True).raw)

print(image.size)
# (532, 432)
Photo of a cat

現在我們可以使用管道進行推理。我們將獲得貓影像的放大版本。

upscaled = pipe(image)
print(upscaled.size)
# (1072, 880)

如果您希望在沒有管道的情況下自行進行推理,可以使用 transformers 的 `Swin2SRForImageSuperResolution` 和 `Swin2SRImageProcessor` 類。我們將使用相同的模型檢查點。讓我們初始化模型和處理器。

from transformers import Swin2SRForImageSuperResolution, Swin2SRImageProcessor 

model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-lightweight-x2-64").to(device)
processor = Swin2SRImageProcessor("caidas/swin2SR-lightweight-x2-64")

`pipeline` 抽象了我們必須自己完成的預處理和後處理步驟,所以讓我們預處理影像。我們將影像傳遞給處理器,然後將畫素值移動到 GPU。

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

pixel_values = pixel_values.to(device)

現在我們可以透過將畫素值傳遞給模型來推斷影像。

import torch

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

輸出是一個型別為 `ImageSuperResolutionOutput` 的物件,如下所示 👇

(loss=None, reconstruction=tensor([[[[0.8270, 0.8269, 0.8275,  ..., 0.7463, 0.7446, 0.7453],
          [0.8287, 0.8278, 0.8283,  ..., 0.7451, 0.7448, 0.7457],
          [0.8280, 0.8273, 0.8269,  ..., 0.7447, 0.7446, 0.7452],
          ...,
          [0.5923, 0.5933, 0.5924,  ..., 0.0697, 0.0695, 0.0706],
          [0.5926, 0.5932, 0.5926,  ..., 0.0673, 0.0687, 0.0705],
          [0.5927, 0.5914, 0.5922,  ..., 0.0664, 0.0694, 0.0718]]]],
       device='cuda:0'), hidden_states=None, attentions=None)

我們需要獲取 `reconstruction` 並對其進行後處理以進行視覺化。讓我們看看它是什麼樣子。

outputs.reconstruction.data.shape
# torch.Size([1, 3, 880, 1072])

我們需要壓縮輸出並去除軸 0,裁剪值,然後將其轉換為 numpy 浮點數。然後我們將排列軸使其具有形狀 [1072, 880],最後將輸出恢復到範圍 [0, 255]。

import numpy as np

# squeeze, take to CPU and clip the values
output = outputs.reconstruction.data.squeeze().cpu().clamp_(0, 1).numpy()
# rearrange the axes
output = np.moveaxis(output, source=0, destination=-1)
# bring values back to pixel values range
output = (output * 255.0).round().astype(np.uint8)
Image.fromarray(output)
Upscaled photo of a cat
< > 在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.