Repository: voxel51/fiftyone
Issue: [BUG] Interactive confusion matrix callbacks don't work for large datasets 

A race condition seems to occur when selecting cells in an interactive confusion matrix for a sufficiently large dataset that prevents the App from updating as expected.

The example below (when run as cells in a Jupyter notebook) reliably causes this error:

```py
import random
import numpy as np

import fiftyone as fo
import fiftyone.zoo as foz
from fiftyone import ViewField as F

tmp_dataset = foz.load_zoo_dataset("coco-2017", split="validation", label_types=["detections"])
```

```py
def jitter(detections, delta, gamma, classes):
    detections = detections.copy()

    for detection in detections.detections:
        alpha = random.random()

        if alpha < gamma:
            detection.label = random.choice(classes)

        bounding_box = np.array(detection.bounding_box)
        bounding_box += alpha * delta * np.array([1, 1, 0, 0])
        detection.bounding_box = bounding_box.tolist()
        detection.confidence = alpha

    return detections
```

```py
counts = tmp_dataset.count_values("ground_truth.detections.label")
classes = sorted(counts, key=counts.get, reverse=True)[:20]

# Create a detection dataset containing 20 classes
dataset = (
    tmp_dataset
    .filter_labels("ground_truth", F("label").is_in(classes))
    .clone()
)

dataset.default_classes = None

# Add predictions
jit = lambda labels: jitter(labels, 0.03, 0.1, classes)
for sample in dataset.iter_samples(progress=True):
    sample["predictions"] = jit(sample["ground_truth"])
    sample.save()

# Perform evaluation
dataset.evaluate_detections("predictions", gt_field="ground_truth", eval_key="eval")

dataset.name = "confusion-matrix-test"
dataset.persistent = True
```

```py
import fiftyone as fo

dataset = fo.load_dataset("confusion-matrix-test")
results = dataset.load_evaluation_results("eval")
```

```py
session = fo.launch_app(dataset)
```

```py
plot = results.plot_confusion_matrix()
plot.show(height=800)
session.plots.attach(plot, name="eval")
```


Changed files:
fiftyone/core/plots/manager.py, fiftyone/core/plots/plotly.py, fiftyone/core/collections.py, fiftyone/core/plots/base.py, fiftyone/core/plots/matplotlib.py

diff --git a/fiftyone/core/collections.py b/fiftyone/core/collections.py
index dbb5df67b..ff3908e85 100644
--- a/fiftyone/core/collections.py
+++ b/fiftyone/core/collections.py
@@ -702,19 +702,52 @@ class SampleCollection(object):
             view.set_values(tags_path, tags)
 
     def _get_selected_labels(self, ids=None, tags=None, fields=None):
-        view = self.select_labels(ids=ids, tags=tags, fields=fields)
+        if ids is not None or tags is not None:
+            view = self.select_labels(ids=ids, tags=tags, fields=fields)
+        else:
+            view = self
 
-        sample_ids = view.values("id")
+        if fields is None:
+            label_fields = view._get_label_fields()
+        elif etau.is_str(fields):
+            label_fields = [fields]
+        else:
+            label_fields = fields
 
-        labels = []
-        for label_field in view._get_label_fields():
+        if not label_fields:
+            return []
+
+        paths = ["id"]
+        is_list_fields = []
+        is_frame_fields = []
+        for label_field in label_fields:
             label_type, id_path = view._get_label_field_path(label_field, "id")
             is_list_field = issubclass(label_type, fol._LABEL_LIST_FIELDS)
+            is_frame_field = view._is_frame_field(label_field)
+
+            paths.append(id_path)
+            is_list_fields.append(is_list_field)
+            is_frame_fields.append(is_frame_field)
 
-            label_ids = view.values(id_path)
+        has_frame_fields = any(is_frame_fields)
 
-            if self._is_frame_field(label_field):
-                frame_numbers = view.values("frames.frame_number")
+        if has_frame_fields:
+            paths.insert(0, "frames.frame_number")
+
+        results = list(view.values(paths))
+
+        if has_frame_fields:
+            frame_numbers = results.pop(0)
+
+        sample_ids = results[0]
+        all_label_ids = results[1:]
+
+        labels = []
+
+        for label_field, label_ids, is_list_field, is_frame_field in zip(
+            label_fields, all_label_ids, is_list_fields, is_frame_fields
+        ):
+            if is_frame_field:
                 for sample_id, sample_frame_numbers, sample_label_ids in zip(
                     sample_ids, frame_numbers, label_ids
                 ):
diff --git a/fiftyone/core/plots/base.py b/fiftyone/core/plots/base.py
index f31166945..7d374ff05 100644
--- a/fiftyone/core/plots/base.py
+++ b/fiftyone/core/plots/base.py
@@ -780,18 +780,20 @@ class InteractivePlot(ResponsivePlot):
     def _register_disconnect_callback(self, callback):
         pass
 
-    def select_ids(self, ids):
+    def select_ids(self, ids, view=None):
         """Selects the points with the given IDs in this plot.
 
         Args:
-            ids: a list of IDs
+            ids: a list of IDs, or None to reset the plot to its default state
+            view (None): the :class:`fiftyone.core.view.DatasetView`
+                corresponding to the given IDs, if available
         """
         if not self.is_connected:
             return
 
