Add labels property to Models

Author: ludwigschubert

lucid/modelzoo/util.py

@@ -20,6 +20,8 @@
 import tensorflow as tf
 from lucid.misc.io.reading import read
 
+def load_text_labels(labels_path):
+  return read(labels_path, encoding='utf-8').splitlines()
 
 def load_graphdef(model_url, reset_device=True):
   """Load GraphDef from a binary proto file."""

lucid/modelzoo/vision_base.py

@@ -16,17 +16,21 @@
 from __future__ import absolute_import, division, print_function
 
 import tensorflow as tf
-from lucid.modelzoo.util import load_graphdef, forget_xy
+from lucid.modelzoo.util import load_text_labels, load_graphdef, forget_xy
 
 class Model(object):
   """Base pretrained model importer."""
 
   model_path = None
+  labels_path = None
+  labels = None
   image_value_range = (-1, 1)
   image_shape = [None, None, 3]
 
   def __init__(self):
     self.graph_def = None
+    if self.labels_path is not None:
+      self.labels = load_text_labels(self.labels_path)
 
   def load_graphdef(self):
     self.graph_def = load_graphdef(self.model_path)

lucid/modelzoo/vision_models.py

@@ -36,6 +36,7 @@ def populate_inception_bottlenecks(scope):
 
 class InceptionV1(Model):
   model_path = 'gs://modelzoo/InceptionV1.pb'
+  labels_path = 'gs://modelzoo/InceptionV1-labels.txt'
   image_shape = [224, 224, 3]
   image_value_range = (-117, 255-117)
   input_name = 'input:0'