|
import json |
|
import os |
|
from functools import lru_cache |
|
from typing import Mapping |
|
|
|
import gradio as gr |
|
from huggingface_hub import HfFileSystem, hf_hub_download |
|
from imgutils.data import ImageTyping, load_image |
|
from natsort import natsorted |
|
|
|
from onnx_ import _open_onnx_model |
|
from preprocess import _img_encode |
|
|
|
hfs = HfFileSystem() |
|
|
|
|
|
@lru_cache() |
|
def open_model_from_repo(repository, model): |
|
runtime = _open_onnx_model(hf_hub_download(repository, f'{model}/model.onnx')) |
|
with open(hf_hub_download(repository, f'{model}/meta.json'), 'r') as f: |
|
labels = json.load(f)['labels'] |
|
|
|
return runtime, labels |
|
|
|
|
|
class Classification: |
|
def __init__(self, title: str, repository: str, default_model=None, imgsize: int = 384): |
|
self.title = title |
|
self.repository = repository |
|
self.models = natsorted([ |
|
os.path.dirname(os.path.relpath(file, self.repository)) |
|
for file in hfs.glob(f'{self.repository}/*/model.onnx') |
|
]) |
|
self.default_model = default_model or self.models[0] |
|
self.imgsize = imgsize |
|
|
|
def _open_onnx_model(self, model): |
|
return open_model_from_repo(self.repository, model) |
|
|
|
def _gr_classification(self, image: ImageTyping, model_name: str, size=384) -> Mapping[str, float]: |
|
image = load_image(image, mode='RGB') |
|
input_ = _img_encode(image, size=(size, size))[None, ...] |
|
model, labels = self._open_onnx_model(model_name) |
|
output, = model.run(['output'], {'input': input_}) |
|
|
|
values = dict(zip(labels, map(lambda x: x.item(), output[0]))) |
|
return values |
|
|
|
def create_gr(self): |
|
with gr.Tab(self.title): |
|
with gr.Row(): |
|
with gr.Column(): |
|
gr_input_image = gr.Image(type='pil', label='Original Image') |
|
gr_model = gr.Dropdown(self.models, value=self.default_model, label='Model') |
|
gr_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size') |
|
gr_submit = gr.Button(value='Submit', variant='primary') |
|
|
|
with gr.Column(): |
|
gr_output = gr.Label(label='Classes') |
|
|
|
gr_submit.click( |
|
self._gr_classification, |
|
inputs=[gr_input_image, gr_model, gr_infer_size], |
|
outputs=[gr_output], |
|
) |
|
|
