オフライン環境＋CPU環境で動く画像キャプション生成モデルが生み出した説明文を元の画像に重ね書きしてファイル保存

###やったこと

__前回の記事__の続編です。

- __（前回の記事）__[オフライン環境 + CPU環境でも動く画像キャプション文生成モデルをMacbookで動かしてみた](https://qiita.com/electronics_diy721/items/2b7c2801b93b6611d5f2)

前回の記事では、指定した画像ファイルの内容を説明するキャプション文を自動作成させて、結果をTerimnalに表示しました。

今回はその続編として、Terminalにされたキャプション文をもとの画像に重ね書きして、ファイルとして書き出して保存する機能を追加しました。

同じ機能を持つ画像キャプション文生成スクリプトは、すでに別の記事で、今回用いたのとは異なるキャプション生成モデルを使って、作成済みです。

- [画像キャプション生成モデル 「CATR」の出力文を入力画像に埋込んでファイル出力するように書き換えてみた](https://qiita.com/electronics_diy721/items/03b8d601da4280c8368b)

その時使ったモデルは、CATRでした。
この画像キャプション生成器には、ある難点がありました。

それは、インターネットに接続中のPC端末でないと、画像のキャプション文生成ができない、という問題でした。

今後、飛行中のドローンの単眼カメラから受信したフレーム画像のキャプション文を生成して、フレーム画像に埋め込んだ画像をPCのウィンドウに表示させいたと考えています。

そのとき、画像を受信するPCは、ドローンとWifi接続している（Wifi親機：ドローン、Wifi子機：PC）ため、このPCはインターネットにつながっていません。

そこで、前回の記事では、オフライン環境でもキャプション生成を実行できるモデルを見つけて、動かしてみました。
今回は、そのモデルを使って、生成した画像キャプション文をもとの画像に埋め込む処理を作り込んでみます。

###作成したスクリプトの実行結果

```bash:Terminal
electron@diynoMacBook-Pro image_captioning % python3 caption_embedded_image_file_save.py --image='png/test001.png' 
 a man is riding a horse in a field . 
electron@diynoMacBook-Pro image_captioning % 
```

__ファイルが出力されたかを確認。__

- __解析結果2021-08-07_23:35:23.827168.jpg」__

```bash:Terminal
electron@diynoMacBook-Pro image_captioning % ls
README.md					download.sh					sample.py
__pycache__					model.py					train.py
build_vocab.py					models						解析結果.jpg
caption_embedded_image_file_save.py		png						解析結果2021-08-07_23:33:03.699471.jpg
data						requirements.txt				解析結果2021-08-07_23:35:23.827168.jpg
data_loader.py					resize.py
electron@diynoMacBook-Pro image_captioning % 
electron@diynoMacBook-Pro image_captioning % open 解析結果2021-08-07_23:35:23.827168.jpg
electron@diynoMacBook-Pro image_captioning % 
```

![スクリーンショット 2021-08-07 23.36.14.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/ef7d7789-14a1-c746-fdfa-80c8bffa5ca2.png)


#####他の画像で試してみる

```bash:Terminal
electron@diynoMacBook-Pro image_captioning % python3 caption_embedded_image_file_save.py --image='png/test003.jpg'
 a group of people standing around a table with a frisbee . 
electron@diynoMacBook-Pro image_captioning % 
```
__ファイルが出力されたかを確認。__

- 解析結果2021-08-07_23:36:34.111694.jpg

```bash:Terminal
electron@diynoMacBook-Pro image_captioning % ls
README.md					download.sh					sample.py
__pycache__					model.py					train.py
build_vocab.py					models						解析結果.jpg
caption_embedded_image_file_save.py		png						解析結果2021-08-07_23:33:03.699471.jpg
data						requirements.txt				解析結果2021-08-07_23:35:23.827168.jpg
data_loader.py					resize.py					解析結果2021-08-07_23:36:34.111694.jpg
electron@diynoMacBook-Pro image_captioning % 
electron@diynoMacBook-Pro image_captioning % open 解析結果2021-08-07_23:36:34.111694.jpg
electron@diynoMacBook-Pro image_captioning % 
```

![スクリーンショット 2021-08-07 23.36.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/69b428de-d048-ca31-9ee1-a25a5741a5ae.png)


#####他の画像でやってみる

```bash:Terminal
electron@diynoMacBook-Pro image_captioning % python3 caption_embedded_image_file_save.py --image='png/test007.jpeg'
 a man is holding a large umbrella in a room . 
 electron@diynoMacBook-Pro image_captioning %
````

__ファイルが出力されたかを確認。__

- 解析結果2021-08-07_23:37:14.609043.jpg
 
```bash:Terminal
electron@diynoMacBook-Pro image_captioning % ls
README.md					model.py					解析結果.jpg
__pycache__					models						解析結果2021-08-07_23:33:03.699471.jpg
build_vocab.py					png						解析結果2021-08-07_23:35:23.827168.jpg
caption_embedded_image_file_save.py		requirements.txt				解析結果2021-08-07_23:36:34.111694.jpg
data						resize.py					解析結果2021-08-07_23:37:14.609043.jpg
data_loader.py					sample.py
download.sh					train.py
electron@diynoMacBook-Pro image_captioning % 
electron@diynoMacBook-Pro image_captioning % open 解析結果2021-08-07_23:37:14.609043.jpg
electron@diynoMacBook-Pro image_captionin
```

####ファイルの大きさが小さいと、文字が溢れてしまう。

- ファイルの縦横サイズを取得して、埋め込む文字数と改行位置やサイズを工夫する必要がある。
 
![スクリーンショット 2021-08-07 23.37.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1754134/c1e9c86d-5a71-a230-ec33-b3bf644c1fe4.png)


###作成したスクリプト・ファイル

```Python3:caption_embedded_image_file_save.py
import torch
#import matplotlib.pyplot as plt
import numpy as np 
import argparse
import pickle 
import os
import cv2
import textwrap
import datetime
import locale
from torchvision import transforms 
from build_vocab import Vocabulary
from model import EncoderCNN, DecoderRNN
from PIL import Image


# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def load_image(image_path, transform=None):
    global img_path
    img_path = image_path
    image = Image.open(image_path).convert('RGB')
    image = image.resize([224, 224], Image.LANCZOS)
    
    if transform is not None:
        image = transform(image).unsqueeze(0)
    
    return image

def main(args):
    # Image preprocessing
    transform = transforms.Compose([
        transforms.ToTensor(), 
        transforms.Normalize((0.485, 0.456, 0.406), 
                             (0.229, 0.224, 0.225))])
    
    # Load vocabulary wrapper
    with open(args.vocab_path, 'rb') as f:
        vocab = pickle.load(f)

    # Build models
    encoder = EncoderCNN(args.embed_size).eval()  # eval mode (batchnorm uses moving mean/variance)
    decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers)
    encoder = encoder.to(device)
    decoder = decoder.to(device)

    # Load the trained model parameters
    encoder.load_state_dict(torch.load(args.encoder_path))
    decoder.load_state_dict(torch.load(args.decoder_path))

    # Prepare an image
    image = load_image(args.image, transform)
    image_tensor = image.to(device)
    
    # Generate an caption from the image
    feature = encoder(image_tensor)
    sampled_ids = decoder.sample(feature)
    sampled_ids = sampled_ids[0].cpu().numpy()          # (1, max_seq_length) -> (max_seq_length)
    
    # Convert word_ids to words
    sampled_caption = []
    for word_id in sampled_ids:
        word = vocab.idx2word[word_id]
        sampled_caption.append(word)
        if word == '<end>':
            break
    sentence = ' '.join(sampled_caption)
    sentence = sentence.replace("<start>", "")
    sentence = sentence.replace("<end>", "")
    
    # Print out the image and the generated caption
    print(sentence)
    caption_text_wrap_list = textwrap.wrap(str(sentence), 40)
    # 入力した画像ファイルに、生成したキャプション文の文字列を埋込む
    img = cv2.imread(img_path)
    
    if len(caption_text_wrap_list) == 1:
        cv2.putText(img, str(caption_text_wrap_list[0]), (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)
    
    datetime_string = str(datetime.datetime.now())
    datetime_string = datetime_string.replace(" ", "_").replace("　", "_")
    output_file_name = "解析結果" + datetime_string + ".jpg"
    cv2.imwrite(output_file_name,img)
    #   image = Image.open(args.image)
    #   plt.imshow(np.asarray(image))
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', type=str, required=True, help='input image for generating caption')
    parser.add_argument('--encoder_path', type=str, default='models/encoder-5-3000.pkl', help='path for trained encoder')
    parser.add_argument('--decoder_path', type=str, default='models/decoder-5-3000.pkl', help='path for trained decoder')
    parser.add_argument('--vocab_path', type=str, default='data/vocab.pkl', help='path for vocabulary wrapper')
    
    # Model parameters (should be same as paramters in train.py)
    parser.add_argument('--embed_size', type=int , default=256, help='dimension of word embedding vectors')
    parser.add_argument('--hidden_size', type=int , default=512, help='dimension of lstm hidden states')
    parser.add_argument('--num_layers', type=int , default=1, help='number of layers in lstm')
    args = parser.parse_args()
    main(args)

```