| import os | |
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| image_directory = '/home/user/app/example_images' | |
| os.makedirs(image_directory, exist_ok=True) | |
| img_files = [file for file in os.listdir( | |
| image_directory) if file.lower().endswith('.jpg') or file.lower().endswith('.png')] | |
| path = [os.path.join(image_directory, filename) for filename in img_files] | |
| model = YOLO('/home/user/app/train_cls_best_small.pt') | |
| inputs_image = [ | |
| gr.components.Image(type="filepath", label="Input Image"), | |
| ] | |
| outputs_text = [ | |
| gr.components.Textbox(type="text", label="Model predict"), | |
| ] | |
| def show_preds_image(image_path): | |
| results = model(image_path) | |
| names_dict = results[0].names | |
| probs = results[0].probs.data.tolist() | |
| return names_dict[np.argmax(probs)], np.max(probs) | |
| interface_image = gr.Interface( | |
| fn=show_preds_image, | |
| inputs=inputs_image, | |
| outputs=outputs_text, | |
| title="Cat vs Dog", | |
| examples=path, | |
| cache_examples=False, | |
| ) | |
| gr.TabbedInterface( | |
| [interface_image], | |
| tab_names=['Image Inference'], | |
| ).queue().launch(debug=True) | |