class DeepImageFeaturizer(Transformer, HasInputCol, HasOutputCol):
"""
Applies the model specified by its popular name, with its prediction layer(s) chopped off,
to the image column in DataFrame. The output is a MLlib Vector so that DeepImageFeaturizer
can be used in a MLlib Pipeline.
The input image column should be 3-channel SpImage.
"""
modelName = Param(
Params._dummy(),
"modelName",
"A deep learning model name",
typeConverter=SparkDLTypeConverters.supportedNameConverter(
SUPPORTED_MODELS))
@keyword_only
def __init__(self, inputCol=None, outputCol=None, modelName=None):
"""
__init__(self, inputCol=None, outputCol=None, modelName=None)
"""
super(DeepImageFeaturizer, self).__init__()
kwargs = self._input_kwargs
self.setParams(**kwargs)
@keyword_only
def setParams(self, inputCol=None, outputCol=None, modelName=None):
"""
setParams(self, inputCol=None, outputCol=None, modelName=None)
"""
kwargs = self._input_kwargs
self._set(**kwargs)
return self
def setModelName(self, value):
return self._set(modelName=value)
def getModelName(self):
return self.getOrDefault(self.modelName)
def _transform(self, dataset):
transformer = _NamedImageTransformer(inputCol=self.getInputCol(),
outputCol=self.getOutputCol(),
modelName=self.getModelName(),
featurize=True)
return transformer.transform(dataset)
class DeepImagePredictor(Transformer, HasInputCol, HasOutputCol):
"""
Applies the model specified by its popular name to the image column in DataFrame.
The output is a MLlib Vector.
"""
modelName = Param(
Params._dummy(),
"modelName",
"A deep learning model name",
typeConverter=SparkDLTypeConverters.supportedNameConverter(
SUPPORTED_MODELS))
decodePredictions = Param(
Params._dummy(),
"decodePredictions",
"If true, output predictions in the (class, description, probability) format",
typeConverter=TypeConverters.toBoolean)
topK = Param(Params._dummy(),
"topK",
"How many classes to return if decodePredictions is True",
typeConverter=TypeConverters.toInt)
@keyword_only
def __init__(self,
inputCol=None,
outputCol=None,
modelName=None,
decodePredictions=False,
topK=5):
"""
__init__(self, inputCol=None, outputCol=None, modelName=None, decodePredictions=False,
topK=5)
"""
super(DeepImagePredictor, self).__init__()
self._setDefault(decodePredictions=False)
self._setDefault(topK=5)
kwargs = self._input_kwargs
self.setParams(**kwargs)
@keyword_only
def setParams(self,
inputCol=None,
outputCol=None,
modelName=None,
decodePredictions=False,
topK=5):
"""
setParams(self, inputCol=None, outputCol=None, modelName=None, decodePredictions=False,
topK=5)
"""
kwargs = self._input_kwargs
self._set(**kwargs)
return self
def setModelName(self, value):
return self._set(modelName=value)
def getModelName(self):
return self.getOrDefault(self.modelName)
def _transform(self, dataset):
transformer = _NamedImageTransformer(
inputCol=self.getInputCol(),
outputCol=self._getIntermediateOutputCol(),
modelName=self.getModelName(),
featurize=False)
transformed = transformer.transform(dataset)
if self.getOrDefault(self.decodePredictions):
return self._decodeOutputAsPredictions(transformed)
else:
return transformed.withColumnRenamed(
self._getIntermediateOutputCol(), self.getOutputCol())
def _decodeOutputAsPredictions(self, df):
# If we start having different weights than imagenet, we'll need to
# move this logic to individual model building in NamedImageTransformer.
# Also, we could put the computation directly in the main computation
# graph or use a scala UDF for potentially better performance.
topK = self.getOrDefault(self.topK)
def decode(predictions):
pred_arr = np.expand_dims(np.array(predictions), axis=0)
decoded = decode_predictions(pred_arr, top=topK)[0]
# convert numpy dtypes to python native types
return [(t[0], t[1], t[2].item()) for t in decoded]
decodedSchema = ArrayType(
StructType([
StructField("class", StringType(), False),
StructField("description", StringType(), False),
StructField("probability", FloatType(), False)
]))
decodeUDF = udf(decode, decodedSchema)
interim_output = self._getIntermediateOutputCol()
return (df.withColumn(self.getOutputCol(), decodeUDF(
df[interim_output])).drop(interim_output))
def _getIntermediateOutputCol(self):
return "__tmp_" + self.getOutputCol()
class _NamedImageTransformer(Transformer, HasInputCol, HasOutputCol):
"""
For internal use only. NamedImagePredictor and NamedImageFeaturizer are the recommended classes
to use.
Applies the model specified by its popular name to the image column in DataFrame. There are
two output modes: predictions or the featurization from the model. In either case the output
is a MLlib Vector.
"""
modelName = Param(
Params._dummy(),
"modelName",
"A deep learning model name",
typeConverter=SparkDLTypeConverters.supportedNameConverter(
SUPPORTED_MODELS))
featurize = Param(
Params._dummy(),
"featurize",
"If true, output features. If false, output predictions. Either way the output is a vector.",
typeConverter=TypeConverters.toBoolean)
@keyword_only
def __init__(self,
inputCol=None,
outputCol=None,
modelName=None,
featurize=False):
"""
__init__(self, inputCol=None, outputCol=None, modelName=None, featurize=False)
"""
super(_NamedImageTransformer, self).__init__()
kwargs = self._input_kwargs
self.setParams(**kwargs)
self._inputTensorName = None
self._outputTensorName = None
self._outputMode = None
@keyword_only
def setParams(self,
inputCol=None,
outputCol=None,
modelName=None,
featurize=False):
"""
setParams(self, inputCol=None, outputCol=None, modelName=None, featurize=False)
"""
kwargs = self._input_kwargs
self._set(**kwargs)
return self
def setModelName(self, value):
return self._set(modelName=value)
def getModelName(self):
return self.getOrDefault(self.modelName)
def setFeaturize(self, value):
return self._set(featurize=value)
def getFeaturize(self):
return self.getOrDefault(self.featurize)
def _transform(self, dataset):
modelGraphSpec = _buildTFGraphForName(self.getModelName(),
self.getFeaturize())
inputCol = self.getInputCol()
resizedCol = "__sdl_imagesResized"
tfTransformer = TFImageTransformer(
inputCol=resizedCol,
outputCol=self.getOutputCol(),
graph=modelGraphSpec["graph"],
inputTensor=modelGraphSpec["inputTensorName"],
outputTensor=modelGraphSpec["outputTensorName"],
outputMode=modelGraphSpec["outputMode"])
resizeUdf = resizeImage(modelGraphSpec["inputTensorSize"])
result = tfTransformer.transform(
dataset.withColumn(resizedCol, resizeUdf(inputCol)))
return result.drop(resizedCol)
class TFImageTransformer(Transformer, HasInputCol, HasOutputCol):
"""
Applies the Tensorflow graph to the image column in DataFrame.
Restrictions of the current API:
* Does not use minibatches, which is a major low-hanging fruit for performance.
* Only one output node can be specified.
* The output is expected to be an image or a 1-d vector.
* All images in the dataframe are expected be of the same numerical data type
(i.e. the dtype of the values in the numpy array representation is the same).
We assume all graphs have a "minibatch" dimension (i.e. an unknown leading
dimension) in the tensor shapes.
Note: The input tensorflow graph should have appropriate weights constantified,
since a new session is created inside this transformer.
"""
USER_GRAPH_NAMESPACE = 'given'
NEW_OUTPUT_PREFIX = 'sdl_flattened'
graph = Param(Params._dummy(), "graph", "A TensorFlow computation graph",
typeConverter=SparkDLTypeConverters.toTFGraph)
inputTensor = Param(Params._dummy(), "inputTensor",
"A TensorFlow tensor object or name representing the input image",
typeConverter=SparkDLTypeConverters.toStringOrTFTensor)
outputTensor = Param(Params._dummy(), "outputTensor",
"A TensorFlow tensor object or name representing the output",
typeConverter=SparkDLTypeConverters.toStringOrTFTensor)
outputMode = Param(Params._dummy(), "outputMode",
"How the output column should be formatted. 'vector' for a 1-d MLlib " +
"Vector of floats. 'image' to format the output to work with the image " +
" tools in this package.",
typeConverter=SparkDLTypeConverters.supportedNameConverter(OUTPUT_MODES))
@keyword_only
def __init__(self, inputCol=None, outputCol=None, graph=None,
inputTensor=utils.IMAGE_INPUT_PLACEHOLDER_NAME, outputTensor=None,
outputMode="vector"):
"""
__init__(self, inputCol=None, outputCol=None, graph=None,
inputTensor=utils.IMAGE_INPUT_PLACEHOLDER_NAME, outputTensor=None,
outputMode="vector")
"""
super(TFImageTransformer, self).__init__()
self._setDefault(inputTensor=utils.IMAGE_INPUT_PLACEHOLDER_NAME)
self._setDefault(outputMode="vector")
kwargs = self._input_kwargs
self.setParams(**kwargs)
@keyword_only
def setParams(self, inputCol=None, outputCol=None, graph=None,
inputTensor=utils.IMAGE_INPUT_PLACEHOLDER_NAME, outputTensor=None,
outputMode="vector"):
"""
setParams(self, inputCol=None, outputCol=None, graph=None,
inputTensor=utils.IMAGE_INPUT_PLACEHOLDER_NAME, outputTensor=None,
outputMode="vector")
"""
kwargs = self._input_kwargs
return self._set(**kwargs)
def setGraph(self, value):
return self._set(graph=value)
def setInputTensor(self, value):
return self._set(inputTensor=value)
def setOutputTensor(self, value):
return self._set(outputTensor=value)
def getGraph(self):
return self.getOrDefault(self.graph)
def getInputTensor(self):
tensor_or_name = self.getOrDefault(self.inputTensor)
if isinstance(tensor_or_name, tf.Tensor):
return tensor_or_name
else:
return self.getGraph().get_tensor_by_name(tensor_or_name)
def getOutputTensor(self):
tensor_or_name = self.getOrDefault(self.outputTensor)
if isinstance(tensor_or_name, tf.Tensor):
return tensor_or_name
else:
return self.getGraph().get_tensor_by_name(tensor_or_name)
def _transform(self, dataset):
graph = self.getGraph()
composed_graph = self._addReshapeLayers(graph, self._getImageDtype(dataset))
final_graph = self._stripGraph(composed_graph)
with final_graph.as_default():
image = dataset[self.getInputCol()]
image_df_exploded = (dataset
.withColumn("__sdl_image_height", image.height)
.withColumn("__sdl_image_width", image.width)
.withColumn("__sdl_image_nchannels", image.nChannels)
.withColumn("__sdl_image_data", image.data)
)
final_output_name = self._getFinalOutputTensorName()
output_tensor = final_graph.get_tensor_by_name(final_output_name)
final_df = (
tfs.map_rows([output_tensor], image_df_exploded,
feed_dict={
"height": "__sdl_image_height",
"width": "__sdl_image_width",
"num_channels": "__sdl_image_nchannels",
"image_buffer": "__sdl_image_data"})
.drop("__sdl_image_height", "__sdl_image_width", "__sdl_image_nchannels",
"__sdl_image_data")
)
tfs_output_name = tfx.op_name(final_graph, output_tensor)
original_output_name = self._getOriginalOutputTensorName()
output_shape = final_graph.get_tensor_by_name(original_output_name).shape
output_mode = self.getOrDefault(self.outputMode)
# TODO: support non-1d tensors (return np.array).
if output_mode == "image":
return self._convertOutputToImage(final_df, tfs_output_name, output_shape)
else:
assert output_mode == "vector", "Unknown output mode: %s" % output_mode
return self._convertOutputToVector(final_df, tfs_output_name)
def _getImageDtype(self, dataset):
# This may not be the best way to get the type of image, but it is one way.
# Assumes that the dtype for all images is the same in the given dataframe.
pdf = dataset.select(self.getInputCol()).take(1)
img = pdf[0][self.getInputCol()]
img_type = sparkModeLookup[img.mode]
return img_type.dtype
def _addReshapeLayers(self, tf_graph, dtype="uint8"):
input_tensor_name = self.getInputTensor().name
gdef = tf_graph.as_graph_def(add_shapes=True)
g = tf.Graph()
with g.as_default():
# Flat image data -> image dimensions
height = tf.placeholder(tf.int32, [], name="height")
width = tf.placeholder(tf.int32, [], name="width")
num_channels = tf.placeholder(tf.int32, [], name="num_channels")
image_buffer = tf.placeholder(tf.string, [], name="image_buffer")
# Note: the shape argument is required for tensorframes as it uses a
# slightly older version of tensorflow.
shape = tf.reshape(tf.stack([height, width, num_channels], axis=0), shape=(3,),
name='shape')
if dtype == "uint8":
image_uint8 = tf.decode_raw(image_buffer, tf.uint8, name="decode_raw")
image_float = tf.to_float(image_uint8)
else:
assert dtype == SparkMode.FLOAT32, "Unsupported dtype for image: %s" % dtype
image_float = tf.decode_raw(image_buffer, tf.float32, name="decode_raw")
image_reshaped = tf.reshape(image_float, shape, name="reshaped")
image_reshaped_expanded = tf.expand_dims(image_reshaped, 0, name="expanded")
# Add on the original graph
tf.import_graph_def(gdef, input_map={input_tensor_name: image_reshaped_expanded},
return_elements=[self.getOutputTensor().name],
name=self.USER_GRAPH_NAMESPACE)
# Flatten the output for tensorframes
output_node = g.get_tensor_by_name(self._getOriginalOutputTensorName())
_ = tf.reshape(output_node[0], # batch-size = 1,
shape=[-1], name=self._getFinalOutputOpName())
return g
# Sometimes the tf graph contains a bunch of stuff that doesn't lead to the
# output. TensorFrames does not like that, so we strip out the parts that
# are not necessary for the computation at hand.
def _stripGraph(self, tf_graph):
gdef = tfx.strip_and_freeze_until([self._getFinalOutputOpName()], tf_graph)
g = tf.Graph()
with g.as_default():
tf.import_graph_def(gdef, name='')
return g
def _getOriginalOutputTensorName(self):
return self.USER_GRAPH_NAMESPACE + '/' + self.getOutputTensor().name
def _getFinalOutputTensorName(self):
return self.NEW_OUTPUT_PREFIX + '_' + self.getOutputTensor().name
def _getFinalOutputOpName(self):
return tfx.as_op_name(self._getFinalOutputTensorName())
def _convertOutputToImage(self, df, tfs_output_col, output_shape):
assert len(output_shape) == 4, str(output_shape) + " does not have 4 dimensions"
height = int(output_shape[1])
width = int(output_shape[2])
def to_image(orig_image, numeric_data):
# Assume the returned image has float pixels but same #channels as input
mode = orig_image.mode if orig_image.mode == "float32" else "RGB-float32"
return [mode, height, width, orig_image.nChannels,
bytearray(np.array(numeric_data).astype(np.float32).tobytes())]
to_image_udf = udf(to_image, imageSchema)
return (
df.withColumn(self.getOutputCol(),
to_image_udf(df[self.getInputCol()], df[tfs_output_col]))
.drop(tfs_output_col)
)
def _convertOutputToVector(self, df, tfs_output_col):
"""
Converts the output python list to MLlib Vector.
"""
return (
df.withColumn(self.getOutputCol(), JVMAPI.list_to_vector_udf(df[tfs_output_col]))
.drop(tfs_output_col)
)
class KerasImageFileTransformer(Transformer, HasInputCol, HasOutputCol):
"""
Applies the Tensorflow-backed Keras model (specified by a file name) to
images (specified by the URI in the inputCol column) in the DataFrame.
Restrictions of the current API:
* see TFImageTransformer.
* Only supports Tensorflow-backed Keras models (no Theano).
"""
modelFile = Param(
Params._dummy(),
"modelFile",
"h5py file containing the Keras model (architecture and weights)",
typeConverter=TypeConverters.toString)
# TODO :add a lambda type converter e.g callable(mylambda)
imageLoader = Param(
Params._dummy(), "imageLoader",
"Function containing the logic for loading and pre-processing images. "
+
"The function should take in a URI string and return a 4-d numpy.array "
+ "with shape (batch_size (1), height, width, num_channels).")
outputMode = Param(
Params._dummy(),
"outputMode",
"How the output column should be formatted. 'vector' for a 1-d MLlib "
+
"Vector of floats. 'image' to format the output to work with the image "
+ "tools in this package.",
typeConverter=SparkDLTypeConverters.supportedNameConverter(
OUTPUT_MODES))
@keyword_only
def __init__(self,
inputCol=None,
outputCol=None,
modelFile=None,
imageLoader=None,
outputMode="vector"):
"""
__init__(self, inputCol=None, outputCol=None, modelFile=None, imageLoader=None,
outputMode="vector")
"""
super(KerasImageFileTransformer, self).__init__()
kwargs = self._input_kwargs
self.setParams(**kwargs)
self._inputTensor = None
self._outputTensor = None
@keyword_only
def setParams(self,
inputCol=None,
outputCol=None,
modelFile=None,
imageLoader=None,
outputMode="vector"):
"""
setParams(self, inputCol=None, outputCol=None, modelFile=None, imageLoader=None,
outputMode="vector")
"""
kwargs = self._input_kwargs
self._set(**kwargs)
return self
def setModelFile(self, value):
return self._set(modelFile=value)
def getModelFile(self):
return self.getOrDefault(self.modelFile)
def _transform(self, dataset):
graph = self._loadTFGraph()
image_df = self._loadImages(dataset)
assert self._inputTensor is not None, "self._inputTensor must be set"
assert self._outputTensor is not None, "self._outputTensor must be set"
transformer = TFImageTransformer(inputCol=self._loadedImageCol(),
outputCol=self.getOutputCol(),
graph=graph,
inputTensor=self._inputTensor,
outputTensor=self._outputTensor,
outputMode=self.getOrDefault(
self.outputMode))
return transformer.transform(image_df).drop(self._loadedImageCol())
def _loadTFGraph(self):
with KSessionWrap() as (sess, g):
assert K.backend() == "tensorflow", \
"Keras backend is not tensorflow but KerasImageTransformer only supports " + \
"tensorflow-backed Keras models."
with g.as_default():
K.set_learning_phase(0) # Testing phase
model = load_model(self.getModelFile())
out_op_name = utils.op_name(model.output)
self._inputTensor = model.input.name
self._outputTensor = model.output.name
return utils.stripAndFreezeGraph(
g.as_graph_def(add_shapes=True), sess, [out_op_name])
def _loadedImageCol(self):
return "__sdl_img"
def _loadImages(self, dataset):
"""
Load image files specified in dataset as image format specified in sparkdl.image.imageIO.
"""
# plan 1: udf(loader() + convert from np.array to imageSchema) -> call TFImageTransformer
# plan 2: udf(loader()) ... we don't support np.array as a dataframe column type...
loader = self.getOrDefault(self.imageLoader)
def load(uri):
img = loader(uri)
return imageIO.imageArrayToStruct(img)
load_udf = udf(load, imageIO.imageSchema)
return dataset.withColumn(self._loadedImageCol(),
load_udf(dataset[self.getInputCol()]))