画像キャプション生成モデル 「CATR」の出力文を入力画像に埋込んでファイル出力するように書き換えてみた


__[前回の記事](https://qiita.com/electronics_diy721/items/10ce68bf25ea239b2395)__では、CATRモデルの実装コードを、githubから落として、動作検証を行ってみました。

###（ 前回の記事 ）

- __[Transformerベースの画像キャプション生成モデル CATRを動かしてみた](https://qiita.com/electronics_diy721/items/10ce68bf25ea239b2395)__

今回は、git cloneした推論工程の実行ファイル __*predict.py*__を改修して、出力されたキャプション文を入力画像に文字列埋込して、結果をファイル保存するように書き換えてみました。

なお、キャプション文が40文字を超える場合は、改行して、改行した結果を入力画像に埋め込むようにしました。4行目以降は、埋め込まずに捨てることにしました。

####（ 入力ファイル ）

![test008.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/02d3c8a0-43a1-8144-63ff-f30b4b6b3b2f.jpeg)

####（ 出力ファイル ）

__*image_files*ディレクトリの直下に、キャプション文を埋め込んだ画像ファイルが出力されます。__

![test008.jpe_解析結果.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/0a6b3c53-261a-d388-f89a-5f6b316dbcef.jpeg)

####（ 実行画面 ）

```bash:Terminal
electron@diynoMacBook-Pro catr % python3 predict_output_file_save.py --path ./image_files/test008.jpeg --v v3
Using cache found in /Users/ocean/.cache/torch/hub/saahiluppal_catr_master
A man sitting on a beach with a surfboard.
2
electron@diynoMacBook-Pro catr % 
```
___

###他の画像ファイルを試してみます。

####（ 入力ファイル ）

__ファイル名： *test005.jpg*__

![test005.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/a1878455-3f58-4bd3-8bfe-c7a09ff31b97.jpeg)


####（ 出力ファイル ）

__*image_files*ディレクトリの直下に、キャプション文を埋め込んだ画像ファイルが出力されます。__
__出力ファイル名： *test005.jp_解析結果.jpeg*__

![test005.jp_解析結果.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/0c66ebad-bc91-186d-3b67-a57dd08e33ff.jpeg)

```bash:Terminal
electron@diynoMacBook-Pro catr %  ls image_files 
dog.jpg				test002.jpg			test005.jpg			test010.jpg
special.jp_解析結果.jpeg	test003.jpg			test007.jpeg
special.jpg			test004.jpg			test008.jpe_解析結果.jpeg
test001.png			test005.jp_解析結果.jpeg	test008.jpeg
electron@diynoMacBook-Pro catr %  
electron@diynoMacBook-Pro catr %  open ./image_files/test005.jp_解析結果.jpeg 
electron@diynoMacBook-Pro catr %
```  

####（ 実行画面 ）

```bash:Terminal
electron@diynoMacBook-Pro catr % python3 predict_output_file_save.py --path ./image_files/test005.jpg --v v3 
Using cache found in /Users/ocean/.cache/torch/hub/saahiluppal_catr_master
A man holding a bottle of beer and a man in a yellow shirt.
2
electron@diynoMacBook-Pro catr % 
```

___

###もう１つ別の画像ファイルでも試してみます。
####（ 入力ファイル ）

__ファイル名： *special.jpg*__

![special.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/12e7b2fe-6b2c-74d0-9a2a-75adc3402586.jpeg)


####（ 出力ファイル ）

__*image_files*ディレクトリの直下に、キャプション文を埋め込んだ画像ファイルが出力されます。__
__出力ファイル名： *special.jp_解析結果.jpeg*__

![special.jp_解析結果.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/e5a046bf-562c-466a-0000-0493f296e579.jpeg)


####（ 実行画面 ）

```bash:Terminal
electron@diynoMacBook-Pro catr %  python3 predict_output_file_save.py --path ./image_files/special.jpg --v v3
Using cache found in /Users/ocean/.cache/torch/hub/saahiluppal_catr_master
A picture of a bunch of books and a teddy bear.
2
electron@diynoMacBook-Pro catr % 
```

___

#実行したのは、次のスクリプト・ファイルです。

* __スクリプト・ファイル名： *predict_output_file_save.py*__

```Python3:predict_output_file_save.py
import torch
from transformers import BertTokenizer
from PIL import Image
import argparse
import cv2
import matplotlib.pyplot as plt
from models import caption
from datasets import coco, utils
from configuration import Config
import os
import textwrap

parser = argparse.ArgumentParser(description='Image Captioning')
parser.add_argument('--path', type=str, help='path to image', required=True)
parser.add_argument('--v', type=str, help='version', default='v3')
parser.add_argument('--checkpoint', type=str, help='checkpoint path', default=None)
args = parser.parse_args()
image_path = args.path
version = args.v
checkpoint_path = args.checkpoint

config = Config()

if version == 'v1':
    model = torch.hub.load('saahiluppal/catr', 'v1', pretrained=True)
elif version == 'v2':
    model = torch.hub.load('saahiluppal/catr', 'v2', pretrained=True)
elif version == 'v3':
    model = torch.hub.load('saahiluppal/catr', 'v3', pretrained=True)
else:
    print("Checking for checkpoint.")
    if checkpoint_path is None:
      raise NotImplementedError('No model to chose from!')
    else:
      if not os.path.exists(checkpoint_path):
        raise NotImplementedError('Give valid checkpoint path')
      print("Found checkpoint! Loading!")
      model,_ = caption.build_model(config)
      print("Loading Checkpoint...")
      checkpoint = torch.load(checkpoint_path, map_location='cpu')
      model.load_state_dict(checkpoint['model'])
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

start_token = tokenizer.convert_tokens_to_ids(tokenizer._cls_token)
end_token = tokenizer.convert_tokens_to_ids(tokenizer._sep_token)

image = Image.open(image_path)
image = coco.val_transform(image)
image = image.unsqueeze(0)


def create_caption_and_mask(start_token, max_length):
    caption_template = torch.zeros((1, max_length), dtype=torch.long)
    mask_template = torch.ones((1, max_length), dtype=torch.bool)

    caption_template[:, 0] = start_token
    mask_template[:, 0] = False

    return caption_template, mask_template


caption, cap_mask = create_caption_and_mask(
    start_token, config.max_position_embeddings)


@torch.no_grad()
def evaluate():
    model.eval()
    for i in range(config.max_position_embeddings - 1):
        predictions = model(image, caption, cap_mask)
        predictions = predictions[:, i, :]
        predicted_id = torch.argmax(predictions, axis=-1)

        if predicted_id[0] == 102:
            return caption

        caption[:, i+1] = predicted_id[0]
        cap_mask[:, i+1] = False

    return caption


output = evaluate()
result = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
#result = tokenizer.decode(output[0], skip_special_tokens=True)
final_caption_text = str(result.capitalize())
#print(result.capitalize())
print(final_caption_text)
#print(result) 上と同じ文字列が出力される

# Pythonで文字列を折り返し・切り詰めして整形する
# https://note.nkmk.me/python-textwrap-wrap-fill-shorten/
# 生成されたキャプション文を40文字で改行する。-> List型が返却される
caption_text_wrap_list = textwrap.wrap(str(final_caption_text), 40)
print(len(caption_text_wrap_list))
# 入力した画像ファイルに、生成したキャプション文の文字列を埋込む
img = cv2.imread(image_path)

#https://qiita.com/Daiki_P/items/91acd5cc208fedd16ee9
img_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #BGRをRGBに変換.

if len(caption_text_wrap_list) == 1:
    cv2.putText(img, str(final_caption_text), (0, 50), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
elif len(caption_text_wrap_list) ==2 :
    cv2.putText(img, str(caption_text_wrap_list[0]), (0, 50), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
    cv2.putText(img, str(caption_text_wrap_list[1]), (0, 90), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
elif len(caption_text_wrap_list) >= 3:
    cv2.putText(img, str(caption_text_wrap_list[0]), (0, 50), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
    cv2.putText(img, str(caption_text_wrap_list[1]), (0, 90), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
    cv2.putText(img, str(caption_text_wrap_list[2]), (0, 130), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
    
#output_file_name = "result.jpeg"
# ファイルパスの最期の.jpg, .png, .jpegを削除
output_file_name = "{0}_解析結果.jpeg".format(image_path[0:-1])
cv2.imwrite(output_file_name,img)

#cv2.putText(img_RGB, str(final_caption_text), (0, 100), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
#plt.imshow(img_RGB)
#plt.show()

```


###（ 注意点 ）

このスクリプト・ファイルを実行する際は、PC端末がインターネットにつながっている必要があります。
内部処理の中で、BERTを使う場面でHuggingfaceのURLに自動アクセスする処理が走るからです。
インターネットにつながっていない端末で実行すると、以下のエラーが吐かれます。

>```bash:Terminal
>electron@diynoMacBook-Pro catr % # インターネット接尾を切る
>zsh: command not found: #
>electron@diynoMacBook-Pro catr % python3 predict_output_file_save.py --path >./image_files/test005.jpg --v v3
>Using cache found in /Users/electron/.cache/torch/hub/saahiluppal_catr_master
>Traceback (most recent call last):
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connection.py", line 169, in >_new_conn
>    conn = connection.create_connection(
>  File "/usr/local/lib/python3.9/site-packages/urllib3/util/connection.py", line 73, in <create_connection
>    for res in socket.getaddrinfo(host, port, family, socket.SOCK_STREAM):
>  File >"/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3>.9/socket.py", line 953, in getaddrinfo
>    for res in _socket.getaddrinfo(host, port, family, type, proto, flags):
>socket.gaierror: [Errno 8] nodename nor servname provided, or not known
>
>During handling of the above exception, another exception occurred:
>
>Traceback (most recent call last):
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connectionpool.py", line 699, in >urlopen
>    httplib_response = self._make_request(
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connectionpool.py", line 382, in >_make_request
>    self._validate_conn(conn)
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connectionpool.py", line 1010, in >_validate_conn
>    conn.connect()
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connection.py", line 353, in >connect
>    conn = self._new_conn()
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connection.py", line 181, in >_new_conn
>    raise NewConnectionError(
>urllib3.exceptions.NewConnectionError: <urllib3.connection.HTTPSConnection object at >0x1665cc7c0>: Failed to establish a new connection: [Errno 8] nodename nor servname >provided, or not known
>
>During handling of the above exception, another exception occurred:
>
>Traceback (most recent call last):
>  File "/usr/local/lib/python3.9/site-packages/requests/adapters.py", line 439, in send
>    resp = conn.urlopen(
>  File "/usr/local/lib/python3.9/site-packages/urllib3/connectionpool.py", line 755, in >urlopen
>    retries = retries.increment(
>  File "/usr/local/lib/python3.9/site-packages/urllib3/util/retry.py", line 573, in >increment
>    raise MaxRetryError(_pool, url, error or ResponseError(cause))
>urllib3.exceptions.MaxRetryError: HTTPSConnectionPool(host='huggingface.co', port=443): >Max retries exceeded with url: /api/models/bert-base-uncased (Caused by >NewConnectionError('<urllib3.connection.HTTPSConnection object at 0x1665cc7c0>: Failed >to establish a new connection: [Errno 8] nodename nor servname provided, or not known'))
>
>During handling of the above exception, another exception occurred:
>
>Traceback (most recent call last):
>  File "/Users/electron/Desktop/catr/predict_output_file_save.py", line 42, in <module>
>    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
>  File "/usr/local/lib/python3.9/site-packages/transformers/tokenization_utils_base.py", >line 1647, in from_pretrained
>    fast_tokenizer_file = get_fast_tokenizer_file(
>  File "/usr/local/lib/python3.9/site-packages/transformers/tokenization_utils_base.py", >line 3408, in get_fast_tokenizer_file
>    all_files = get_list_of_files(path_or_repo, revision=revision, >use_auth_token=use_auth_token)
>  File "/usr/local/lib/python3.9/site-packages/transformers/file_utils.py", line 1686, >in get_list_of_files
>    model_info = HfApi(endpoint=HUGGINGFACE_CO_RESOLVE_ENDPOINT).model_info(
>  File "/usr/local/lib/python3.9/site-packages/huggingface_hub/hf_api.py", line 247, in >model_info
>    r = requests.get(path, headers=headers)
>  File "/usr/local/lib/python3.9/site-packages/requests/api.py", line 76, in get
>    return request('get', url, params=params, **kwargs)
>  File "/usr/local/lib/python3.9/site-packages/requests/api.py", line 61, in request
>    return session.request(method=method, url=url, **kwargs)
>  File "/usr/local/lib/python3.9/site-packages/requests/sessions.py", line 542, in >request
>    resp = self.send(prep, **send_kwargs)
>  File "/usr/local/lib/python3.9/site-packages/requests/sessions.py", line 655, in send
>    r = adapter.send(request, **kwargs)
>  File "/usr/local/lib/python3.9/site-packages/requests/adapters.py", line 516, in send
>    raise ConnectionError(e, request=request)
>requests.exceptions.ConnectionError: HTTPSConnectionPool(host='huggingface.co', >port=443): Max retries exceeded with url: /api/models/bert-base-uncased (Caused by >NewConnectionError('<urllib3.connection.HTTPSConnection object at 0x1665cc7c0>: Failed >to establish a new connection: [Errno 8] nodename nor servname provided, or not known'))
>electron@diynoMacBook-Pro catr % 
>```