hotcoco 1.0.0

Perception evaluation in pure Rust — a pycocotools-compatible COCO/LVIS/Open Images engine with diagnostics and dataset tools
Documentation
use std::collections::{HashMap, HashSet};

use rayon::prelude::*;

use super::matching::{EvalImgContext, IouMatrix};
use super::{COCOeval, EvalMode};

impl COCOeval {
    /// Populate `params.img_ids` and `params.cat_ids` from the GT dataset if not already set.
    ///
    /// The writing half of [`COCOeval::resolved_ids`], which owns the derivation:
    /// "which ids does this dataset cover" is one question, and `confusion_matrix`
    /// asks it too without being allowed to mutate. Answering it separately here
    /// meant two places to keep agreeing about it.
    fn resolve_params(&mut self) {
        let (img_ids, cat_ids) = self.resolved_ids();
        let (img_ids, cat_ids) = (img_ids.into_owned(), cat_ids.into_owned());
        self.params.img_ids = img_ids;
        self.params.cat_ids = cat_ids;
    }

    /// Build the sorted list of (img_id, cat_id) pairs to evaluate.
    ///
    /// Takes the union of non-empty GT and DT pairs, filters to the active img_ids/cat_ids
    /// from params, and returns them sorted for deterministic output order.
    ///
    /// In LVIS mode, DT-only pairs are dropped unless the category appears in `neg_cats`
    /// for that image — that is, it was confirmed absent and unmatched DTs should count as FP.
    fn collect_sparse_pairs(
        &self,
        cat_ids: &[u64],
        neg_cats: &HashMap<u64, HashSet<u64>>,
    ) -> Vec<(u64, u64)> {
        let allowed_imgs: HashSet<u64> = self.params.img_ids.iter().copied().collect();
        let allowed_cats: HashSet<u64> = cat_ids.iter().copied().collect();

        // At large-scale (e.g. Objects365: 365 cats × 80K imgs = 29M pairs), ~96% of pairs
        // are empty. Driving evaluation from the index instead reduces pairs by ~35x.
        let mut sparse_set: HashSet<(u64, u64)> = HashSet::new();
        if self.params.use_cats {
            // Collect GT pairs first (needed for LVIS DT filtering).
            let mut gt_pairs: HashSet<(u64, u64)> = HashSet::new();
            for pair in self.coco_gt.nonempty_img_cat_pairs() {
                if allowed_imgs.contains(&pair.0) && allowed_cats.contains(&pair.1) {
                    gt_pairs.insert(pair);
                    sparse_set.insert(pair);
                }
            }
            for pair in self.coco_dt.nonempty_img_cat_pairs() {
                if allowed_imgs.contains(&pair.0) && allowed_cats.contains(&pair.1) {
                    if self.eval_mode == EvalMode::Lvis {
                        // Keep DT pair only if GT exists OR cat is explicitly neg for this image.
                        if gt_pairs.contains(&pair)
                            || neg_cats.get(&pair.0).is_some_and(|s| s.contains(&pair.1))
                        {
                            sparse_set.insert(pair);
                        }
                    } else {
                        sparse_set.insert(pair);
                    }
                }
            }
        } else {
            for img_id in self.coco_gt.nonempty_img_ids() {
                if allowed_imgs.contains(&img_id) {
                    sparse_set.insert((img_id, u64::MAX));
                }
            }
            for img_id in self.coco_dt.nonempty_img_ids() {
                if allowed_imgs.contains(&img_id) {
                    sparse_set.insert((img_id, u64::MAX));
                }
            }
        }

        let mut pairs: Vec<(u64, u64)> = sparse_set.into_iter().collect();
        pairs.sort_unstable();
        pairs
    }

