|
import sys, os |
|
|
|
import gradio as gr |
|
import plotly.express as px |
|
import numpy as np |
|
import random |
|
from ultralytics import YOLO |
|
|
|
from src.sahi_onnx import * |
|
from sahi.predict import get_sliced_prediction |
|
from sahi.utils.cv import visualize_object_predictions |
|
import PIL |
|
|
|
|
|
model_base = 'onnx_models' |
|
|
|
def inference( |
|
im:gr.inputs.Image=None, |
|
model_path:gr.inputs.Dropdown='YOLOv8n', |
|
conf_thr:gr.inputs.Slider=0.25 |
|
): |
|
|
|
model = Yolov8onnxDetectionModel(model_path=f'{model_base}/{model_path}/best.onnx', |
|
config_path=f'{model_base}/{model_path}/args.yaml', |
|
device='cpu', |
|
confidence_threshold=conf_thr, |
|
category_mapping={'0': 'Boat'}, |
|
image_size=640) |
|
|
|
res = get_sliced_prediction(im, model, slice_width=320, |
|
slice_height=320, overlap_height_ratio=0.2, |
|
overlap_width_ratio=0.2, verbose=0) |
|
img = PIL.Image.open(im) |
|
visual_result = visualize_object_predictions(image=np.array(img), |
|
object_prediction_list=res.object_prediction_list, |
|
text_size=0.4, |
|
rect_th=1) |
|
fig = px.imshow(visual_result['image']) |
|
fig.update_layout(showlegend=False, hovermode=False) |
|
fig.update_xaxes(visible=False) |
|
fig.update_yaxes(visible=False) |
|
return fig |
|
|
|
inputs = [ |
|
gr.Image(type='filepath', label='Input'), |
|
gr.components.Dropdown([ |
|
'YOLOv8n', |
|
'YOLOv8s', |
|
'YOLOv8m', |
|
'YOLOv8l', |
|
'YOLOv8x' |
|
], |
|
value='YOLOv8n', label='Model'), |
|
gr.components.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label='Confidence Threshold'), |
|
] |
|
|
|
outputs = [ |
|
gr.Plot(label='Predictions') |
|
] |
|
|
|
example_images = [[f'examples/{f}'] for f in os.listdir('examples')] |
|
|
|
|
|
|
|
gr.Interface( |
|
fn=inference, |
|
inputs=inputs, |
|
outputs=outputs, |
|
allow_flagging='never', |
|
examples=example_images, |
|
cache_examples=False, |
|
examples_per_page=10, |
|
title='Marine vessel detection from Sentinel 2 images', |
|
description="""Models detect potential marine vessels from Sentinel 2 imagery. |
|
Each example image covers 7.68x7.68 km (768x768 pixels).""" |
|
).launch() |
