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.CenterCrop(SIZE), 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("\\", "") # must clean slashes using filename 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))