Skip to content

Instantly share code, notes, and snippets.

@TakesxiSximada
Last active August 29, 2015 14:18
Show Gist options
  • Save TakesxiSximada/ca59974659dbc5c284e2 to your computer and use it in GitHub Desktop.
Save TakesxiSximada/ca59974659dbc5c284e2 to your computer and use it in GitHub Desktop.
WATSONの画像認識の結果を1画像ずつのjsonファイルに出力するscript
# -*- coding: utf-8 -*-
import os
import sys
import json
import argparse
import itertools
import requests
from pit import Pit
base_url = 'https://gateway.watsonplatform.net/visual-recognition-beta/api'
recog_url = base_url + '/v1/tag/recognize'
def main(argv=sys.argv[1:]):
parser = argparse.ArgumentParser()
parser.add_argument('target')
parser.add_argument('-c', '--count', default=5, type=int)
parser.add_argument('-o', '--output', default='output')
args = parser.parse_args(argv)
os.makedirs(args.output, exist_ok=True)
output = args.output
setting = Pit.get('iwdcat',
{'require': {'username': '',
'password': '',
}})
username = setting['username']
password = setting['password']
auth_token = (username, password)
for root, dirs, files in os.walk(args.target):
for filenames in itertools.zip_longest(*[iter(files)]*args.count):
filenames = filter(lambda n: n, filenames) # filter None
filename_filepath = dict(
(filename, os.path.join(root, filename))
for filename in filenames
)
files_params = dict(
('imgFile{}'.format(ii), ('{}'.format(filename), open(filepath, 'rb')))
for ii, (filename, filepath) in enumerate(filename_filepath.items())
)
res = requests.post(
recog_url, auth=auth_token, files=files_params)
for name, fp in files_params.values():
fp.close()
data = json.loads(res.text)
for image in data['images']:
output_filepath = os.path.join(output, image['image_name'] + '.json')
with open(output_filepath, 'w+t') as fp:
json.dump(image['labels'], fp)
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment