unitrefine 0.2.0__py3-none-any.whl

unitrefine/manage_models.py

@@ -0,0 +1,193 @@
+import json
+import shutil
+from functools import partial
+from pathlib import Path
+
+import pandas as pd
+import PyQt5.QtWidgets as QtWidgets
+from PyQt5.QtCore import pyqtSignal
+
+from unitrefine.models_hub import hfh_model_groups
+
+
+class HFHModelsDialog(QtWidgets.QDialog):
+    """Lets the user pick one or more Hugging Face Hub models to download into the project."""
+
+    def __init__(self, existing_repo_ids, parent=None):
+
+        super().__init__(parent)
+
+        self.setWindowTitle("Load models from Hugging Face Hub")
+        self.setMinimumWidth(550)
+
+        layout = QtWidgets.QVBoxLayout(self)
+        layout.addWidget(QtWidgets.QLabel("Select the UnitRefine models to download:"))
+
+        self.checkboxes = {}
+        for group_index, group in enumerate(hfh_model_groups):
+            if group_index > 0:
+                layout.addSpacing(15)
+            if group["title"] is not None:
+                group_title = QtWidgets.QLabel(group["title"])
+                group_title.setStyleSheet("font-weight: bold;")
+                layout.addWidget(group_title)
+
+            for repo_id, note in group["models"]:
+                checkbox = QtWidgets.QCheckBox(f"{repo_id}  {note}".strip())
+                if repo_id in existing_repo_ids:
+                    checkbox.setText(f"{repo_id}  (already in project)")
+                    checkbox.setChecked(True)
+                    checkbox.setEnabled(False)
+                self.checkboxes[repo_id] = checkbox
+                layout.addWidget(checkbox)
+
+        other_label = QtWidgets.QLabel("Other Hugging Face repos (one per line, e.g. username/model-name):")
+        other_label.setStyleSheet("font-weight: bold;")
+        layout.addWidget(other_label)
+
+        self.other_repos_form = QtWidgets.QPlainTextEdit()
+        self.other_repos_form.setFixedHeight(70)
+        layout.addWidget(self.other_repos_form)
+
+        button_box = QtWidgets.QDialogButtonBox(
+            QtWidgets.QDialogButtonBox.StandardButton.Ok | QtWidgets.QDialogButtonBox.StandardButton.Cancel
+        )
+        button_box.accepted.connect(self.accept)
+        button_box.rejected.connect(self.reject)
+        layout.addWidget(button_box)
+
+    def get_repo_ids(self):
+        """Returns the selected repo ids which are not already in the project, without duplicates."""
+
+        repo_ids = [repo_id for repo_id, checkbox in self.checkboxes.items() if checkbox.isChecked() and checkbox.isEnabled()]
+        for line in self.other_repos_form.toPlainText().replace(",", "\n").splitlines():
+            repo_id = line.strip()
+            if repo_id != "" and repo_id not in repo_ids:
+                repo_ids.append(repo_id)
+        return repo_ids
+
+
+def get_model_summary(model_folder):
+    """
+    Returns the classes a model predicts, the number of training units per class and the
+    best balanced accuracy, read from the files saved next to the model.
+    """
+
+    model_folder = Path(model_folder)
+    summary = {"classes": [], "class_counts": {}, "balanced_accuracy": None}
+
+    model_info_path = model_folder / "model_info.json"
+    if model_info_path.is_file():
+        with open(model_info_path) as f:
+            label_conversion = json.load(f).get("label_conversion", {})
+        summary["classes"] = [label_conversion[key] for key in sorted(label_conversion, key=int)]
+
+        labels_path = model_folder / "labels.csv"
+        if labels_path.is_file():
+            labels = pd.read_csv(labels_path)
+            label_column = [column for column in labels.columns if column != "unit_index"][0]
+            # labels are stored as class indices; retrained models can store them as floats
+            class_indices = pd.to_numeric(labels[label_column], errors="coerce").dropna().astype(int).astype(str)
+            counts = class_indices.map(label_conversion).value_counts()
+            summary["class_counts"] = {label: int(counts.get(label, 0)) for label in summary["classes"]}
+
+    accuracies_path = model_folder / "model_accuracies.csv"
+    if accuracies_path.is_file():
+        accuracies = pd.read_csv(accuracies_path)
+        # older models saved the balanced accuracy in an `accuracy` column
+        accuracy_column = "balanced_accuracy" if "balanced_accuracy" in accuracies.columns else "accuracy"
+        if accuracy_column in accuracies.columns and len(accuracies) > 0:
+            summary["balanced_accuracy"] = float(accuracies[accuracy_column].max())
+
+    return summary
+
+
+class ManageModelsWindow(QtWidgets.QDialog):
+    """Lists the models in a project with the classes they were trained on, and lets the user delete them."""
+
+    models_changed = pyqtSignal()
+
+    def __init__(self, project, parent=None):
+
+        super().__init__(parent)
+
+        self.project = project
+
+        self.setWindowTitle("UNITREFINE: Manage models")
+        self.resize(900, 400)
+
+        layout = QtWidgets.QVBoxLayout(self)
+        layout.addWidget(QtWidgets.QLabel("Models in this project. Balanced accuracy is on the held-out test set."))
+
+        self.table = QtWidgets.QTableWidget()
+        self.table.setColumnCount(5)
+        self.table.setHorizontalHeaderLabels(["Model", "Source", "Classes (training units)", "Balanced accuracy", ""])
+        self.table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
+        self.table.verticalHeader().setVisible(False)
+        layout.addWidget(self.table)
+
+        self.make_table()
+
+    def make_table(self):
+
+        self.table.setRowCount(len(self.project.model_paths))
+
+        for row, (model_folder, hfh_or_local) in enumerate(self.project.model_paths):
+
+            summary = get_model_summary(model_folder)
+
+            if len(summary["class_counts"]) > 0:
+                classes_text = ", ".join(f"{label}: {count}" for label, count in summary["class_counts"].items())
+            else:
+                classes_text = ", ".join(summary["classes"]) if len(summary["classes"]) > 0 else "unknown"
+
+            accuracy = summary["balanced_accuracy"]
+            accuracy_text = f"{accuracy:.3f}" if accuracy is not None else "---"
+
+            source_text = "Hugging Face" if hfh_or_local == "hfh" else "local"
+
+            model_item = QtWidgets.QTableWidgetItem(Path(model_folder).name)
+            model_item.setToolTip(str(model_folder))
+            self.table.setItem(row, 0, model_item)
+            self.table.setItem(row, 1, QtWidgets.QTableWidgetItem(source_text))
+            self.table.setItem(row, 2, QtWidgets.QTableWidgetItem(classes_text))
+            self.table.setItem(row, 3, QtWidgets.QTableWidgetItem(accuracy_text))
+
+            delete_button = QtWidgets.QToolButton()
+            delete_button.setIcon(self.style().standardIcon(QtWidgets.QStyle.SP_TrashIcon))
+            delete_button.setToolTip(f"Delete {Path(model_folder).name}")
+            delete_button.clicked.connect(partial(self.delete_model, model_folder))
+            self.table.setCellWidget(row, 4, delete_button)
+
+        self.table.resizeColumnsToContents()
+
+    def model_is_in_project_folder(self, model_folder):
+        return Path(model_folder).resolve().parent == (Path(self.project.folder_name) / "models").resolve()
+
+    def delete_model(self, model_folder):
+
+        model_name = Path(model_folder).name
+        in_project = self.model_is_in_project_folder(model_folder)
+
+        if in_project:
+            message = f"Delete the model '{model_name}'?\n\nThis permanently deletes the folder:\n{model_folder}"
+        else:
+            message = (
+                f"Remove the model '{model_name}' from this project?\n\n"
+                f"It was loaded from outside the project, so its folder is not deleted:\n{model_folder}"
+            )
+
+        answer = QtWidgets.QMessageBox.question(self, "Delete model", message)
+        if answer != QtWidgets.QMessageBox.Yes:
+            return
+
+        if in_project:
+            shutil.rmtree(model_folder)
+            print(f"Deleted model at {model_folder}.")
+        else:
+            print(f"Removed model at {model_folder} from project.")
+
+        self.project.model_paths = [model for model in self.project.model_paths if Path(model[0]) != Path(model_folder)]
+
+        self.make_table()
+        self.models_changed.emit()

