mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
"""Functions for generating image captions"""
|
|
|
|
import os
|
|
import os.path
|
|
from functools import lru_cache
|
|
|
|
import torch
|
|
from torchvision import transforms
|
|
from torchvision.transforms.functional import InterpolationMode
|
|
|
|
from imaginairy.utils import get_device
|
|
from imaginairy.utils.downloads import get_cached_url_path
|
|
from imaginairy.vendored.blip.blip import BLIP_Decoder, load_checkpoint
|
|
|
|
device = get_device()
|
|
if "mps" in device:
|
|
device = "cpu"
|
|
|
|
BLIP_EVAL_SIZE = 384
|
|
|
|
|
|
@lru_cache
|
|
def blip_model():
|
|
from imaginairy.utils.paths import PKG_ROOT
|
|
|
|
config_path = os.path.join(
|
|
PKG_ROOT, "vendored", "blip", "configs", "med_config.json"
|
|
)
|
|
url = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_base_caption.pth"
|
|
|
|
model = BLIP_Decoder(image_size=BLIP_EVAL_SIZE, vit="base", med_config=config_path)
|
|
cached_url_path = get_cached_url_path(url)
|
|
model, msg = load_checkpoint(model, cached_url_path)
|
|
model.eval()
|
|
model = model.to(device)
|
|
return model
|
|
|
|
|
|
def generate_caption(image, min_length=30):
|
|
"""Given an image, return a caption."""
|
|
image = image.convert("RGB")
|
|
gpu_image = (
|
|
transforms.Compose(
|
|
[
|
|
transforms.Resize(
|
|
(BLIP_EVAL_SIZE, BLIP_EVAL_SIZE),
|
|
interpolation=InterpolationMode.BICUBIC,
|
|
),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(
|
|
(0.48145466, 0.4578275, 0.40821073),
|
|
(0.26862954, 0.26130258, 0.27577711),
|
|
),
|
|
]
|
|
)(image)
|
|
.unsqueeze(0)
|
|
.to(device)
|
|
)
|
|
|
|
with torch.no_grad():
|
|
caption = blip_model().generate(
|
|
gpu_image, sample=True, num_beams=3, max_length=80, min_length=min_length
|
|
)
|
|
return caption[0]
|