Add MTRCNN-DG weights and evaluation

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1eeda46ea42e0dfe59f2724fc204728bfef7fe38ae3c3c5800cf00a2aa13712e
+size 918678

mtrcnn.py

@@ -0,0 +1,1537 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "huggingface-hub",
+#     "jinja2",
+#     "numpy",
+#     "scipy",
+#     "soundfile",
+#     "torch",
+# ]
+# ///
+"""Train a compact domain-generalized mosquito species classifier."""
+
+# pylint: disable=too-few-public-methods,too-many-arguments,too-many-instance-attributes
+# pylint: disable=too-many-lines,too-many-locals,too-many-positional-arguments
+
+import argparse
+import json
+import math
+import os
+import random
+import re
+import shutil
+import statistics
+import urllib.request
+import zipfile
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+import soundfile
+import torch
+from huggingface_hub import (
+    EvalResult,
+    HfApi,
+    ModelCard,
+    ModelCardData,
+    hf_hub_download,
+    snapshot_download,
+)
+from huggingface_hub.errors import RepositoryNotFoundError
+from scipy.signal import resample_poly
+from torch import nn
+from torch.nn import functional
+from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
+
+DATASET = "aptemvs/mosquitoes-biodcase2026-task5"
+REPO = "jgalego/mosquito-mtrcnn-dg"
+ZENODO = "https://zenodo.org/api/records/20478577/files/Development_data.zip/content"
+SAMPLE_RATE = 8_000
+N_FFT = 512
+HOP_LENGTH = 80
+N_MELS = 64
+# Mains hum and handling noise sit below 200 Hz; wingbeat fundamentals are higher.
+MIN_HZ = 200
+LEVEL_RMS = 0.05
+# Every D5 clip lasts 0.625 s; longer crops would pad them with telltale silence.
+SECONDS = 0.625
+SPECIES = [
+    "Aedes aegypti",
+    "Aedes albopictus",
+    "Culex quinquefasciatus",
+    "Anopheles gambiae",
+    "Anopheles arabiensis",
+    "Anopheles dirus",
+    "Culex pipiens",
+    "Anopheles minimus",
+    "Anopheles stephensi",
+]
+DOMAINS = ["D1", "D2", "D3", "D4", "D5"]
+FILE_ID = re.compile(r"^S_(\d+)_D_(\d+)_(\d+)$")
+SYNTHETIC_DOMAINS = [
+    [0, 3, 4],
+    [0, 2, 4],
+    [3, 4],
+    [4],
+    [4],
+    [0],
+    [0, 1, 4],
+    [0, 2, 3],
+    [0, 1, 2],
+]
+
+
+def parse_file_id(file_id):
+    """Return zero-based species and domain indices from an official file ID."""
+    match = FILE_ID.fullmatch(file_id)
+    if not match:
+        raise ValueError(f"invalid BioDCASE file ID: {file_id}")
+    species, domain, _ = map(int, match.groups())
+    if not (1 <= species <= len(SPECIES) and 1 <= domain <= len(DOMAINS)):
+        raise ValueError(f"out-of-range BioDCASE file ID: {file_id}")
+    return species - 1, domain - 1
+
+
+def held_domain_split(ids, fold, seed=42, singleton_fraction=0.1):
+    """Hold out a rare domain, or rows when only one domain exists, per species."""
+    groups = {}
+    for file_id in ids:
+        groups.setdefault(parse_file_id(file_id), []).append(file_id)
+    held_domains = {}
+    validation_ids = []
+    for species, species_name in enumerate(SPECIES):
+        available = [domain for domain in range(len(DOMAINS)) if groups.get((species, domain))]
+        candidates = sorted(available)
+        if len(candidates) == 1:
+            held_domains[species] = None
+            rows = groups[species, candidates[0]].copy()
+            random.Random(seed + species).shuffle(rows)
+            validation_size = min(
+                len(rows) - 1,
+                max(1, round(singleton_fraction * len(rows))),
+            )
+            validation_ids.extend(rows[:validation_size])
+            continue
+        if fold >= len(candidates):
+            raise ValueError(
+                f"fold {fold} unavailable for {species_name}: "
+                f"only {len(candidates)} domains"
+            )
+        held_domains[species] = candidates[fold]
+        validation_ids.extend(groups[species, candidates[fold]])
+    validation_set = set(validation_ids)
+    training_ids = [file_id for file_id in ids if file_id not in validation_set]
+    training_species = {parse_file_id(file_id)[0] for file_id in training_ids}
+    validation_species = {parse_file_id(file_id)[0] for file_id in validation_ids}
+    expected_species = set(range(len(SPECIES)))
+    if training_species != expected_species or validation_species != expected_species:
+        raise ValueError("held-domain split does not cover every species")
+    return training_ids, validation_ids, held_domains
+
+
+def load_ids(split="Training"):
+    """Download and read one official ID list."""
+    path = hf_hub_download(DATASET, f"data/metadata/{split}_ids.txt", repo_type="dataset")
+    with open(path, encoding="utf-8") as handle:
+        return [line.strip() for line in handle if line.strip()]
+
+
+def load_unseen_domains():
+    """Return the official unseen domain index for each species index."""
+    path = hf_hub_download(DATASET, "data/metadata/split_summary.json", repo_type="dataset")
+    mapping = json.loads(Path(path).read_text(encoding="utf-8"))["unseen_domain_by_species"]
+    return {SPECIES.index(species): DOMAINS.index(domain) for species, domain in mapping.items()}
+
+
+def download(url, path):
+    """Download a large file atomically with coarse progress output."""
+    partial = path.with_suffix(path.suffix + ".part")
+    with urllib.request.urlopen(url) as response, open(partial, "wb") as output:
+        total = int(response.headers.get("Content-Length", 0))
+        downloaded = 0
+        report_at = 0
+        while chunk := response.read(1 << 20):
+            output.write(chunk)
+            downloaded += len(chunk)
+            if downloaded >= report_at:
+                print(
+                    f"downloaded {downloaded / 1e9:.1f}/{total / 1e9:.1f} GB",
+                    flush=True,
+                )
+                report_at += 250_000_000
+    partial.rename(path)
+
+
+def prepare_data(data):
+    """Download and flatten the official Zenodo development archive."""
+    destination = Path(data)
+    marker = destination / ".complete"
+    if marker.exists():
+        print(f"audio already prepared at {destination}")
+        return
+    destination.mkdir(parents=True, exist_ok=True)
+    archive = destination.parent / "Development_data.zip"
+    if not archive.exists():
+        download(ZENODO, archive)
+    with zipfile.ZipFile(archive) as source:
+        members = [
+            name
+            for name in source.namelist()
+            if name.lower().endswith(".wav")
+            and not name.startswith("__MACOSX")
+            and not Path(name).name.startswith("._")
+        ]
+        print(f"extracting {len(members)} waveforms", flush=True)
+        for index, name in enumerate(members, 1):
+            with source.open(name) as src, open(
+                destination / Path(name).name, "wb"
+            ) as dst:
+                shutil.copyfileobj(src, dst)
+            if index % 25_000 == 0:
+                print(f"extracted {index}/{len(members)}", flush=True)
+    marker.write_text(str(len(members)), encoding="utf-8")
+    archive.unlink()
+    print(f"prepared {len(members)} waveforms at {destination}")
+
+
+def select_ids(ids, size, seed):
+    """Select a seeded species-domain-stratified subset."""
+    if not size or size >= len(ids):
+        return ids
+    groups = {}
+    for file_id in ids:
+        groups.setdefault(parse_file_id(file_id), []).append(file_id)
+    rng = random.Random(seed)
+    for rows in groups.values():
+        rng.shuffle(rows)
+    selected = []
+    while len(selected) < size:
+        added = False
+        for key in sorted(groups):
+            if groups[key] and len(selected) < size:
+                selected.append(groups[key].pop())
+                added = True
+        if not added:
+            break
+    rng.shuffle(selected)
+    return selected
+
+
+def cpus():
+    """Return the CPU quota visible to this process."""
+    try:
+        quota, period = Path("/sys/fs/cgroup/cpu.max").read_text(
+            encoding="utf-8"
+        ).split()
+        if quota != "max":
+            return max(1, int(quota) // int(period))
+    except OSError:
+        pass
+    return len(os.sched_getaffinity(0))
+
+
+def load_waveform(path, seconds=None, training=False):
+    """Read a mono waveform at the model sample rate, or only a random or central crop."""
+    start, frames = 0, -1
+    if seconds is not None:
+        info = soundfile.info(path)
+        frames = math.ceil(seconds * info.samplerate)
+        gap = max(0, info.frames - frames)
+        start = random.randrange(gap + 1) if training else gap // 2
+    waveform, sample_rate = soundfile.read(
+        path, start=start, frames=frames, dtype="float32", always_2d=False
+    )
+    if waveform.ndim > 1:
+        waveform = waveform.mean(axis=1)
+    if sample_rate != SAMPLE_RATE:
+        waveform = resample_poly(waveform, SAMPLE_RATE, sample_rate).astype(np.float32)
+    return waveform
+
+
+def normalize_level(waveforms, level_rms):
+    """Scale each waveform to one RMS level so recording gain cannot reveal its domain."""
+    if level_rms <= 0:
+        return waveforms.astype(np.float32)
+    rms = np.sqrt(np.mean(np.square(waveforms), axis=-1, keepdims=True))
+    return (waveforms * (level_rms / np.maximum(rms, 1e-6))).astype(np.float32)
+
+
+def split_windows(waveform, samples):
+    """Cover a complete clip with fixed-length windows, the last aligned to its end."""
+    starts = list(range(0, max(len(waveform) - samples + 1, 1), samples))
+    if starts[-1] + samples < len(waveform):
+        starts.append(len(waveform) - samples)
+    return np.stack(
+        [fixed_length(waveform[start : start + samples], samples, False) for start in starts]
+    )
+
+
+def fixed_length(waveform, samples, training):
+    """Randomly or centrally crop and zero-pad a waveform."""
+    if len(waveform) < samples:
+        gap = samples - len(waveform)
+        left = random.randrange(gap + 1) if training else gap // 2
+        waveform = np.pad(waveform, (left, gap - left))
+    if len(waveform) > samples:
+        gap = len(waveform) - samples
+        start = random.randrange(gap + 1) if training else gap // 2
+        waveform = waveform[start : start + samples]
+    return waveform.astype(np.float32)
+
+
+class MosquitoDataset(Dataset):
+    """Fold-filtered 8 kHz audio with species-domain sampling weights."""
+
+    def __init__(
+        self, root, ids, seconds, training, balance_power, synthetic=False, level_rms=LEVEL_RMS
+    ):
+        self.root = Path(root) if root else None
+        self.ids = ids
+        self.samples = round(seconds * SAMPLE_RATE)
+        self.training = training
+        self.synthetic = synthetic
+        self.level_rms = level_rms
+        self.labels = [parse_file_id(file_id) for file_id in ids]
+        counts = Counter(self.labels)
+        self.sample_weights = [counts[label] ** -balance_power for label in self.labels]
+
+    def __len__(self):
+        return len(self.ids)
+
+    def _synthetic_waveform(self, index):
+        species, domain = self.labels[index]
+        row = int(self.ids[index].rsplit("_", maxsplit=1)[1])
+        time = np.arange(self.samples, dtype=np.float32) / SAMPLE_RATE
+        phase = (row % 7) * math.pi / 7
+        return 0.2 * np.sin(2 * np.pi * (180 + 35 * species + domain) * time + phase)
+
+    def __getitem__(self, index):
+        if self.synthetic:
+            waveform = self._synthetic_waveform(index)
+        else:
+            waveform = load_waveform(
+                self.root / f"{self.ids[index]}.wav",
+                self.samples / SAMPLE_RATE,
+                self.training,
+            )
+            waveform = fixed_length(waveform, self.samples, self.training)
+        waveform = normalize_level(waveform, self.level_rms)
+        if self.training:
+            waveform = np.roll(waveform, random.randrange(len(waveform)))
+            waveform = waveform * 10 ** random.uniform(-0.3, 0.3)
+            waveform = waveform + np.random.normal(
+                0, random.uniform(0, 0.005), len(waveform)
+            )
+        species, domain = self.labels[index]
+        return {
+            "file_id": self.ids[index],
+            "waveform": waveform.astype(np.float32),
+            "species": species,
+            "domain": domain,
+        }
+
+
+class AudioCollator:
+    """Stack fixed-length waveforms and labels."""
+
+    def __call__(self, rows):
+        waveforms = torch.from_numpy(np.stack([row["waveform"] for row in rows]))
+        return {
+            "file_ids": [row["file_id"] for row in rows],
+            "waveforms": waveforms,
+            "lengths": torch.full((len(rows),), waveforms.size(1), dtype=torch.long),
+            "species": torch.tensor([row["species"] for row in rows]),
+            "domains": torch.tensor([row["domain"] for row in rows]),
+        }
+
+
+class ClipDataset(MosquitoDataset):
+    """Complete clips split into training-length evaluation windows."""
+
+    def __getitem__(self, index):
+        if self.synthetic:
+            waveform = self._synthetic_waveform(index)
+        else:
+            waveform = load_waveform(self.root / f"{self.ids[index]}.wav")
+        species, domain = self.labels[index]
+        return {
+            "windows": normalize_level(split_windows(waveform, self.samples), self.level_rms),
+            "species": species,
+            "domain": domain,
+        }
+
+
+class WindowCollator:
+    """Concatenate every clip's windows and record which clip owns each."""
+
+    def __call__(self, rows):
+        counts = torch.tensor([len(row["windows"]) for row in rows])
+        waveforms = torch.from_numpy(np.concatenate([row["windows"] for row in rows]))
+        return {
+            "waveforms": waveforms,
+            "lengths": torch.full((len(waveforms),), waveforms.size(1), dtype=torch.long),
+            "owners": torch.repeat_interleave(torch.arange(len(rows)), counts),
+            "species": torch.tensor([row["species"] for row in rows]),
+            "domains": torch.tensor([row["domain"] for row in rows]),
+        }
+
+
+def synthetic_ids(rows_per_cell):
+    """Build official-shaped IDs with the real species-domain availability pattern."""
+    return [
+        f"S_{species + 1}_D_{domain + 1}_{row + 1}"
+        for species, domains in enumerate(SYNTHETIC_DOMAINS)
+        for domain in domains
+        for row in range(rows_per_cell)
+    ]
+
+
+def make_datasets(args):
+    """Create one pseudo-unseen fold, or for fold -1 all training rows and the official
+    validation split."""
+    ids = synthetic_ids(args.synthetic_rows) if args.synthetic else load_ids()
+    if args.fold < 0:
+        training_ids = ids
+        validation_ids = ids if args.synthetic else load_ids("Validation")
+        held_domains = dict.fromkeys(range(len(SPECIES)))
+    else:
+        training_ids, validation_ids, held_domains = held_domain_split(
+            ids, args.fold, seed=args.seed
+        )
+    training_ids = select_ids(training_ids, args.train_size, args.seed)
+    validation_ids = select_ids(validation_ids, args.validation_size, args.seed)
+    training = MosquitoDataset(
+        args.data,
+        training_ids,
+        args.seconds,
+        True,
+        args.balance_power,
+        args.synthetic,
+        args.level_rms,
+    )
+    validation = MosquitoDataset(
+        args.data,
+        validation_ids,
+        args.seconds,
+        False,
+        args.balance_power,
+        args.synthetic,
+        args.level_rms,
+    )
+    expected = set(range(len(SPECIES)))
+    if {label[0] for label in training.labels} != expected:
+        raise ValueError("training size limit removed a species")
+    if {label[0] for label in validation.labels} != expected:
+        raise ValueError("validation size limit removed a species")
+    return training, validation, held_domains
+
+
+def seed_worker(_worker_id):
+    """Seed NumPy and Python from a DataLoader worker's PyTorch seed."""
+    worker_seed = torch.initial_seed() % (2**32)
+    random.seed(worker_seed)
+    np.random.seed(worker_seed)
+
+
+def make_loader(dataset, batch_size, workers, training, seed, collator=None):
+    """Create a deterministic loader, balanced by species-domain cell for training."""
+    generator = torch.Generator().manual_seed(seed)
+    sampler = None
+    if training:
+        sampler = WeightedRandomSampler(
+            dataset.sample_weights,
+            num_samples=len(dataset),
+            replacement=True,
+            generator=generator,
+        )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        shuffle=False,
+        num_workers=workers,
+        pin_memory=torch.cuda.is_available(),
+        persistent_workers=workers > 0,
+        worker_init_fn=seed_worker,
+        generator=generator,
+        collate_fn=collator or AudioCollator(),
+    )
+
+
+def deterministic_view(dataset, ids=None):
+    """Return an unaugmented view over a dataset's IDs."""
+    return MosquitoDataset(
+        dataset.root,
+        dataset.ids if ids is None else ids,
+        dataset.samples / SAMPLE_RATE,
+        training=False,
+        balance_power=0.0,
+        synthetic=dataset.synthetic,
+        level_rms=dataset.level_rms,
+    )
+
+
+@torch.no_grad()
+def compute_normalization(dataset, size, batch_size, workers, device, seed, min_hz):
+    """Estimate per-mel mean and variance from fold-training rows only."""
+    ids = select_ids(dataset.ids, size, seed)
+    loader = make_loader(
+        deterministic_view(dataset, ids), batch_size, workers, False, seed
+    )
+    frontend = LogMelFrontend(min_hz=min_hz).to(device)
+    total = torch.zeros(N_MELS, dtype=torch.float64, device=device)
+    squared = torch.zeros_like(total)
+    frames = 0
+    for batch in loader:
+        waveforms = batch["waveforms"].to(device, non_blocking=True)
+        lengths = batch["lengths"].to(device, non_blocking=True)
+        features, frame_lengths = frontend(waveforms, lengths)
+        positions = torch.arange(features.size(1), device=device)[None, :]
+        mask = positions < frame_lengths[:, None]
+        selected = features[mask].double()
+        total += selected.sum(0)
+        squared += selected.square().sum(0)
+        frames += selected.size(0)
+    mean = total / frames
+    variance = squared / frames - mean.square()
+    return mean.float().cpu(), variance.clamp_min(1e-6).sqrt().float().cpu()
+
+
+def balanced_accuracy(gold, predicted):
+    """Return mean per-species recall over labels present in gold."""
+    gold = np.asarray(gold)
+    predicted = np.asarray(predicted)
+    recalls = [
+        np.mean(predicted[gold == label] == label)
+        for label in sorted(set(gold.tolist()))
+    ]
+    return float(np.mean(recalls)) if recalls else 0.0
+
+
+def prediction_metrics(labels, predictions):
+    """Return balanced accuracy, per-species recall and a gold-by-predicted confusion."""
+    labels = labels.cpu().numpy()
+    predictions = predictions.cpu().numpy()
+    confusion = np.bincount(
+        labels * len(SPECIES) + predictions, minlength=len(SPECIES) ** 2
+    ).reshape(len(SPECIES), len(SPECIES))
+    return {
+        "balanced_accuracy": balanced_accuracy(labels, predictions),
+        "rows": len(labels),
+        "per_species": {
+            SPECIES[label]: float(np.mean(predictions[labels == label] == label))
+            for label in sorted(set(labels.tolist()))
+        },
+        "confusion": confusion.tolist(),
+    }
+
+
+def official_metrics(labels, predictions, domains, unseen_domains):
+    """Score the BioDCASE seen and unseen species-domain partitions."""
+    unseen = torch.tensor(
+        [
+            unseen_domains[label.item()] == domain.item()
+            for label, domain in zip(labels, domains)
+        ],
+        dtype=torch.bool,
+    )
+    return {
+        "seen": prediction_metrics(labels[~unseen], predictions[~unseen]),
+        "unseen": prediction_metrics(labels[unseen], predictions[unseen]),
+    }
+
+
+def fold_prediction_metrics(labels, predictions, domains, held_domains):
+    """Split fold metrics into true domain holdouts and singleton reserves."""
+    metrics = prediction_metrics(labels, predictions)
+    domain_holdout = torch.tensor(
+        [held_domains[label.item()] is not None for label in labels],
+        dtype=torch.bool,
+    )
+    metrics["domain_holdout"] = prediction_metrics(
+        labels[domain_holdout], predictions[domain_holdout]
+    )
+    metrics["same_domain"] = prediction_metrics(
+        labels[~domain_holdout], predictions[~domain_holdout]
+    )
+    expected_domains = torch.tensor(
+        [
+            held_domains[label.item()]
+            if held_domains[label.item()] is not None
+            else domain.item()
+            for label, domain in zip(labels, domains)
+        ]
+    )
+    if not torch.equal(domains, expected_domains):
+        raise ValueError("validation rows do not match the declared fold policy")
+    return metrics
+
+
+@torch.no_grad()
+def collect_outputs(model, frontend, dataset, batch_size, workers, device, seed):
+    """Collect deterministic logits and embeddings for one dataset."""
+    model.eval()
+    loader = make_loader(dataset, batch_size, workers, False, seed)
+    collected = {"labels": [], "domains": [], "logits": [], "embeddings": []}
+    for batch in loader:
+        waveforms = batch["waveforms"].to(device, non_blocking=True)
+        lengths = batch["lengths"].to(device, non_blocking=True)
+        domains = batch["domains"].to(device, non_blocking=True)
+        features, frame_lengths = frontend(waveforms, lengths)
+        outputs = model(features, frame_lengths, domains)
+        collected["labels"].append(batch["species"])
+        collected["domains"].append(batch["domains"])
+        collected["logits"].append(outputs["species_logits"].cpu())
+        collected["embeddings"].append(outputs["embedding"].cpu())
+    return {name: torch.cat(values) for name, values in collected.items()}
+
+
+@torch.no_grad()
+def collect_clip_outputs(model, frontend, dataset, batch_size, workers, device, seed):
+    """Average window probabilities and embeddings over each complete clip."""
+    model.eval()
+    loader = make_loader(dataset, batch_size, workers, False, seed, WindowCollator())
+    collected = {"labels": [], "domains": [], "probabilities": [], "embeddings": []}
+    for batch in loader:
+        owners = batch["owners"].to(device, non_blocking=True)
+        domains = batch["domains"].to(device, non_blocking=True)
+        window_outputs = {"probabilities": [], "embeddings": []}
+        for start in range(0, len(owners), batch_size):
+            window = slice(start, start + batch_size)
+            features, frame_lengths = frontend(
+                batch["waveforms"][window].to(device, non_blocking=True),
+                batch["lengths"][window].to(device, non_blocking=True),
+            )
+            outputs = model(features, frame_lengths, domains[owners[window]])
+            window_outputs["probabilities"].append(
+                functional.softmax(outputs["species_logits"], dim=1)
+            )
+            window_outputs["embeddings"].append(outputs["embedding"])
+        counts = torch.bincount(owners, minlength=len(domains)).unsqueeze(1)
+        for name, values in window_outputs.items():
+            values = torch.cat(values)
+            totals = values.new_zeros(len(domains), values.size(1))
+            collected[name].append((totals.index_add_(0, owners, values) / counts).cpu())
+        collected["labels"].append(batch["species"])
+        collected["domains"].append(batch["domains"])
+    return {name: torch.cat(values) for name, values in collected.items()}
+
+
+def mahalanobis_fit(reference_embeddings, reference_labels):
+    """Return class means and the precision of one regularized shared covariance matrix."""
+    if set(reference_labels.tolist()) != set(range(len(SPECIES))):
+        raise ValueError("Mahalanobis reference data must cover every species")
+    means = torch.stack(
+        [
+            reference_embeddings[reference_labels == label].mean(0)
+            for label in range(len(SPECIES))
+        ]
+    )
+    residuals = reference_embeddings - means[reference_labels]
+    covariance = residuals.T @ residuals / max(1, len(residuals) - len(SPECIES))
+    regularization = covariance.diagonal().mean().clamp_min(1e-6) * 1e-3
+    covariance = covariance + regularization * torch.eye(covariance.size(0))
+    return means, torch.linalg.pinv(covariance)
+
+
+def mahalanobis_predict(means, precision, embeddings):
+    """Assign each embedding to the nearest class mean in Mahalanobis distance."""
+    differences = embeddings[:, None, :] - means[None, :, :]
+    distances = torch.einsum("ncd,de,nce->nc", differences, precision, differences)
+    return distances.argmin(1)
+
+
+def sampled_priors(dataset, device):
+    """Return species priors under the balanced training sampler for logit adjustment."""
+    totals = torch.zeros(len(SPECIES), dtype=torch.float64).index_add_(
+        0,
+        torch.tensor([species for species, _domain in dataset.labels]),
+        torch.tensor(dataset.sample_weights, dtype=torch.float64),
+    )
+    return (totals / totals.sum()).float().to(device)
+
+
+def grl_strength(step, total_steps):
+    """Return the standard sigmoid gradient-reversal warmup."""
+    progress = step / max(1, total_steps - 1)
+    return 2 / (1 + math.exp(-10 * progress)) - 1
+
+
+def learning_rate_schedule(total_steps, warmup_ratio):
+    """Return a linear-warmup cosine-decay multiplier."""
+    warmup_steps = max(1, round(total_steps * warmup_ratio))
+
+    def multiplier(step):
+        if step < warmup_steps:
+            return (step + 1) / warmup_steps
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))
+
+    return multiplier
+
+
+def train_epoch(model, frontend, loader, optimizer, scheduler, priors, epoch, state):
+    """Run one balanced training epoch or stop at the configured step limit."""
+    model.train()
+    totals = {}
+    batches = 0
+    for batch in loader:
+        if 0 < state["max_steps"] <= state["step"]:
+            break
+        waveforms = batch["waveforms"].to(state["device"], non_blocking=True)
+        lengths = batch["lengths"].to(state["device"], non_blocking=True)
+        species = batch["species"].to(state["device"], non_blocking=True)
+        domains = batch["domains"].to(state["device"], non_blocking=True)
+        with torch.no_grad():
+            features, frame_lengths = frontend(waveforms, lengths)
+        optimizer.zero_grad(set_to_none=True)
+        outputs = model(
+            features,
+            frame_lengths,
+            domains,
+            domain_strength=grl_strength(state["step"], state["total_steps"]),
+            augment=True,
+        )
+        losses = training_loss(outputs, species, domains, priors, epoch + 1)
+        losses["loss"].backward()
+        nn.utils.clip_grad_norm_(model.parameters(), state["gradient_clip"])
+        optimizer.step()
+        scheduler.step()
+        state["step"] += 1
+        batches += 1
+        for name, value in losses.items():
+            totals[name] = totals.get(name, 0.0) + value.detach().item()
+    if not batches:
+        return None
+    return {name: value / batches for name, value in totals.items()}
+
+
+def checkpoint_state(model, frontend, optimizer, scheduler, epoch, state, metrics, args):
+    """Build a resumable best-checkpoint payload."""
+    return {
+        "model": model.state_dict(),
+        "frontend": frontend.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scheduler": scheduler.state_dict(),
+        "epoch": epoch,
+        "step": state["step"],
+        "metrics": metrics,
+        "arguments": {key: value for key, value in vars(args).items() if key != "run"},
+    }
+
+
+def final_metrics(model, frontend, training, validation, held_domains, args, device):
+    """Score softmax and shared-covariance Mahalanobis inference."""
+    reference_ids = select_ids(training.ids, args.reference_size, args.seed + 1)
+    reference = deterministic_view(training, reference_ids)
+    reference_outputs = collect_outputs(
+        model,
+        frontend,
+        reference,
+        args.eval_batch_size,
+        args.workers,
+        device,
+        args.seed,
+    )
+    validation_outputs = collect_outputs(
+        model,
+        frontend,
+        validation,
+        args.eval_batch_size,
+        args.workers,
+        device,
+        args.seed,
+    )
+    softmax = fold_prediction_metrics(
+        validation_outputs["labels"],
+        validation_outputs["logits"].argmax(1),
+        validation_outputs["domains"],
+        held_domains,
+    )
+    mahalanobis = fold_prediction_metrics(
+        validation_outputs["labels"],
+        mahalanobis_predict(
+            *mahalanobis_fit(reference_outputs["embeddings"], reference_outputs["labels"]),
+            validation_outputs["embeddings"],
+        ),
+        validation_outputs["domains"],
+        held_domains,
+    )
+    return softmax, mahalanobis, reference_outputs
+
+
+def train(args):
+    """Train and select MTRCNN-DG on one training-only pseudo-unseen fold."""
+    set_seed(args.seed)
+    if not args.synthetic and not (Path(args.data) / ".complete").exists():
+        prepare_data(args.data)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    training, validation, held_domains = make_datasets(args)
+    mean, std = compute_normalization(
+        training,
+        args.stats_size,
+        args.eval_batch_size,
+        args.workers,
+        device,
+        args.seed,
+        args.min_hz,
+    )
+    frontend = LogMelFrontend(mean, std, args.min_hz).to(device)
+    model = MTRCNNDG(dropout=args.dropout).to(device)
+    loader = make_loader(
+        training, args.batch_size, args.workers, True, args.seed
+    )
+    natural_steps = args.epochs * len(loader)
+    total_steps = min(natural_steps, args.max_steps) if args.max_steps > 0 else natural_steps
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay
+    )
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer, learning_rate_schedule(total_steps, args.warmup_ratio)
+    )
+    priors = sampled_priors(training, device)
+    output = Path(args.output)
+    output.mkdir(parents=True, exist_ok=True)
+    checkpoint = output / "checkpoint.pt"
+    state = {
+        "device": device,
+        "step": 0,
+        "total_steps": total_steps,
+        "max_steps": args.max_steps,
+        "gradient_clip": args.gradient_clip,
+    }
+    history = []
+    best_score = -1.0
+    best_epoch = 0
+    stale_epochs = 0
+    for epoch in range(args.epochs):
+        losses = train_epoch(
+            model, frontend, loader, optimizer, scheduler, priors, epoch, state
+        )
+        if losses is None:
+            break
+        validation_outputs = collect_outputs(
+            model,
+            frontend,
+            validation,
+            args.eval_batch_size,
+            args.workers,
+            device,
+            args.seed,
+        )
+        metrics = fold_prediction_metrics(
+            validation_outputs["labels"],
+            validation_outputs["logits"].argmax(1),
+            validation_outputs["domains"],
+            held_domains,
+        )
+        record = {"epoch": epoch + 1, "step": state["step"], "losses": losses, **metrics}
+        history.append(record)
+        print(json.dumps(record), flush=True)
+        # Without a held-out domain (fold -1) the last epoch is kept.
+        score = metrics["domain_holdout"]["balanced_accuracy"] if args.fold >= 0 else epoch
+        if score > best_score:
+            best_score = score
+            best_epoch = epoch + 1
+            stale_epochs = 0
+            torch.save(
+                checkpoint_state(
+                    model,
+                    frontend,
+                    optimizer,
+                    scheduler,
+                    epoch + 1,
+                    state,
+                    metrics,
+                    args,
+                ),
+                checkpoint,
+            )
+        else:
+            stale_epochs += 1
+        reached_limit = 0 < args.max_steps <= state["step"]
+        if reached_limit or stale_epochs >= args.patience:
+            break
+    saved = torch.load(checkpoint, map_location=device, weights_only=True)
+    model.load_state_dict(saved["model"])
+    frontend.load_state_dict(saved["frontend"])
+    softmax, mahalanobis, reference = final_metrics(
+        model, frontend, training, validation, held_domains, args, device
+    )
+    result = {
+        "device": str(device),
+        "fold": args.fold,
+        "seed": args.seed,
+        "training_rows": len(training),
+        "validation_rows": len(validation),
+        "normalization_rows": len(select_ids(training.ids, args.stats_size, args.seed)),
+        "reference_rows": len(reference["labels"]),
+        "held_domains": {
+            SPECIES[species]: DOMAINS[domain] if domain is not None else "same-domain"
+            for species, domain in held_domains.items()
+        },
+        "best_epoch": best_epoch,
+        "steps": state["step"],
+        "softmax": softmax,
+        "mahalanobis": mahalanobis,
+        "history": history,
+        "checkpoint": str(checkpoint),
+        "test": test_metrics(model, frontend, reference, args, device) if args.test else None,
+    }
+    (output / "results.json").write_text(
+        json.dumps(result, indent=1), encoding="utf-8"
+    )
+    print(json.dumps(result, indent=1))
+
+
+def test_metrics(model, frontend, reference, args, device):
+    """Score the official development test split with softmax and Mahalanobis inference."""
+    test_ids = (
+        [
+            f"S_{species + 1}_D_{domain + 1}_{row + 1}"
+            for species in range(len(SPECIES))
+            for domain in range(len(DOMAINS))
+            for row in range(args.synthetic_rows)
+        ]
+        if args.synthetic
+        else load_ids("Test")
+    )
+    test = ClipDataset(
+        args.data,
+        select_ids(test_ids, args.test_size, args.seed),
+        args.seconds,
+        False,
+        0.0,
+        args.synthetic,
+        args.level_rms,
+    )
+    outputs = collect_clip_outputs(
+        model, frontend, test, args.eval_batch_size, args.workers, device, args.seed
+    )
+    unseen_domains = load_unseen_domains()
+    mahalanobis = mahalanobis_predict(
+        *mahalanobis_fit(reference["embeddings"], reference["labels"]), outputs["embeddings"]
+    )
+    return {
+        "rows": len(test),
+        "softmax": official_metrics(
+            outputs["labels"],
+            outputs["probabilities"].argmax(1),
+            outputs["domains"],
+            unseen_domains,
+        ),
+        "mahalanobis": official_metrics(
+            outputs["labels"], mahalanobis, outputs["domains"], unseen_domains
+        ),
+    }
+
+
+def evaluate(args):
+    """Score a checkpoint on its validation rows and on the official development test split."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    saved = torch.load(args.checkpoint, map_location=device, weights_only=True)
+    options = argparse.Namespace(**saved["arguments"])
+    options.data = args.data
+    options.eval_batch_size = args.eval_batch_size
+    options.workers = args.workers
+    options.test_size = args.test_size
+    set_seed(options.seed)
+    if not options.synthetic and not (Path(args.data) / ".complete").exists():
+        prepare_data(args.data)
+    training, validation, held_domains = make_datasets(options)
+    frontend = LogMelFrontend(min_hz=options.min_hz).to(device)
+    frontend.load_state_dict(saved["frontend"])
+    model = MTRCNNDG(dropout=options.dropout).to(device)
+    model.load_state_dict(saved["model"])
+    fold_softmax, fold_mahalanobis, reference = final_metrics(
+        model, frontend, training, validation, held_domains, options, device
+    )
+    result = {
+        "checkpoint": args.checkpoint,
+        "epoch": saved["epoch"],
+        "fold": options.fold,
+        "seed": options.seed,
+        "arguments": saved["arguments"],
+        "test": test_metrics(model, frontend, reference, options, device),
+        "fold_softmax": fold_softmax,
+        "fold_mahalanobis": fold_mahalanobis,
+    }
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(result, indent=1), encoding="utf-8")
+    print(json.dumps(result, indent=1))
+    if args.export:
+        export(model, frontend, reference, saved["arguments"], result, Path(args.export))
+        if args.repo:
+            HfApi().upload_folder(
+                repo_id=args.repo,
+                folder_path=args.export,
+                commit_message="Add MTRCNN-DG weights and evaluation",
+            )
+
+
+def export(model, frontend, reference, arguments, result, folder):
+    """Write weights, Mahalanobis statistics, this script and the evaluation to folder."""
+    means, precision = mahalanobis_fit(reference["embeddings"], reference["labels"])
+    (folder / "results").mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "frontend": frontend.state_dict(),
+            "arguments": arguments,
+            "means": means,
+            "precision": precision,
+        },
+        folder / "model.pt",
+    )
+    shutil.copy(__file__, folder / "mtrcnn.py")
+    (folder / "results" / "eval.json").write_text(json.dumps(result, indent=1), encoding="utf-8")
+
+
+def predict(args):
+    """Print the Mahalanobis species prediction for each WAV file."""
+    path = Path(args.model)
+    if not path.is_file():
+        path = hf_hub_download(args.model, "model.pt")
+    saved = torch.load(path, map_location="cpu", weights_only=True)
+    options = saved["arguments"]
+    frontend = LogMelFrontend(min_hz=options["min_hz"])
+    frontend.load_state_dict(saved["frontend"])
+    model = MTRCNNDG().eval()
+    model.load_state_dict(saved["model"])
+    samples = round(options["seconds"] * SAMPLE_RATE)
+    for file in args.files:
+        windows = split_windows(load_waveform(file), samples)
+        waveforms = torch.from_numpy(normalize_level(windows, options["level_rms"]))
+        with torch.no_grad():
+            features, lengths = frontend(waveforms, torch.full((len(waveforms),), samples))
+            embedding = model(features, lengths, torch.zeros(len(waveforms), dtype=torch.long))[
+                "embedding"
+            ].mean(0, keepdim=True)
+        species = mahalanobis_predict(saved["means"], saved["precision"], embedding).item()
+        print(json.dumps({"file": file, "species": SPECIES[species]}))
+
+
+def card(args):
+    """Render card.jinja into card/README.md with the results stored in the model repo."""
+    try:
+        folder = Path(snapshot_download(args.repo, allow_patterns="results/*.json"))
+        paths = folder.glob("results/*.json")
+    except RepositoryNotFoundError:
+        paths = []
+    results = {path.stem: json.loads(path.read_text(encoding="utf-8")) for path in paths}
+    runs = {}
+    for name, result in sorted(results.items()):
+        if name.startswith("train-"):
+            runs.setdefault(name.split("-")[1], []).append(result["test"])
+    seeds = {
+        config: {
+            inference: {
+                partition: [run[inference][partition]["balanced_accuracy"] * 100 for run in tests]
+                for partition in ("seen", "unseen")
+            }
+            for inference in ("softmax", "mahalanobis")
+        }
+        for config, tests in runs.items()
+    }
+    evaluation = results.get("eval")
+    data = ModelCardData(
+        model_name=args.repo.split("/")[1],
+        datasets=[DATASET],
+        license="cc-by-4.0",
+        tags=["audio-classification", "bioacoustics", "mosquito", "domain-generalization"],
+        eval_results=[
+            EvalResult(
+                task_type="audio-classification",
+                dataset_type=DATASET,
+                dataset_name="BioDCASE 2026 Task 5 development test, unseen domains",
+                metric_type="balanced_accuracy",
+                metric_name="Unseen-domain balanced accuracy",
+                metric_value=round(
+                    evaluation["test"]["mahalanobis"]["unseen"]["balanced_accuracy"] * 100, 2
+                ),
+            )
+        ]
+        if evaluation
+        else None,
+    )
+    rendered = ModelCard.from_template(
+        data,
+        template_path=Path(__file__).parent / "card.jinja",
+        repo=args.repo,
+        dataset=DATASET,
+        species=SPECIES,
+        evaluation=evaluation,
+        seeds=seeds,
+        mean=statistics.mean,
+        stdev=statistics.stdev,
+    )
+    output = Path(__file__).parent / "card"
+    output.mkdir(exist_ok=True)
+    rendered.save(output / "README.md")
+
+
+def set_seed(seed):
+    """Seed Python, NumPy and PyTorch."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def mel_filter_bank(min_hz):
+    """Return a triangular Slaney-style mel filter bank from min_hz to Nyquist."""
+    frequencies = torch.linspace(0, SAMPLE_RATE / 2, N_FFT // 2 + 1)
+    mel_min = 2595 * math.log10(1 + min_hz / 700)
+    mel_max = 2595 * math.log10(1 + (SAMPLE_RATE / 2) / 700)
+    mel_points = torch.linspace(mel_min, mel_max, N_MELS + 2)
+    hz_points = 700 * (torch.pow(10, mel_points / 2595) - 1)
+    lower = (frequencies[:, None] - hz_points[:-2]) / (
+        hz_points[1:-1] - hz_points[:-2]
+    )
+    upper = (hz_points[2:] - frequencies[:, None]) / (
+        hz_points[2:] - hz_points[1:-1]
+    )
+    return torch.minimum(lower, upper).clamp_min(0)
+
+
+class LogMelFrontend(nn.Module):
+    """Fixed 8 kHz, 64-bin log-mel frontend with global normalization."""
+
+    def __init__(self, mean=None, std=None, min_hz=MIN_HZ):
+        super().__init__()
+        self.register_buffer("window", torch.hann_window(N_FFT), persistent=False)
+        self.register_buffer("mel_filters", mel_filter_bank(min_hz), persistent=False)
+        self.register_buffer(
+            "mean",
+            torch.zeros(N_MELS) if mean is None else torch.as_tensor(mean).float(),
+        )
+        self.register_buffer(
+            "std",
+            torch.ones(N_MELS) if std is None else torch.as_tensor(std).float(),
+        )
+
+    def forward(self, waveforms, sample_lengths):
+        """Convert padded waveforms to normalized log-mel frames."""
+        spectrum = torch.stft(
+            waveforms,
+            n_fft=N_FFT,
+            hop_length=HOP_LENGTH,
+            win_length=N_FFT,
+            window=self.window,
+            center=True,
+            pad_mode="reflect",
+            return_complex=True,
+        )
+        power = spectrum.abs().square().transpose(1, 2)
+        features = 10 * torch.log10((power @ self.mel_filters).clamp_min(1e-10))
+        features = (features - self.mean) / self.std.clamp_min(1e-6)
+        frame_lengths = torch.div(sample_lengths, HOP_LENGTH, rounding_mode="floor") + 1
+        return features, frame_lengths
+
+
+def different_domain_partners(domains):
+    """Choose a random different-domain partner for every available sample."""
+    partners = []
+    for index, domain in enumerate(domains):
+        candidates = torch.nonzero(domains != domain, as_tuple=False).flatten()
+        if candidates.numel():
+            choice = torch.randint(candidates.numel(), (), device=domains.device)
+            partners.append(candidates[choice])
+        else:
+            partners.append(torch.tensor(index, device=domains.device))
+    return torch.stack(partners)
+
+
+class FourierMix(nn.Module):
+    """Mix log-mel Fourier amplitudes between recording domains."""
+
+    def __init__(self, probability=0.5):
+        super().__init__()
+        self.probability = probability
+
+    def forward(self, features, domains):
+        """Optionally replace each sample's amplitude style across domains."""
+        if not self.training or torch.rand((), device=features.device) >= self.probability:
+            return features
+        partners = different_domain_partners(domains)
+        spectrum = torch.fft.rfft2(features, dim=(-2, -1), norm="ortho")
+        amplitude = spectrum.abs()
+        phase = spectrum / amplitude.clamp_min(1e-8)
+        weight = torch.rand(features.size(0), 1, 1, device=features.device)
+        mixed_amplitude = (1 - weight) * amplitude + weight * amplitude[partners]
+        return torch.fft.irfft2(
+            mixed_amplitude * phase,
+            s=features.shape[-2:],
+            dim=(-2, -1),
+            norm="ortho",
+        )
+
+
+class MixStyle(nn.Module):
+    """Mix channel statistics between domains after a convolutional stage."""
+
+    def __init__(self, probability=0.5, alpha=0.1):
+        super().__init__()
+        self.probability = probability
+        self.beta = torch.distributions.Beta(alpha, alpha)
+
+    def forward(self, features, domains):
+        """Optionally mix feature statistics across recording domains."""
+        if not self.training or torch.rand((), device=features.device) >= self.probability:
+            return features
+        mean = features.mean(dim=(2, 3), keepdim=True).detach()
+        std = (features.var(dim=(2, 3), keepdim=True, unbiased=False) + 1e-6).sqrt().detach()
+        normalized = (features - mean) / std
+        partners = different_domain_partners(domains)
+        weight = self.beta.sample((features.size(0), 1, 1, 1)).to(features.device)
+        mixed_mean = weight * mean + (1 - weight) * mean[partners]
+        mixed_std = weight * std + (1 - weight) * std[partners]
+        return normalized * mixed_std + mixed_mean
+
+
+class ConvStage(nn.Module):
+    """One convolution, normalization, pooling and dropout stage."""
+
+    def __init__(self, in_channels, out_channels, kernel, dilation, padding, dropout):
+        super().__init__()
+        self.kernel = kernel
+        self.dilation = dilation
+        self.padding = padding
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel,
+            dilation=dilation,
+            padding=padding,
+            bias=False,
+        )
+        self.batch_norm = nn.BatchNorm2d(out_channels)
+        self.pool = nn.AvgPool2d(2)
+        self.dropout = nn.Dropout2d(dropout)
+
+    def forward(self, features):
+        """Transform and downsample one feature map."""
+        features = functional.relu(self.batch_norm(self.conv(features)), inplace=True)
+        return self.dropout(self.pool(features))
+
+    def output_lengths(self, lengths):
+        """Propagate valid time lengths through convolution and pooling."""
+        convolved = lengths + 2 * self.padding[0] - self.dilation[0] * (
+            self.kernel[0] - 1
+        )
+        return torch.div(convolved.clamp_min(0), 2, rounding_mode="floor")
+
+
+def masked_mean_max(features, lengths):
+    """Pool valid time frames by adding their mean and maximum."""
+    lengths = lengths.clamp(min=1, max=features.size(2))
+    positions = torch.arange(features.size(2), device=features.device).view(1, 1, -1, 1)
+    mask = positions < lengths.view(-1, 1, 1, 1)
+    mean = (features * mask).sum(2) / lengths.view(-1, 1, 1)
+    maximum = features.masked_fill(~mask, float("-inf")).max(2).values
+    return mean + torch.where(torch.isfinite(maximum), maximum, torch.zeros_like(maximum))
+
+
+class MTRCNNBranch(nn.Module):
+    """One temporal-resolution branch of MTRCNN."""
+
+    def __init__(self, stage_specs, dropout):
+        super().__init__()
+        channels = [1, 16, 32, 64]
+        self.stages = nn.ModuleList(
+            ConvStage(channels[index], channels[index + 1], *specification, dropout)
+            for index, specification in enumerate(stage_specs)
+        )
+        self.mix_style = MixStyle()
+        self.frequency_projection = nn.Linear(self._frequency_bins(), 1)
+
+    def _frequency_bins(self):
+        with torch.no_grad():
+            features = torch.zeros(1, 1, 128, N_MELS)
+            for stage in self.stages:
+                features = stage(features)
+        return features.shape[-1]
+
+    def forward(self, features, lengths, domains, augment):
+        """Encode one temporal-resolution branch."""
+        for index, stage in enumerate(self.stages):
+            features = stage(features)
+            lengths = stage.output_lengths(lengths)
+            if index == 0 and augment:
+                features = self.mix_style(features, domains)
+        pooled = masked_mean_max(features, lengths)
+        return functional.relu(self.frequency_projection(pooled).squeeze(-1), inplace=True)
+
+
+class GradientReversal(torch.autograd.Function):
+    """Identity forward pass with a sign-reversed backward pass."""
+
+    @staticmethod
+    def forward(ctx, features, strength):
+        """Return features unchanged and retain the reversal strength."""
+        ctx.strength = strength
+        return features.view_as(features)
+
+    @staticmethod
+    def backward(ctx, gradient):
+        """Reverse and scale the upstream embedding gradient."""
+        return -ctx.strength * gradient, None
+
+
+class MTRCNNDG(nn.Module):
+    """MTRCNN with spectral style, adversarial, and conditional domain heads."""
+
+    def __init__(self, dropout=0.2):
+        super().__init__()
+        self.fourier_mix = FourierMix()
+        self.input_batch_norm = nn.BatchNorm2d(N_MELS)
+        specifications = [
+            [((3, 3), (1, 1), (1, 1)), ((3, 3), (2, 1), (2, 0)), ((3, 3), (3, 1), (3, 0))],
+            [((5, 5), (1, 1), (2, 2)), ((5, 5), (2, 1), (4, 1)), ((5, 5), (3, 1), (6, 1))],
+            [((7, 7), (1, 1), (3, 3)), ((7, 7), (2, 1), (6, 2)), ((7, 7), (3, 1), (9, 2))],
+        ]
+        self.branches = nn.ModuleList(MTRCNNBranch(specs, dropout) for specs in specifications)
+        self.embedding = nn.Linear(64 * 3, 32)
+        self.species_classifier = nn.Linear(32, len(SPECIES))
+        self.domain_classifier = nn.Linear(32, len(DOMAINS))
+        self.aux_domain_classifier = nn.Linear(32, len(DOMAINS))
+
+    def forward(self, features, lengths, domains, domain_strength=0.0, augment=False):
+        """Return species, domain, auxiliary, and embedding outputs."""
+        if augment:
+            features = self.fourier_mix(features, domains)
+        features = features.unsqueeze(1).transpose(1, 3)
+        features = self.input_batch_norm(features).transpose(1, 3)
+        branch_outputs = [
+            branch(features, lengths, domains, augment) for branch in self.branches
+        ]
+        embedding = functional.gelu(self.embedding(torch.cat(branch_outputs, dim=1)))
+        reversed_embedding = GradientReversal.apply(embedding, domain_strength)
+        auxiliary = self.aux_domain_classifier
+        return {
+            "embedding": embedding,
+            "species_logits": self.species_classifier(embedding),
+            "domain_logits": self.domain_classifier(reversed_embedding),
+            "aux_domain_logits": auxiliary(embedding),
+            "frozen_aux_domain_logits": functional.linear(
+                embedding,
+                auxiliary.weight.detach(),
+                auxiliary.bias.detach(),
+            ),
+        }
+
+
+def supervised_contrastive(embedding, labels, temperature=0.01):
+    """Pull embeddings of the same species together."""
+    normalized = functional.normalize(embedding, dim=1)
+    similarity = normalized @ normalized.T / temperature
+    identity = torch.eye(len(labels), dtype=torch.bool, device=labels.device)
+    positives = labels[:, None].eq(labels[None, :]) & ~identity
+    similarity = similarity.masked_fill(identity, float("-inf"))
+    log_probability = similarity - torch.logsumexp(similarity, dim=1, keepdim=True)
+    log_probability = torch.where(positives, log_probability, torch.zeros_like(log_probability))
+    counts = positives.sum(1)
+    valid = counts > 0
+    if not valid.any():
+        return embedding.sum() * 0
+    return -(log_probability.sum(1)[valid] / counts[valid]).mean()
+
+
+def conditional_domain_losses(logits, species_labels):
+    """Return class-conditional domain entropy and balance losses."""
+    probabilities = functional.softmax(logits, dim=1).clamp_min(1e-8)
+    entropy_losses = []
+    balance_losses = []
+    for species in species_labels.unique():
+        class_probabilities = probabilities[species_labels == species]
+        entropy_losses.append((class_probabilities * class_probabilities.log()).sum(1).mean())
+        mean_probability = class_probabilities.mean(0)
+        balance_losses.append(
+            (mean_probability * (mean_probability.log() + math.log(len(DOMAINS)))).sum()
+        )
+    return torch.stack(entropy_losses).mean(), torch.stack(balance_losses).mean()
+
+
+def training_loss(outputs, species_labels, domain_labels, priors, epoch):
+    """Compute the full MTRCNN-DG objective and its components."""
+    adjusted_logits = outputs["species_logits"] + priors.clamp_min(1e-8).log()
+    species_loss = functional.cross_entropy(adjusted_logits, species_labels)
+    domain_loss = functional.cross_entropy(outputs["domain_logits"], domain_labels)
+    auxiliary_loss = functional.cross_entropy(outputs["aux_domain_logits"], domain_labels)
+    contrastive_loss = supervised_contrastive(outputs["embedding"], species_labels)
+    ccde_loss, cdb_loss = conditional_domain_losses(
+        outputs["frozen_aux_domain_logits"], species_labels
+    )
+    conditional_weight = 0.05 if epoch >= 5 else 0.0
+    total = (
+        species_loss
+        + domain_loss
+        + auxiliary_loss
+        + 0.5 * contrastive_loss
+        + conditional_weight * (ccde_loss + cdb_loss)
+    )
+    return {
+        "loss": total,
+        "species": species_loss,
+        "domain": domain_loss,
+        "aux_domain": auxiliary_loss,
+        "contrastive": contrastive_loss,
+        "ccde": ccde_loss,
+        "cdb": cdb_loss,
+    }
+
+
+def smoke(args):
+    """Exercise frontend, augmentations, all heads, losses, and backward pass."""
+    set_seed(args.seed)
+    multi_domain_species = {0, 1, 2, 6, 7, 8}
+    fold_ids = [
+        f"S_{species + 1}_D_{domain + 1}_{row + 1}"
+        for species in range(len(SPECIES))
+        for domain in range(2 if species in multi_domain_species else 1)
+        for row in range(2)
+    ]
+    training_ids, validation_ids, held_domains = held_domain_split(
+        fold_ids, 0, seed=args.seed
+    )
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dataset = MosquitoDataset(
+        None,
+        synthetic_ids(4),
+        args.seconds,
+        training=True,
+        balance_power=0.5,
+        synthetic=True,
+    )
+    sampler = WeightedRandomSampler(
+        dataset.sample_weights,
+        num_samples=args.batch_size,
+        replacement=True,
+        generator=torch.Generator().manual_seed(args.seed),
+    )
+    batch = next(
+        iter(
+            DataLoader(
+                dataset,
+                batch_size=args.batch_size,
+                sampler=sampler,
+                collate_fn=AudioCollator(),
+            )
+        )
+    )
+    waveforms = batch["waveforms"].to(device)
+    lengths = batch["lengths"].to(device)
+    labels = batch["species"].to(device)
+    domains = batch["domains"].to(device)
+    mean, std = compute_normalization(
+        dataset,
+        size=min(20, len(dataset)),
+        batch_size=args.batch_size,
+        workers=0,
+        device=device,
+        seed=args.seed,
+        min_hz=MIN_HZ,
+    )
+    frontend = LogMelFrontend(mean, std).to(device)
+    model = MTRCNNDG().to(device).train()
+    features, frame_lengths = frontend(waveforms, lengths)
+    outputs = model(features, frame_lengths, domains, domain_strength=0.3, augment=True)
+    priors = torch.full((len(SPECIES),), 1 / len(SPECIES), device=device)
+    losses = training_loss(outputs, labels, domains, priors, epoch=5)
+    losses["loss"].backward()
+    assert outputs["species_logits"].shape == (args.batch_size, len(SPECIES))
+    assert outputs["domain_logits"].shape == (args.batch_size, len(DOMAINS))
+    assert all(torch.isfinite(value) for value in losses.values())
+    assert all(
+        parameter.grad is None or torch.isfinite(parameter.grad).all()
+        for parameter in model.parameters()
+    )
+    result = {
+        "device": str(device),
+        "features": list(features.shape),
+        "embedding": list(outputs["embedding"].shape),
+        "parameters": sum(parameter.numel() for parameter in model.parameters()),
+        "normalization": {
+            "mean": [round(mean.min().item(), 4), round(mean.max().item(), 4)],
+            "std": [round(std.min().item(), 4), round(std.max().item(), 4)],
+        },
+        "fold": {
+            "train": len(training_ids),
+            "validation": len(validation_ids),
+            "held_domains": {
+                SPECIES[species]: DOMAINS[domain] if domain is not None else "same-domain"
+                for species, domain in held_domains.items()
+            },
+        },
+        "losses": {name: round(value.item(), 6) for name, value in losses.items()},
+    }
+    print(json.dumps(result, indent=1))
+
+
+def add_train_parser(commands):
+    """Register the train command and its options."""
+    train_parser = commands.add_parser("train")
+    train_parser.add_argument("--data", default="/tmp/mosquito-audio")
+    train_parser.add_argument("--output", default="/tmp/mosquito-mtrcnn")
+    train_parser.add_argument("--fold", type=int, default=0)
+    train_parser.add_argument("--epochs", type=int, default=20)
+    train_parser.add_argument("--max-steps", type=int, default=-1)
+    train_parser.add_argument("--train-size", type=int, default=None)
+    train_parser.add_argument("--validation-size", type=int, default=None)
+    train_parser.add_argument("--stats-size", type=int, default=10_000)
+    train_parser.add_argument("--reference-size", type=int, default=20_000)
+    train_parser.add_argument("--batch-size", type=int, default=128)
+    train_parser.add_argument("--eval-batch-size", type=int, default=256)
+    train_parser.add_argument("--learning-rate", type=float, default=1e-3)
+    train_parser.add_argument("--weight-decay", type=float, default=1e-4)
+    train_parser.add_argument("--warmup-ratio", type=float, default=0.05)
+    train_parser.add_argument("--gradient-clip", type=float, default=5.0)
+    train_parser.add_argument("--dropout", type=float, default=0.2)
+    train_parser.add_argument("--balance-power", type=float, default=1.0)
+    train_parser.add_argument("--seconds", type=float, default=SECONDS)
+    train_parser.add_argument("--min-hz", type=float, default=MIN_HZ)
+    train_parser.add_argument("--level-rms", type=float, default=LEVEL_RMS)
+    train_parser.add_argument("--patience", type=int, default=4)
+    train_parser.add_argument("--workers", type=int, default=max(0, cpus() - 1))
+    train_parser.add_argument("--seed", type=int, default=42)
+    train_parser.add_argument("--synthetic", action="store_true")
+    train_parser.add_argument("--synthetic-rows", type=int, default=4)
+    train_parser.add_argument("--test", action="store_true")
+    train_parser.add_argument("--test-size", type=int, default=None)
+    train_parser.set_defaults(run=train)
+
+
+def main():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    smoke_parser = commands.add_parser("smoke")
+    smoke_parser.add_argument("--batch-size", type=int, default=10)
+    smoke_parser.add_argument("--seconds", type=float, default=SECONDS)
+    smoke_parser.add_argument("--seed", type=int, default=42)
+    smoke_parser.set_defaults(run=smoke)
+    add_train_parser(commands)
+    evaluate_parser = commands.add_parser("evaluate")
+    evaluate_parser.add_argument("--checkpoint", required=True)
+    evaluate_parser.add_argument("--data", default="/tmp/mosquito-audio")
+    evaluate_parser.add_argument("--output", default="/tmp/mosquito-mtrcnn-eval.json")
+    evaluate_parser.add_argument("--test-size", type=int, default=None)
+    evaluate_parser.add_argument("--eval-batch-size", type=int, default=128)
+    evaluate_parser.add_argument("--workers", type=int, default=max(0, cpus() - 1))
+    evaluate_parser.add_argument("--export", default=None)
+    evaluate_parser.add_argument("--repo", default=None)
+    evaluate_parser.set_defaults(run=evaluate)
+    predict_parser = commands.add_parser("predict")
+    predict_parser.add_argument("--model", default=REPO)
+    predict_parser.add_argument("files", nargs="+")
+    predict_parser.set_defaults(run=predict)
+    card_parser = commands.add_parser("card")
+    card_parser.add_argument("--repo", default=REPO)
+    card_parser.set_defaults(run=card)
+    args = parser.parse_args()
+    args.run(args)
+
+
+if __name__ == "__main__":
+    main()

