Start dataset.makeTranslator pattern (#547)

This pattern builds translators that match the behavior of a dataset. It works best inside the abstract datasets that represent an Application. Change-Id: Ic52d6d21e66aec102c88ece3a3c99666453b5413
deepjavalibrary · Jan 20, 2021 · 2994b95 · 2994b95
1 parent 9077d35
commit 2994b95
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 9 deletions.
diff --git a/...taset/src/main/java/ai/djl/basicdataset/cv/classification/ImageClassificationDataset.java b/...taset/src/main/java/ai/djl/basicdataset/cv/classification/ImageClassificationDataset.java
@@ -12,13 +12,19 @@
  */
 package ai.djl.basicdataset.cv.classification;
 
+import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.transform.Resize;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.modality.cv.translator.ImageClassificationTranslator;
 import ai.djl.modality.cv.util.NDImageUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.training.dataset.RandomAccessDataset;
 import ai.djl.training.dataset.Record;
+import ai.djl.translate.Pipeline;
+import ai.djl.translate.Translator;
 import java.io.IOException;
 import java.util.List;
 import java.util.Optional;
@@ -77,6 +83,28 @@ public Record get(NDManager manager, long index) throws IOException {
         return new Record(data, label);
     }
 
+    /**
+     * Returns the {@link ImageClassificationTranslator} matching the format of this dataset.
+     *
+     * @return the {@link ImageClassificationTranslator} matching the format of this dataset
+     */
+    public Translator<Image, Classifications> makeTranslator() {
+        Pipeline pipeline = new Pipeline();
+
+        // Resize the image if the image size is fixed
+        Optional<Integer> width = getImageWidth();
+        Optional<Integer> height = getImageHeight();
+        if (width.isPresent() && height.isPresent()) {
+            pipeline.add(new Resize(width.get(), height.get()));
+        }
+        pipeline.add(new ToTensor());
+
+        return ImageClassificationTranslator.builder()
+                .optSynset(getClasses())
+                .setPipeline(pipeline)
+                .build();
+    }
+
     /**
      * Returns the number of channels in the images in the dataset.
      *

diff --git a/djl-zero/src/main/java/ai/djl/zero/cv/ImageClassification.java b/djl-zero/src/main/java/ai/djl/zero/cv/ImageClassification.java
@@ -21,9 +21,6 @@
 import ai.djl.basicmodelzoo.cv.classification.ResNetV1;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.transform.Resize;
-import ai.djl.modality.cv.transform.ToTensor;
-import ai.djl.modality.cv.translator.ImageClassificationTranslator;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
 import ai.djl.repository.zoo.Criteria;
@@ -151,12 +148,7 @@ public static ZooModel<Image, Classifications> train(
             EasyTrain.fit(trainer, 35, trainDataset, validateDataset);
         }
 
-        Translator<Image, Classifications> translator =
-                ImageClassificationTranslator.builder()
-                        .optSynset(classes)
-                        .addTransform(new Resize(width, height))
-                        .addTransform(new ToTensor())
-                        .build();
+        Translator<Image, Classifications> translator = dataset.makeTranslator();
         return new ZooModel<>(model, translator);
     }