Add box plots and unittests (#137)

* Add box plots and unittests

* improve unit tests

* fix positional args

* fix positional args

* change default title to none

Co-authored-by: Bowen Li <[email protected]>

Author: BlairLee

datasetinsights/stats/__init__.py

@@ -3,6 +3,8 @@
     bar_plot,
     grid_plot,
     histogram_plot,
+    model_performance_box_plot,
+    model_performance_comparison_box_plot,
     plot_bboxes,
     rotation_plot,
 )
@@ -12,6 +14,8 @@
     "grid_plot",
     "histogram_plot",
     "plot_bboxes",
+    "model_performance_box_plot",
+    "model_performance_comparison_box_plot",
     "rotation_plot",
     "RenderedObjectInfo",
 ]

datasetinsights/stats/visualization/plots.py

@@ -22,6 +22,8 @@
 COLORS = list(ImageColor.colormap.values())
 FONT_SCALE = 35
 LINE_WIDTH_SCALE = 250
+ERROR_BAR_BASE_COLOR = "indianred"
+ERROR_BAR_COMPARE_COLOR = "lightseagreen"
 
 
 def decode_segmap(labels, dataset="cityscapes"):
@@ -370,3 +372,127 @@ def rotation_plot(df, x, y, z=None, max_samples=None, title=None, **kwargs):
         )
     )
     return fig
+
+
+def model_performance_comparison_box_plot(
+    title=None,
+    mean_ap_base=None,
+    mean_ap_50_base=None,
+    mean_ar_base=None,
+    mean_ap_new=None,
+    mean_ap_50_new=None,
+    mean_ar_new=None,
+    range=[0, 1.0],
+    **kwargs,
+):
+    """Create a box plot for a base and new model performance
+    Args:
+        title (str): title of the plot
+        mean_ap_base (list): a list of base mAP
+        mean_ap_50_base (list): a list of base mAP
+        mean_ar_base (list): a list of base mAP
+        mean_ap_new (list): a list of base mAP
+        mean_ap_50_new (list): a list of base mAP
+        mean_ar_new (list): a list of base mAP
+        range (list): the range of y axis. Defaults to [0, 1.0]
+
+    Returns:
+        A plotly.graph_objects.Figure containing the box plot
+    """
+    fig = go.Figure(
+        layout=go.Layout(title=go.layout.Title(text=title), **kwargs)
+    )
+    fig.update_yaxes(range=range)
+    _fig_add_trace(
+        fig,
+        mean_ap_base,
+        name="baes mAP",
+        base=True,
+        color=ERROR_BAR_BASE_COLOR,
+    )
+    _fig_add_trace(
+        fig,
+        mean_ap_new,
+        name="new mAP",
+        base=False,
+        color=ERROR_BAR_COMPARE_COLOR,
+    )
+    _fig_add_trace(
+        fig,
+        mean_ap_50_base,
+        name="base mAP50",
+        base=True,
+        color=ERROR_BAR_BASE_COLOR,
+    )
+    _fig_add_trace(
+        fig,
+        mean_ap_50_new,
+        name="new mAP50",
+        base=False,
+        color=ERROR_BAR_COMPARE_COLOR,
+    )
+    _fig_add_trace(
+        fig,
+        mean_ar_base,
+        name="base mAR",
+        base=True,
+        color=ERROR_BAR_BASE_COLOR,
+    )
+    _fig_add_trace(
+        fig,
+        mean_ar_new,
+        name="new mAR",
+        base=False,
+        color=ERROR_BAR_COMPARE_COLOR,
+    )
+
+    return fig
+
+
+def model_performance_box_plot(
+    title=None,
+    mean_ap=None,
+    mean_ap_50=None,
+    mean_ar=None,
+    range=[0, 1.0],
+    **kwargs,
+):
+    """Create a box plot for one model performance
+    Args:
+        title (str): title of the plot
+        mean_ap (list): a list of base mAP
+        mean_ap_50 (list): a list of base mAP
+        mean_ar (list): a list of base mAP
+        range (list): the range of y axis. Defaults to [0, 1.0]
+
+    Returns:
+        A plotly.graph_objects.Figure containing the box plot
+    """
+    fig = go.Figure(
+        layout=go.Layout(title=go.layout.Title(text=title), **kwargs)
+    )
+    fig.update_yaxes(range=range)
+    _fig_add_trace(fig, mean_ap, name="mAP")
+    _fig_add_trace(fig, mean_ap_50, name="mAP@IOU50")
+    _fig_add_trace(fig, mean_ar, name="mAR")
+
+    return fig
+
+
+def _fig_add_trace(
+    fig=None, data=[], name="", base=True, color=ERROR_BAR_BASE_COLOR
+):
+    """Add box plot in figure
+    Args:
+        fig (go.Figure): figure you want to add box plot in
+        data (list): metric values
+        name (str): name of the added box plot
+        base (bool): whether is a base metric. Defaults to true
+        color (str): color for box plot
+
+    Returns:
+        A plotly.graph_objects.Figure containing the box plot
+    """
+    if not data:
+        return
+    fig.add_trace(go.Box(y=data, name=name, marker_color=color))

