Windows上のPythonでCOCO APIを使ってみる

Microsoft COCO

はじめに

pycocotoolsの導入はこちらを参照。

データのダウンロード

今回は試しに「captions_train-val2014.zip」をダウンロードする。
http://msvocds.blob.core.windows.net/annotations-1-0-3/captions_train-val2014.zip
解凍すると「annotations」フォルダのなかに「captions_train2014.json」ファイル、「captions_val2014.json」ファイルがある。

jsonファイルの読み込み

from pycocotools.coco import COCO
coco = COCO('captions_train2014.json')

読み込んだファイルのサイズを見てみる

>>> len(coco.anns)
414113
>>> len(coco.imgs)
82783

82783の画像に対して説明文が414113個。
だいたい1画像に5個の説明文ということがわかる。

一つの画像を見てみる

「57870」は適当な数字。

coco.imgs[57870]['flickr_url']

http://farm4.staticflickr.com/3153/2970773875_164f0c0b83_z.jpg

http://farm4.staticflickr.com/3153/2970773875_164f0c0b83_z.jpg

説明文を見てみる

anns_keys = coco.anns.keys()

for key in anns_keys:
    if coco.anns[key]['image_id']==57870:
        print(coco.anns[key]['caption'])

このように出力される。

A restaurant has modern wooden tables and chairs.
A long restaurant table with rattan rounded back chairs.
a long table with a plant on top of it surrounded with wooden chairs
A long table with a flower arrangement in the middle for meetings
A table is adorned with wooden chairs with blue accents.

大文字から始まったりそうでなかったり。
ピリオドがあったりなかったり。

画像のファイル名を知る

画像ファイルをダウンロードしたと仮定して対応するファイル名を知るには1行でOK。

coco.imgs[57870]['file_name']

このように出力される。

'COCO_train2014_000000057870.jpg'

最も長い説明文

>>> coco.anns[776153]['caption']
'A black and white shot shows a  bus interior with windows, hand rail,  and a woman, her
eyes closed, as she holds a bag on her lap, an empty seat, and a second, male passenger,
holding a cell and reaching under his jacket, while holding a suitcase.  '

http://farm4.staticflickr.com/3465/3203708166_dc3078e2ff_z.jpg

最も短い説明文

>>> coco.anns[27468]['caption']
'Person snowboarding down a hill.   '

http://farm8.staticflickr.com/7068/6932931159_14236c1dab_z.jpg

長い説明文を除外する

import nltk

original_token = []

for key in anns_keys:
    caption = coco.anns[key]['caption']
    tokens = nltk.tokenize.word_tokenize(caption.lower())
    if tokens[-1]=='.':
        tokens = tokens[:-1]
    if len(tokens)<=13:
        img_id = coco.anns[key]['image_id']
        a = [coco.loadImgs(img_id)[0]['file_name'], tokens]
        original_token.append(a)

文の最後のピリオドを除いて13単語以下の文章のみ抽出した。
開始文字と終了文字を足してさらにパディングで15単語に統一することを想定している。

>>> len(original_token)
375551

ほとんどの文章が残っているので、もともと説明文は一部を除いて短いことがわかる。

最後に

画像ファイル名とID化した説明文書をリストにして保存。
ついでに「id_to_word」、「word_to_id」も保存。

from pycocotools.coco import COCO
import nltk
import pickle

coco = COCO('annotations/captions_train2014.json')

anns_keys = coco.anns.keys()

original_token = []

for key in anns_keys:
    caption = coco.anns[key]['caption']
    tokens = nltk.tokenize.word_tokenize(caption.lower())
    if tokens[-1]=='.':
        tokens = tokens[:-1]
    if len(tokens)<=13:
        img_id = coco.anns[key]['image_id']
        a = [coco.loadImgs(img_id)[0]['file_name'], tokens]
        original_token.append(a)

from collections import Counter
freq = Counter()

for i in range(len(original_token)):
    freq.update(set(original_token[i][1]))

common = freq.most_common()
vocab = sorted([t for t, c in common if c>=3])
vocab.append('<pad>')
vocab.append('<start>')
vocab.append('<end>')
vocab.append('<unk>')

id_to_word = {i+1:t for i, t in enumerate(vocab)}
word_to_id = {t:i+1 for i, t in enumerate(vocab)}

original_token_id = []

for i in range(len(original_token)):
    sent = ['<start>'] + original_token[i][1] + ['<end>']
    sent_id = [word_to_id[t] if (t in word_to_id) else word_to_id['<unk>'] for t in sent]
    if(len(sent_id)<15):
        sent_id = sent_id + [word_to_id['<pad>']] * (15-len(sent_id))
    original_token_id.append([original_token[i][0], sent_id])

with open('id_to_word.pkl', 'wb') as f:
    pickle.dump(id_to_word, f)

with open('word_to_id.pkl', 'wb') as f:
    pickle.dump(word_to_id, f)

with open('filename_token.pkl', 'wb') as f:
    pickle.dump(original_token_id, f)

>>> len(word_to_id)
10593
>>> len(filename_token)
375551
>>> len(list(set([c[0] for c in filename_token])))
82775

10593単語
375551文章
82775画像

#ファイル名の拡張子をjpgからpngに変更する
import os

png_token = []

for i in range(len(filename_token)):
    name, ext = os.path.splitext(filename_token[i][0])
    new_filename = name + '.png'
    png_token.append([new_filename, filename_token[i][1]])

#ファイル名をID化する
temp_list = list(set([c[0] for c in png_token]))
id_to_image = {i+1:t for i, t in enumerate(temp_list)}
image_to_id = {t:i+1 for i, t in enumerate(temp_list)}

imageid_token = [[image_to_id[c[0]], c[1]] for c in png_token]

#pickleで保存する
with open('id_to_image.pkl', 'wb') as f:
    pickle.dump(id_to_image, f)

with open('image_to_id.pkl', 'wb') as f:
    pickle.dump(image_to_id, f)

with open('imageid_token.pkl', 'wb') as f:
    pickle.dump(imageid_token, f)

#確認
random_num = 300000
id_to_image[imageid_token[random_num][0]]
' '.join([id_to_word[c] for c in png_token[random_num][1]])

これでImage Captioningのデータ準備は終了。