i2vを使用して、指定ディレクトリのすべての画像ファイルに対して画像種類のタグを付けて結果をファイルに保存。タグの確からしさ(0~1)も保存

taggingImages_withProbability.py

#coding:utf-8
"""
i2vを使用して、指定ディレクトリのすべての画像ファイルに対して画像種類のタグを付けて結果をファイルに保存
タグの確からしさ(0~1)も保存

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

"""
import argparse
import csv
import decimal
import glob
import json

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("--tag_list_file", default='tag_list.json', help="tag list file 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.tag_list_file, 'r') as f:
    #     # content = f.read()
    #     labels = json.load(f)

    #header
    # labelStr = ""
    # for label in labels:
    #     labelStr += label+","
    # f.write("name,"+labels+"\r\n")

    header = set()
    content=[]
    #content
    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'
        record = {}
        record['name']=os.path.basename(filepath)
        for tagForImage in tagsForImages:
            tagStr=""
            for tag in tagForImage['general']:  #tag = ('tagname',probability)
                tagname, prob =tag
                header.add(tagname)
                record[tagname]=round(prob,2)
                # tagStr += tagname+","
            content += [record,]
        # print(str(record))


        # line = os.path.basename(filepath) + ","+tagStr+"\r\n"
        # # print(line)
        # f.write(line)

    with open(a.resultFile, 'w') as f:
        headerList = ['name']
        headerList.extend(list(header))
        writer = csv.DictWriter(f, headerList)
        writer.writeheader()
        for row in content:
            writer.writerow(row)

if __name__ == main():
    main()