yield_estimation.py

"""Yield estimation from tomato ripeness detections.

Two building blocks:

1. `estimate_yield_from_counts` -- turns a {class_name: count} dict into a weight estimate using
   an average-fruit-weight table. This is the same method commercial per-plant yield-counting
   vision systems use as a baseline: total mass = sum(count_c * avg_weight_c).

2. `VideoYieldTracker` -- for video/flythrough input, wraps Ultralytics' built-in ByteTrack
   integration (`model.track(..., persist=True)`) so the *same* physical tomato seen in several
   overlapping frames is counted once (by track ID), not once per frame. Per-frame detection
   counts alone systematically overcount yield on any video with frame-to-frame FOV overlap --
   which is exactly the case here (these datasets are frames extracted from a source video, most
   densely the mobile capture at a fixed 10-frame stride).

Weight-table caveat (read before trusting an absolute kg number)
------------------------------------------------------------------
`DEFAULT_AVG_FRUIT_WEIGHT_G` is a literature-typical greenhouse round-tomato fruit weight, not a
measurement calibrated to this specific crop/cultivar/camera setup. Tomatoes reach most of their
final size by the Breakers stage and change mostly in color (not much in size) while ripening, so
a single weight-per-fruit assumption across ripeness classes is reasonable -- but the *absolute*
number should be calibrated (e.g. weighing a sample of counted fruit from one plant/plot and
adjusting `AVG_FRUIT_WEIGHT_G`) before being reported as ground truth yield for a real deployment.
"""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field
from typing import Any

CLASSES = ["Green", "Breakers", "Turning", "Pink", "Light Red", "Red"]

DEFAULT_AVG_FRUIT_WEIGHT_G: dict[str, float] = {name: 120.0 for name in CLASSES}

HARVESTABLE_CLASSES = {"Light Red", "Red"}


@dataclass
class YieldEstimate:
    counts: dict[str, int]
    weight_g_by_class: dict[str, float]
    total_weight_g: float
    total_weight_kg: float
    harvestable_weight_kg: float
    total_fruit_count: int

    def as_dict(self) -> dict[str, Any]:
        return {
            "counts": self.counts,
            "weight_g_by_class": {k: round(v, 1) for k, v in self.weight_g_by_class.items()},
            "total_weight_g": round(self.total_weight_g, 1),
            "total_weight_kg": round(self.total_weight_kg, 3),
            "harvestable_weight_kg": round(self.harvestable_weight_kg, 3),
            "total_fruit_count": self.total_fruit_count,
        }


def estimate_yield_from_counts(
    counts: dict[str, int], weight_table: dict[str, float] | None = None
) -> YieldEstimate:
    weight_table = weight_table or DEFAULT_AVG_FRUIT_WEIGHT_G
    weight_g_by_class = {
        name: counts.get(name, 0) * weight_table.get(name, 0.0) for name in CLASSES
    }
    total_g = sum(weight_g_by_class.values())
    harvestable_g = sum(
        weight_g_by_class[name] for name in CLASSES if name in HARVESTABLE_CLASSES
    )
    return YieldEstimate(
        counts={name: counts.get(name, 0) for name in CLASSES},
        weight_g_by_class=weight_g_by_class,
        total_weight_g=total_g,
        total_weight_kg=total_g / 1000.0,
        harvestable_weight_kg=harvestable_g / 1000.0,
        total_fruit_count=sum(counts.get(name, 0) for name in CLASSES),
    )


def counts_from_detections(class_names: list[str]) -> dict[str, int]:
    counter = Counter(class_names)
    return {name: counter.get(name, 0) for name in CLASSES}


@dataclass
class VideoYieldResult:
    frames: int
    naive_per_frame_total: int
    unique_track_total: int
    overcount_ratio: float
    unique_counts_by_class: dict[str, int]
    yield_estimate: YieldEstimate
    per_frame_track_ids: list[list[int]] = field(default_factory=list)


def summarize_tracked_video(
    frame_results: list[Any], class_names_by_id: dict[int, str], weight_table: dict[str, float] | None = None
) -> VideoYieldResult:
    naive_total = 0
    first_seen_class: dict[int, str] = {}
    per_frame_ids: list[list[int]] = []

    for result in frame_results:
        boxes = result.boxes
        frame_ids: list[int] = []
        if boxes is None or boxes.id is None:
            per_frame_ids.append(frame_ids)
            continue
        ids = boxes.id.int().tolist()
        classes = boxes.cls.int().tolist()
        naive_total += len(ids)
        for track_id, class_idx in zip(ids, classes):
            frame_ids.append(track_id)
            if track_id not in first_seen_class:
                first_seen_class[track_id] = class_names_by_id.get(class_idx, str(class_idx))
        per_frame_ids.append(frame_ids)

    unique_counts = counts_from_detections(list(first_seen_class.values()))
    unique_total = sum(unique_counts.values())
    yield_est = estimate_yield_from_counts(unique_counts, weight_table)

    return VideoYieldResult(
        frames=len(frame_results),
        naive_per_frame_total=naive_total,
        unique_track_total=unique_total,
        overcount_ratio=(naive_total / unique_total) if unique_total else 0.0,
        unique_counts_by_class=unique_counts,
        yield_estimate=yield_est,
        per_frame_track_ids=per_frame_ids,
    )