onnx/CaffeNet

1import mxnet as mx2import matplotlib.pyplot as plt3import numpy as np4from collections import namedtuple5from mxnet.gluon.data.vision import transforms6import os7import gradio as gr8 9from PIL import Image10import imageio11import onnxruntime as ort12 13def get_image(path):14    '''15        Using path to image, return the RGB load image16    '''17    img = imageio.imread(path, pilmode='RGB')18    return img19 20# Pre-processing function for ImageNet models using numpy21def preprocess(img):22    '''23    Preprocessing required on the images for inference with mxnet gluon24    The function takes loaded image and returns processed tensor25    '''26    img = np.array(Image.fromarray(img).resize((224, 224))).astype(np.float32)27    img[:, :, 0] -= 123.6828    img[:, :, 1] -= 116.77929    img[:, :, 2] -= 103.93930    img[:,:,[0,1,2]] = img[:,:,[2,1,0]]31    img = img.transpose((2, 0, 1))32    img = np.expand_dims(img, axis=0)33 34    return img35 36mx.test_utils.download('https://s3.amazonaws.com/model-server/inputs/kitten.jpg')37 38mx.test_utils.download('https://s3.amazonaws.com/onnx-model-zoo/synset.txt')39with open('synset.txt', 'r') as f:40    labels = [l.rstrip() for l in f]41    42os.system("wget https://github.com/AK391/models/raw/main/vision/classification/caffenet/model/caffenet-12.onnx")43 44ort_session = ort.InferenceSession("caffenet-12.onnx")45 46    47def predict(path):48    img_batch = preprocess(get_image(path))49 50    outputs = ort_session.run(51        None,52        {"data_0": img_batch.astype(np.float32)},53    )54 55    a = np.argsort(-outputs[0].flatten())56    results = {}57    for i in a[0:5]:58        results[labels[i]]=float(outputs[0][0][i])59    return results60       61 62title="CaffeNet"63description="CaffeNet a variant of AlexNet. AlexNet is the name of a convolutional neural network for classification, which competed in the ImageNet Large Scale Visual Recognition Challenge in 2012."64 65examples=[['catonnx.jpg']]66gr.Interface(predict,gr.inputs.Image(type='filepath'),"label",title=title,description=description,examples=examples).launch(enable_queue=True,debug=True)