Many improvements and fixes

app.py

@@ -10,7 +10,6 @@ import streamlit as st  # Import last
 import pandas as pd
 from pandas.api.types import is_float_dtype
 from itertools import combinations
-from collections import OrderedDict
 
 from utils import call_subprocess, clean_string_for_non_alphanumerics, clean_str_column
 from data import read_data, read_data_cached, DownloadHeader, generate_data
@@ -157,7 +156,6 @@ elif input_choice == "Upload counts":
                 st.session_state["step"] = 2
 
         if st.session_state["step"] >= 2:
-            print(st.session_state["count_data"])
             # Ensure targets and predictions are clean strings
             st.session_state["count_data"][target_col] = clean_str_column(
                 st.session_state["count_data"][target_col]
@@ -242,47 +240,43 @@ elif input_choice == "Enter counts":
             st.session_state["count_data"] = pd.DataFrame(
                 all_pairs, columns=["Target", "Prediction"]
             )
+            st.session_state["count_data"]["N"] = 0
 
             st.session_state["step"] = 1
 
     if st.session_state["step"] >= 1:
         with st.form(key="enter_counts_form"):
-            st.write("Fill in the counts for `N(Target, Prediction)` pairs.")
-            count_input_fields = OrderedDict()
-
-            num_cols = 3
-            cols = st.columns(num_cols)
-            for i, (targ, pred) in enumerate(
-                zip(
-                    st.session_state["count_data"]["Target"],
-                    st.session_state["count_data"]["Prediction"],
-                )
-            ):
-                count_input_fields[f"{targ}____{pred}"] = cols[
-                    i % num_cols
-                ].number_input(f"N({targ}, {pred})", step=1)
+            st.write(
+                "Fill in the counts by pressing each cell in the `N` column and inputting the counts."
+            )
+
+            new_counts = st.data_editor(
+                st.session_state["count_data"],
+                hide_index=True,
+                column_config={
+                    "Target": st.column_config.TextColumn(disabled=True),
+                    "Prediction": st.column_config.TextColumn(disabled=True),
+                    "N": st.column_config.NumberColumn(
+                        disabled=False, min_value=0, step=1
+                    ),
+                },
+            )
 
             if st.form_submit_button(
                 label="Generate data",
             ):
-                st.session_state["count_data"]["N"] = [
-                    int(val) for val in count_input_fields.values()
-                ]
+                st.session_state["count_data"] = new_counts
                 st.session_state["step"] = 2
 
     if st.session_state["step"] >= 2:
         DownloadHeader.header_and_data_download(
-            "Entered counts",
+            "",
             data=st.session_state["count_data"],
             file_name="confusion_matrix_counts.csv",
+            label="Download counts",
             help="Download counts",
-            col_sizes=[10, 2],
+            col_sizes=[10, 3],
         )
-        col1, col2, col3 = st.columns([4, 5, 4])
-        with col2:
-            st.write(st.session_state["count_data"])
-            st.write(f"{st.session_state['count_data'].shape}")
-
         target_col = "Target"
         prediction_col = "Prediction"
         n_col = "N"
@@ -318,7 +312,7 @@ if st.session_state["step"] >= 2:
             st.subheader("The Data")
             col1, col2, col3 = st.columns([2, 2, 2])
             with col2:
-                st.write(df.head(5))
+                st.dataframe(df.head(5), hide_index=True)
                 st.write(f"{df.shape} (Showing first 5 rows)")
 
     else:

design.py

@@ -78,16 +78,14 @@ def design_section(
                 "Any observation with either a target or prediction "
                 "of another class is excluded.",
             )
-            # TODO: Once the arrow bug in cvms is fixed, enable reversing!
             # Reverse by default
-            # selected_classes.reverse()
+            selected_classes.reverse()
         with col2:
-            # st.write(" ")
-            # st.write(" ")
-            # reverse_class_order = st.checkbox(
-            #     "Reverse order", value=False, help="Reverse the order of the classes."
-            # )
-            pass
+            st.write(" ")
+            st.write(" ")
+            reverse_class_order = st.checkbox(
+                "Reverse order", value=False, help="Reverse the order of the classes."
+            )
 
         # Color palette
         output["palette"] = _add_select_box(
@@ -151,6 +149,45 @@ def design_section(
         st.markdown("""---""")
         st.markdown("**Advanced**:")
 
+        with st.expander("Labels"):
+            col1, col2 = st.columns(2)
+            with col1:
+                output["x_label"] = st.text_input(
+                    "x-axis",
+                    value=get_uploaded_setting(
+                        key="x_label", default="True Class", type_=str
+                    ),
+                )
+            with col2:
+                output["y_label"] = st.text_input(
+                    "y-axis",
+                    value=get_uploaded_setting(
+                        key="y_label", default="Predicted Class", type_=str
+                    ),
+                )
+
+            st.markdown("---")
+            col1, col2 = st.columns(2)
+            with col1:
+                output["title_label"] = st.text_input(
+                    "Title",
+                    value=get_uploaded_setting(
+                        key="title_label", default="", type_=str
+                    ),
+                )
+            with col2:
+                output["caption_label"] = st.text_input(
+                    "Caption",
+                    value=get_uploaded_setting(
+                        key="caption_label", default="", type_=str
+                    ),
+                )
+            st.info(
+                "Note: When adding a title or caption, "
+                "you may need to adjust the height and "
+                "width of the plot as well."
+            )
+
         with st.expander("Elements"):
             col1, col2 = st.columns(2)
             with col1:
@@ -454,8 +491,10 @@ def design_section(
                 # Save settings as json
                 with open(design_settings_store_path, "w") as f:
                     json.dump(output, f)
-            # if reverse_class_order:
-            #     selected_classes.reverse()
+            if not output["place_x_axis_above"]:
+                selected_classes.reverse()
+            if reverse_class_order:
+                selected_classes.reverse()
 
         design_ready = False
         if st.session_state["step"] >= 3:

plot.R

@@ -267,6 +267,16 @@ if (isTRUE(design_settings$counts_on_top) ||
     )
 }
 
+tile_border_color <- NA
+if (isTRUE(design_settings$show_tile_border)) {
+    tile_border_color <- design_settings$tile_border_color
+}
+
+intensity_by <- ifelse(
+    tolower(design_settings$intensity_by) == "counts",
+    "counts",
+    "normalized"
+)
 
 confusion_matrix_plot <- tryCatch(
     {
@@ -282,7 +292,13 @@ confusion_matrix_plot <- tryCatch(
             rm_zero_text = !design_settings$show_zero_text,
             add_zero_shading = design_settings$show_zero_shading,
             add_arrows = design_settings$show_arrows,
+            arrow_size = design_settings$arrow_size,
+            arrow_nudge_from_text = design_settings$arrow_nudge_from_text,
+            intensity_by = intensity_by,
+            darkness = design_settings$darkness,
             counts_on_top = design_settings$counts_on_top,
+            place_x_axis_above = design_settings$place_x_axis_above,
+            rotate_y_text = design_settings$rotate_y_text,
             diag_percentages_only = design_settings$diag_percentages_only,
             digits = as.integer(design_settings$num_digits),
             palette = design_settings$palette,
@@ -290,7 +306,10 @@ confusion_matrix_plot <- tryCatch(
             font_counts = do.call("font", counts_font_args),
             font_normalized = do.call("font", normalized_font_args),
             font_row_percentages = do.call("font", percentages_font_args),
-            font_col_percentages = do.call("font", percentages_font_args)
+            font_col_percentages = do.call("font", percentages_font_args),
+            tile_border_color = tile_border_color,
+            tile_border_size = design_settings$tile_border_size,
+            tile_border_linetype = design_settings$tile_border_linetype
         )
     },
     error = function(e) {
@@ -301,6 +320,30 @@ confusion_matrix_plot <- tryCatch(
     }
 )
 
+# Add labels on x and y axes
+confusion_matrix_plot <- confusion_matrix_plot +
+    ggplot2::labs(
+        x = design_settings$x_label,
+        y = design_settings$y_label
+    )
+
+# Add title
+if (nchar(design_settings$title_label) > 0) {
+    confusion_matrix_plot <- confusion_matrix_plot +
+        ggplot2::labs(
+            title = design_settings$title_label
+        )
+}
+
+# Add caption
+if (nchar(design_settings$caption_label) > 0) {
+    confusion_matrix_plot <- confusion_matrix_plot +
+        ggplot2::labs(
+            caption = design_settings$caption_label
+        )
+}
+
+
 tryCatch(
     {
         ggplot2::ggsave(