results/eval.json

@@ -0,0 +1,1177 @@
+{
+ "checkpoint": "/checkpoint/checkpoint.pt",
+ "epoch": 12,
+ "fold": -1,
+ "seed": 42,
+ "arguments": {
+  "command": "train",
+  "data": "/tmp/mosquito-audio",
+  "output": "/data/mosquito-mtrcnn-dg-v1-full-seed42",
+  "fold": -1,
+  "epochs": 12,
+  "max_steps": -1,
+  "train_size": null,
+  "validation_size": null,
+  "stats_size": 10000,
+  "reference_size": 20000,
+  "batch_size": 128,
+  "eval_batch_size": 256,
+  "learning_rate": 0.001,
+  "weight_decay": 0.0001,
+  "warmup_ratio": 0.05,
+  "gradient_clip": 5.0,
+  "dropout": 0.2,
+  "balance_power": 1.0,
+  "seconds": 2.0,
+  "min_hz": 0.0,
+  "level_rms": 0.0,
+  "patience": 4,
+  "workers": 2,
+  "seed": 42,
+  "synthetic": false,
+  "synthetic_rows": 4,
+  "test": true,
+  "test_size": null
+ },
+ "test": {
+  "rows": 27217,
+  "softmax": {
+   "seen": {
+    "balanced_accuracy": 0.6932280918167799,
+    "rows": 23015,
+    "per_species": {
+     "Aedes aegypti": 0.6460399146479227,
+     "Aedes albopictus": 0.44731332868108864,
+     "Culex quinquefasciatus": 0.8741965105601469,
+     "Anopheles gambiae": 0.7882534775888718,
+     "Anopheles arabiensis": 0.6027397260273972,
+     "Culex pipiens": 0.8008255933952528
+    },
+    "confusion": [
+     [
+      5147,
+      139,
+      765,
+      598,
+      736,
+      24,
+      555,
+      0,
+      3
+     ],
+     [
+      82,
+      641,
+      6,
+      533,
+      146,
+      5,
+      19,
+      0,
+      1
+     ],
+     [
+      146,
+      2,
+      5712,
+      12,
+      623,
+      10,
+      29,
+      0,
+      0
+     ],
+     [
+      80,
+      28,
+      8,
+      3060,
+      639,
+      4,
+      63,
+      0,
+      0
+     ],
+     [
+      2,
+      1,
+      13,
+      93,
+      176,
+      0,
+      7,
+      0,
+      0
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ],
+     [
+      275,
+      3,
+      49,
+      95,
+      151,
+      4,
+      2328,
+      0,
+      2
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ]
+    ]
+   },
+   "unseen": {
+    "balanced_accuracy": 0.2546596781809005,
+    "rows": 4202,
+    "per_species": {
+     "Aedes aegypti": 0.0,
+     "Aedes albopictus": 0.29832935560859186,
+     "Culex quinquefasciatus": 0.02976190476190476,
+     "Anopheles gambiae": 0.0,
+     "Anopheles arabiensis": 0.0016483516483516484,
+     "Anopheles dirus": 0.0,
+     "Culex pipiens": 0.9411764705882353,
+     "Anopheles minimus": 0.7777777777777778,
+     "Anopheles stephensi": 0.24324324324324326
+    },
+    "confusion": [
+     [
+      0,
+      192,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ],
+     [
+      81,
+      125,
+      0,
+      9,
+      0,
+      114,
+      28,
+      41,
+      21
+     ],
+     [
+      121,
+      5,
+      20,
+      0,
+      0,
+      281,
+      6,
+      102,
+      137
+     ],
+     [
+      61,
+      152,
+      90,
+      0,
+      0,
+      13,
+      0,
+      474,
+      28
+     ],
+     [
+      544,
+      253,
+      262,
+      0,
+      3,
+      363,
+      12,
+      307,
+      76
+     ],
+     [
+      1,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      39,
+      0
+     ],
+     [
+      2,
+      0,
+      1,
+      0,
+      0,
+      0,
+      64,
+      0,
+      1
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      77,
+      22
+     ],
+     [
+      0,
+      0,
+      25,
+      0,
+      0,
+      0,
+      0,
+      31,
+      18
+     ]
+    ]
+   }
+  },
+  "mahalanobis": {
+   "seen": {
+    "balanced_accuracy": 0.749029147696192,
+    "rows": 23015,
+    "per_species": {
+     "Aedes aegypti": 0.7094263838333125,
+     "Aedes albopictus": 0.7997208653175157,
+     "Culex quinquefasciatus": 0.8790939700030609,
+     "Anopheles gambiae": 0.6254507985574446,
+     "Anopheles arabiensis": 0.6335616438356164,
+     "Culex pipiens": 0.8469212246302029
+    },
+    "confusion": [
+     [
+      5652,
+      269,
+      630,
+      155,
+      513,
+      1,
+      747,
+      0,
+      0
+     ],
+     [
+      99,
+      1146,
+      5,
+      75,
+      83,
+      0,
+      25,
+      0,
+      0
+     ],
+     [
+      268,
+      25,
+      5744,
+      0,
+      443,
+      0,
+      54,
+      0,
+      0
+     ],
+     [
+      189,
+      317,
+      9,
+      2428,
+      830,
+      0,
+      109,
+      0,
+      0
+     ],
+     [
+      4,
+      16,
+      18,
+      58,
+      185,
+      0,
+      11,
+      0,
+      0
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ],
+     [
+      267,
+      15,
+      40,
+      26,
+      96,
+      0,
+      2462,
+      0,
+      1
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ],
+     [
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ]
+    ]
+   },
+   "unseen": {
+    "balanced_accuracy": 0.2824912329416518,
+    "rows": 4202,
+    "per_species": {
+     "Aedes aegypti": 0.0,
+     "Aedes albopictus": 0.5918854415274463,
+     "Culex quinquefasciatus": 0.05357142857142857,
+     "Anopheles gambiae": 0.0,
+     "Anopheles arabiensis": 0.002197802197802198,
+     "Anopheles dirus": 0.0,
+     "Culex pipiens": 0.9411764705882353,
+     "Anopheles minimus": 0.7373737373737373,
+     "Anopheles stephensi": 0.21621621621621623
+    },
+    "confusion": [
+     [
+      0,
+      192,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      0
+     ],
+     [
+      85,
+      248,
+      0,
+      4,
+      0,
+      53,
+      25,
+      4,
+      0
+     ],
+     [
+      194,
+      65,
+      36,
+      0,
+      0,
+      302,
+      43,
+      3,
+      29
+     ],
+     [
+      56,
+      309,
+      221,
+      0,
+      1,
+      15,
+      26,
+      179,
+      11
+     ],
+     [
+      601,
+      457,
+      330,
+      0,
+      4,
+      327,
+      92,
+      8,
+      1
+     ],
+     [
+      3,
+      1,
+      2,
+      0,
+      0,
+      0,
+      1,
+      33,
+      0
+     ],
+     [
+      1,
+      0,
+      2,
+      0,
+      0,
+      0,
+      64,
+      0,
+      1
+     ],
+     [
+      1,
+      1,
+      1,
+      0,
+      0,
+      0,
+      3,
+      73,
+      20
+     ],
+     [
+      4,
+      0,
+      25,
+      0,
+      0,
+      0,
+      1,
+      28,
+      16
+     ]
+    ]
+   }
+  }
+ },
+ "fold_softmax": {
+  "balanced_accuracy": 0.7731446928816349,
+  "rows": 30516,
+  "per_species": {
+   "Aedes aegypti": 0.6406232973738695,
+   "Aedes albopictus": 0.4512722035525684,
+   "Culex quinquefasciatus": 0.8714373843306601,
+   "Anopheles gambiae": 0.7983733686400605,
+   "Anopheles arabiensis": 0.5806315789473684,
+   "Anopheles dirus": 1.0,
+   "Culex pipiens": 0.8031072602330445,
+   "Anopheles minimus": 0.8928571428571429,
+   "Anopheles stephensi": 0.92
+  },
+  "confusion": [
+   [
+    5879,
+    136,
+    873,
+    653,
+    928,
+    31,
+    669,
+    1,
+    7
+   ],
+   [
+    126,
+    940,
+    12,
+    785,
+    176,
+    3,
+    40,
+    0,
+    1
+   ],
+   [
+    200,
+    2,
+    7063,
+    13,
+    787,
+    13,
+    26,
+    0,
+    1
+   ],
+   [
+    86,
+    27,
+    11,
+    4221,
+    872,
+    7,
+    63,
+    0,
+    0
+   ],
+   [
+    41,
+    17,
+    99,
+    746,
+    1379,
+    2,
+    90,
+    0,
+    1
+   ],
+   [
+    0,
+    0,
+    0,
+    0,
+    0,
+    11,
+    0,
+    0,
+    0
+   ],
+   [
+    345,
+    1,
+    36,
+    96,
+    179,
+    2,
+    2688,
+    0,
+    0
+   ],
+   [
+    0,
+    0,
+    0,
+    0,
+    0,
+    3,
+    0,
+    50,
+    3
+   ],
+   [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    6,
+    69
+   ]
+  ],
+  "domain_holdout": {
+   "balanced_accuracy": 0.0,
+   "rows": 0,
+   "per_species": {},
+   "confusion": [
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ]
+   ]
+  },
+  "same_domain": {
+   "balanced_accuracy": 0.7731446928816349,
+   "rows": 30516,
+   "per_species": {
+    "Aedes aegypti": 0.6406232973738695,
+    "Aedes albopictus": 0.4512722035525684,
+    "Culex quinquefasciatus": 0.8714373843306601,
+    "Anopheles gambiae": 0.7983733686400605,
+    "Anopheles arabiensis": 0.5806315789473684,
+    "Anopheles dirus": 1.0,
+    "Culex pipiens": 0.8031072602330445,
+    "Anopheles minimus": 0.8928571428571429,
+    "Anopheles stephensi": 0.92
+   },
+   "confusion": [
+    [
+     5879,
+     136,
+     873,
+     653,
+     928,
+     31,
+     669,
+     1,
+     7
+    ],
+    [
+     126,
+     940,
+     12,
+     785,
+     176,
+     3,
+     40,
+     0,
+     1
+    ],
+    [
+     200,
+     2,
+     7063,
+     13,
+     787,
+     13,
+     26,
+     0,
+     1
+    ],
+    [
+     86,
+     27,
+     11,
+     4221,
+     872,
+     7,
+     63,
+     0,
+     0
+    ],
+    [
+     41,
+     17,
+     99,
+     746,
+     1379,
+     2,
+     90,
+     0,
+     1
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     11,
+     0,
+     0,
+     0
+    ],
+    [
+     345,
+     1,
+     36,
+     96,
+     179,
+     2,
+     2688,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     3,
+     0,
+     50,
+     3
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     6,
+     69
+    ]
+   ]
+  }
+ },
+ "fold_mahalanobis": {
+  "balanced_accuracy": 0.8076237330803099,
+  "rows": 30516,
+  "per_species": {
+   "Aedes aegypti": 0.6983763757219135,
+   "Aedes albopictus": 0.7820451272203552,
+   "Culex quinquefasciatus": 0.8768661320172733,
+   "Anopheles gambiae": 0.6372233780972196,
+   "Anopheles arabiensis": 0.6058947368421053,
+   "Anopheles dirus": 1.0,
+   "Culex pipiens": 0.8598745144905886,
+   "Anopheles minimus": 0.875,
+   "Anopheles stephensi": 0.9333333333333333
+  },
+  "confusion": [
+   [
+    6409,
+    297,
+    744,
+    165,
+    672,
+    0,
+    890,
+    0,
+    0
+   ],
+   [
+    146,
+    1629,
+    13,
+    122,
+    120,
+    0,
+    53,
+    0,
+    0
+   ],
+   [
+    354,
+    28,
+    7107,
+    1,
+    545,
+    0,
+    70,
+    0,
+    0
+   ],
+   [
+    205,
+    413,
+    13,
+    3369,
+    1160,
+    0,
+    127,
+    0,
+    0
+   ],
+   [
+    78,
+    125,
+    125,
+    469,
+    1439,
+    0,
+    139,
+    0,
+    0
+   ],
+   [
+    0,
+    0,
+    0,
+    0,
+    0,
+    11,
+    0,
+    0,
+    0
+   ],
+   [
+    300,
+    7,
+    37,
+    27,
+    98,
+    0,
+    2878,
+    0,
+    0
+   ],
+   [
+    0,
+    0,
+    0,
+    0,
+    0,
+    3,
+    1,
+    49,
+    3
+   ],
+   [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    5,
+    70
+   ]
+  ],
+  "domain_holdout": {
+   "balanced_accuracy": 0.0,
+   "rows": 0,
+   "per_species": {},
+   "confusion": [
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0
+    ]
+   ]
+  },
+  "same_domain": {
+   "balanced_accuracy": 0.8076237330803099,
+   "rows": 30516,
+   "per_species": {
+    "Aedes aegypti": 0.6983763757219135,
+    "Aedes albopictus": 0.7820451272203552,
+    "Culex quinquefasciatus": 0.8768661320172733,
+    "Anopheles gambiae": 0.6372233780972196,
+    "Anopheles arabiensis": 0.6058947368421053,
+    "Anopheles dirus": 1.0,
+    "Culex pipiens": 0.8598745144905886,
+    "Anopheles minimus": 0.875,
+    "Anopheles stephensi": 0.9333333333333333
+   },
+   "confusion": [
+    [
+     6409,
+     297,
+     744,
+     165,
+     672,
+     0,
+     890,
+     0,
+     0
+    ],
+    [
+     146,
+     1629,
+     13,
+     122,
+     120,
+     0,
+     53,
+     0,
+     0
+    ],
+    [
+     354,
+     28,
+     7107,
+     1,
+     545,
+     0,
+     70,
+     0,
+     0
+    ],
+    [
+     205,
+     413,
+     13,
+     3369,
+     1160,
+     0,
+     127,
+     0,
+     0
+    ],
+    [
+     78,
+     125,
+     125,
+     469,
+     1439,
+     0,
+     139,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     11,
+     0,
+     0,
+     0
+    ],
+    [
+     300,
+     7,
+     37,
+     27,
+     98,
+     0,
+     2878,
+     0,
+     0
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     3,
+     1,
+     49,
+     3
+    ],
+    [
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     0,
+     5,
+     70
+    ]
+   ]
+  }
+ }
+}