Repository: deepjavalibrary/djl (java)

Issue:
[examples] ImageClassification gives invalid probability.
## Description
As discussion #1566, `examples.inference.ImageClassification` gives probability >0 or <1.

## What have you tried to solve it?

Set `.optApplySoftmax(true)` to set a Softmax classification layer.


Pull request description:
[examples] Fix ImageClassification invalid probability
## Description ##

- Set a softmax classification layer in `examples.inference.ImageClassification`.
- Write a test for Class `ImageClassification`
- close #1574 


Patch:
diff --git a/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java b/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java
index 9712ecb82..69ae0f442 100644
--- a/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java
+++ b/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java
@@ -68,6 +68,7 @@ public final class ImageClassification {
                     ImageClassificationTranslator.builder()
                             .addTransform(new ToTensor())
                             .optSynset(classes)
+                            .optApplySoftmax(true)
                             .build();
 
             try (Predictor<Image, Classifications> predictor = model.newPredictor(translator)) {


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
index 511e780ec..5c29059be 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
@@ -28,6 +28,7 @@ public class TrainMnistTest {
     public void testTrainMnist() throws ModelException, TranslateException, IOException {
         TestRequirements.engine("MXNet", "PyTorch");
 
+        double expectedProb;
         if (Boolean.getBoolean("nightly")) {
             String[] args = new String[] {"-g", "1"};
 
@@ -39,14 +40,18 @@ public class TrainMnistTest {
             Assert.assertTrue(accuracy > 0.9f, "Accuracy: " + accuracy);
             Assert.assertTrue(loss < 0.35f, "Loss: " + loss);
 
-            Classifications classifications = ImageClassification.predict();
-            Classifications.Classification best = classifications.best();
-            Assert.assertEquals(best.getClassName(), "0");
-            Assert.assertTrue(Double.compare(best.getProbability(), 0.9) > 0);
+            expectedProb = 0.9;
         } else {
             String[] args = new String[] {"-g", "1", "-m", "2"};
 
             TrainMnist.runExample(args);
+            expectedProb = 0;
         }
+
+        Classifications classifications = ImageClassification.predict();
+        Classifications.Classification best = classifications.best();
+        Assert.assertEquals(best.getClassName(), "0");
+        double probability = best.getProbability();
+        Assert.assertTrue(probability > expectedProb && probability <= 1);
     }
 }