152 lines
6.2 KiB
Python
152 lines
6.2 KiB
Python
import argparse
|
||
import os
|
||
import re
|
||
|
||
from PIL import Image
|
||
from tqdm import tqdm
|
||
import torch
|
||
from transformers import AutoProcessor, AutoModelForCausalLM
|
||
from transformers.generation.utils import GenerationMixin
|
||
|
||
import library.train_util as train_util
|
||
|
||
|
||
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
PATTERN_REPLACE = [
|
||
re.compile(r'(has|with|and) the (words?|letters?|name) (" ?[^"]*"|\w+)( ?(is )?(on|in) (the |her |their |him )?\w+)?'),
|
||
re.compile(r'(with a sign )?that says ?(" ?[^"]*"|\w+)( ?on it)?'),
|
||
re.compile(r"(with a sign )?that says ?(' ?(i'm)?[^']*'|\w+)( ?on it)?"),
|
||
re.compile(r'with the number \d+ on (it|\w+ \w+)'),
|
||
re.compile(r'with the words "'),
|
||
re.compile(r'word \w+ on it'),
|
||
re.compile(r'that says the word \w+ on it'),
|
||
re.compile('that says\'the word "( on it)?'),
|
||
]
|
||
|
||
# 誤検知しまくりの with the word xxxx を消す
|
||
|
||
|
||
def remove_words(captions, debug):
|
||
removed_caps = []
|
||
for caption in captions:
|
||
cap = caption
|
||
for pat in PATTERN_REPLACE:
|
||
cap = pat.sub("", cap)
|
||
if debug and cap != caption:
|
||
print(caption)
|
||
print(cap)
|
||
removed_caps.append(cap)
|
||
return removed_caps
|
||
|
||
|
||
def collate_fn_remove_corrupted(batch):
|
||
"""Collate function that allows to remove corrupted examples in the
|
||
dataloader. It expects that the dataloader returns 'None' when that occurs.
|
||
The 'None's in the batch are removed.
|
||
"""
|
||
# Filter out all the Nones (corrupted examples)
|
||
batch = list(filter(lambda x: x is not None, batch))
|
||
return batch
|
||
|
||
|
||
def main(args):
|
||
# GITにバッチサイズが1より大きくても動くようにパッチを当てる: transformers 4.26.0用
|
||
org_prepare_input_ids_for_generation = GenerationMixin._prepare_input_ids_for_generation
|
||
curr_batch_size = [args.batch_size] # ループの最後で件数がbatch_size未満になるので入れ替えられるように
|
||
|
||
# input_idsがバッチサイズと同じ件数である必要がある：バッチサイズはこの関数から参照できないので外から渡す
|
||
# ここより上で置き換えようとするとすごく大変
|
||
def _prepare_input_ids_for_generation_patch(self, bos_token_id, encoder_outputs):
|
||
input_ids = org_prepare_input_ids_for_generation(self, bos_token_id, encoder_outputs)
|
||
if input_ids.size()[0] != curr_batch_size[0]:
|
||
input_ids = input_ids.repeat(curr_batch_size[0], 1)
|
||
return input_ids
|
||
GenerationMixin._prepare_input_ids_for_generation = _prepare_input_ids_for_generation_patch
|
||
|
||
print(f"load images from {args.train_data_dir}")
|
||
image_paths = train_util.glob_images(args.train_data_dir)
|
||
print(f"found {len(image_paths)} images.")
|
||
|
||
# できればcacheに依存せず明示的にダウンロードしたい
|
||
print(f"loading GIT: {args.model_id}")
|
||
git_processor = AutoProcessor.from_pretrained(args.model_id)
|
||
git_model = AutoModelForCausalLM.from_pretrained(args.model_id).to(DEVICE)
|
||
print("GIT loaded")
|
||
|
||
# captioningする
|
||
def run_batch(path_imgs):
|
||
imgs = [im for _, im in path_imgs]
|
||
|
||
curr_batch_size[0] = len(path_imgs)
|
||
inputs = git_processor(images=imgs, return_tensors="pt").to(DEVICE) # 画像はpil形式
|
||
generated_ids = git_model.generate(pixel_values=inputs.pixel_values, max_length=args.max_length)
|
||
captions = git_processor.batch_decode(generated_ids, skip_special_tokens=True)
|
||
|
||
if args.remove_words:
|
||
captions = remove_words(captions, args.debug)
|
||
|
||
for (image_path, _), caption in zip(path_imgs, captions):
|
||
with open(os.path.splitext(image_path)[0] + args.caption_extension, "wt", encoding='utf-8') as f:
|
||
f.write(caption + "\n")
|
||
if args.debug:
|
||
print(image_path, caption)
|
||
|
||
# 読み込みの高速化のためにDataLoaderを使うオプション
|
||
if args.max_data_loader_n_workers is not None:
|
||
dataset = train_util.ImageLoadingDataset(image_paths)
|
||
data = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
|
||
num_workers=args.max_data_loader_n_workers, collate_fn=collate_fn_remove_corrupted, drop_last=False)
|
||
else:
|
||
data = [[(None, ip)] for ip in image_paths]
|
||
|
||
b_imgs = []
|
||
for data_entry in tqdm(data, smoothing=0.0):
|
||
for data in data_entry:
|
||
if data is None:
|
||
continue
|
||
|
||
image, image_path = data
|
||
if image is None:
|
||
try:
|
||
image = Image.open(image_path)
|
||
if image.mode != 'RGB':
|
||
image = image.convert("RGB")
|
||
except Exception as e:
|
||
print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
|
||
continue
|
||
|
||
b_imgs.append((image_path, image))
|
||
if len(b_imgs) >= args.batch_size:
|
||
run_batch(b_imgs)
|
||
b_imgs.clear()
|
||
|
||
if len(b_imgs) > 0:
|
||
run_batch(b_imgs)
|
||
|
||
print("done!")
|
||
|
||
|
||
def setup_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
|
||
parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption file / 出力されるキャプションファイルの拡張子")
|
||
parser.add_argument("--model_id", type=str, default="microsoft/git-large-textcaps",
|
||
help="model id for GIT in Hugging Face / 使用するGITのHugging FaceのモデルID")
|
||
parser.add_argument("--batch_size", type=int, default=1, help="batch size in inference / 推論時のバッチサイズ")
|
||
parser.add_argument("--max_data_loader_n_workers", type=int, default=None,
|
||
help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）")
|
||
parser.add_argument("--max_length", type=int, default=50, help="max length of caption / captionの最大長")
|
||
parser.add_argument("--remove_words", action="store_true",
|
||
help="remove like `with the words xxx` from caption / `with the words xxx`のような部分をキャプションから削除する")
|
||
parser.add_argument("--debug", action="store_true", help="debug mode")
|
||
|
||
return parser
|
||
|
||
|
||
if __name__ == '__main__':
|
||
parser = setup_parser()
|
||
|
||
args = parser.parse_args()
|
||
main(args)
|