    /// Run per-image evaluation.
    ///
    /// # Open Images replaces `coco_gt` (and possibly `coco_dt`)
    ///
    /// In [`EvalMode::OpenImages`], this method **overwrites the public
    /// `coco_gt` field** — and, when `params.expand_dt` is set, `coco_dt` —
    /// with hierarchy-expanded copies: every annotation is duplicated at each
    /// ancestor category and virtual categories are added for hierarchy-only
    /// nodes (see [`super::expand::expand_annotations`]). Any read of those
    /// fields after `evaluate()` sees the expanded datasets, not the ones the
    /// evaluator was constructed with. The expansion deduplicates, so calling
    /// `evaluate()` again does not expand further. The eval paths must see the
    /// expanded data through the same fields the analysis surfaces read (TIDE,
    /// diagnostics, category names), which is why the originals are replaced
    /// rather than shadowed by private copies.
    pub fn evaluate(&mut self) {
        // OID: expand GT (and optionally DT) using hierarchy — this replaces
        // the public `coco_gt`/`coco_dt` fields; see the method docs above.
        if self.eval_mode == EvalMode::OpenImages {
            let hierarchy = self.hierarchy.clone().unwrap_or_else(|| {
                crate::detection::hierarchy::Hierarchy::from_categories(
                    &self.coco_gt.dataset.categories,
                )
            });
            self.coco_gt = super::expand::expand_annotations(&self.coco_gt, &hierarchy);
            if self.params.expand_dt {
                self.coco_dt = super::expand::expand_annotations(&self.coco_dt, &hierarchy);
            }
            self.hierarchy = Some(hierarchy);
        }

        self.resolve_params();

        let cat_ids = if self.params.use_cats {
            self.params.cat_ids.clone()
        } else {
            vec![u64::MAX] // placeholder single category (avoids collision with real category_id=0)
        };

        // LVIS: scan GT image metadata to build per-image category sets.
        // Deferred from construction so the scan only happens when evaluate() is called.
        // neg_cats:       img_id → categories confirmed absent (unmatched DTs count as FP).
        // not_exhaustive: img_id → categories not fully checked (unmatched DTs are ignored).
        let (neg_cats, not_exhaustive) = if self.eval_mode == EvalMode::Lvis {
            let mut neg: HashMap<u64, HashSet<u64>> = HashMap::new();
            let mut not_ex: HashMap<u64, HashSet<u64>> = HashMap::new();
            for img in &self.coco_gt.dataset.images {
                if !img.neg_category_ids.is_empty() {
                    neg.insert(img.id, img.neg_category_ids.iter().copied().collect());
                }
                if !img.not_exhaustive_category_ids.is_empty() {
                    not_ex.insert(
                        img.id,
                        img.not_exhaustive_category_ids.iter().copied().collect(),
                    );
                }
            }
            (neg, not_ex)
        } else {
            (HashMap::new(), HashMap::new())
        };

        // LVIS: build freq_groups now that cat_ids are established.
        if self.eval_mode == EvalMode::Lvis {
            let cat_id_to_k_idx: HashMap<u64, usize> =
                cat_ids.iter().enumerate().map(|(i, &id)| (id, i)).collect();
            let mut freq_groups = super::mode::FreqGroups::default();
            for cat in &self.coco_gt.dataset.categories {
                if let Some(&k_idx) = cat_id_to_k_idx.get(&cat.id) {
                    match cat.frequency.as_deref() {
                        Some("r") => freq_groups.rare.push(k_idx),
                        Some("c") => freq_groups.common.push(k_idx),
                        Some("f") => freq_groups.frequent.push(k_idx),
                        _ => {}
                    }
                }
            }
            self.freq_groups = freq_groups;
        }

        let sparse_pairs = self.collect_sparse_pairs(&cat_ids, &neg_cats);

        // Segm only: convert every in-scope mask to RLE once, up front —
        // pycocotools' `_prepare` step. The per-cell IoU computation below and
        // the cross-category matrices in `confusion_matrix`/`tide` all read
        // this instead of re-rasterizing polygons per call site.
        use crate::primitives::sim::SimKind;
        self.segm_rles = (SimKind::from(self.params.iou_type) == SimKind::Mask)
            .then(|| super::iou::SegmRles::prepare(&self.coco_gt, &self.coco_dt, &self.params));

        // Compute IoUs only for pairs where both GT and DT are non-empty.
        // Pairs with only GT or only DT produce empty IoU matrices — skip storing them.
        let iou_results: Vec<((u64, u64), IouMatrix)> = sparse_pairs
            .par_iter()
            .filter_map(|&(img_id, cat_id)| {
                let iou_matrix = Self::compute_iou_static(
                    &self.coco_gt,
                    &self.coco_dt,
                    &self.params,
                    img_id,
                    cat_id,
                    self.eval_mode,
                    self.segm_rles.as_ref(),
                );
                if iou_matrix.is_empty() {
                    None
                } else {
                    Some(((img_id, cat_id), iou_matrix))
                }
            })
            .collect();

        // Replaces the cache wholesale, so `collect` sizes the map from the vec's
        // exact length rather than inheriting the previous run's capacity.
        self.ious = iou_results.into_iter().collect();

        // Evaluate each (image, category, area_range) combination in parallel,
        // over sparse_pairs × area_ranges rather than the full
        // cat_ids × area_ranges × img_ids product.
        //
        // Empty `max_dets` is degraded, not panicked on: `Params::max_det()`
        // owns the fallback cap (100), matching how every other degenerate
        // configuration on this path (missing area label, absent threshold)
        // degrades to the `-1.0` sentinel downstream instead of aborting —
        // `evaluate()` has no `Result` channel, and its siblings do not panic.
        let max_det = self.params.max_det();

        // pycocotools searches from `min(t, 1-1e-10)`, not from `t`. Inert below
        // 1.0, so the default 0.50:0.95 sweep is untouched; at t == 1.0 it admits
        // near-identical pairs, which is the drop-in behavior. Resolved once here
        // and shared — see `EvalImgContext::match_floors`.
        let match_floors: Vec<f64> = self
            .params
            .iou_thrs
            .iter()
            .map(|&t| crate::primitives::greedy::coco_match_floor(t))
            .collect();

        // Build shared context (borrows self after self.ious is fully populated).
        let ctx = EvalImgContext {
            coco_gt: &self.coco_gt,
            coco_dt: &self.coco_dt,
            params: &self.params,
            ious: &self.ious,
            eval_mode: self.eval_mode,
            match_floors: &match_floors,
        };

        // Fan out over pairs, not (pair, area range) cells: `gather_pair` resolves
        // everything the ranges share once per pair. Cells are written in place,
        // one `area_ranges.len()` chunk per pair — collect-then-flatten would move
        // ~800 MB of `EvalImg`s single-threaded on Objects365 (measured 2.0 s vs
        // 1.5 s). Every pair gets its full chunk, including empty gathers, so
        // `eval_imgs` keeps exactly the length, order, and `None` positions that
        // `accumulate`'s grouping walk and the public `eval_imgs()` accessor read.
        let is_lvis = self.eval_mode == EvalMode::Lvis;
        let area_ranges = &ctx.params.area_ranges;

        // `par_iter().map(..).collect()`, not `resize_with`: rayon's indexed
        // collect writes straight into the vector's uninitialized capacity across
        // all threads, while a sequential fill single-threads the first touch of
        // every page in that ~800 MB buffer. Measured at 270 ms on Objects365 —
        // more than the fan-out below saves.
        let mut eval_imgs: Vec<Option<super::matching::EvalImg>> = (0..sparse_pairs.len()
            * area_ranges.len())
            .into_par_iter()
            .map(|_| None)
            .collect();

        eval_imgs
            .par_chunks_mut(area_ranges.len())
            .zip(sparse_pairs.par_iter())
            .for_each(|(chunk, &(img_id, cat_id))| {
                let Some(pair) = super::matching::gather_pair(&ctx, img_id, cat_id, max_det) else {
                    return;
                };
                let not_exhaustive_cat = is_lvis
                    && not_exhaustive
                        .get(&img_id)
                        .is_some_and(|s| s.contains(&cat_id));

                for (slot, ar) in chunk.iter_mut().zip(area_ranges) {
                    *slot =
                        super::matching::evaluate_cell(&ctx, &pair, ar.range, not_exhaustive_cat);
                }
            });

        self.eval_imgs = eval_imgs;
    }
}