Merge pull request #4 from simonellistonball/master

Timothy Spann · web-flow · commit c95f7a6bc578 · 2017-10-02T17:25:05.000-04:00
Added more output to attributes
diff --git a/build.sh b/build.sh
@@ -1 +1,3 @@
-mvn install -DskipTests
+#!/bin/sh
+
+mvn clean package
diff --git a/nifi-tensorflow-nar/.gitignore b/nifi-tensorflow-nar/.gitignore
diff --git a/nifi-tensorflow-nar/pom.xml b/nifi-tensorflow-nar/pom.xml
@@ -19,7 +19,7 @@
     <parent>
         <groupId>com.dataflowdeveloper</groupId>
         <artifactId>tensorflow-processor</artifactId>
-        <version>2.0</version>
+        <version>2.1</version>
     </parent>
 
     <artifactId>nifi-tensorflow-nar</artifactId>
@@ -33,7 +33,7 @@
         <dependency>
             <groupId>com.dataflowdeveloper</groupId>
             <artifactId>nifi-tensorflow-processors</artifactId>
-            <version>2.0</version>
+            <version>2.1</version>
         </dependency>
     </dependencies>
 
diff --git a/nifi-tensorflow-processors/.gitignore b/nifi-tensorflow-processors/.gitignore
diff --git a/nifi-tensorflow-processors/pom.xml b/nifi-tensorflow-processors/pom.xml
@@ -19,7 +19,7 @@
     <parent>
         <groupId>com.dataflowdeveloper</groupId>
         <artifactId>tensorflow-processor</artifactId>
-        <version>2.0</version>
+        <version>2.1</version>
     </parent>
   
     <artifactId>nifi-tensorflow-processors</artifactId>
diff --git a/nifi-tensorflow-processors/src/main/java/com/dataflowdeveloper/processors/process/TensorFlowProcessor.java b/nifi-tensorflow-processors/src/main/java/com/dataflowdeveloper/processors/process/TensorFlowProcessor.java
@@ -23,9 +23,18 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Set;
+import java.util.SortedMap;
+import java.util.TreeMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
 
 import org.apache.commons.io.IOUtils;
 import org.apache.nifi.annotation.behavior.EventDriven;
