Working with Images

Ray Data를 이용하여 대용량 이미지 데이터셋을 쉽게 읽고 변형할 수 있다.

Reading images

Ray Data는 다양한 포맷의 이미지를 읽어올 수 있다.

import ray

ds = ray.data.read_images("s3://anonymous@ray-example-data/batoidea/JPEGImages")

print(ds.schema())

Column  Type
------  ----
image   numpy.ndarray(shape=(32, 32, 3), dtype=uint8)

Transforming images

이미지를 변형시키기 위해서는 map 또는 map_batches를 사용한다

from typing import Any, Dict
import numpy as np
import ray

def increase_brightness(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
    batch["image"] = np.clip(batch["image"] + 4, 0, 255)
    return batch

ds = (
    ray.data.read_images("s3://anonymous@ray-example-data/batoidea/JPEGImages")
    .map_batches(increase_brightness)
)

Performing inference on images

pre-trained model을 사용하여 inference를 수행하기 위해서는 먼저 데이터를 로드하고 입력에 맞게 변형하는 것이 필요하다

from typing import Any, Dict
from torchvision import transforms
import ray

def transform_image(row: Dict[str, Any]) -> Dict[str, Any]:
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((32, 32))
    ])
    row["image"] = transform(row["image"])
    return row

ds = (
    ray.data.read_images("s3://anonymous@ray-example-data/batoidea/JPEGImages")
    .map(transform_image)
)

다음으로 모델을 설정하고 호출 가능한 클래스를 적용한다

import torch
from torchvision import models

class ImageClassifier:
    def __init__(self):
        weights = models.ResNet18_Weights.DEFAULT
        self.model = models.resnet18(weights=weights)
        self.model.eval()

    def __call__(self, batch):
        inputs = torch.from_numpy(batch["image"])
        with torch.inference_mode():
            outputs = self.model(inputs)
        return {"class": outputs.argmax(dim=1)}

마지막으로 Dataset.map_batches를 호출한다

predictions = ds.map_batches(
    ImageClassifier,
    concurrency=2,
    batch_size=4
)
predictions.show(3)

{'class': 118}
{'class': 153}
{'class': 296}