Update app.py

app.py

@@ -10,29 +10,6 @@ from PIL import Image
 import imageio
 import onnxruntime as ort
 
-def get_image(path):
-    '''
-        Using path to image, return the RGB load image
-    '''
-    img = imageio.imread(path, pilmode='RGB')
-    return img
-
-# Pre-processing function for ImageNet models using numpy
-def preprocess(img):
-    '''
-    Preprocessing required on the images for inference with mxnet gluon
-    The function takes loaded image and returns processed tensor
-    '''
-    img = np.array(Image.fromarray(img).resize((224, 224))).astype(np.float32)
-    img[:, :, 0] -= 123.68
-    img[:, :, 1] -= 116.779
-    img[:, :, 2] -= 103.939
-    img[:,:,[0,1,2]] = img[:,:,[2,1,0]]
-    img = img.transpose((2, 0, 1))
-    img = np.expand_dims(img, axis=0)
-
-    return img
-
 mx.test_utils.download('https://s3.amazonaws.com/model-server/inputs/kitten.jpg')
 
 mx.test_utils.download('https://s3.amazonaws.com/onnx-model-zoo/synset.txt')
@@ -45,11 +22,18 @@ ort_session = ort.InferenceSession("shufflenet-v2-10.onnx")
 
     
 def predict(path):
-    img_batch = preprocess(get_image(path))
-
+    input_image = Image.open(path)
+    preprocess = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    input_tensor = preprocess(input_image)
+    input_batch = input_tensor.unsqueeze(0)
     outputs = ort_session.run(
         None,
-        {"input": img_batch.astype(np.float32)},
+        {"input": input_batch.astype(np.float32)},
     )
 
     a = np.argsort(outputs[0].flatten())