Update app.py

app.py

@@ -7,6 +7,33 @@ from mxnet.contrib.onnx.onnx2mx.import_model import import_model
 import os
 import gradio as gr
 
+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')
@@ -15,56 +42,23 @@ with open('synset.txt', 'r') as f:
     
 os.system("wget https://github.com/AK391/models/raw/main/vision/classification/alexnet/model/bvlcalexnet-7.onnx")
 
-# Enter path to the ONNX model file
-
-sym, arg_params, aux_params = import_model('bvlcalexnet-7.onnx')
+ort_session = ort.InferenceSession("bvlcalexnet-7.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 = preprocess(img)
-    mod.forward(Batch([img]))
-    # Take softmax to generate probabilities
-    scores = mx.ndarray.softmax(mod.get_outputs()[0]).asnumpy()
-    # print the top-5 inferences class
-    scores = np.squeeze(scores)
-    a = np.argsort(scores)[::-1]
+    img_batch = preprocess(get_image(path))
+
+    outputs = ort_session.run(
+        None,
+        {"data_0": img_batch.astype(np.float32)},
+    )
+
+    a = np.argsort(-outputs[0].flatten())
     results = {}
     for i in a[0:5]:
-        results[labels[i]] = float(scores[i])
+        results[label[i]]=outputs[0][0][i]
     return results
-        
- # Determine and set context
-if len(mx.test_utils.list_gpus())==0:
-    ctx = mx.cpu()
-else:
-    ctx = mx.gpu(0)
-# Load module
-
-mod = mx.mod.Module(symbol=sym, context=ctx,data_names=['data_0'],label_names=None)
-mod.bind(for_training=False, data_shapes=[('data_0', (1,3,224,224))], 
-         label_shapes=mod._label_shapes)
-mod.set_params(arg_params, aux_params, allow_missing=True, allow_extra=True)
+       
 
 title="AlexNet"
 description="AlexNet is the name of a convolutional neural network for classification, which competed in the ImageNet Large Scale Visual Recognition Challenge in 2012."