Update app.py

app.py

@@ -7,6 +7,32 @@ from mxnet.contrib.onnx.onnx2mx.import_model import import_model
 import os
 import gradio as gr
 
+from PIL import Image
+import imageio
+
+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')
@@ -20,28 +46,10 @@ os.system("wget https://github.com/onnx/models/raw/main/vision/classification/in
 sym, arg_params, aux_params = import_model('googlenet-9.onnx')
 
 Batch = namedtuple('Batch', ['data'])
-def get_image(path, show=False):
-    img = mx.image.imread(path)
-    if img is None:
-        return None
-    if show:
-        plt.imshow(img.asnumpy())
-        plt.axis('off')
-    return img
-    
-def preprocess(img):   
-    transform_fn = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-    img = transform_fn(img)
-    img = img.expand_dims(axis=0)
-    return img
+
     
 def predict(path):
-    img = get_image(path, show=True)
+    img = get_image(path)
     img = preprocess(img)
     mod.forward(Batch([img]))
     # Take softmax to generate probabilities