@@ -64,7 +73,7 @@
  */
 public class TensorFlowProcessor extends AbstractProcessor {
 
-	public static final String ATTRIBUTE_OUTPUT_NAME = "probabilities";
+	public static final String ATTRIBUTE_OUTPUT_NAME = "tf.probabilities";
 	public static final String MODEL_DIR_NAME = "modeldir";
 	public static final String PROPERTY_NAME_EXTRA = "Extra Resources";
 
@@ -134,21 +143,28 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
 			service = new TensorFlowService();
 			// read all bytes of the flowfile (tensor requires whole image)
 			InputStream is = session.read(flowFile);
-			String value; 
+			List<Entry<Float, String>> results;
 			try {
 				byte[] byteArray = IOUtils.toByteArray(is);
-				value = service.getInception(byteArray, modelDir);
-
+				results = service.getInception(byteArray, modelDir).limit(10).collect(Collectors.toList());
 			} catch(Exception e) {
 				throw new ProcessException(e);
 			} finally {
 				is.close();
 			}
 			
-			if (value == null) {
+			if (results == null) {
 				session.transfer(flowFile, REL_UNMATCHED);
 			} else {
-				flowFile = session.putAttribute(flowFile, ATTRIBUTE_OUTPUT_NAME, value);
+				HashMap<String,String> attributes = new HashMap<String,String>(results.size() * 2);
+				for(int i = 0; i < results.size(); i++) {
+					Object[] key = new Object[] { ATTRIBUTE_OUTPUT_NAME, i };
+					
+					Entry<Float, String> entry = results.get(i);
+					attributes.put(String.format("%s.%d.label", key), entry.getValue());
+					attributes.put(String.format("%s.%d.probability", key), entry.getKey().toString());
+				}					
+				flowFile = session.putAllAttributes(flowFile, attributes);
 				session.transfer(flowFile, REL_SUCCESS);
 			}
 			session.commit();
diff --git a/nifi-tensorflow-processors/src/main/java/com/dataflowdeveloper/processors/process/TensorFlowService.java b/nifi-tensorflow-processors/src/main/java/com/dataflowdeveloper/processors/process/TensorFlowService.java
@@ -8,10 +8,15 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.tensorflow.DataType;
 import org.tensorflow.Graph;
 import org.tensorflow.Output;
@@ -26,31 +31,44 @@
  */
 public class TensorFlowService {
 
+	private final Logger logger = LoggerFactory.getLogger(TensorFlowService.class);
+
 	private Map<Path, Graph> modelCache = new HashMap<Path, Graph>();
 	private Map<Path, List<String>> labelCache = new HashMap<Path, List<String>>();
 
-	public String getInception(byte[] imageBytes, String modelDir) {
+	public List<Entry<Float, String>> getInception(byte[] imageBytes, String modelDir) {
+		logger.info(String.format("getInception: %d bytes %s", new Object[] { imageBytes.length, Paths.get(modelDir, "graph.pb") }));
 		Graph g = getOrCreate(Paths.get(modelDir, "graph.pb"));
 		try (Session s = new Session(g)) {
 			List<String> labels = getOrCreateLabels(Paths.get(modelDir, "label.txt"));
-			try {
-				Tensor image = constructAndExecuteGraphToNormalizeImage(imageBytes);
-				Tensor result = s.runner().feed("input", image).fetch("output").run().get(0);
-				final long[] rshape = result.shape();
-				if (result.numDimensions() != 2 || rshape[0] != 1) {
-					throw new RuntimeException(String.format(
-							"Expected model to produce a [1 N] shaped tensor where N is the number of labels, instead it produced one with shape %s",
-							Arrays.toString(rshape)));
-				}
-				int nlabels = (int) rshape[1];
-				float[] labelProbabilities = result.copyTo(new float[1][nlabels])[0];
-				int bestLabelIdx = maxIndex(labelProbabilities);
-				return String.format("BEST MATCH: %s (%.2f%% likely)", labels.get(bestLabelIdx),
-						labelProbabilities[bestLabelIdx] * 100f);
-			} catch (Exception x) {
-				x.printStackTrace();
+			Tensor image = constructAndExecuteGraphToNormalizeImage(imageBytes);
+			Tensor result = s.runner().feed("input", image).fetch("output").run().get(0);
+			logger.debug("found results");
+
+			final long[] rshape = result.shape();
+			if (result.numDimensions() != 2 || rshape[0] != 1) {
+				throw new RuntimeException(String.format(
+						"Expected model to produce a [1 N] shaped tensor where N is the number of labels, instead it produced one with shape %s",
+						Arrays.toString(rshape)));
+			}
+			int nlabels = (int) rshape[1];
+
+			logger.debug(String.format("number of labels %d, %d", new Object[] { labels.size(), nlabels }));
+			int mLabeled = Math.min(labels.size(), nlabels);
+
+			float[] labelProbabilities = result.copyTo(new float[1][nlabels])[0];
+
+			HashMap<Float, String> results = new HashMap<Float, String>();
+			for (int i = 0; i < mLabeled; i++) {
+				results.put(labelProbabilities[i], labels.get(i));
 			}
-			return "Unknown";
+
+			return Collections.synchronizedList(
+					results.entrySet().stream().sorted(Collections.reverseOrder(Map.Entry.comparingByKey())).limit(10)
+							.collect(Collectors.toList()));
+		} catch (Exception e) {
+			logger.error("Failed in tensorflow", e);
+			throw(e);
 		}
 	}
 
@@ -104,16 +122,6 @@ private static Tensor constructAndExecuteGraphToNormalizeImage(byte[] imageBytes
 		}
 	}
 
-	private static int maxIndex(float[] probabilities) {
-		int best = 0;
-		for (int i = 1; i < probabilities.length; ++i) {
-			if (probabilities[i] > probabilities[best]) {
-				best = i;
-			}
-		}
-		return best;
-	}
-
 	private static byte[] readAllBytesOrExit(Path path) {
 		try {
 			return Files.readAllBytes(path);
diff --git a/nifi-tensorflow-processors/src/test/java/com/dataflowdeveloper/processors/process/TensorFlowProcessorTest.java b/nifi-tensorflow-processors/src/test/java/com/dataflowdeveloper/processors/process/TensorFlowProcessorTest.java
@@ -16,20 +16,14 @@
  */
 package com.dataflowdeveloper.processors.process;
 
-import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertEquals;
 
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
-import java.io.UnsupportedEncodingException;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.file.Paths;
 import java.util.List;
 
-import org.apache.nifi.components.PropertyDescriptor;
-import org.apache.nifi.processor.util.StandardValidators;
 import org.apache.nifi.util.MockFlowFile;
 import org.apache.nifi.util.TestRunner;
 import org.apache.nifi.util.TestRunners;
@@ -76,8 +70,7 @@ private void runAndAssertHappy() {
 		List<MockFlowFile> successFiles = testRunner.getFlowFilesForRelationship(TensorFlowProcessor.REL_SUCCESS);
 
 		for (MockFlowFile mockFile : successFiles) {
-			System.out.println("Attribute: " + mockFile.getAttribute(TensorFlowProcessor.ATTRIBUTE_OUTPUT_NAME));
-			assertNotNull(mockFile.getAttribute(TensorFlowProcessor.ATTRIBUTE_OUTPUT_NAME));
+			assertEquals("giant panda", mockFile.getAttribute("tf.probabilities.0.label"));
 		}
 	}
 }
diff --git a/pom.xml b/pom.xml
@@ -19,12 +19,12 @@
     <parent>
         <groupId>org.apache.nifi</groupId>
         <artifactId>nifi-nar-bundles</artifactId>
-        <version>1.0.0</version>
+        <version>1.3.0</version>
     </parent>
 
     <groupId>com.dataflowdeveloper</groupId>
     <artifactId>tensorflow-processor</artifactId>
-    <version>2.0</version>
+    <version>2.1</version>
     <packaging>pom</packaging>
 
     <modules>
