Fix backwards-incompatible changes to Brain.add_annotation (#14083)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: wmvanvliet

mne/label.py

@@ -2124,6 +2124,87 @@ def _read_annot_cands(dir_name, raise_error=True):
     return cands
 
 
+def _read_annot(fname):
+    """Read a Freesurfer annotation from a .annot file.
+
+    Note : Copied from PySurfer
+
+    Parameters
+    ----------
+    fname : str
+        Path to annotation file
+
+    Returns
+    -------
+    annot : numpy array, shape=(n_verts)
+        Annotation id at each vertex
+    ctab : numpy array, shape=(n_entries, 5)
+        RGBA + label id colortable array
+    names : list of str
+        List of region names as stored in the annot file
+
+    """
+    if not op.isfile(fname):
+        dir_name = op.split(fname)[0]
+        cands = _read_annot_cands(dir_name)
+        if len(cands) == 0:
+            raise OSError(
+                f"No such file {fname}, no candidate parcellations found in directory"
+            )
+        else:
+            raise OSError(
+                f"No such file {fname}, candidate parcellations in "
+                "that directory:\n" + "\n".join(cands)
+            )
+    with open(fname, "rb") as fid:
+        n_verts = np.fromfile(fid, ">i4", 1)[0]
+        data = np.fromfile(fid, ">i4", n_verts * 2).reshape(n_verts, 2)
+        annot = data[data[:, 0], 1]
+        ctab_exists = np.fromfile(fid, ">i4", 1)[0]
+        if not ctab_exists:
+            raise Exception("Color table not found in annotation file")
+        n_entries = np.fromfile(fid, ">i4", 1)[0]
+        if n_entries > 0:
+            length = np.fromfile(fid, ">i4", 1)[0]
+            np.fromfile(fid, ">c", length)  # discard orig_tab
+
+            names = list()
+            ctab = np.zeros((n_entries, 5), np.int64)
+            for i in range(n_entries):
+                name_length = np.fromfile(fid, ">i4", 1)[0]
+                name = np.fromfile(fid, f"|S{name_length}", 1)[0]
+                names.append(name)
+                ctab[i, :4] = np.fromfile(fid, ">i4", 4)
+                ctab[i, 4] = (
+                    ctab[i, 0]
+                    + ctab[i, 1] * (2**8)
+                    + ctab[i, 2] * (2**16)
+                    + ctab[i, 3] * (2**24)
+                )
+        else:
+            ctab_version = -n_entries
+            if ctab_version != 2:
+                raise Exception("Color table version not supported")
+            n_entries = np.fromfile(fid, ">i4", 1)[0]
+            ctab = np.zeros((n_entries, 5), np.int64)
+            length = np.fromfile(fid, ">i4", 1)[0]
+            np.fromfile(fid, f"|S{length}", 1)  # Orig table path
+            entries_to_read = np.fromfile(fid, ">i4", 1)[0]
+            names = list()
+            for i in range(entries_to_read):
+                np.fromfile(fid, ">i4", 1)  # Structure
+                name_length = np.fromfile(fid, ">i4", 1)[0]
+                name = np.fromfile(fid, f"|S{name_length}", 1)[0]
+                names.append(name)
+                ctab[i, :4] = np.fromfile(fid, ">i4", 4)
+                ctab[i, 4] = ctab[i, 0] + ctab[i, 1] * (2**8) + ctab[i, 2] * (2**16)
+
+        # convert to more common alpha value
+        ctab[:, 3] = 255 - ctab[:, 3]
+
+    return annot, ctab, names
+
+
 def _get_annot_fname(annot_fname, subject, hemi, parc, subjects_dir):
     """Get the .annot filenames and hemispheres."""
     if annot_fname is not None:

mne/viz/_brain/_brain.py

@@ -46,8 +46,10 @@
 from ...utils import (
     Bunch,
     _auto_weakref,
+    _check_fname,
     _check_option,
     _ensure_int,
+    _path_like,
     _ReuseCycle,
     _to_rgb,
     _validate_type,
@@ -3078,21 +3080,27 @@ def add_annotation(
         color=None,
         hover=True,
     ):
-        """Add an annotation file.
+        """Add an annotation (i.e. an atlas of many labels) to the brain figure.
 
         Parameters
         ----------
-        annot : str
-            Either path to annotation file or annotation name.
+        annot : path-like | str | list of Label
+            Either path to annotation file, an annotation name, or a list of
+            :class:`mne.Label` objects.
+
+            DEPRECATED: The annotation can be specified as a ``(labels, ctab)`` tuple
+            per hemisphere, i.e. ``annot=(labels, ctab)`` for a single hemisphere or
+            ``annot=((lh_labels, lh_ctab), (rh_labels, rh_ctab))`` for both hemispheres.
+
+            .. versionadded:: 1.13
+               The ability to supply a list of :class:`~mne.Label` objects.
         borders : bool | int
             Show only label borders. If int, specify the number of steps
             (away from the true border) along the cortical mesh to include
             as part of the border definition.
         %(alpha)s Default is 1.
         hemi : str | None
-            If None, it is assumed to belong to the hemisphere being
-            shown. If two hemispheres are being shown, data must exist
-            for both hemispheres.
+            Optionally restrict the annotation to the given hemisphere.
         remove_existing : bool
             If True (default), remove old annotations.
         color : matplotlib-style color code
@@ -3105,19 +3113,49 @@ def add_annotation(
         """
         from ...label import read_labels_from_annot
 
+        if (isinstance(annot, tuple) and isinstance(annot[0], np.ndarray)) or (
+            isinstance(annot, (tuple, list)) and isinstance(annot[0], tuple)
+        ):
+            # Deprecated old style of passing a (labels, cmap) pair per hemisphere.
+            # Shortcut to old code that can be removed in version 1.14.
+            warn(
+                "Passing the annotation as a `(label, cmap)` tuple is deprecated and "
+                "will be removed in MNE-Python version 1.14.",
+                FutureWarning,
+            )
+            self._old_add_annotation(
+                annot,
+                borders=borders,
+                alpha=alpha,
+                hemi=hemi,
+                remove_existing=remove_existing,
+                color=color,
+            )
+            return
+
+        _validate_type(annot, ("path-like", str, list), "annot")
+
         hemis = self._check_hemis(hemi)
         kwargs = dict()
-        if os.path.isfile(annot):
-            kwargs["annot_fname"] = annot
-        else:
-            kwargs["parc"] = annot
 
-        for hemi in hemis:
-            labels = read_labels_from_annot(
-                self._subject, hemi=hemi, subjects_dir=self._subjects_dir, **kwargs
-            )
+        for hemi_idx, hemi in enumerate(hemis):
+            if _path_like(annot):
+                if os.path.isfile(annot):
+                    kwargs["annot_fname"] = annot
+                else:
+                    kwargs["parc"] = annot
+                labels = read_labels_from_annot(
+                    self._subject,
+                    hemi=hemi,
+                    subjects_dir=self._subjects_dir,
+                    **kwargs,
+                )
+                name = annot
+            else:
+                labels = [label for label in annot if label.hemi == hemi]
+                name = "annotation"  # placeholder name for the annotation
             n_labels = len(labels)
-            ids = np.zeros(self.geo[hemi].coords.shape[0], dtype=int)
+            ids = np.full(self.geo[hemi].coords.shape[0], -1, dtype=int)
             cmap = np.zeros((len(labels) + 1, 4))
             cmap[:, 3] = 1
             cmap[0] = np.array(self._brain_color)
@@ -3131,7 +3169,7 @@ def add_annotation(
                     label.center_of_mass(subjects_dir=self._subjects_dir)
                 ]
             self._annots[hemi].append(
-                dict(name=annot, labels=labels, ids=ids, centroids=centroids)
+                dict(name=name, labels=labels, ids=ids, centroids=centroids)
             )
             del labels
 
@@ -3151,7 +3189,7 @@ def add_annotation(
                     colormap=ctable * 255,
                     rng=[0, n_labels],
                     opacity=alpha,
-                    name=annot,
+                    name=name,
                 )
 
         if hover:
@@ -3175,6 +3213,99 @@ def on_annotation_hover(iren, event):
                     )
         self._renderer._update()
 
+    # DEPRECATED: Can be removed in version 1.14. Also remove _read_annot from
+    # mne/labels.py.
+    def _old_add_annotation(
+        self, annot, borders=True, alpha=1, hemi=None, remove_existing=True, color=None
+    ):
+        from ...label import _read_annot
+
+        hemis = self._check_hemis(hemi)
+
+        # Figure out where the data is coming from
+        if _path_like(annot):
+            if os.path.isfile(annot):
+                filepath = _check_fname(annot, overwrite="read")
+                file_hemi, annot = filepath.name.split(".", 1)
+                if len(hemis) > 1:
+                    if file_hemi == "lh":
+                        filepaths = [filepath, filepath.parent / ("rh." + annot)]
+                    elif file_hemi == "rh":
+                        filepaths = [filepath.parent / ("lh." + annot), filepath]
+                    else:
+                        raise RuntimeError(
+                            "To add both hemispheres simultaneously, filename must "
+                            'begin with "lh." or "rh."'
+                        )
+                else:
+                    filepaths = [filepath]
+            else:
+                filepaths = []
+                for hemi in hemis:
+                    filepath = op.join(
+                        self._subjects_dir,
+                        self._subject,
+                        "label",
+                        ".".join([hemi, annot, "annot"]),
+                    )
+                    if not os.path.exists(filepath):
+                        raise ValueError(f"Annotation file {filepath} does not exist")
+                    filepaths += [filepath]
+            annots = []
+            for hemi, filepath in zip(hemis, filepaths):
+                # Read in the data
+                labels, cmap, _ = _read_annot(filepath)
+                annots.append((labels, cmap))
+        else:
+            annots = [annot] if len(hemis) == 1 else annot
+            annot = "annotation"
+
+        for hemi, (labels, cmap) in zip(hemis, annots):
+            # Maybe zero-out the non-border vertices
+            self._to_borders(labels, hemi, borders)
+
+            # Handle null labels properly
+            cmap[:, 3] = 255
+            bgcolor = np.round(np.array(self._brain_color) * 255).astype(int)
+            bgcolor[-1] = 0
+            cmap[cmap[:, 4] < 0, 4] += 2**24  # wrap to positive
+            cmap[cmap[:, 4] <= 0, :4] = bgcolor
+            if np.any(labels == 0) and not np.any(cmap[:, -1] <= 0):
+                cmap = np.vstack((cmap, np.concatenate([bgcolor, [0]])))
+
+            # Set label ids sensibly
+            order = np.argsort(cmap[:, -1])
+            cmap = cmap[order]
+            ids = np.searchsorted(cmap[:, -1], labels)
+            cmap = cmap[:, :4]
+
+            #  Set the alpha level
+            alpha_vec = cmap[:, 3]
+            alpha_vec[alpha_vec > 0] = alpha * 255
+
+            # Override the cmap when a single color is used
+            if color is not None:
+                rgb = np.round(np.multiply(_to_rgb(color), 255))
+                cmap[:, :3] = rgb.astype(cmap.dtype)
+
+            ctable = cmap.astype(np.float64)
+            for _ in self._iter_views(hemi):
+                mesh = self.layered_meshes[hemi]
+                mesh.add_overlay(
+                    scalars=ids,
+                    colormap=ctable,
+                    rng=[np.min(ids), np.max(ids)],
+                    opacity=alpha,
+                    name=annot,
+                )
+                self._annots[hemi].append(annot)
+                if not self.time_viewer or self.traces_mode == "vertex":
+                    self._renderer._set_colormap_range(
+                        mesh._actor, cmap.astype(np.uint8), None
+                    )
+
+        self._renderer._update()
+
     def _create_caption(self):
         from vtkmodules.vtkRenderingAnnotation import vtkCaptionActor2D
 

mne/viz/_brain/tests/test_brain.py

@@ -52,6 +52,7 @@
 fname_trans = sample_dir / "sample_audvis_trunc-trans.fif"
 fname_stc = sample_dir / "sample_audvis_trunc-meg"
 fname_label = sample_dir / "labels" / "Vis-lh.label"
+fname_label2 = sample_dir / "labels" / "Vis-rh.label"
 fname_cov = sample_dir / "sample_audvis_trunc-cov.fif"
 fname_evoked = sample_dir / "sample_audvis_trunc-ave.fif"
 fname_fwd = sample_dir / "sample_audvis_trunc-meg-eeg-oct-4-fwd.fif"
@@ -590,10 +591,11 @@ def test_add_annotation(renderer_interactive_pyvistaqt, brain_gc):
     annots = [
         "aparc",
         subjects_dir / "fsaverage" / "label" / "lh.PALS_B12_Lobes.annot",
+        [read_label(fname_label, "fsaverage"), read_label(fname_label2, "fsaverage")],
     ]
-    borders = [True, 2]
-    alphas = [1, 0.5]
-    colors = [None, "r"]
+    borders = [True, 1, 2]
+    alphas = [1, 0.5, 0]
+    colors = [None, "r", "b"]
     size = (100, 100)
     brain = Brain(
         subject="fsaverage",
@@ -655,7 +657,7 @@ def GetPickPosition(self):
         subjects_dir=subjects_dir,
     )
     for a, b, p, color in zip(annots, borders, alphas, colors):
-        brain.add_annotation(str(a), b, p, color=color)
+        brain.add_annotation(a, b, p, color=color)
     brain.close()