Skip to content

Commit 47d61e2

Browse files
kohya-sskohya-ss
committed
format by black
1 parent 8f6fc8d commit 47d61e2

5 files changed

+783
-606
lines changed

‎finetune/make_captions.py

+159-132
Original file line numberDiff line numberDiff line change
@@ -14,158 +14,185 @@
1414
from blip.blip import blip_decoder
1515
import library.train_util as train_util
1616

17-
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
17+
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
1818

1919

2020
IMAGE_SIZE = 384
2121

2222
# 正方形でいいのか？　という気がするがソースがそうなので
23-
IMAGE_TRANSFORM = transforms.Compose([
24-
transforms.Resize((IMAGE_SIZE, IMAGE_SIZE), interpolation=InterpolationMode.BICUBIC),
25-
transforms.ToTensor(),
26-
transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
27-
])
23+
IMAGE_TRANSFORM = transforms.Compose(
24+
[
25+
transforms.Resize((IMAGE_SIZE, IMAGE_SIZE), interpolation=InterpolationMode.BICUBIC),
26+
transforms.ToTensor(),
27+
transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
28+
]
29+
)
30+
2831

2932
# 共通化したいが微妙に処理が異なる……
3033
class ImageLoadingTransformDataset(torch.utils.data.Dataset):
31-
def __init__(self, image_paths):
32-
self.images = image_paths
34+
def __init__(self, image_paths):
35+
self.images = image_paths
3336

34-
def __len__(self):
35-
return len(self.images)
37+
def __len__(self):
38+
return len(self.images)
3639

37-
def __getitem__(self, idx):
38-
img_path = self.images[idx]
40+
def __getitem__(self, idx):
41+
img_path = self.images[idx]
3942

40-
try:
41-
image = Image.open(img_path).convert("RGB")
42-
# convert to tensor temporarily so dataloader will accept it
43-
tensor = IMAGE_TRANSFORM(image)
44-
except Exception as e:
45-
print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
46-
return None
43+
try:
44+
image = Image.open(img_path).convert("RGB")
45+
# convert to tensor temporarily so dataloader will accept it
46+
tensor = IMAGE_TRANSFORM(image)
47+
except Exception as e:
48+
print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
49+
return None
4750

48-
return (tensor, img_path)
51+
return (tensor, img_path)
4952

5053

5154
def collate_fn_remove_corrupted(batch):
52-
"""Collate function that allows to remove corrupted examples in the
53-
dataloader. It expects that the dataloader returns 'None' when that occurs.
54-
The 'None's in the batch are removed.
55-
"""
56-
# Filter out all the Nones (corrupted examples)
57-
batch = list(filter(lambda x: x is not None, batch))
58-
return batch
55+
"""Collate function that allows to remove corrupted examples in the
56+
dataloader. It expects that the dataloader returns 'None' when that occurs.
57+
The 'None's in the batch are removed.
58+
"""
59+
# Filter out all the Nones (corrupted examples)
60+
batch = list(filter(lambda x: x is not None, batch))
61+
return batch
5962

6063

