| import argparse | |
| import glob | |
| import os | |
| import cv2 | |
| import numpy | |
| import torch | |
| from PIL import Image | |
| from Model import TRCaptionNet, clip_transform | |
| def demo(opt): | |
| preprocess = clip_transform(224) | |
| model = TRCaptionNet({ | |
| "max_length": 35, | |
| "clip": "ViT-L/14", | |
| "bert": "dbmdz/bert-base-turkish-cased", | |
| "proj": True, | |
| "proj_num_head": 16 | |
| }) | |
| device = torch.device(opt.device) | |
| model.load_state_dict(torch.load(opt.model_ckpt, map_location=device)["model"], strict=True) | |
| model = model.to(device) | |
| model.eval() | |
| image_paths = glob.glob(os.path.join(opt.input_dir, '*.jpg')) | |
| for image_path in sorted(image_paths): | |
| img_name = image_path.split('/')[-1] | |
| img0 = Image.open(image_path) | |
| batch = preprocess(img0).unsqueeze(0).to(device) | |
| caption = model.generate(batch, min_length=11, repetition_penalty=1.6)[0] | |
| print(f"{img_name} :", caption) | |
| orj_img = numpy.array(img0)[:, :, ::-1] | |
| h, w, _ = orj_img.shape | |
| new_h = 800 | |
| new_w = int(new_h * (w / h)) | |
| orj_img = cv2.resize(orj_img, (new_w, new_h)) | |
| cv2.imshow("image", orj_img) | |
| cv2.waitKey(0) | |
| return | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='Turkish-Image-Captioning!') | |
| parser.add_argument('--model-ckpt', type=str, default='./checkpoints/TRCaptionNet_L14_berturk.pth') | |
| parser.add_argument('--input-dir', type=str, default='./images/') | |
| parser.add_argument('--device', type=str, default='cuda:0') | |
| args = parser.parse_args() | |
| demo(args) | |