named_image.py 文件源码

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)
        )