이미지에서 원하는 텍스트 뽑아내기, 이미지 캡셔닝 BLIP-2(Colab 가능)

LAVIS/projects/blip2 at main · salesforce/LAVIS

이 글은 위 Github을 참고하여 작성했습니다.

---

• 본 글에서는 이미지 캡셔닝(Image2 Text) BLIP-2 모델 사용법을 알려드리려고 합니다.
• BLIP-2은 입력 이미지에 대한 Text를 출력해주는데, 사용자가 원하는 형태의 답변을 지정해 줄 수 있습니다.
• Google Colab T4(15GB) 환경에서 실행가능하며, 약 12GB 정도의 GPU memory를 사용합니다.

---

1. Install

BLIP-2 패키지를 설치해줍니다. BLIP-2은 salesforce AI research에서 만들어졌고 salesforce-lavis는 BLIP-2의 소스코드가 있는 github입니다.

pip install salesforce-lavis

 

2. BLIP-2 실행

 

BLIP-2 모델을 실행하기 전에 이미지를 불러옵니다. 6번 라인 이미지 경로를 알맞게 수정해 주세요.

import torch
from PIL import Image
# setup device to use
device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
# load sample image
raw_image = Image.open("../../docs/_static/merlion.png").convert("RGB")
display(raw_image.resize((596, 437)))

 

그 후, load_model_and_preprocess 메서드로 model과 전처리 메서드를 불러옵니다. 이미지와 모델을 device에 할당해 줍니다.

import torch
from lavis.models import load_model_and_preprocess
# loads BLIP-2 pre-trained model
model, vis_processors, _ = load_model_and_preprocess(name="blip2_t5", model_type="pretrain_flant5xxl", is_eval=True, device=device)
# prepare the image
image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)

 

모델에게 image와 prompt를 전달해 줄 수 있는데, 이때 prompt로 자신이 원하는 답변의 형식을 지정해 줄 수 있습니다.

아래 코드는 이미지에 대해 어느 도시인지 물어보는 prompt를 추가한 예시입니다.

model.generate({"image": image, "prompt": "Question: which city is this? Answer:"})
# 'singapore'

3. 실행 결과

BLIP-2에게 전달한 강아지 이미지

 

위 이미지를 BLIP-2에게 prompt 없이 전달하면 다음과 같은 결과를 출력합니다.

 

BLIP-2의 출력 결과

 

"a small white dog running through a field of grass"라는 텍스트를 출력한 것을 확인할 수 있습니다.