Update app.py

app.py

@@ -9,29 +9,15 @@ import gradio as gr
 from PIL import Image
 import imageio
 import onnxruntime as ort
+from torchvision import transforms
 
-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
+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]),
+])
 
 mx.test_utils.download('https://s3.amazonaws.com/model-server/inputs/kitten.jpg')
 
@@ -44,8 +30,9 @@ os.system("wget https://github.com/AK391/models/raw/main/vision/classification/d
 ort_session = ort.InferenceSession("densenet-9.onnx")
 
 
-def predict(path):
-    img_batch = preprocess(get_image(path))
+def predict(pil):
+    input_tensor = preprocess(pil)
+    img_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
 
     outputs = ort_session.run(
         None,
@@ -63,4 +50,4 @@ title="DenseNet-121"
 description="DenseNet-121 is a convolutional neural network for classification."
 
 examples=[['apple.jpg']]
-gr.Interface(predict,gr.inputs.Image(type='filepath'),"label",title=title,description=description,examples=examples).launch(enable_queue=True,debug=True)
+gr.Interface(predict,gr.inputs.Image(type='pil'),"label",title=title,description=description,examples=examples).launch(enable_queue=True,debug=True)