tests/test_visual.py

@@ -21,6 +21,8 @@
     decode_segmap,
     histogram_plot,
     match_boxes,
+    model_performance_box_plot,
+    model_performance_comparison_box_plot,
     plot_bboxes,
 )
 
@@ -33,6 +35,14 @@ def get_image_and_bbox():
     return image, bbox
 
 
+@pytest.fixture
+def get_evaluation_metrics():
+    mean_ap = [0.1, 0.2, 0.3]
+    mean_ap_50 = [0.3, 0.4, 0.5]
+    mean_ar = [0.2, 0.3, 0.4]
+    return [mean_ap, mean_ap_50, mean_ar]
+
+
 def test_decode_segmap():
     ids = list(CITYSCAPES_COLOR_MAPPING.keys())
     colors = list(CITYSCAPES_COLOR_MAPPING.values())
@@ -57,6 +67,41 @@ def test_histogram_plot():
         assert fig == mock_layout
 
 
+@patch("datasetinsights.stats.visualization.plots.go.Figure.add_trace")
+@patch("datasetinsights.stats.visualization.plots.go.Figure.update_yaxes")
+def test_model_performance_box_plot(
+    mock_update, mock_add_trace, get_evaluation_metrics
+):
+    mean_ap, mean_ap_50, mean_ar = get_evaluation_metrics
+    title = "test plot"
+    model_performance_box_plot(
+        title=title, mean_ap=mean_ap, mean_ap_50=mean_ap_50, mean_ar=mean_ar
+    )
+    assert mock_add_trace.call_count == 3
+    assert mock_update.call_count == 1
+
+
+@patch("datasetinsights.stats.visualization.plots.go.Figure.add_trace")
+@patch("datasetinsights.stats.visualization.plots.go.Figure.update_yaxes")
+def test_model_performance_comparison_box_plot(
+    mock_update, mock_add_trace, get_evaluation_metrics
+):
+    mean_ap_base, mean_ap_50_base, mean_ar_base = get_evaluation_metrics
+    mean_ap_new, mean_ap_50_new, mean_ar_new = get_evaluation_metrics
+    title = "test plot"
+    model_performance_comparison_box_plot(
+        title=title,
+        mean_ap_base=mean_ap_base,
+        mean_ap_50_base=mean_ap_50_base,
+        mean_ar_base=mean_ar_base,
+        mean_ap_new=mean_ap_new,
+        mean_ap_50_new=mean_ap_50_new,
+        mean_ar_new=mean_ar_new,
+    )
+    assert mock_add_trace.call_count == 6
+    mock_update.assert_called_once()
+
+
 def test_bar_plot():
     df = pd.DataFrame({"x": ["a", "b", "c"], "y": [1, 2, 3]})
     mock_figure = Mock()