import plotly.express as px import plotly.graph_objects as go from utils import DEEPLITE_LIGHT_BLUE_HEX, load_yolobench_data df, pareto_indices = load_yolobench_data() METRIC_NAME_MAPPING = { 'mAP@0.5': 'mAP_0.5', 'mAP@0.5:0.95': 'mAP_0.5:0.95', 'Precision': 'precision', 'Recall': 'recall', } METRIC_KEYS_TO_NAMES = {v: k for k, v in METRIC_NAME_MAPPING.items()} LATENCY_KEYS = { 'Raspberry Pi 4 Model B (CPU, TFLite, FP32)': 'raspi4_tflite_latency', 'Jetson Nano (GPU, ONNX Runtime, FP32)': 'nano_gpu_latency', 'Intel® Core™i7-10875H (CPU, OpenVINO, FP32)': 'openvino_latency', 'Khadas VIM3 (NPU, INT16)': 'vim3_latency', 'Orange Pi 5 (NPU, FP16)': 'orange_pi_latency', 'NVIDIA A40 (GPU, TensorRT, FP32)': 'a40_trt_latency', } LATENCY_KEYS_TO_NAMES = {v: k for k, v in LATENCY_KEYS.items()} DATASET_TAGS = { 'PASCAL VOC': 'voc', 'SKU-110K': 'sku', 'WIDERFACE': 'wider', 'COCO': 'coco', } DATASET_TAGS_TO_NAMES = {v: k for k, v in DATASET_TAGS.items()} def get_scatter_plot( dataset_tag, metric_tag, latency_key, model_family_coloring=True, add_pareto_frontier=False, plot_pareto_only=False, log_axis=False, ): fig_opts, layout_opts = {'opacity': 0.5, 'color_discrete_sequence': [DEEPLITE_LIGHT_BLUE_HEX]}, {} if model_family_coloring: fig_opts = { 'color': 'model_family', 'opacity': 0.75, 'color_discrete_sequence': px.colors.qualitative.Plotly, } layout_opts = { 'legend': dict( title='Model family
(click to toggle)', ) } frontier = None if plot_pareto_only: metric_key = f'{metric_tag}_{dataset_tag}' frontier = pareto_indices[metric_key][latency_key] fig = px.scatter( df if frontier is None else df.iloc[frontier, :], x=latency_key, y=f'{metric_tag}_{dataset_tag}', title=f'{METRIC_KEYS_TO_NAMES[metric_tag]}-latency scatter plot', hover_data={ 'model_name': True, 'model_family': False, latency_key: ':.2f', f'{metric_tag}_{dataset_tag}': ':.2f', }, labels={ 'model_name': 'Model name', latency_key: 'Latency', f'{metric_tag}_{dataset_tag}': METRIC_KEYS_TO_NAMES[metric_tag], }, template='plotly_white', **fig_opts, ) if log_axis: fig.update_xaxes(type='log') fig.update_layout( height=600, modebar_remove=['lasso', 'autoscale', 'zoomin', 'zoomout', 'select2d', 'select'], xaxis_title=f'{LATENCY_KEYS_TO_NAMES[latency_key]} latency, ms', yaxis_title=f"{METRIC_KEYS_TO_NAMES[metric_tag]}", xaxis=dict( rangeslider=dict( visible=True, bgcolor=DEEPLITE_LIGHT_BLUE_HEX, thickness=0.02, ), ), yaxis=dict( fixedrange=False, ), hoverlabel=dict( # bgcolor="white", font_size=14, font_family='Source Sans Pro' ), **layout_opts, ) if add_pareto_frontier: fig = pareto_frontier_layer(fig, dataset_tag, metric_tag, latency_key) return fig def create_yolobench_plots( dataset_name, hardware_name, metric_name, vis_options, table_mode, ): model_family_coloring = 'Model family' in vis_options add_pareto_frontier = 'Highlight Pareto' in vis_options plot_pareto_only = 'Show Pareto only' in vis_options log_axis = 'Log x-axis' in vis_options fig = get_scatter_plot( DATASET_TAGS[dataset_name], METRIC_NAME_MAPPING[metric_name], LATENCY_KEYS[hardware_name], model_family_coloring, add_pareto_frontier, plot_pareto_only, log_axis, ) pareto_table = get_pareto_table( dataset_name, hardware_name, metric_name, expand_table='Show all' in table_mode ) return fig, pareto_table def pareto_frontier_layer( fig, dataset_tag, metric_tag, latency_key, ): metric_key = f'{metric_tag}_{dataset_tag}' frontier = pareto_indices[metric_key][latency_key] fig.add_trace( go.Scatter( x=df.iloc[frontier, :][latency_key], y=df.iloc[frontier, :][metric_key], mode='lines', opacity=0.5, line=go.scatter.Line(color='grey'), showlegend=False, name=metric_key, ) ) return fig def get_pareto_table( dataset_name, hardware_name, metric_name, expand_table=False, ): dataset_tag = DATASET_TAGS[dataset_name] metric_tag = METRIC_NAME_MAPPING[metric_name] latency_key = LATENCY_KEYS[hardware_name] metric_key = f'{metric_tag}_{dataset_tag}' latency_key_final = f'{LATENCY_KEYS_TO_NAMES[latency_key]} latency, ms' metric_key_final = METRIC_KEYS_TO_NAMES[metric_tag] frontier = pareto_indices[metric_key][latency_key] table_df = df.iloc[frontier, :][['model_name', metric_key, latency_key]] table_df['Input resolution (px)'] = table_df['model_name'].apply(lambda name: name.split('_')[-1]) table_df['Model name'] = table_df['model_name'].apply(lambda name: name.split('_')[0]) table_df[metric_key_final] = table_df[metric_key].apply(lambda val: round(val, 3)) table_df[latency_key_final] = table_df[latency_key].apply(lambda val: round(val, 2)) def make_clickable(url, name): return f'{name}' if dataset_name == 'COCO': table_df['Download link'] = table_df['model_name'].apply( lambda name: f'https://download.deeplite.ai/zoo/models/YOLOBench/{name.split("_")[0]}_640.pt' ) table_df['Download link'] = table_df.apply(lambda x: make_clickable(x['Download link'], 'Weights download'), axis=1) else: table_df['Download link'] = table_df['model_name'].apply(lambda s: 'Coming soon') table_df = table_df[['Model name', 'Input resolution (px)', metric_key_final, latency_key_final, 'Download link']].sort_values(by=metric_key_final, ascending=False) if not expand_table: table_df = table_df.iloc[:10, :] table_df = table_df.to_html( classes='table', escape=False, render_links=True, index=False ) return table_df