|
from custom_torch_module.deploy_utils import Onnx_deploy_model |
|
import gradio as gr |
|
import time |
|
from PIL import Image |
|
import os |
|
|
|
model_path = "deploying model/" + "vit_xsmall_patch16_clip_224(trainble_0.15) (Acc 98.44%, Loss 0.168152).onnx" |
|
input_size = [1, 3, 224, 224] |
|
img_size = input_size[-1] |
|
|
|
title = "Gender Vision mini" |
|
description = "An ViT(xsmall_clip) based model(fine tuned with Custom dataset : around 800 train images & 200 test iamges) Accuracy : around 98.4% with the custom test dataset. Optimized with ONNX(around 1.7 times faster than PyTorch version on cpu)" |
|
article = "Through bunch of fine tuning and experiments. !REMEMBER! This model can be wrong." |
|
|
|
def predict(img): |
|
start_time = time.time() |
|
output = onnx_model.run(img, return_prob=True) |
|
end_time = time.time() |
|
elapsed_time = end_time - start_time |
|
prediction_fps = 1 / elapsed_time |
|
|
|
pred_label_and_probs = {"Men" : output[0],"Women" : output[1]} |
|
|
|
return pred_label_and_probs, prediction_fps |
|
|
|
onnx_model = Onnx_deploy_model(model_path=model_path, img_size=img_size) |
|
example_list = [["examples/" + example] for example in os.listdir("examples")] |
|
|
|
|
|
demo = gr.Interface(fn=predict, |
|
inputs=gr.Image(type="pil"), |
|
outputs=[gr.Label(num_top_classes=2, label="Predictions"), |
|
gr.Number(label="Prediction speed(FPS)")], |
|
examples=example_list, |
|
title=title, |
|
description=description, |
|
article=article) |
|
|
|
|
|
demo.launch() |
