| import argparse |
| import glob |
| import os |
| from PIL import Image |
| import sys |
| from torchvision import transforms |
| from torchvision.transforms.functional import InterpolationMode |
| import torch |
| import aiohttp |
| import asyncio |
| import subprocess |
| import numpy as np |
| import io |
| import aiofiles |
|
|
| SIZE = 384 |
| BLIP_MODEL_URL = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth' |
|
|
| def get_parser(**parser_kwargs): |
| parser = argparse.ArgumentParser(**parser_kwargs) |
| parser.add_argument( |
| "--img_dir", |
| type=str, |
| nargs="?", |
| const=True, |
| default="input", |
| help="directory with images to be captioned", |
| ), |
| parser.add_argument( |
| "--out_dir", |
| type=str, |
| nargs="?", |
| const=True, |
| default="output", |
| help="directory to put captioned images", |
| ), |
| parser.add_argument( |
| "--format", |
| type=str, |
| nargs="?", |
| const=True, |
| default="filename", |
| help="'filename', 'mrwho', 'txt', or 'caption'", |
| ), |
| parser.add_argument( |
| "--nucleus", |
| type=bool, |
| nargs="?", |
| const=True, |
| default=False, |
| help="use nucleus sampling instead of beam", |
| ), |
| parser.add_argument( |
| "--q_factor", |
| type=float, |
| nargs="?", |
| const=True, |
| default=1.0, |
| help="adjusts the likelihood of a word being repeated", |
| ), |
| parser.add_argument( |
| "--min_length", |
| type=int, |
| nargs="?", |
| const=True, |
| default=22, |
| help="adjusts the likelihood of a word being repeated", |
| ), |
| parser.add_argument( |
| "--torch_device", |
| type=str, |
| nargs="?", |
| const=False, |
| default="cuda", |
| help="specify a different torch device, e.g. 'cpu'", |
| ), |
|
|
| return parser |
|
|
| def load_image(raw_image, device): |
| transform = transforms.Compose([ |
| |
| transforms.Resize((SIZE, SIZE), interpolation=InterpolationMode.BICUBIC), |
| transforms.ToTensor(), |
| transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) |
| ]) |
| image = transform(raw_image).unsqueeze(0).to(device) |
| return image |
|
|
| def get_out_file_name(out_dir, base_name, ext): |
| return os.path.join(out_dir, f"{base_name}{ext}") |
|
|
| async def main(opt): |
| print("starting") |
| import models.blip |
|
|
| sample = False |
| if opt.nucleus: |
| sample = True |
|
|
| input_dir = opt.img_dir |
| print("input_dir: ", input_dir) |
|
|
| config_path = "scripts/BLIP/configs/med_config.json" |
|
|
| cache_folder = ".cache" |
| model_cache_path = ".cache/model_base_caption_capfilt_large.pth" |
|
|
| if not os.path.exists(cache_folder): |
| os.makedirs(cache_folder) |
|
|
| if not os.path.exists(opt.out_dir): |
| os.makedirs(opt.out_dir) |
|
|
| if not os.path.exists(model_cache_path): |
| print(f"Downloading model to {model_cache_path}... please wait") |
|
|
| async with aiohttp.ClientSession() as session: |
| async with session.get(BLIP_MODEL_URL) as res: |
| with open(model_cache_path, 'wb') as f: |
| async for chunk in res.content.iter_chunked(1024): |
| f.write(chunk) |
| print(f"Model cached to: {model_cache_path}") |
| else: |
| print(f"Model already cached to: {model_cache_path}") |
|
|
| blip_decoder = models.blip.blip_decoder(pretrained=model_cache_path, image_size=SIZE, vit='base', med_config=config_path) |
| blip_decoder.eval() |
|
|
| print(f"loading model to {opt.torch_device}") |
|
|
| blip_decoder = blip_decoder.to(torch.device(opt.torch_device)) |
|
|
| ext = ('.jpg', '.jpeg', '.png', '.webp', '.tif', '.tga', '.tiff', '.bmp', '.gif') |
|
|
| i = 0 |
|
|
| for idx, img_file_name in enumerate(glob.iglob(os.path.join(opt.img_dir, "*.*"))): |
| if img_file_name.endswith(ext): |
| caption = None |
| file_ext = os.path.splitext(img_file_name)[1] |
| if (file_ext in ext): |
| async with aiofiles.open(img_file_name, "rb") as input_file: |
| print("working image: ", img_file_name) |
|
|
| image_bin = await input_file.read() |
| image = Image.open(io.BytesIO(image_bin)) |
|
|
| if not image.mode == "RGB": |
| image = image.convert("RGB") |
|
|
| image = load_image(image, device=torch.device(opt.torch_device)) |
|
|
| if opt.nucleus: |
| captions = blip_decoder.generate(image, sample=True, top_p=opt.q_factor) |
| else: |
| captions = blip_decoder.generate(image, sample=sample, num_beams=16, min_length=opt.min_length, \ |
| max_length=48, repetition_penalty=opt.q_factor) |
|
|
| caption = captions[0] |
|
|
| if opt.format in ["mrwho","joepenna"]: |
| prefix = f"{i:05}@" |
| i += 1 |
| caption = prefix+caption |
| elif opt.format == "filename": |
| postfix = f"_{i}" |
| i += 1 |
| caption = caption+postfix |
|
|
| if opt.format in ["txt","text","caption"]: |
| out_base_name = os.path.splitext(os.path.basename(img_file_name))[0] |
|
|
| if opt.format in ["txt","text"]: |
| out_file = get_out_file_name(opt.out_dir, out_base_name, ".txt") |
|
|
| if opt.format in ["caption"]: |
| out_file = get_out_file_name(opt.out_dir, out_base_name, ".caption") |
|
|
| if opt.format in ["txt","text","caption"]: |
| print("writing caption to: ", out_file) |
| async with aiofiles.open(out_file, "w") as out_file: |
| await out_file.write(caption) |
|
|
| if opt.format in ["filename", "mrwho", "joepenna"]: |
| caption = caption.replace("/", "").replace("\\", "") |
| out_file = get_out_file_name(opt.out_dir, caption, file_ext) |
| async with aiofiles.open(out_file, "wb") as out_file: |
| await out_file.write(image_bin) |
| elif opt.format == "json": |
| raise NotImplementedError |
| elif opt.format == "parquet": |
| raise NotImplementedError |
|
|
| def isWindows(): |
| return sys.platform.startswith("win") |
|
|
| if __name__ == "__main__": |
| parser = get_parser() |
| opt = parser.parse_args() |
|
|
| if opt.format not in ["filename", "mrwho", "joepenna", "txt", "text", "caption"]: |
| raise ValueError("format must be 'filename', 'mrwho', 'txt', or 'caption'") |
|
|
| if (isWindows()): |
| print("Windows detected, using asyncio.WindowsSelectorEventLoopPolicy") |
| asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) |
| else: |
| print("Unix detected, using default asyncio event loop policy") |
|
|
| if not os.path.exists("scripts/BLIP"): |
| print("BLIP not found, cloning BLIP repo") |
| subprocess.run(["git", "clone", "https://github.com/salesforce/BLIP", "scripts/BLIP"]) |
| blip_path = "scripts/BLIP" |
| sys.path.append(blip_path) |
|
|
| asyncio.run(main(opt)) |
|
|