-        self._select_ids(ids)
+        self._select_ids(ids, view=view)
 
-    def _select_ids(self, ids):
+    def _select_ids(self, ids, view=None):
         raise ValueError("Subclass must implement _select_ids()")
 
     def reset(self):
diff --git a/fiftyone/core/plots/manager.py b/fiftyone/core/plots/manager.py
index 3687fa33f..6b758531c 100644
--- a/fiftyone/core/plots/manager.py
+++ b/fiftyone/core/plots/manager.py
@@ -47,7 +47,7 @@ class PlotManager(object):
         session: a :class:`fiftyone.core.session.Session`
     """
 
-    _MIN_UPDATE_DELTA_SECONDS = 0.5
+    _MIN_UPDATE_DELTA_SECONDS = 1
     _LISTENER_NAME = "plots"
 
     def __init__(self, session):
@@ -56,6 +56,7 @@ class PlotManager(object):
         self._aggs = {}
         self._current_sample_ids = None
         self._current_labels = None
+        self._last_session_view = None
         self._last_update = None
         self._last_updates = {}
         self._connected = False
@@ -374,6 +375,9 @@ class PlotManager(object):
         if not self._ready_for_update("session"):
             return
 
+        if self._session.view == self._last_session_view:
+            return
+
         self._update_ids_from_session()
         self._update_plots_from_session()
 
@@ -451,7 +455,7 @@ class PlotManager(object):
             warnings.warn(msg)
 
         self._update_session(plot_view)
-        self._update_plots_from_session()
+        self._update_plots_from_session(exclude=[name])
 
     def _update_ids_from_session(self):
         session = self._session
@@ -489,14 +493,16 @@ class PlotManager(object):
         if not self._needs_update("session"):
             return
 
+        self._last_session_view = view
         with self._session.no_show():
             self._session.view = view
 
-    def _update_plots_from_session(self):
+    def _update_plots_from_session(self, exclude=None):
+        exclude = set(exclude or [])
         names = [
             name
             for name, plot in self._plots.items()
-            if plot.supports_session_updates
+            if plot.supports_session_updates and name not in exclude
         ]
         self._update_plots(names)
 
@@ -526,18 +532,18 @@ class PlotManager(object):
                     % (name, plot.link_type)
                 )
 
+        view = self._session._collection.view()
+
         if view_plot_names:
-            self._update_view_plots(view_plot_names)
+            self._update_view_plots(view_plot_names, view)
 
         for name in interactive_plot_names:
-            self._update_interactive_plot(name)
+            self._update_interactive_plot(name, view)
 
-    def _update_view_plots(self, names):
+    def _update_view_plots(self, names, view):
         # For efficiency, aggregations for all supported `ViewPlot`s are
         # computed in a single batch
 
-        view = self._session._collection.view()
-
         # Build flat list of all aggregations
         aggregations = []
         agg_lengths = []
@@ -571,14 +577,14 @@ class PlotManager(object):
             plot = self._plots[name]
             plot.update_view(view, agg_results=_agg_results)
 
-    def _update_interactive_plot(self, name):
+    def _update_interactive_plot(self, name, view):
         plot = self._plots[name]
 
         if plot.link_type == "samples":
-            plot.select_ids(self._current_sample_ids)
+            plot.select_ids(self._current_sample_ids, view=view)
         elif plot.link_type == "labels":
             label_ids = self._get_current_label_ids_for_plot(plot)
-            plot.select_ids(label_ids)
+            plot.select_ids(label_ids, view=view)
         else:
             raise ValueError(
                 "InteractivePlot '%s' has unsupported link type '%s'"
diff --git a/fiftyone/core/plots/matplotlib.py b/fiftyone/core/plots/matplotlib.py
index 6c5ff3b15..4859610d0 100644
--- a/fiftyone/core/plots/matplotlib.py
+++ b/fiftyone/core/plots/matplotlib.py
@@ -917,7 +917,7 @@ class InteractiveCollection(InteractiveMatplotlibPlot):
 
         self._select_inds(inds)
 
-    def _select_ids(self, ids):
+    def _select_ids(self, ids, view=None):
         if ids is not None:
             inds = [self._ids_to_inds[_id] for _id in ids]
         else:
diff --git a/fiftyone/core/plots/plotly.py b/fiftyone/core/plots/plotly.py
index fb48a43a7..2bb85833b 100644
--- a/fiftyone/core/plots/plotly.py
+++ b/fiftyone/core/plots/plotly.py
@@ -1279,7 +1279,7 @@ class InteractiveScatter(PlotlyInteractivePlot):
     def _reopen(self):
         self._widget = self._make_widget()
 
-    def _select_ids(self, ids):
+    def _select_ids(self, ids, view=None):
         deselect = ids is None
         if deselect:
             ids = []
@@ -1404,8 +1404,11 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
         self._bgw = None
         self._selected_cells = []
 
-        # Lower bound at 1 to avoid zero division errors
-        self._ids_counts = np.vectorize(lambda a: max(1, len(a)))(ids)
+        self._curr_view = None
+        self._curr_ids = None
+        self._curr_Z = None
+        self._curr_zlim = None
+
         self._cells_map = {}
 
         widget = self._make_widget()
@@ -1417,14 +1420,26 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
     def supports_session_updates(self):
         return True
 
+    @property
+    def init_view(self):
+        if self._curr_view is not None:
+            return self._curr_view
+
+        return super().init_view
+
     @property
     def _selected_ids(self):
         if not self._selected_cells:
             return None
 
+        if self._curr_ids is not None:
+            ids = self._curr_ids
+        else:
+            ids = self.ids
+
         return list(
             itertools.chain.from_iterable(
-                self.ids[y, x] for x, y in self._selected_cells
+                ids[y, x] for x, y in self._selected_cells
             )
         )
 
@@ -1456,27 +1471,35 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
     def _reopen(self):
         self._widget = self._make_widget()
 
-    def _select_ids(self, ids):
+    def _select_ids(self, ids, view=None):
         if ids is None:
             self._deselect()
             return
 
-        counter = defaultdict(int)
+        num_rows, num_cols = self.Z.shape
+        curr_ids = np.empty((num_rows, num_cols), dtype=object)
+        for i in range(num_rows):
+            for j in range(num_cols):
+                curr_ids[i, j] = []
+
         for _id in ids:
             cell = self._cells_map.get(_id, None)
             if cell is not None:
-                counter[cell] += 1
+                x, y = cell
+                curr_ids[y, x].append(_id)
 
-        if counter:
-            cells, counts = zip(*counter.items())
-        else:
-            cells, counts = [], []
+        Z = np.vectorize(lambda a: len(a))(curr_ids)
+        zlim = [0, Z.max()]
+
+        self._curr_view = view
+        self._curr_ids = curr_ids
+        self._curr_Z = Z
+        self._curr_zlim = zlim
 
-        self._select_fractional_cells(cells, counts)
+        self._update_heatmap([], Z, Z, zlim)
 
     def _on_click(self, point_inds):
         # `point_inds` is a list of `(y, x)` coordinates of selected cells
-
         if not point_inds:
             self._deselect()
             return
@@ -1498,7 +1521,6 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
     def _on_selection(self, point_inds):
         # `point_inds` are linear indices into the flattened meshgrid of points
         # from `_make_heatmap()`
-
         if point_inds:
             y, x = np.unravel_index(point_inds, self.Z.shape)
             cells = list(zip(x, y))
@@ -1507,7 +1529,12 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
             self._deselect()
 
     def _deselect(self):
-        self._update_heatmap([], self.Z, self.Z)
+        self._curr_view = None
+        self._curr_ids = None
+        self._curr_Z = None
+        self._curr_zlim = None
+
+        self._update_heatmap([], self.Z, self.Z, self.zlim)
 
     def _select_cells(self, cells):
         if cells is None:
@@ -1515,42 +1542,21 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
             return
 
         Z_active = np.full(self.Z.shape, None)
-        Z_bg = self.Z.copy()
 
-        if cells:
-            x, y = zip(*cells)
-            Z_active[y, x] = self.Z[y, x]
-
-        self._update_heatmap(cells, Z_active, Z_bg)
-
-    def _select_fractional_cells(self, cells, counts):
-        #
-        # Compute "effective" heatmap values by comparing the number of
-        # observed IDs to the total possible IDs. This is important because
-        # some heatmap cells may have two IDs (confusion matrix cells with GT
-        # and predicted labels) while other cells have have only one associated
-        # will have both GT and predicted IDs, while cells corresponding to
-        # false positive/negative predictions will only have one ID
-        #
-
-        Z_active = np.zeros_like(self.Z)
-        Z_bg = np.zeros_like(self.Z)
+        if self._curr_ids is not None:
+            Z_bg = self._curr_Z
+            zlim = self._curr_zlim
+        else:
+            Z_bg = self.Z
+            zlim = self.zlim
 
         if cells:
             x, y = zip(*cells)
-            observed_frac = np.array(counts) / self._ids_counts[y, x]
-            z_observed = observed_frac * self.Z[y, x]
-            Z_active[y, x] = z_observed
-            Z_bg[y, x] = z_observed
-
-        zlim = [0, Z_active.max()]
+            Z_active[y, x] = Z_bg[y, x]
 
-        self._update_heatmap(cells, Z_active, Z_bg, zlim=zlim)
-
-    def _update_heatmap(self, cells, Z_active, Z_bg, zlim=None):
-        if zlim is None:
-            zlim = self.zlim
+        self._update_heatmap(cells, Z_active, Z_bg, zlim)
 
+    def _update_heatmap(self, cells, Z_active, Z_bg, zlim):
         self._selected_cells = cells
 
         with self._widget.batch_update():
@@ -1576,7 +1582,7 @@ class InteractiveHeatmap(PlotlyInteractivePlot):
                     self._cells_map[_id] = (x, y)
 
     def _make_heatmap(self):
-        Z = self.Z.copy()
+        Z = self.Z
 
         num_rows, num_cols = Z.shape
         xticks = np.arange(num_cols)