Update default image preprocessing (#1268)

Change-Id: I79144cd8b1353a08b8232ee0372641abb6c6b55e
deepjavalibrary · Oct 5, 2021 · b94f0da · b94f0da
1 parent 34c02a2
commit b94f0da
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
@@ -184,11 +184,20 @@ protected void configPreProcess(Map<String, ?> arguments) {
             if (arguments.containsKey("flag")) {
                 flag = Image.Flag.valueOf(arguments.get("flag").toString());
             }
-            if (getBooleanValue(arguments, "centerCrop", false)) {
-                addTransform(new CenterCrop());
-            }
-            if (getBooleanValue(arguments, "resize", false)) {
+            String resize = getStringValue(arguments, "resize", "false");
+            if ("true".equals(resize)) {
                 addTransform(new Resize(width, height));
+            } else if (!"false".equals(resize)) {
+                String[] tokens = resize.split("\\s*,\\s*");
+                if (tokens.length > 1) {
+                    addTransform(
+                            new Resize(Integer.parseInt(tokens[0]), Integer.parseInt(tokens[1])));
+                } else {
+                    addTransform(new Resize(Integer.parseInt(tokens[0])));
+                }
+            }
+            if (getBooleanValue(arguments, "centerCrop", false)) {
+                addTransform(new CenterCrop(width, height));
             }
             if (getBooleanValue(arguments, "toTensor", true)) {
                 addTransform(new ToTensor());