queen_clf.py

"""
Binary queen-vs-worker classifier inference.

A dedicated EfficientNet-B0 binary classifier trained on cropped bee
images. Given a cropped bee, returns a probability that it is a queen.

This is much more focused than asking a multi-class YOLO to classify
queens (where localization + classification compete), or a VLM cascade
(where the generalist model has no bee-specific training).

Usage:
    import queen_clf
    if queen_clf.is_available():
        results = queen_clf.classify_crops([crop1, crop2, ...])
        # results = [{'queen_prob': 0.92, 'is_queen': True}, ...]
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

from PIL import Image


_HERE = Path(os.path.dirname(os.path.abspath(__file__)))
WEIGHTS_PATH = _HERE / "weights" / "queen_classifier.pt"

# Confidence threshold for calling a crop a queen. Tuned empirically -
# if false positives happen on the live Space, raise this.
# Set conservatively to avoid false queens on borderline crops.
QUEEN_PROB_THRESHOLD = 0.92
# At most one queen per frame (real frames almost never have more).
MAX_QUEENS_PER_FRAME = 1

_model = None
_meta = None
_tf = None
_failed = False
_logged = False


def is_available() -> bool:
    return WEIGHTS_PATH.exists() and WEIGHTS_PATH.stat().st_size > 1024 and not _failed


def _log_once():
    global _logged
    if _logged:
        return
    _logged = True
    print(f"[queen_clf] weights path: {WEIGHTS_PATH}", file=sys.stderr)
    print(f"[queen_clf] weights exist: {WEIGHTS_PATH.exists()}", file=sys.stderr)
    if WEIGHTS_PATH.exists():
        size = WEIGHTS_PATH.stat().st_size
        print(f"[queen_clf] weights size: {size} bytes ({size/1024/1024:.1f} MB)",
              file=sys.stderr)


def _load():
    global _model, _meta, _tf, _failed
    _log_once()
    if _model is not None or _failed:
        return _model
    if not WEIGHTS_PATH.exists():
        _failed = True
        return None
    try:
        import torch
        import timm
        from torchvision import transforms

        ckpt = torch.load(str(WEIGHTS_PATH), map_location="cpu", weights_only=False)
        arch = ckpt.get("arch", "efficientnet_b0")
        img_size = ckpt.get("img_size", 224)
        class_to_idx = ckpt.get("class_to_idx", {"queen": 0, "worker": 1})

        model = timm.create_model(arch, pretrained=False, num_classes=2)
        model.load_state_dict(ckpt["state_dict"])
        model.eval()

        tf = transforms.Compose([
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225]),
        ])

        _model = model
        _meta = {"class_to_idx": class_to_idx, "img_size": img_size, "arch": arch,
                 "queen_idx": class_to_idx.get("queen", 0)}
        _tf = tf
        print(f"[queen_clf] loaded {arch}, classes={class_to_idx}", file=sys.stderr)
    except Exception as e:
        print(f"[queen_clf] load failed: {type(e).__name__}: {e}", file=sys.stderr)
        _failed = True
    return _model


def classify_crops(crops: list[Image.Image]) -> list[dict]:
    """Given a list of PIL crops, return per-crop queen probabilities
    averaged across 4 test-time augmentations (original, h-flip, v-flip,
    180-rotation). False positives often score high on one orientation
    but collapse under augmentation; true queens stay high across all."""
    model = _load()
    if model is None or not crops:
        return [{"queen_prob": 0.0, "is_queen": False} for _ in crops]

    import torch

    queen_idx = _meta["queen_idx"]
    per_crop_probs = []

    for c in crops:
        c = c.convert("RGB")
        # Build 4 augmented views
        views = [
            c,
            c.transpose(Image.FLIP_LEFT_RIGHT),
            c.transpose(Image.FLIP_TOP_BOTTOM),
            c.transpose(Image.ROTATE_180),
        ]
        batch = torch.stack([_tf(v) for v in views])
        with torch.no_grad():
            probs = torch.softmax(model(batch), dim=1)
        # Average queen probability across the 4 views
        avg_qp = float(probs[:, queen_idx].mean().item())
        per_crop_probs.append(avg_qp)

    out = []
    for qp in per_crop_probs:
        out.append({
            "queen_prob": round(qp, 3),
            "is_queen": qp >= QUEEN_PROB_THRESHOLD,
        })
    return out