[파이썬] fastai 캡션 생성 및 이미지 설명

많은 웹사이트 및 소셜 미디어 플랫폼은 이미지와 함께 설명(캡션)을 제공하여 사용자에게 더 나은 경험을 제공합니다. 이러한 기능은 주로 기계 학습 및 자연어 처리(Natural Language Processing, NLP) 기술을 사용하여 구현됩니다.

fastai는 파이썬에서 이미지 분류(classification), 객체 감지(object detection), 세분화(segmentation), 캡션 생성(captioning) 등 다양한 컴퓨터 비전 및 자연어 처리 작업을 수행할 수 있는 높은 수준의 API를 제공합니다. 이번 블로그 포스트에서는 fastai를 사용하여 캡션 생성 및 이미지 설명을 제공하는 방법에 대해 알아보겠습니다.

데이터 준비

먼저, 캡션 생성을 위해 필요한 데이터를 수집하고 준비해야 합니다. 일반적으로 이미지와 해당 이미지에 대한 캡션으로 구성된 데이터셋을 필요로 합니다. Fastai는 다양한 데이터셋을 불러오고 전처리할 수 있는 도구를 제공하므로, 캡션 데이터셋을 불러온 후 필요한 전처리 작업을 수행할 수 있습니다.

# 데이터셋 로드
from fastai.vision.all import *
from fastai.data.transforms import *
from torchvision.models import resnet34

path = untar_data(URLs.COCO_SAMPLE)

# 이미지 데이터 불러오기
data = DataBlock(blocks=(ImageBlock, TextBlock),
                 get_items=get_image_files,
                 get_y=parent_label,
                 splitter=RandomSplitter(valid_pct=0.2, seed=42),
                 item_tfms=Resize(460),
                 batch_tfms=aug_transforms(size=224, min_scale=0.75))

dls = data.dataloaders(path)

위의 코드는 fastai를 사용하여 COCO 데이터셋의 이미지와 캡션을 불러오는 예시입니다. 실제 사용 시에는 데이터셋에 맞게 경로와 데이터 전처리 방법을 설정해주어야 합니다.

캡션 모델 구축

데이터를 준비했다면, 다음은 캡션을 생성하기 위한 모델을 구축해야 합니다. Fastai는 이미지 및 텍스트 작업에 대한 사전 학습된 모델을 제공하므로, 이를 사용하여 캡션 모델을 구축할 수 있습니다.

# 캡션 생성 모델 구축
encoder = create_body(resnet34, pretrained=True, cut=-2)
encoder = nn.Sequential(*children(encoder), nn.AdaptiveAvgPool2d(1), Flatten())
encoder_output_size = 512

class CaptionModel(Module):
    def __init__(self, encoder, vocab_size, hidden_size, encoder_output_size):
        self.encoder = encoder
        self.attention = SelfAttention(hidden_size)
        self.decoder = LinearDecoder(hidden_size, vocab_size)
        self.rnn = LSTMCell(encoder_output_size + hidden_size, hidden_size)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.encoder_output_size = encoder_output_size

    def forward(self, x):
        # 이미지 인코딩
        encoded = self.encoder(x)

        # 캡션 생성
        hidden_state, cell_state = self.init_hidden_states(x.shape[0], self.hidden_size, device=x.device)
        output = torch.zeros(x.shape[0], self.vocab_size).to(x.device)
        for i in range(x.shape[1]):
            context_vector, attention_weights = self.attention(encoded, hidden_state)
            input_embedding = self.decoder.embedding(output.argmax(dim=-1))
            lstm_input = torch.cat([context_vector, input_embedding], dim=1)
            hidden_state, cell_state = self.rnn(lstm_input, (hidden_state, cell_state))
            output = self.decoder(hidden_state)

        return output

model = CaptionModel(encoder, vocab_size, hidden_size, encoder_output_size)

위의 코드는 fastai를 사용하여 이미지와 캡션을 입력받아 캡션을 생성하는 모델을 구축하는 예시입니다. 모델의 구조와 피처 추출을 위한 encoder, 어텐션 모듈, 디코더, LSTM 셀 등을 포함하여 캡션을 생성하는 방법을 정의합니다.

모델 학습

모델이 준비되었다면, 이제 모델을 학습해야 합니다. Fastai는 다양한 학습 기술 및 학습률 최적화 알고리즘을 제공하므로, 이를 사용하여 모델을 학습할 수 있습니다.

# 모델 학습
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), metrics=[accuracy])
learn.fit_one_cycle(10)

위의 코드는 모델을 학습시키는 예시입니다. 실제 사용 시에는 데이터셋과 학습 설정에 맞게 변경해주어야 합니다.

이미지 설명 생성

캡션 모델이 학습되었다면, 이제 이미지에 대한 설명을 생성하는 것이 가능해집니다. 다음은 새로운 이미지에 대한 캡션을 생성하는 예시 코드입니다.

# 이미지 설명 생성
def generate_caption(image):
    image = PILImage.create(image)
    pred, _, _ = learn.predict(image)
    caption = ' '.join(dls.vocab[pred.cpu().argmax(dim=-1)])
    return caption

image_path = 'path/to/image.jpg'
caption = generate_caption(image_path)
print(caption)

위의 코드는 새로운 이미지에 대해 학습된 모델을 사용하여 캡션을 생성하는 예시입니다. 실제 사용 시에는 이미지 경로를 적절한 값으로 변경해주어야 합니다.

이와 같이 fastai를 사용하여 캡션 생성 및 이미지 설명을 제공할 수 있습니다. fastai의 강력한 기능과 컴퓨터 비전, 자연어 처리 모델을 통해 다양한 이미지와 대응하는 텍스트 데이터셋을 활용하여 더 나은 사용자 경험을 제공할 수 있습니다.