code review items.

Signed-off-by: Robert Altena <Rob@Ra-ai.com>

Author: RobAltena

dl4j-examples/src/main/java/org/deeplearning4j/examples/dataexamples/ImageDrawer.java

@@ -34,19 +34,21 @@
 import org.nd4j.linalg.indexing.conditions.Conditions;
 import org.nd4j.linalg.learning.config.Adam;
 import org.nd4j.linalg.lossfunctions.LossFunctions;
+import org.nd4j.shade.guava.collect.Streams;
 
 import javax.imageio.ImageIO;
 import javax.swing.*;
 import java.awt.image.BufferedImage;
+import java.awt.image.DataBufferByte;
 import java.io.File;
+import java.util.Random;
 
 /**
  * Application to show a neural network learning to draw an image.
  * Demonstrates how to feed an NN with externally originated data.
  *
  *  Updates from previous versions:
  *   - Now uses swing. No longer uses JavaFX which caused problems with the OpenJDK.
- *   - All slow java loops in the dataset creation and image drawing are replaced with fast vectorized code.
  *
  * @author Robert Altena
  * Many thanks to @tmanthey for constructive feedback and suggestions.
@@ -59,17 +61,11 @@ public class ImageDrawer {
     private BufferedImage originalImage;
     private JLabel generatedLabel;
 
-    private INDArray blueMat; // color channels of he original image.
-    private INDArray greenMat;
-    private INDArray redMat;
-
-    private INDArray xPixels; // x coordinates of the pixels for the NN.
-    private INDArray yPixels; // y coordinates of the pixels for the NN.
-
     private INDArray xyOut; //x,y grid to calculate the output image. Needs to be calculated once, then re-used.
 
     private Java2DNativeImageLoader j2dNil; //Datavec class used to read and write images to /from INDArrays.
-
+    private FastRGB rgb; // helper class for fast access to the image pixels.
+    private Random random;
 
     private void init() throws Exception {
 
@@ -78,6 +74,7 @@ private void init() throws Exception {
 
         String localDataPath = DownloaderUtility.DATAEXAMPLES.Download();
         originalImage = ImageIO.read(new File(localDataPath, "Mona_Lisa.png"));
+
         //start with a blank image of the same size as the original.
         BufferedImage generatedImage = new BufferedImage(originalImage.getWidth(), originalImage.getHeight(), originalImage.getType());
 
@@ -98,15 +95,13 @@ private void init() throws Exception {
         mainFrame.setVisible(true);  // Show UI
 
 
-        j2dNil = new Java2DNativeImageLoader(); //Datavec class used to read and write images.
+        j2dNil = new Java2DNativeImageLoader(); //Datavec class used to write images.
+        random = new Random();
         nn = createNN(); // Create the neural network.
         xyOut = calcGrid(); //Create a mesh used to generate the image.
 
         // read the color channels from the original image.
-        INDArray imageMat = j2dNil.asMatrix(originalImage).castTo(DataType.DOUBLE).div(255.0);
-        blueMat = imageMat .tensorAlongDimension(1, 0, 2, 3).reshape(width * height, 1);
-        greenMat = imageMat .tensorAlongDimension(2, 0, 2, 3).reshape(width * height, 1);
-        redMat = imageMat .tensorAlongDimension(3, 0, 2, 3).reshape(width * height, 1);
+        rgb = new FastRGB(originalImage);
 
         SwingUtilities.invokeLater(this::onCalc);
     }
@@ -127,30 +122,30 @@ private static MultiLayerNetwork createNN() {
         int numOutputs = 3 ; //R, G and B value.
 
         MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
-                .seed(seed)
-                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
-                .weightInit(WeightInit.XAVIER)
-                .updater(new Adam(learningRate))
-                .list()
-                .layer(new DenseLayer.Builder().nIn(numInputs).nOut(numHiddenNodes)
-                        .activation(Activation.LEAKYRELU)
-                        .build())
-                .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
-                        .activation(Activation.LEAKYRELU)
-                        .build())
-                .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
-                        .activation(Activation.LEAKYRELU)
-                        .build())
-                .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
-                        .activation(Activation.LEAKYRELU)
-                        .build())
-                .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
-                        .activation(Activation.LEAKYRELU)
-                        .build())
-                .layer( new OutputLayer.Builder(LossFunctions.LossFunction.L2)
-                        .activation(Activation.IDENTITY)
-                        .nOut(numOutputs).build())
-                .build();
+            .seed(seed)
+            .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
+            .weightInit(WeightInit.XAVIER)
+            .updater(new Adam(learningRate))
+            .list()
+            .layer(new DenseLayer.Builder().nIn(numInputs).nOut(numHiddenNodes)
+                .activation(Activation.LEAKYRELU)
+                .build())
+            .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
+                .activation(Activation.LEAKYRELU)
+                .build())
+            .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
+                .activation(Activation.LEAKYRELU)
+                .build())
+            .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
+                .activation(Activation.LEAKYRELU)
+                .build())
+            .layer(new DenseLayer.Builder().nOut(numHiddenNodes )
+                .activation(Activation.LEAKYRELU)
+                .build())
+            .layer( new OutputLayer.Builder(LossFunctions.LossFunction.L2)
+                .activation(Activation.IDENTITY)
+                .nOut(numOutputs).build())
+            .build();
 
         MultiLayerNetwork net = new MultiLayerNetwork(conf);
         net.init();
@@ -162,8 +157,9 @@ private static MultiLayerNetwork createNN() {
      * Training the NN and updating the current graphical output.
      */
     private void onCalc(){
-        int batchSize = 1000;
-        int numBatches = 10;
+        // Find a reasonable balance between batch size and number of batches per generated redraw.
+        int batchSize = 1000; //larger batch size slows the calculation but speeds up the learning per batch
+        int numBatches = 10; // Drawing the generated image is slow. Doing multiple batches before redrawing increases speed.
         for (int i =0; i< numBatches; i++){
             DataSet ds = generateDataSet(batchSize);
             nn.fit(ds);
@@ -172,11 +168,13 @@ private void onCalc(){
         mainFrame.invalidate();
         mainFrame.repaint();
 
-        SwingUtilities.invokeLater(this::onCalc);
+        SwingUtilities.invokeLater(this::onCalc); //TODO: move training to a worker thread,
     }
 
     /**
      * Take a batchsize of random samples from the source image.
+     * This illustrates how to generate a custom dataset. The normal way of doing this would be to generate a dataset
+     * of the entire source image, train om shuffled batches from there.
      *
      * @param batchSize number of sample points to take out of the image.
      * @return DeepLearning4J DataSet.
@@ -185,22 +183,22 @@ private DataSet generateDataSet(int batchSize) {
         int w = originalImage.getWidth();
         int h = originalImage.getHeight();
 
-        INDArray xindex = Nd4j.rand(batchSize).muli(w-1).castTo(DataType.UINT32);
-        INDArray yindex = Nd4j.rand(batchSize).muli(h-1).castTo(DataType.UINT32);
-
-        INDArray xPos = xPixels.get(xindex).reshape(batchSize); // Look up the normalized positions pf the pixels.
-        INDArray yPos = yPixels.get(yindex).reshape(batchSize);
-
-        INDArray xy =  Nd4j.vstack(xPos, yPos).transpose(); // Create the array that can be fed into the NN.
-
-        //Look up the correct colors fot our random pixels.
-        INDArray xyIndex = yindex.mul(w).add(xindex); //TODO: figure out the 2D version of INDArray.get.
-        INDArray b = blueMat.get(xyIndex).reshape(batchSize);
-        INDArray g = greenMat.get(xyIndex).reshape(batchSize);
-        INDArray r = redMat.get(xyIndex).reshape(batchSize);
-        INDArray out = Nd4j.vstack(r, g, b).transpose(); // Create the array that can be used for NN training.
-
-        return new DataSet(xy, out);
+        float[][] in = new float[batchSize][2];
+        float[][] out = new float[batchSize][3];
+        final int[] i = {0};
+        Streams.forEachPair(
+            random.ints(batchSize, 0, w).boxed(),
+            random.ints(batchSize, 0, h).boxed(),
+            (a, b) -> {
+                final short[] parts = rgb.getRGB(a, b);
+                in[i[0]] = new float[]{((a / (float)w) - 0.5f) * 2f, ((b / (float)h) - 0.5f) * 2f};
+                out[i[0]] = new float[]{parts[0], parts[1], parts[2]};
+                i[0]++;
+            }
+        );
+        final INDArray input = Nd4j.create(in);
+        final INDArray labels = Nd4j.create(out).divi(255);
+        return new DataSet(input, labels);
     }
 
     /**
@@ -211,7 +209,7 @@ private void drawImage() {
         int h = originalImage.getHeight();
 
         INDArray out = nn.output(xyOut); // The raw NN output.
-        BooleanIndexing.replaceWhere(out, 0.0, Conditions.lessThan(0.0)); // Cjip between 0 and 1.
+        BooleanIndexing.replaceWhere(out, 0.0, Conditions.lessThan(0.0)); // Clip between 0 and 1.
         BooleanIndexing.replaceWhere(out, 1.0, Conditions.greaterThan(1.0));
         out = out.mul(255).castTo(DataType.BYTE); //convert to bytes.
 
@@ -231,13 +229,40 @@ private void drawImage() {
     private INDArray calcGrid(){
         int w = originalImage.getWidth();
         int h = originalImage.getHeight();
-        xPixels = Nd4j.linspace(-1.0, 1.0, w, DataType.DOUBLE);
-        yPixels = Nd4j.linspace(-1.0, 1.0, h, DataType.DOUBLE);
+        INDArray xPixels = Nd4j.linspace(-1.0, 1.0, w, DataType.DOUBLE);
+        INDArray yPixels = Nd4j.linspace(-1.0, 1.0, h, DataType.DOUBLE);
         INDArray [] mesh = Nd4j.meshgrid(xPixels, yPixels);
 
-        xPixels = xPixels.reshape(w, 1); // This is a hack to work around a bug in INDArray.get()
-        yPixels = yPixels.reshape(h, 1); // in the dataset generation.
-
         return Nd4j.vstack(mesh[0].ravel(), mesh[1].ravel()).transpose();
     }
+
+
+    public class FastRGB {
+        int width;
+        int height;
+        private boolean hasAlphaChannel;
+        private int pixelLength;
+        private byte[] pixels;
+
+        FastRGB(BufferedImage image) {
+            pixels = ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
+            width = image.getWidth();
+            height = image.getHeight();
+            hasAlphaChannel = image.getAlphaRaster() != null;
+            pixelLength = 3;
+            if (hasAlphaChannel)
+                pixelLength = 4;
+        }
+
+        short[] getRGB(int x, int y) {
+            int pos = (y * pixelLength * width) + (x * pixelLength);
+            short rgb[] = new short[4];
+            if (hasAlphaChannel)
+                rgb[3] = (short) (pixels[pos++] & 0xFF); // Alpha
+            rgb[2] = (short) (pixels[pos++] & 0xFF); // Blue
+            rgb[1] = (short) (pixels[pos++] & 0xFF); // Green
+            rgb[0] = (short) (pixels[pos] & 0xFF); // Red
+            return rgb;
+        }
+    }
 }