unitrefine/metrics.py

@@ -0,0 +1,146 @@
+"""
+Computes the metrics a model needs which an analyzer doesn't have yet, e.g. the drift metrics or the multi-channel
+template metrics (spread, velocity_above, ...) used by the models on the Hugging Face Hub. Used by
+unitrefine-apply --compute_missing_metrics (see cli.py), and by Apply in the GUI, which runs this file in its own
+process (python -m unitrefine.metrics, see main below) so the window keeps responding.
+
+The metrics are computed with SpikeInterface's default parameters, as the published models were trained with, and
+saved in the analyzer. The metrics it already has are kept.
+"""
+
+import argparse
+import json
+import sys
+import traceback
+
+from unitrefine.import_labels import compute_all_metrics
+from unitrefine.predictions import RENAMED_TEMPLATE_METRICS
+
+# printed by main, so the GUI can show which analyzer is being done, and read the results
+METRICS_STEP_MARKER = "UNITREFINE_METRICS_STEP:"
+METRICS_RESULT_MARKER = "UNITREFINE_METRICS_RESULT:"
+
+# the extensions metrics can depend on, in the order they can be computed in
+PARENT_EXTENSIONS = ["random_spikes", "waveforms", "templates", "noise_levels", "spike_amplitudes", "spike_locations",
+                     "principal_components"]
+
+
+def required_metric_columns(model_folders):
+    """The metric columns the models use, with the names SpikeInterface uses now (e.g. peak_to_valley is now
+    peak_to_trough_duration)."""
+    from spikeinterface.curation.model_based_curation import get_required_metrics_from_model
+
+    columns = set()
+    for model_folder in model_folders:
+        columns |= set(get_required_metrics_from_model(model_folder=model_folder, trust_model=True))
+    return {RENAMED_TEMPLATE_METRICS.get(column, (column,))[0] for column in columns}
+
+
+def metrics_to_compute(analyzer, model_folders):
+    """
+    The metrics to compute so the models can be applied: {"quality_metrics" or "template_metrics": [metric names]},
+    empty if the analyzer has them all.
+    """
+    from spikeinterface.metrics import ComputeQualityMetrics, ComputeTemplateMetrics
+
+    columns = set(compute_all_metrics(analyzer).columns)
+    # metrics computed before SpikeInterface 0.103.2 have the old names (e.g. peak_to_valley), which models accept
+    columns |= {new_name for old_name, (new_name, _) in RENAMED_TEMPLATE_METRICS.items() if old_name in columns}
+    missing = required_metric_columns(model_folders) - columns
+    to_compute = {}
+    for extension in [ComputeQualityMetrics, ComputeTemplateMetrics]:
+        for metric in extension.metric_list:
+            if len(missing & set(metric.metric_columns)) > 0:
+                to_compute.setdefault(extension.extension_name, []).append(metric.metric_name)
+                missing -= set(metric.metric_columns)
+    if len(missing) > 0:
+        raise ValueError(f"SpikeInterface can't compute the metrics {sorted(missing)}, which the models use.")
+    return to_compute
+
+
+def compute_missing_metrics(analyzer, model_folders, job_kwargs=None):
+    """
+    Computes and saves the metrics the models need which `analyzer` doesn't have, and the extensions they
+    depend on. Returns {extension: [metric names]} of what was computed.
+    """
+    from spikeinterface.metrics import ComputeQualityMetrics, ComputeTemplateMetrics, get_multi_channel_template_metric_names
+
+    to_compute = metrics_to_compute(analyzer, model_folders)
+    if len(to_compute) == 0:
+        return to_compute
+
+    needed = set()
+    if "template_metrics" in to_compute:
+        needed |= {"random_spikes", "templates"}
+    if "quality_metrics" in to_compute:
+        needed |= set(ComputeQualityMetrics.get_optional_dependencies(metric_names=to_compute["quality_metrics"]))
+        # e.g. "amplitude_scalings" or "spike_amplitudes": spike_amplitudes is used
+        needed -= {"amplitude_scalings"}
+    if "principal_components" in needed:
+        needed |= {"waveforms"}
+    if len(needed) > 0:
+        needed |= {"random_spikes"}
+    parents = [name for name in PARENT_EXTENSIONS if name in needed and not analyzer.has_extension(name)]
+    if len(parents) > 0 and not (analyzer.has_recording() or analyzer.has_temporary_recording()):
+        raise ValueError(f"The analyzer has no recording, which is needed to compute {', '.join(parents)} for the "
+                         f"metrics {to_compute}. Set its recording, or compute these metrics where the recording is.")
+
+    for name in parents:
+        compute_keeping_children(analyzer, name, job_kwargs)
+    for extension_name, metric_names in to_compute.items():
+        print(f"    Computing {', '.join(metric_names)}...", flush=True)
+        params = dict(metric_names=list(metric_names), delete_existing_metrics=False)
+        if extension_name == "template_metrics":
+            params["include_multi_channel_metrics"] = len(set(metric_names) & set(get_multi_channel_template_metric_names())) > 0
+        analyzer.compute(extension_name, **params, **(job_kwargs or {}))
+    return to_compute
+
+
+def compute_keeping_children(analyzer, extension_name, job_kwargs=None):
+    """
+    Computes an extension. SpikeInterface deletes the extensions which depend on it (e.g. computing waveforms, for
+    the PCA metrics, deletes the templates and the template metrics), so these are computed again, with the
+    parameters they had. Except for unit_locations, whose fit can change a little, they come back the same.
+    """
+    from spikeinterface.core.sortinganalyzer import _get_children_dependencies
+
+    children = [child for child in dict.fromkeys(_get_children_dependencies(extension_name)) if analyzer.has_extension(child)]
+    children_params = {child: analyzer.get_extension(child).params for child in children}
+    print(f"    Computing {extension_name}...", flush=True)
+    analyzer.compute(extension_name, **(job_kwargs or {}))
+    for child in children:
+        if not analyzer.has_extension(child):
+            print(f"    Computing {child} again (computing {extension_name} deleted it), with the same parameters...", flush=True)
+            analyzer.compute(child, **children_params[child], **(job_kwargs or {}))
+
+
+def main(argv=None):
+    """
+    Computes the metrics the models need for each analyzer (see compute_missing_metrics). Before each analyzer it
+    prints METRICS_STEP_MARKER and "analyzer_{index}"; at the end, METRICS_RESULT_MARKER and a json list with, per
+    analyzer, its "path" and the "computed" metrics or the "error".
+    """
+    from spikeinterface.core import load_sorting_analyzer
+
+    parser = argparse.ArgumentParser(description="Computes the metrics models need which analyzers don't have.")
+    parser.add_argument("analyzers", nargs="+")
+    parser.add_argument("--model", action="append", required=True, dest="models", help="a model folder")
+    parser.add_argument("--n_jobs", type=int, default=-1, help="processes to use (-1: all cores)")
+    args = parser.parse_args(argv)
+
+    results = []
+    for index, path in enumerate(args.analyzers):
+        print(f"{METRICS_STEP_MARKER} analyzer_{index}", flush=True)
+        print(f"\nComputing the missing metrics of {path}...", flush=True)
+        try:
+            analyzer = load_sorting_analyzer(path, load_extensions=False)
+            results.append(dict(path=path, computed=compute_missing_metrics(analyzer, args.models, dict(n_jobs=args.n_jobs))))
+        except Exception as error:
+            traceback.print_exc()
+            results.append(dict(path=path, error=str(error)))
+    print(f"{METRICS_RESULT_MARKER}{json.dumps(results)}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

unitrefine/model_selection.py

@@ -0,0 +1,159 @@
+from pathlib import Path
+
+import PyQt5.QtWidgets as QtWidgets
+from PyQt5.QtCore import Qt, pyqtSignal
+
+from unitrefine.predictions import get_model_classes, is_layer1_model, is_layer2_model, selection_key
+
+
+class ModelSelectionWidget(QtWidgets.QWidget):
+    """
+    Choose the models to inspect, relabel and retrain:
+        - Binary: one or more models, each applied on its own
+        - Cascade: a noise vs neural model (layer 1), then a SUA vs MUA model (layer 2) on the neural units
+    """
+
+    selection_changed = pyqtSignal()
+
+    def __init__(self, parent=None):
+
+        super().__init__(parent)
+
+        self.model_paths = []
+
+        layout = QtWidgets.QVBoxLayout(self)
+        layout.setContentsMargins(0, 0, 0, 0)
+
+        mode_layout = QtWidgets.QHBoxLayout()
+        self.binary_radio = QtWidgets.QRadioButton("Each model separately")
+        self.cascade_radio = QtWidgets.QRadioButton("As a cascade (noise vs neural, then SUA vs MUA on the neural units)")
+        self.binary_radio.setChecked(True)
+        self.binary_radio.toggled.connect(self.on_mode_changed)
+        mode_layout.addWidget(self.binary_radio)
+        mode_layout.addWidget(self.cascade_radio)
+        mode_layout.addStretch()
+        layout.addLayout(mode_layout)
+
+        self.binary_list = QtWidgets.QListWidget()
+        # fixed, so that on short screens the page scrolls instead of the list being squeezed
+        self.binary_list.setFixedHeight(140)
+        self.binary_list.itemChanged.connect(lambda _: self.selection_changed.emit())
+        layout.addWidget(self.binary_list)
+
+        self.cascade_widget = QtWidgets.QWidget()
+        cascade_layout = QtWidgets.QFormLayout(self.cascade_widget)
+        cascade_layout.setContentsMargins(0, 0, 0, 0)
+        self.layer1_combo = QtWidgets.QComboBox()
+        self.layer2_combo = QtWidgets.QComboBox()
+        self.layer1_combo.currentIndexChanged.connect(lambda _: self.selection_changed.emit())
+        self.layer2_combo.currentIndexChanged.connect(lambda _: self.selection_changed.emit())
+        cascade_layout.addRow("Layer 1 (noise vs neural):", self.layer1_combo)
+        cascade_layout.addRow("Layer 2 (SUA vs MUA):", self.layer2_combo)
+        layout.addWidget(self.cascade_widget)
+        self.cascade_widget.setVisible(False)
+
+    def set_models(self, model_paths):
+        """
+        Rebuilds the lists from the project's models. Which binary models are ticked:
+            - first time: the last model
+            - after models were trained or loaded: the new models
+            - otherwise: the same models as before
+        New models which fit a cascade layer are also chosen for that layer.
+        """
+
+        previously_checked = set(self.checked_binary_names())
+        previous_names = [Path(model_folder).name for model_folder, _ in self.model_paths]
+        previous_layers = [self.layer1_combo.currentText(), self.layer2_combo.currentText()]
+
+        self.model_paths = list(model_paths)
+        names = [Path(model_folder).name for model_folder, _ in self.model_paths]
+        new_names = [name for name in names if name not in previous_names]
+
+        if len(previous_names) == 0:
+            names_to_check = set(names[-1:])
+        elif len(new_names) > 0:
+            names_to_check = set(new_names)
+        else:
+            names_to_check = previously_checked
+
+        for widget in [self, self.binary_list, self.layer1_combo, self.layer2_combo]:
+            widget.blockSignals(True)
+
+        self.binary_list.clear()
+        self.layer1_combo.clear()
+        self.layer2_combo.clear()
+
+        for model_folder, _ in self.model_paths:
+            name = Path(model_folder).name
+            classes = get_model_classes(model_folder)
+
+            item = QtWidgets.QListWidgetItem(f"{name}   [{', '.join(classes)}]")
+            item.setData(Qt.UserRole, name)
+            item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
+            item.setCheckState(Qt.Checked if name in names_to_check else Qt.Unchecked)
+            self.binary_list.addItem(item)
+
+            if is_layer1_model(classes):
+                self.layer1_combo.addItem(name)
+            if is_layer2_model(classes):
+                self.layer2_combo.addItem(name)
+
+        for combo, previous in zip([self.layer1_combo, self.layer2_combo], previous_layers):
+            combo_names = [combo.itemText(index) for index in range(combo.count())]
+            new_for_layer = [name for name in new_names if name in combo_names]
+            choice = new_for_layer[-1] if len(new_for_layer) > 0 else previous
+            if choice in combo_names:
+                combo.setCurrentIndex(combo_names.index(choice))
+
+        for widget in [self, self.binary_list, self.layer1_combo, self.layer2_combo]:
+            widget.blockSignals(False)
+
+        self.selection_changed.emit()
+
+    def on_mode_changed(self):
+        is_binary = self.binary_radio.isChecked()
+        self.binary_list.setVisible(is_binary)
+        self.cascade_widget.setVisible(not is_binary)
+        self.selection_changed.emit()
+
+    def mode(self):
+        return "binary" if self.binary_radio.isChecked() else "cascade"
+
+    def checked_binary_names(self):
+        return [
+            self.binary_list.item(row).data(Qt.UserRole)
+            for row in range(self.binary_list.count())
+            if self.binary_list.item(row).checkState() == Qt.Checked
+        ]
+
+    def selected_model_names(self):
+        if self.mode() == "binary":
+            return self.checked_binary_names()
+        layer1, layer2 = self.layer1_combo.currentText(), self.layer2_combo.currentText()
+        if layer1 == "" or layer2 == "":
+            return []
+        return [layer1, layer2]
+
+    def selected_models(self):
+        """Returns [(model_folder, hfh_or_local), ...]; for a cascade, [layer 1, layer 2]."""
+        models_by_name = {Path(model_folder).name: (Path(model_folder), hfh_or_local) for model_folder, hfh_or_local in self.model_paths}
+        return [models_by_name[name] for name in self.selected_model_names()]
+
+    def selection_key(self):
+        names = self.selected_model_names()
+        if len(names) == 0:
+            return None
+        return selection_key(self.mode(), names)
+
+    def describe_problem(self):
+        """Returns why the selection can't be used, or None if it can."""
+        if len(self.model_paths) == 0:
+            return "No models trained or loaded into project! Please train or load one first."
+        if self.mode() == "binary" and len(self.checked_binary_names()) == 0:
+            return "No model selected. Tick at least one model."
+        if self.mode() == "cascade":
+            if self.layer1_combo.count() == 0:
+                return "No noise vs neural model in the project for cascade layer 1."
+            if self.layer2_combo.count() == 0:
+                return "No SUA vs MUA model in the project for cascade layer 2."
+        return None

unitrefine/models_hub.py

@@ -0,0 +1,118 @@
+"""
+Downloads UnitRefine models from the Hugging Face Hub. Used by the GUI (Load models from Hugging Face Hub) and
+by unitrefine-download-models (see cli.py), so has no Qt.
+"""
+
+import json
+import re
+import shutil
+from pathlib import Path
+
+# UnitRefine models published on the Hugging Face Hub, shown as options when loading from HFH.
+# Each group is shown with a gap before it and an optional heading; models are (repo id, note shown after the name).
+CASCADE_NOTE = "(can be combined in a cascade for good / MUA / noise labels)"
+LIGHTWEIGHT_NOTE = "(based on fast-to-compute metrics, no PCA metrics; can be combined in a cascade for good / MUA / noise labels)"
+hfh_model_groups = [
+    dict(title=None, models=[
+        ("SpikeInterface/UnitRefine_noise_neural_classifier", CASCADE_NOTE),
+        ("SpikeInterface/UnitRefine_sua_mua_classifier", CASCADE_NOTE),
+    ]),
+    dict(title=None, models=[
+        ("SpikeInterface/UnitRefine_noise_neural_classifier_lightweight", LIGHTWEIGHT_NOTE),
+        ("SpikeInterface/UnitRefine_sua_mua_classifier_lightweight", LIGHTWEIGHT_NOTE),
+    ]),
+    dict(title="SUA only models", models=[
+        ("AnoushkaJain3/UnitRefine-generalized-sua-classifier", ""),
+        ("AnoushkaJain3/UnitRefine-mice-sua-classifier", ""),
+        ("AnoushkaJain3/UnitRefine-rat-sua-classifier", ""),
+        ("AnoushkaJain3/UnitRefine-mole-rat-sua-classifier", ""),
+        ("AnoushkaJain3/UnitRefine-monkey-sua-classifier", ""),
+        ("AnoushkaJain3/UnitRefine-human-sua-classifier", ""),
+    ]),
+]
+
+
+# downloaded into every new project (unless unitrefine is started with --no_pretrained_models), so models can be
+# applied straight away: the first two as a cascade (good / MUA / noise), the third on its own (SUA or not)
+DEFAULT_MODELS = [
+    "SpikeInterface/UnitRefine_noise_neural_classifier",
+    "SpikeInterface/UnitRefine_sua_mua_classifier",
+    "AnoushkaJain3/UnitRefine-generalized-sua-classifier",
+]
+
+# the file in a downloaded model's folder which says which Hugging Face repo it came from
+HFH_PATH_FILE = "hfh_path.txt"
+
+# class names used by models on the Hugging Face Hub, as the GUI labels them
+LABEL_TRANSLATION = {
+    'MUA': ['mua'],
+    'good': ['sua', 'SUA', 'neural'],
+    'noise': [],
+}
+
+
+def looks_like_repo_id(text):
+    """True for a Hugging Face repo id, e.g. SpikeInterface/UnitRefine_sua_mua_classifier."""
+    return re.fullmatch(r"[\w.-]+/[\w.-]+", str(text)) is not None
+
+
+def downloaded_repo_id(model_folder):
+    """The Hugging Face repo a model folder was downloaded from, or None if it wasn't."""
+    hfh_path = Path(model_folder) / HFH_PATH_FILE
+    return hfh_path.read_text().strip() if hfh_path.is_file() else None
+
+
+def translate_label_conversion(model_info_path):
+    """Renames the classes in a downloaded model's model_info.json to the GUI's labels, e.g. sua to good."""
+    with open(model_info_path, 'r') as f:
+        model_info = json.load(f)
+
+    updated_model_info = {}
+    for label_index, label in model_info['label_conversion'].items():
+        new_label = label
+        for sigui_label, user_labels in LABEL_TRANSLATION.items():
+            if label in user_labels:
+                new_label = sigui_label
+        updated_model_info[label_index] = new_label
+    model_info['label_conversion'] = updated_model_info
+
+    with open(model_info_path, 'w') as f:
+        json.dump(model_info, f, indent=4)
+
+
+def download_model(repo_id, model_directory):
+    """
+    Downloads the model in the Hugging Face repo `repo_id` into `model_directory`, which must not exist yet,
+    and renames its classes to the GUI's labels. Nothing is left behind if the download fails.
+    """
+    from huggingface_hub import snapshot_download
+
+    model_directory = Path(model_directory)
+    print(f"Downloading model from HuggingFaceHub repo {repo_id}.\n")
+    try:
+        snapshot_download(repo_id, local_dir=model_directory)
+        (model_directory / HFH_PATH_FILE).write_text(repo_id)
+        translate_label_conversion(model_directory / "model_info.json")
+    except BaseException:
+        # don't leave a half-downloaded model behind
+        shutil.rmtree(model_directory, ignore_errors=True)
+        raise
+
+
+def download_default_models(models_folder, repo_ids=DEFAULT_MODELS):
+    """
+    Downloads the default models into `models_folder` (a new project's models folder), each into a folder named
+    after its repo. Models already there are kept. Returns the repo ids which could not be downloaded, with the
+    reason, e.g. no internet: {repo id: message}.
+    """
+    failed = {}
+    for repo_id in repo_ids:
+        model_directory = Path(models_folder) / repo_id.split("/")[-1]
+        if model_directory.is_dir():
+            continue
+        try:
+            download_model(repo_id, model_directory)
+        except Exception as error:
+            print(f"Could not download {repo_id}: {error}")
+            failed[repo_id] = str(error)
+    return failed