Update app.py

app.py

@@ -3,69 +3,64 @@ import matplotlib.pyplot as plt
 import numpy as np
 from collections import namedtuple
 from mxnet.gluon.data.vision import transforms
-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')
 with open('synset.txt', 'r') as f:
     labels = [l.rstrip() for l in f]
     
-os.system("wget https://github.com/AK391/models/raw/main/vision/classification/shufflenet/model/shufflenet-v2-10.onnx")
+os.system("wget https://github.com/onnx/models/raw/main/vision/classification/inception_and_googlenet/googlenet/model/googlenet-9.onnx")
 
-# Enter path to the ONNX model file
+ort_session = ort.InferenceSession("googlenet-9.onnx")
 
-sym, arg_params, aux_params = import_model('shufflenet-v2-10.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[labels[i]]=float(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'], label_names=None)
-mod.bind(for_training=False, data_shapes=[('data', (1,3,224,224))],label_shapes=mod._label_shapes)
-mod.set_params(arg_params, aux_params, allow_missing=True, allow_extra=True)
+       
 
-title="ResNet"
-description="ResNet models perform image classification - they take images as input and classify the major object in the image into a set of pre-defined classes. They are trained on ImageNet dataset which contains images from 1000 classes. ResNet models provide very high accuracies with affordable model sizes. They are ideal for cases when high accuracy of classification is required."
+title="GoogleNet"
+description="GoogLeNet is the name of a convolutional neural network for classification, which competed in the ImageNet Large Scale Visual Recognition Challenge in 2014."
 
 examples=[['catonnx.jpg']]
-gr.Interface(predict,gr.inputs.Image(type='filepath'),"label",title=title,description=description,examples=examples).launch(enable_queue=True)
+gr.Interface(predict,gr.inputs.Image(type='filepath'),"label",title=title,description=description,examples=examples).launch(enable_queue=True,debug=True)