WATSONの画像認識の結果を1画像ずつのjsonファイルに出力するscript

gistfile1.py

# -*- 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()