6164
def main(args):
62-
# fix the seed for reproducibility
63-
seed = args.seed # + utils.get_rank()
64-
torch.manual_seed(seed)
65-
np.random.seed(seed)
66-
random.seed(seed)
67-
68-
if not os.path.exists("blip"):
69-
args.train_data_dir = os.path.abspath(args.train_data_dir) # convert to absolute path
70-
71-
cwd = os.getcwd()
72-
print('Current Working Directory is: ', cwd)
73-
os.chdir('finetune')
74-
75-
print(f"load images from {args.train_data_dir}")
76-
train_data_dir_path = Path(args.train_data_dir)
77-
image_paths = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
78-
print(f"found {len(image_paths)} images.")
79-
80-
print(f"loading BLIP caption: {args.caption_weights}")
81-
model = blip_decoder(pretrained=args.caption_weights, image_size=IMAGE_SIZE, vit='large', med_config="./blip/med_config.json")
82-
model.eval()
83-
model = model.to(DEVICE)
84-
print("BLIP loaded")
85-
86-
# captioningする
87-
def run_batch(path_imgs):
88-
imgs = torch.stack([im for _, im in path_imgs]).to(DEVICE)
89-
90-
with torch.no_grad():
91-
if args.beam_search:
92-
captions = model.generate(imgs, sample=False, num_beams=args.num_beams,
93-
max_length=args.max_length, min_length=args.min_length)
94-
else:
95-
captions = model.generate(imgs, sample=True, top_p=args.top_p, max_length=args.max_length, min_length=args.min_length)
96-
97-
for (image_path, _), caption in zip(path_imgs, captions):
98-
with open(os.path.splitext(image_path)[0] + args.caption_extension, "wt", encoding='utf-8') as f:
99-
f.write(caption + "\n")
100-
if args.debug:
101-
print(image_path, caption)
102-
103-
# 読み込みの高速化のためにDataLoaderを使うオプション
104-
if args.max_data_loader_n_workers is not None:
105-
dataset = ImageLoadingTransformDataset(image_paths)
106-
data = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
107-
num_workers=args.max_data_loader_n_workers, collate_fn=collate_fn_remove_corrupted, drop_last=False)
108-
else:
109-
data = [[(None, ip)] for ip in image_paths]
110-
111-
b_imgs = []
112-
for data_entry in tqdm(data, smoothing=0.0):
113-
for data in data_entry:
114-
if data is None:
115-
continue
116-
117-
img_tensor, image_path = data
118-
if img_tensor is None:
119-
try:
120-
raw_image = Image.open(image_path)
121-
if raw_image.mode != 'RGB':
122-
raw_image = raw_image.convert("RGB")
123-
img_tensor = IMAGE_TRANSFORM(raw_image)
124-
except Exception as e:
125-
print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
126-
continue
127-
128-
b_imgs.append((image_path, img_tensor))
129-
if len(b_imgs) >= args.batch_size:
65+
# fix the seed for reproducibility
66+
seed = args.seed # + utils.get_rank()
67+
torch.manual_seed(seed)
68+
np.random.seed(seed)
69+
random.seed(seed)
70+
71+
if not os.path.exists("blip"):
72+
args.train_data_dir = os.path.abspath(args.train_data_dir) # convert to absolute path
73+
74+
cwd = os.getcwd()
75+
print("Current Working Directory is: ", cwd)
76+
os.chdir("finetune")
77+
78+
print(f"load images from {args.train_data_dir}")
79+
train_data_dir_path = Path(args.train_data_dir)
80+
image_paths = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
81+
print(f"found {len(image_paths)} images.")
82+
83+
print(f"loading BLIP caption: {args.caption_weights}")
84+
model = blip_decoder(pretrained=args.caption_weights, image_size=IMAGE_SIZE, vit="large", med_config="./blip/med_config.json")
85+
model.eval()
86+
model = model.to(DEVICE)
87+
print("BLIP loaded")
88+
89+
# captioningする
90+
def run_batch(path_imgs):
91+
imgs = torch.stack([im for _, im in path_imgs]).to(DEVICE)
92+
93+
with torch.no_grad():
94+
if args.beam_search:
95+
captions = model.generate(
96+
imgs, sample=False, num_beams=args.num_beams, max_length=args.max_length, min_length=args.min_length
97+
)
98+
else:
99+
captions = model.generate(
100+
imgs, sample=True, top_p=args.top_p, max_length=args.max_length, min_length=args.min_length
101+
)
102+
103+
for (image_path, _), caption in zip(path_imgs, captions):
104+
with open(os.path.splitext(image_path)[0] + args.caption_extension, "wt", encoding="utf-8") as f:
105+
f.write(caption + "\n")
106+
if args.debug:
107+
print(image_path, caption)
108+
109+
# 読み込みの高速化のためにDataLoaderを使うオプション
110+
if args.max_data_loader_n_workers is not None:
111+
dataset = ImageLoadingTransformDataset(image_paths)
112+
data = torch.utils.data.DataLoader(
113+
dataset,
114+
batch_size=args.batch_size,
115+
shuffle=False,
116+
num_workers=args.max_data_loader_n_workers,
117+
collate_fn=collate_fn_remove_corrupted,
118+
drop_last=False,
119+
)
120+
else:
121+
data = [[(None, ip)] for ip in image_paths]
122+
123+
b_imgs = []
124+
for data_entry in tqdm(data, smoothing=0.0):
125+
for data in data_entry:
126+
if data is None:
127+
continue
128+
129+
img_tensor, image_path = data
130+
if img_tensor is None:
131+
try:
132+
raw_image = Image.open(image_path)
133+
if raw_image.mode != "RGB":
134+
raw_image = raw_image.convert("RGB")
135+
img_tensor = IMAGE_TRANSFORM(raw_image)
136+
except Exception as e:
137+
print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
138+
continue
139+
140+
b_imgs.append((image_path, img_tensor))
141+
if len(b_imgs) >= args.batch_size:
142+
run_batch(b_imgs)
143+
b_imgs.clear()
144+
if len(b_imgs) > 0:
130145
run_batch(b_imgs)
131-
b_imgs.clear()
132-
if len(b_imgs) > 0:
133-
run_batch(b_imgs)
134146

