Update app.py

app.py

@@ -7,7 +7,7 @@ import os
 import gradio as gr
 
 import mxnet
-from mxnet.gluon.data.vision import transforms
+from torchvision import transforms
 
 os.system("wget https://s3.amazonaws.com/onnx-model-zoo/synset.txt")
 
@@ -30,28 +30,21 @@ def get_image(path):
         img = np.array(img.convert('RGB'))
     return img
     
-def preprocess(img):
-    '''
-    Preprocessing required on the images for inference with mxnet gluon
-    The function takes path to an image and returns processed tensor
-    '''
-    transform_fn = transforms.Compose([
-    transforms.Resize(224),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-    img = mxnet.ndarray.array(img)
-    img = transform_fn(img)
-    img = img.expand_dims(axis=0) # batchify
+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]),
+])
+        
 
-    return img.asnumpy()
     
 
 def predict(path):
     img = get_image(path)
-    img = preprocess(img)
-    ort_inputs = {session.get_inputs()[0].name: img}
+    input_tensor = preprocess(img)
+    img = input_tensor.unsqueeze(0)
+    ort_inputs = {session.get_inputs()[0].name: img.cpu().detach().numpy()}
     preds = session.run(None, ort_inputs)[0]
     preds = np.squeeze(preds)
     a = np.argsort(preds)