MNISTDemo UI tweak

asotona · asotona · commit 5a30fac048d7 · 2025-02-25T07:13:22.000Z
Reviewed-by: psandoz
diff --git a/cr-examples/onnx/src/test/java/oracle/code/onnx/MNISTDemo.java b/cr-examples/onnx/src/test/java/oracle/code/onnx/MNISTDemo.java
@@ -29,6 +29,7 @@
 import java.io.*;
 import jdk.incubator.code.CodeReflection;
 import java.lang.foreign.MemorySegment;
+import java.lang.foreign.ValueLayout;
 import java.lang.invoke.MethodHandles;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -47,10 +48,9 @@
 import static oracle.code.onnx.Tensor.ElementType.*;
 
 public class MNISTDemo {
-
     private static float[] loadConstant(String resource) throws IOException {
-        var bb = ByteBuffer.wrap(MNISTDemo.class.getResourceAsStream(resource).readAllBytes()).order(ByteOrder.LITTLE_ENDIAN);
-        return FloatBuffer.allocate(bb.capacity() / 4).put(bb.asFloatBuffer()).array();
+        return MemorySegment.ofArray(MNISTDemo.class.getResourceAsStream(resource).readAllBytes())
+                .toArray(ValueLayout.JAVA_FLOAT_UNALIGNED);
     }
 
     @CodeReflection
@@ -110,11 +110,13 @@ public static Tensor<Float> cnn(Tensor<Float> inputImage) throws IOException {
     static final int IMAGE_SIZE = 28;
     static final int DRAW_AREA_SIZE = 600;
     static final int PEN_SIZE = 20;
+    static final String[] COLORS = {"1034a6", "412f88", "722b6a", "a2264b", "d3212d", "f62d2d"};
 
     public static void main(String[] args) throws Exception {
         var frame = new JFrame("CNN MNIST Demo - Handwritten Digit Classification");
         var drawPane = new JPanel(false);
         var statusBar = new JLabel("   Hold SHIFT key to draw with trackpad or mouse, click ENTER to run digit classification.");
+        var results = new JLabel();
         var cleanFlag = new AtomicBoolean(true);
         var modelRuntimeSession = OnnxRuntime.getInstance().createSession(
                 OnnxProtoBuilder.buildFuncModel(
@@ -128,6 +130,8 @@ public static void main(String[] args) throws Exception {
         var inputArguments = List.of(new Tensor(MemorySegment.ofBuffer(scaledImageDataBuffer), FLOAT, 1, 1, IMAGE_SIZE, IMAGE_SIZE).tensorAddr);
         var sampleArray = new float[IMAGE_SIZE * IMAGE_SIZE];
 
+        results.setPreferredSize(new Dimension(100, 0));
+
         drawPane.setPreferredSize(new Dimension(DRAW_AREA_SIZE, DRAW_AREA_SIZE));
         drawPane.addMouseMotionListener(new MouseAdapter() {
             @Override
@@ -145,6 +149,7 @@ public void mouseMoved(MouseEvent e) {
 
         frame.setLayout(new BorderLayout());
         frame.add(drawPane, BorderLayout.CENTER);
+        frame.add(results, BorderLayout.EAST);
         frame.add(statusBar, BorderLayout.SOUTH);
         frame.pack();
         frame.setResizable(false);
@@ -155,17 +160,13 @@ public void keyPressed(KeyEvent e) {
                     scaledGraphics.drawImage(drawAreaImage.getScaledInstance(IMAGE_SIZE, IMAGE_SIZE, Image.SCALE_SMOOTH), 0, 0, null);
                     scaledImageDataBuffer.put(0, scaledImage.getData().getSamples(0, 0, IMAGE_SIZE, IMAGE_SIZE, 0, sampleArray));
                     FloatBuffer result = OnnxRuntime.getInstance().tensorBuffer(modelRuntimeSession.run(inputArguments).getFirst()).asFloatBuffer();
-                    int max = 0;
-                    for (int i = 1; i < 10; i++) {
-                        if (result.get(i) > result.get(max)) max = i;
-                    }
-                    var msg = new StringBuilder("<html>&nbsp;");
+                    var msg = new StringBuilder("<html>");
                     for (int i = 0; i < 10; i++) {
-                        msg.append((max == i ? "&nbsp;&nbsp;<b><u>%d:&nbsp;%.1f%%</u></b>"
-                                             : "&nbsp;&nbsp;%d:&nbsp;%.1f%%")
-                                .formatted(i, 100 * result.get(i)));
+                        var w = result.get(i);
+                        msg.append("&nbsp;<font size=\"%d\" color=\"#%s\">%d</font>&nbsp;(%.1f%%)&nbsp;<br><br><br>"
+                                .formatted((int)(20 * w) + 3, COLORS[(int)(5.99 * w)], i, 100 * w));
                     }
-                    statusBar.setText(msg.toString());
+                    results.setText(msg.toString());
                     cleanFlag.set(true);
                 }
             }
