Load index->mid and mid->display name maps. Args: labelmap_path: path to the file with the list of mids, describing predictions. dict_path: path to the dict.csv that translates from mids to display names. Returns: labelmap: an index to mid list label_dict: mid to display name di
(num_classes, labelmap_path, dict_path)
| 81 | |
| 82 | |
| 83 | def LoadLabelMaps(num_classes, labelmap_path, dict_path): |
| 84 | """Load index->mid and mid->display name maps. |
| 85 | |
| 86 | Args: |
| 87 | labelmap_path: path to the file with the list of mids, describing predictions. |
| 88 | dict_path: path to the dict.csv that translates from mids to display names. |
| 89 | Returns: |
| 90 | labelmap: an index to mid list |
| 91 | label_dict: mid to display name dictionary |
| 92 | """ |
| 93 | labelmap = [line.rstrip() for line in tf.gfile.GFile(labelmap_path).readlines()] |
| 94 | if len(labelmap) != num_classes: |
| 95 | tf.logging.fatal( |
| 96 | "Label map loaded from {} contains {} lines while the number of classes is {}".format( |
| 97 | labelmap_path, len(labelmap), num_classes)) |
| 98 | sys.exit(1) |
| 99 | |
| 100 | label_dict = {} |
| 101 | for line in tf.gfile.GFile(dict_path).readlines(): |
| 102 | words = [word.strip(' "\n') for word in line.split(',', 1)] |
| 103 | label_dict[words[0]] = words[1] |
| 104 | |
| 105 | return labelmap, label_dict |
| 106 | |
| 107 | |
| 108 | def main(args): |