illustration2vecを使用して、指定フォルダ内の全画像をタグ付け

taggingImages.py

"""
指定ディレクトリのすべての画像ファイルに対してタグ付けして結果をファイルに保存

Usage:
python taggingImages.py --image_dir path/to/images

"""
import argparse
import glob

import os
import pyparsing as pp
from tqdm import tqdm

import i2v
from PIL import Image


argParser = argparse.ArgumentParser()
argParser.add_argument("--image_dir", help="path to folder containing images")
argParser.add_argument("--probability_th", default=0.3, help="probability threashold to record tag name.")
argParser.add_argument("--resultFile", default='result_tag.csv', help="result file name.")
a = argParser.parse_args()

illust2vec = i2v.make_i2v_with_chainer(
    "illust2vec_tag_ver200.caffemodel", "tag_list.json")

# In the case of caffe, please use i2v.make_i2v_with_caffe instead:
# illust2vec = i2v.make_i2v_with_caffe(
#     "illust2vec_tag.prototxt", "illust2vec_tag_ver200.caffemodel",
#     "tag_list.json")


# general=pp.Literal("'general':").suppress()
# tagParser = nemo + c + numMsg + c + msgNum + c + numSV + blocks + cs
# ret = tagParser.parseString(buf)



def main():
    with open(a.resultFile, 'w') as f:
        for filepath in tqdm(glob.glob(a.image_dir+"/*")):
            img = Image.open(filepath)
            tagsForImages = illust2vec.estimate_plausible_tags([img], threshold=a.probability_th)
            #tagsForImages[0]['general'][0][0]
            #Out[21]: 'thighhighs'
            for tagForImage in tagsForImages:
                tagStr=""
                for tag in tagForImage['general']:  #tag = ('tagname',probability)
                    tagname, _=tag
                    tagStr += tagname+","
            line = os.path.basename(filepath) + ","+tagStr+"\r\n"
            # print(line)
            f.write(line)

if __name__ == main():
    main()