label_to_c_strings.py

#! /usr/bin/env python3

import argparse
import os


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-l", "--label_file", type=str, default='labels.txt',
                        help="input image size")
    parser.add_argument("-o", "--header_file", type=str, default='image_classes.h',
                        help="output c strings file")
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()

    header_filename = args.header_file
    labels_filename = args.label_file

    label_strings = 'char * imagenet_classes[] ={'
    lines = ['']
    for line in open(labels_filename, 'r'):
        label_strings += '"%s",\n' % lines[-1] if lines[-1] else ''
        label = line.strip().split(',')[0]
        print(label)
        lines.append(label)
    label_strings += '"%s"' % lines[-1]
    label_strings += '};'
    with open(header_filename, 'w') as f:
        f.write(label_strings)
    print('done.')