Gradio and ONNX 图像分类 Demo

参考：https://www.gradio.app/guides/Gradio-and-ONNX-on-Hugging-Face

import numpy as np
import onnxruntime
import copy
import cv2

def letter_bbox(img, tar_size=(512, 512), color=(0, 0, 0)):
    height = tar_size[0]
    shape = img.shape  # shape = [height,width]
    ratio = float(height) / max(shape)  # ratio = old / new 得到缩放比例
    new_shape = (round(shape[1] * ratio), round(shape[0] * ratio))
    dw = (height - new_shape[0]) / 2  # width padding
    dh = (height - new_shape[1]) / 2  # height padding
    top, bottom = round(dh - 0.1), round(dh + 0.1)
    left, right = round(dw - 0.1), round(dw + 0.1)
    img = cv2.resize(img, new_shape, interpolation=cv2.INTER_LINEAR)
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
    return img, {'dw':dw, 'dh':dh, 'ratio':ratio, 'orin_shape':shape}

class Classifier(object):
    def __init__(self, model_path: str, cls_dict: dict, provider: list = None,
                 letter_box_flag: bool = True, infer_size: tuple = (224, 224)):
        if provider is None:
            self.provider = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        self.onnx_session = onnxruntime.InferenceSession(model_path, providers=self.provider)
        self.input_name = self.get_input_name()
        self.output_name = self.get_output_name()
        self.letter_box_flag = letter_box_flag
        self.infer_size = infer_size
        self.cls_dict = cls_dict
        self.mean = np.array([0.485, 0.456, 0.406])
        self.std = np.array([0.229, 0.224, 0.225])

    def get_input_name(self):
        input_name = []
        for node in self.onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name

    def get_output_name(self):
        output_name = []
        for node in self.onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name

    def get_input_feed(self, img_tensor):
        input_feed = {}
        for name in self.input_name:
            input_feed[name] = img_tensor.astype(np.float32)
        return input_feed

    def pre_process(self, img):
        pre_img = copy.deepcopy(img)
        pre_img = cv2.cvtColor(pre_img, cv2.COLOR_BGR2RGB)
        if self.letter_box_flag:
            pre_img, resize_dict = letter_bbox(pre_img, tar_size=self.infer_size)
            # pre_img = pre_img / 255
            pre_img = (pre_img / 255 - self.mean) / self.std
            pre_img = pre_img.transpose(2, 0, 1)
            pre_img = np.expand_dims(pre_img, axis=0)
            return pre_img, resize_dict
        else:
            ori_shape = img.shape
            pre_img, __ = resize(pre_img, tar_size=self.infer_size)
            # pre_img = pre_img / 255
            pre_img = (pre_img/255 - self.mean) / self.std
            pre_img = pre_img.transpose(2, 0, 1)
            pre_img = np.expand_dims(pre_img, axis=0)
            return pre_img, {'orin_shape': ori_shape}

    def infer(self, img: np.ndarray):
        i_img = copy.deepcopy(img)
        i_img, param_dict = self.pre_process(i_img)
        input_feed = self.get_input_feed(i_img)
        pred = self.onnx_session.run(None, input_feed)[0]
        return pred

cloth_cls_dict = {0: 'HoodieL', 1: 'ShirtL', 2: 'ShirtS', 3: 'Suit', 4: 'DJacket', 5:'Jacket', 6: 'dressL', 
          7: 'dressS', 8: 'dressN', 9: 'TshirtL', 10: 'TshirtS', 11: 'vest'}
model = Classifier(model_path='cloth-clsv3.onnx', cls_dict=cloth_cls_dict, letter_box_flag=True, infer_size=(224, 224))

C:\Users\Administrator\AppData\Roaming\Python\Python38\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py:54: UserWarning: Specified provider 'CUDAExecutionProvider' is not in available provider names.Available providers: 'CPUExecutionProvider'
  warnings.warn(

import gradio as gr
import torch
import requests
from torchvision import transforms

def read_txt_file(filename):
    with open(filename, 'r') as file:
        lines = file.readlines()
        lines = [line.strip() for line in lines]
        return lines

# model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()
filename = 'labels-cloth.txt'  # 替换为你的txt文件路径
labels = read_txt_file(filename)
print(len(labels))

def predict(inp):
  inp = cv2.imread(inp)
  # inp = transforms.ToTensor()(inp).unsqueeze(0)
  with torch.no_grad():
    prediction = torch.nn.functional.softmax(torch.from_numpy(model.infer(inp)[0]), dim=0)
    confidences = {labels[i]: float(prediction[i]) for i in range(12)}
  return confidences

demo = gr.Interface(fn=predict,
             inputs=gr.inputs.Image(type="filepath"),
             outputs=gr.outputs.Label(num_top_classes=3),
             examples=[["_crop/bd_000101_3.jpg"]],
             )

demo.launch()

C:\Users\Administrator\AppData\Local\Temp\ipykernel_11000\3164783745.py:26: GradioDeprecationWarning: Usage of gradio.inputs is deprecated, and will not be supported in the future, please import your component from gradio.components
  inputs=gr.inputs.Image(type="filepath"),
C:\Users\Administrator\AppData\Local\Temp\ipykernel_11000\3164783745.py:26: GradioDeprecationWarning: `optional` parameter is deprecated, and it has no effect
  inputs=gr.inputs.Image(type="filepath"),
C:\Users\Administrator\AppData\Local\Temp\ipykernel_11000\3164783745.py:27: GradioDeprecationWarning: Usage of gradio.outputs is deprecated, and will not be supported in the future, please import your components from gradio.components
  outputs=gr.outputs.Label(num_top_classes=3),
C:\Users\Administrator\AppData\Local\Temp\ipykernel_11000\3164783745.py:27: GradioUnusedKwargWarning: You have unused kwarg parameters in Label, please remove them: {'type': 'auto'}
  outputs=gr.outputs.Label(num_top_classes=3),


12
Running on local URL:  http://127.0.0.1:7872

To create a public link, set `share=True` in `launch()`.