135-
print("done!")
147+
print("done!")
136148

137149

138150
def setup_parser() -> argparse.ArgumentParser:
139-
parser = argparse.ArgumentParser()
140-
parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
141-
parser.add_argument("--caption_weights", type=str, default="https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth",
142-
help="BLIP caption weights (model_large_caption.pth) / BLIP captionの重みファイル(model_large_caption.pth)")
143-
parser.add_argument("--caption_extention", type=str, default=None,
144-
help="extension of caption file (for backward compatibility) / 出力されるキャプションファイルの拡張子（スペルミスしていたのを残してあります）")
145-
parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption file / 出力されるキャプションファイルの拡張子")
146-
parser.add_argument("--beam_search", action="store_true",
147-
help="use beam search (default Nucleus sampling) / beam searchを使う（このオプション未指定時はNucleus sampling）")
148-
parser.add_argument("--batch_size", type=int, default=1, help="batch size in inference / 推論時のバッチサイズ")
149-
parser.add_argument("--max_data_loader_n_workers", type=int, default=None,
150-
help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）")
151-
parser.add_argument("--num_beams", type=int, default=1, help="num of beams in beam search /beam search時のビーム数（多いと精度が上がるが時間がかかる）")
152-
parser.add_argument("--top_p", type=float, default=0.9, help="top_p in Nucleus sampling / Nucleus sampling時のtop_p")
153-
parser.add_argument("--max_length", type=int, default=75, help="max length of caption / captionの最大長")
154-
parser.add_argument("--min_length", type=int, default=5, help="min length of caption / captionの最小長")
155-
parser.add_argument('--seed', default=42, type=int, help='seed for reproducibility / 再現性を確保するための乱数seed')
156-
parser.add_argument("--debug", action="store_true", help="debug mode")
157-
parser.add_argument("--recursive", action="store_true", help="search for images in subfolders recursively / サブフォルダを再帰的に検索する")
158-
159-
return parser
160-
161-
162-
if __name__ == '__main__':
163-
parser = setup_parser()
164-
165-
args = parser.parse_args()
166-
167-
# スペルミスしていたオプションを復元する
168-
if args.caption_extention is not None:
169-
args.caption_extension = args.caption_extention
170-
171-
main(args)
151+
parser = argparse.ArgumentParser()
152+
parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
153+
parser.add_argument(
154+
"--caption_weights",
155+
type=str,
156+
default="https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth",
157+
help="BLIP caption weights (model_large_caption.pth) / BLIP captionの重みファイル(model_large_caption.pth)",
158+
)
159+
parser.add_argument(
160+
"--caption_extention",
161+
type=str,
162+
default=None,
163+
help="extension of caption file (for backward compatibility) / 出力されるキャプションファイルの拡張子（スペルミスしていたのを残してあります）",
164+
)
165+
parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption file / 出力されるキャプションファイルの拡張子")
166+
parser.add_argument(
167+
"--beam_search",
168+
action="store_true",
169+
help="use beam search (default Nucleus sampling) / beam searchを使う（このオプション未指定時はNucleus sampling）",
170+
)
171+
parser.add_argument("--batch_size", type=int, default=1, help="batch size in inference / 推論時のバッチサイズ")
172+
parser.add_argument(
173+
"--max_data_loader_n_workers",
174+
type=int,
175+
default=None,
176+
help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）",
177+
)
178+
parser.add_argument("--num_beams", type=int, default=1, help="num of beams in beam search /beam search時のビーム数（多いと精度が上がるが時間がかかる）")
179+
parser.add_argument("--top_p", type=float, default=0.9, help="top_p in Nucleus sampling / Nucleus sampling時のtop_p")
180+
parser.add_argument("--max_length", type=int, default=75, help="max length of caption / captionの最大長")
181+
parser.add_argument("--min_length", type=int, default=5, help="min length of caption / captionの最小長")
182+
parser.add_argument("--seed", default=42, type=int, help="seed for reproducibility / 再現性を確保するための乱数seed")
183+
parser.add_argument("--debug", action="store_true", help="debug mode")
184+
parser.add_argument("--recursive", action="store_true", help="search for images in subfolders recursively / サブフォルダを再帰的に検索する")
185+
186+
return parser
187+
188+
189+
if __name__ == "__main__":
190+
parser = setup_parser()
191+
192+
args = parser.parse_args()
193+
194+
# スペルミスしていたオプションを復元する
195+
if args.caption_extention is not None:
196+
args.caption_extension = args.caption_extention
197+
198+
main(args)

0 commit comments

Comments
 (0)