Batched vgg preprocessing working for ndarray loaded images in [0,1].

Author: Iainmon

bridge/lib/bridge.cpp

@@ -253,10 +253,10 @@ extern "C" bridge_tensor_t imagenet_normalize(bridge_tensor_t input) {
     auto mean = torch::tensor(kMean).reshape({3, 1, 1});  // (3,1,1)
     auto std  = torch::tensor(kStd).reshape({3, 1, 1});
 
-    // if (image.dim() == 4) {
-    //     mean = mean.unsqueeze(0); // (1,3,1,1)
-    //     std = std.unsqueeze(0);
-    // }
+    if (image.dim() == 4) {
+        mean = mean.unsqueeze(0); // (1,3,1,1)
+        std = std.unsqueeze(0);
+    }
 
     auto output = (image - mean) / std;
     return torch_to_bridge(output);

examples/torch_model_loading/torch_load.chpl

@@ -8,27 +8,18 @@ proc main(args: [] string) {
     // writeln("b sum: ", b.sum());
 
     var image = ndarray.loadFrom(args[1],3,real(32));
-    writeln("Loaded image: ", args[1]);
-    writeln("Image shape: ", image.shape);
 
-    writeln("image : ", max reduce image.data);
+    image = image.resize(224,224).imageNetNormalize();
+    writeln("Resized image: ", image.shape);
 
-    image = image.imageNetNormalize();
-    writeln("image : ", max reduce image.data);
+    var batchedImage = ndarray.loadFrom(args[1],3,real(32)).unsqueeze(0);
+    writeln("Batched image: ", batchedImage.shape);
 
+    batchedImage = batchedImage.resize(224,224);
+    writeln("Batched image resized: ", batchedImage.shape);
 
-
-    // image = image.resize(224,224).imageNetNormalize();
-    // writeln("Resized image: ", image.shape);
-
-    // var batchedImage = ndarray.loadFrom(args[1],3,real(32)).unsqueeze(0);
-    // writeln("Batched image: ", batchedImage.shape);
-
-    // batchedImage = batchedImage.resize(224,224);
-    // writeln("Batched image resized: ", batchedImage.shape);
-
-    // image = batchedImage.squeeze(3).imageNetNormalize();
-    // writeln("Squeezed image: ", image.shape);
+    image = batchedImage.squeeze(3).imageNetNormalize();
+    writeln("Squeezed image: ", image.shape);
 
     image.saveImage("test.jpg");
 }