lilacs2039 · December 10, 2017 02:41
diff --git a/taggingImages_withProbability.py b/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()
	#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()