Keep the PRIMA architecture and verify prediction resume provenance

06b25e4 verified 1 day ago

22.3 kB

#!/usr/bin/env python3
"""Score PRIMA predictions with explicit annotations and metric definitions."""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import tempfile
from pathlib import Path

import numpy as np
import torch
from pycocotools import mask as mask_utils
from tools.predictions import index_prediction_files, sample_identifier

DEFAULT_TEXT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_TEXT_MODEL_REVISION = "1110a243fdf4706b3f48f1d95db1a4f5529b4d41"



def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("prediction_dirs", nargs="+", type=Path)
    parser.add_argument("--gt-json", type=Path, required=True)
    parser.add_argument("--ann-data", type=Path,
                        help="Optional legacy image-size metadata; new predictions store image sizes.")
    parser.add_argument("--output", type=Path)
    parser.add_argument(
        "--subset-mode",
        choices=("available", "intersection"),
        default="available",
        help=(
            "Evaluate each directory's available files, or use the common file "
            "intersection for every directory."
        ),
    )
    parser.add_argument("--with-text", action="store_true")
    parser.add_argument(
        "--with-meteor",
        action="store_true",
        help="Also run the paper's PTB-tokenized COCO METEOR computation.",
    )
    parser.add_argument("--text-model", default=DEFAULT_TEXT_MODEL)
    parser.add_argument(
        "--penalize-missing",
        action=argparse.BooleanOptionalAction,
        default=True,
        help=(
            "Count missing prediction files as empty outputs. The published API "
            "baseline rows used --no-penalize-missing (available outputs only)."
        ),
    )
    return parser.parse_args()


def decode_mask(segmentation, height: int, width: int):
    if isinstance(segmentation, list):
        rle = mask_utils.merge(mask_utils.frPyObjects(segmentation, height, width))
    elif isinstance(segmentation["counts"], list):
        rle = mask_utils.frPyObjects(segmentation, height, width)
    else:
        rle = segmentation
    return mask_utils.decode(rle).astype(bool)


def binary_iou(left: np.ndarray, right: np.ndarray) -> float:
    # Counterfactual irrelevant-image replacement can put a distractor mask at
    # the original image slot with a different native resolution.  Such masks
    # cannot describe the original target and are therefore exact non-matches.
    # Ordinary evaluation masks have matching shapes, so this leaves the paper
    # metric unchanged while making the replacement diagnostic well-defined.
    if left.shape != right.shape:
        return 0.0
    intersection = np.logical_and(left, right).sum()
    union = np.logical_or(left, right).sum()
    return float(intersection / union) if union else 0.0


def greedy_mask_ious(pred_masks, gt_masks):
    matrix = np.zeros((len(pred_masks), len(gt_masks)), dtype=np.float64)
    for pred_index, pred_mask in enumerate(pred_masks):
        for gt_index, gt_mask in enumerate(gt_masks):
            matrix[pred_index, gt_index] = binary_iou(pred_mask, gt_mask)
    paired = []
    while matrix.size and matrix.max() > 0:
        pred_index, gt_index = np.unravel_index(matrix.argmax(), matrix.shape)
        paired.append(float(matrix[pred_index, gt_index]))
        matrix = np.delete(matrix, pred_index, axis=0)
        matrix = np.delete(matrix, gt_index, axis=1)
    return paired


def greedy_mask_miou(pred_masks, gt_masks):
    paired = greedy_mask_ious(pred_masks, gt_masks)
    return float(np.mean(paired)) if paired else 0.0


def extract_caption(saved) -> str:
    output = saved.get("output", [])
    raw = output if isinstance(output, str) else (output[0] if output else "")
    raw = "" if raw is None else str(raw)
    if "ASSISTANT:" in raw:
        raw = raw.split("ASSISTANT:", 1)[1]
    return " ".join(raw.strip().replace("</s>", "").split())


def clean_caption(caption: str) -> str:
    for markup in (
        "<p> ",
        " </p> ",
        "[SEG]",
        "(IMAGE1)",
        "(IMAGE2)",
        "(IMAGE3)",
    ):
        caption = caption.replace(markup, "")
    return re.sub(r"_\d+\b", "", caption).replace("_", " ").strip()


def clean_reference_caption(caption: str) -> str:
    return re.sub(r"_\d+\b", "", caption).replace("_", " ").strip()


def ptb_tokenize(captions):
    """Run the same PTB tokenizer as COCOEvalCap with temp files in /tmp.
    The upstream package writes beside its JAR, which fails on read-only
    installations.  Changing only the temporary-file location leaves the
    tokenizer command and post-processing identical.
    """
    from pycocoevalcap.tokenizer import ptbtokenizer

    image_ids = [key for key, values in captions.items() for _ in values]
    sentences = "\n".join(
        value["caption"].replace("\n", " ")
        for values in captions.values()
        for value in values
    )
    jar_dir = Path(ptbtokenizer.__file__).resolve().parent
    jar_path = jar_dir / ptbtokenizer.STANFORD_CORENLP_3_4_1_JAR
    with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as handle:
        handle.write(sentences.encode())
        temporary_path = Path(handle.name)
    try:
        command = [
            "java",
            "-cp",
            str(jar_path),
            "edu.stanford.nlp.process.PTBTokenizer",
            "-preserveLines",
            "-lowerCase",
            str(temporary_path),
        ]
        completed = subprocess.run(
            command,
            cwd=jar_dir,
            check=True,
            stdout=subprocess.PIPE,
        )
    finally:
        temporary_path.unlink(missing_ok=True)
    lines = completed.stdout.decode().split("\n")
    tokenized = {}
    for image_id, line in zip(image_ids, lines):
        words = [
            word
            for word in line.rstrip().split(" ")
            if word not in ptbtokenizer.PUNCTUATIONS
        ]
        tokenized.setdefault(image_id, []).append(" ".join(words))
    return tokenized


def coco_meteor(references, predictions) -> float:
    from pycocoevalcap.meteor.meteor import Meteor

    ground_truth = ptb_tokenize(references)
    results = ptb_tokenize(predictions)
    scorer = Meteor()
    score, _ = scorer.compute_score(ground_truth, results)
    return float(score)


def convert_name(name: str) -> str:
    """Normalize grounded labels using the paper's evaluation rules."""
    name = re.sub(r"\d+", "", name)
    name = name.replace("_", " ").strip()
    name = re.sub(r"_\d+\b", "", name)
    possessive = re.search(r"(?:'s\s*|’s\s*)(.*)$", name, flags=re.IGNORECASE)
    if possessive:
        return possessive.group(1).strip()
    name = re.sub(r"\d+", "", name)
    return name.replace("_", "").strip()


def extract_labels(caption: str, num_images: int):
    labels = [[] for _ in range(num_images)]
    pattern = re.compile(
        r"<p>\s*(.*?)\s*</p>((?:\s*\[SEG\]\s*\(IMAGE\d+\))+)",
    )
    for phrase, seg_block in pattern.findall(caption):
        phrase = convert_name(phrase.strip().lower())
        for image_number in re.findall(
            r"\[SEG\]\s*\(IMAGE(\d+)\)", seg_block
        ):
            index = int(image_number) - 1
            if 0 <= index < num_images:
                labels[index].append(phrase)
    return labels


def normalize_pred_masks(saved, num_images: int):
    raw = list(saved.get("pred_masks", []))
    # The model's saved structure has a sentinel at index zero.
    raw = raw[1 : num_images + 1]
    raw.extend([None] * (num_images - len(raw)))
    normalized = []
    for masks in raw:
        if masks is None or len(masks) == 0:
            normalized.append([])
            continue
        # Older cached runs wrap the per-image [N,H,W] tensor in a singleton
        # batch list, whereas newer evaluation saves that tensor directly.
        if isinstance(masks, (list, tuple)) and len(masks) == 1:
            masks = masks[0]
        if torch.is_tensor(masks):
            masks = masks.detach().cpu().numpy()
        else:
            masks = np.stack(
                [
                    mask.detach().cpu().numpy() if torch.is_tensor(mask) else np.asarray(mask)
                    for mask in masks
                ]
            )
        normalized.append([(mask > 0) for mask in masks])
    return normalized


def text_siou(left: str, right: str) -> float:
    left_words, right_words = set(left.split()), set(right.split())
    union = left_words | right_words
    return len(left_words & right_words) / len(union) if union else 0.0


class TextEmbedder:
    def __init__(self, name: str):
        from sentence_transformers import SentenceTransformer

        revision = DEFAULT_TEXT_MODEL_REVISION if name == DEFAULT_TEXT_MODEL else None
        self.model = SentenceTransformer(name, revision=revision)
        self.revision = getattr(self.model[0].auto_model.config, '_commit_hash', None)
        self.cache = {}

    def similarity(self, left: str, right: str) -> float:
        missing = [value for value in (left, right) if value not in self.cache]
        if missing:
            encoded = self.model.encode(missing, convert_to_numpy=True, normalize_embeddings=True)
            self.cache.update(zip(missing, encoded))
        return float(np.dot(self.cache[left], self.cache[right]))


def grounded_matches(gt_masks, gt_labels, pred_masks, pred_labels, embedder):
    ious = np.zeros((len(gt_masks), len(pred_masks)), dtype=np.float64)
    for gt_index, gt_mask in enumerate(gt_masks):
        for pred_index, pred_mask in enumerate(pred_masks):
            ious[gt_index, pred_index] = binary_iou(gt_mask, pred_mask)
    similarities = np.zeros((len(gt_labels), len(pred_labels)), dtype=np.float64)
    sious = np.zeros_like(similarities)
    for gt_index, gt_label in enumerate(gt_labels):
        for pred_index, pred_label in enumerate(pred_labels):
            similarities[gt_index, pred_index] = embedder.similarity(gt_label, pred_label)
            sious[gt_index, pred_index] = text_siou(gt_label, pred_label)
    matches = 0
    similarity_sum = 0.0
    siou_sum = 0.0
    while ious.size:
        gt_index, pred_index = np.unravel_index(ious.argmax(), ious.shape)
        # The current evaluator treats a mask/label count mismatch as an
        # unmatched remainder instead of indexing outside the text matrix.
        if gt_index >= similarities.shape[0] or pred_index >= similarities.shape[1]:
            break
        if ious[gt_index, pred_index] < 0.5 or similarities[gt_index, pred_index] < 0.5:
            break
        matches += 1
        similarity_sum += similarities[gt_index, pred_index]
        siou_sum += sious[gt_index, pred_index]
        ious[gt_index, :] = 0
        ious[:, pred_index] = 0
        similarities[gt_index, :] = 0
        similarities[:, pred_index] = 0
    denom = max(matches, 1)
    return matches, similarity_sum / denom, siou_sum / denom


def text_only_matches(gt_labels, pred_labels, embedder):
    similarities = np.zeros((len(gt_labels), len(pred_labels)), dtype=np.float64)
    sious = np.zeros_like(similarities)
    for gt_index, gt_label in enumerate(gt_labels):
        for pred_index, pred_label in enumerate(pred_labels):
            similarities[gt_index, pred_index] = embedder.similarity(gt_label, pred_label)
            sious[gt_index, pred_index] = text_siou(gt_label, pred_label)
    matches = 0
    similarity_sum = 0.0
    siou_sum = 0.0
    while similarities.size:
        gt_index, pred_index = np.unravel_index(similarities.argmax(), similarities.shape)
        if similarities[gt_index, pred_index] < 0.5:
            break
        matches += 1
        similarity_sum += similarities[gt_index, pred_index]
        siou_sum += sious[gt_index, pred_index]
        similarities[gt_index, :] = 0
        similarities[:, pred_index] = 0
    denom = max(matches, 1)
    return matches, similarity_sum / denom, siou_sum / denom


def image_size(saved, image_index, groundings, image_info=None):
    """Resolve native dimensions without requiring a separate annotation cache."""
    if image_info is not None:
        return int(image_info['height']), int(image_info['width'])
    sizes = saved.get('image_sizes', [])
    if image_index < len(sizes) and sizes[image_index] is not None:
        return tuple(map(int, sizes[image_index]))
    for grounding in groundings:
        segmentation = grounding['segmentation']
        if isinstance(segmentation, dict) and 'size' in segmentation:
            return tuple(map(int, segmentation['size']))
    paths = saved.get('image_paths', [])
    if image_index < len(paths) and Path(paths[image_index]).is_file():
        from PIL import Image
        with Image.open(paths[image_index]) as image:
            width, height = image.size
        return height, width
    return None


def evaluate(
    prediction_dir: Path,
    prediction_files,
    records,
    selected_ids,
    ann_data=None,
    embedder=None,
    penalize_missing=False,
    with_meteor=False,
):
    ann_data = ann_data or {}
    image_mious = []
    all_matched_pair_ious = []
    actual_positives = 0
    true_positives = 0
    image_matched_similarity = 0.0
    image_matched_siou = 0.0
    image_text_similarity = 0.0
    image_text_siou = 0.0
    samples = 0
    annotated_images = 0
    images_with_predictions = 0
    predicted_masks = 0
    load_errors = 0
    meteor_references = {}
    meteor_predictions = {}
    for record in records:
        pair_id = sample_identifier(record)
        if pair_id not in selected_ids:
            continue
        path = prediction_files.get(pair_id)
        if path is None and not penalize_missing:
            continue
        if path is None:
            saved = {"output": ["ASSISTANT: No prediction available. </s>"], "pred_masks": []}
        else:
            try:
                saved = torch.load(path, map_location="cpu")
            except (EOFError, OSError, RuntimeError):
                load_errors += 1
                if not penalize_missing:
                    continue
                saved = {"output": ["ASSISTANT: No prediction available. </s>"], "pred_masks": []}
        num_images = len(record["img_ids"])
        caption = extract_caption(saved)
        labels_by_image = extract_labels(caption, num_images)
        masks_by_image = normalize_pred_masks(saved, num_images)
        samples += 1
        for image_index in range(num_images):
            groundings = record["groundings"].get(f"img_{image_index + 1}", [])
            if not groundings:
                continue
            pred_masks = masks_by_image[image_index]
            pred_labels = labels_by_image[image_index]
            gt_masks = []
            gt_labels = []
            dataset_name = str(record["dataset_name"])
            image_id = str(record["img_ids"][image_index])
            image_info = ann_data.get(dataset_name, {}).get(image_id)
            native_size = image_size(saved, image_index, groundings, image_info)
            if native_size is None:
                if pred_masks:
                    raise ValueError(
                        f'Native image size is missing for {pair_id}, image {image_index + 1}. '
                        'Use predictions with image_sizes, restore the source images, or provide --ann-data.'
                    )
                # An absent prediction scores zero without decoding polygon masks.
                image_mious.append(0.0)
                actual_positives += len(groundings)
                annotated_images += 1
                continue
            native_height, native_width = native_size
            for grounding in groundings:
                segmentation = grounding["segmentation"]
                gt_masks.append(decode_mask(segmentation, native_height, native_width))
                gt_label = str(grounding["name"]).lower()
                gt_labels.append(convert_name(gt_label))
            if len(pred_labels) < len(pred_masks):
                pred_labels.extend([""] * (len(pred_masks) - len(pred_labels)))
            paired_ious = greedy_mask_ious(pred_masks, gt_masks)
            image_mious.append(float(np.mean(paired_ious)) if paired_ious else 0.0)
            all_matched_pair_ious.extend(paired_ious)
            actual_positives += len(gt_labels)
            predicted_masks += len(pred_masks)
            annotated_images += 1
            images_with_predictions += int(bool(pred_masks))
            if with_meteor and pred_masks:
                caption_id = f"{pair_id}_{image_index + 1}"
                meteor_references[caption_id] = [
                    {
                        "caption": clean_reference_caption(
                            str(record["a"])
                        )
                    }
                ]
                meteor_predictions[caption_id] = [
                    {"caption": clean_caption(caption)}
                ]
            if embedder is not None and pred_masks and pred_labels:
                matched, similarity, siou = grounded_matches(
                    gt_masks, gt_labels, pred_masks, pred_labels, embedder
                )
                true_positives += matched
                image_matched_similarity += similarity
                image_matched_siou += siou
                _, similarity, siou = text_only_matches(gt_labels, pred_labels, embedder)
                image_text_similarity += similarity
                image_text_siou += siou
    result = {
        "prediction_dir": str(prediction_dir),
        "samples": samples,
        "load_errors": load_errors,
        "annotated_images": annotated_images,
        "ground_truth_masks": actual_positives,
        "predicted_masks": predicted_masks,
        "miou": float(np.mean(image_mious)) if image_mious else 0.0,
        "matched_pairs": len(all_matched_pair_ious),
        "miou_over_all_matched_pairs": (
            float(np.mean(all_matched_pair_ious)) if all_matched_pair_ious else 0.0
        ),
        "images_with_predictions": images_with_predictions,
        "miou_over_images_with_predictions": (
            float(np.sum(image_mious) / images_with_predictions)
            if images_with_predictions
            else 0.0
        ),
        "images_with_positive_iou": sum(value > 0 for value in image_mious),
        "miou_over_images_with_positive_iou": (
            float(np.mean([value for value in image_mious if value > 0]))
            if any(value > 0 for value in image_mious)
            else 0.0
        ),
    }
    if embedder is not None:
        result.update(
            {
                "recall": true_positives / max(actual_positives, 1),
                "iou_matched_semantic_similarity": image_matched_similarity
                / max(annotated_images, 1),
                "iou_matched_text_siou": image_matched_siou / max(annotated_images, 1),
                "semantic_similarity": image_text_similarity / max(annotated_images, 1),
                "text_siou": image_text_siou / max(annotated_images, 1),
            }
        )
    if with_meteor:
        result["meteor_evaluated_images"] = len(meteor_predictions)
        result["METEOR"] = (
            coco_meteor(meteor_references, meteor_predictions)
            if meteor_predictions
            else 0.0
        )
    return result


def score_predictions(prediction_dirs, records, ann_data=None, *,
                      subset_mode="available", penalize_missing=True, with_text=False,
                      with_meteor=False, text_model=DEFAULT_TEXT_MODEL, ground_truth=None):
    """Score M4SEG predictions for the standalone CLI or post-generation workflow."""
    manifest_ids = {sample_identifier(record) for record in records}
    if len(manifest_ids) != len(records):
        raise ValueError('Evaluation annotations must have unique example identifiers.')
    prediction_files = []
    available = []
    for directory in prediction_dirs:
        files = index_prediction_files(directory)
        prediction_files.append(files)
        available.append(set(files) & manifest_ids)
    if subset_mode == "intersection":
        selected_sets = [set.intersection(*available)] * len(prediction_dirs)
    elif penalize_missing:
        selected_sets = [manifest_ids] * len(prediction_dirs)
    else:
        selected_sets = available
    embedder = TextEmbedder(text_model) if with_text else None
    report = {
        "ground_truth": str(ground_truth),
        "subset_mode": subset_mode,
        "text_model": text_model if with_text else None,
        "text_model_revision": getattr(embedder, 'revision', None),
        "penalize_missing": penalize_missing,
        "metric_definition": {
            "miou": "mean over annotated images of greedy matched-pair mean IoU",
            "missing_or_extra_masks": "unmatched masks are not added to the matched-pair denominator",
            "text_threshold": 0.5,
            "mask_iou_threshold": 0.5,
        },
        "results": [
            evaluate(
                directory,
                files,
                records,
                selected,
                ann_data,
                embedder,
                penalize_missing,
                with_meteor,
            )
            for directory, files, selected in zip(
                prediction_dirs, prediction_files, selected_sets
            )
        ],
    }
    return report


def main():
    args = parse_args()
    import gzip
    opener = gzip.open if str(args.gt_json).endswith('.gz') else open
    with opener(args.gt_json, 'rt', encoding='utf-8') as handle:
        records = json.load(handle)
    ann_data = torch.load(args.ann_data, map_location="cpu", weights_only=False) if args.ann_data else None
    report = score_predictions(
        args.prediction_dirs, records, ann_data,
        subset_mode=args.subset_mode, penalize_missing=args.penalize_missing,
        with_text=args.with_text, with_meteor=args.with_meteor,
        text_model=args.text_model, ground_truth=args.gt_json,
    )
    output = json.dumps(report, indent=2, sort_keys=True)
    print(output)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(output + "\n")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Score PRIMA predictions with explicit annotations and metric definitions."""

	from __future__ import annotations

	import argparse
	import json
	import re
	import subprocess
	import tempfile
	from pathlib import Path

	import numpy as np
	import torch
	from pycocotools import mask as mask_utils
	from tools.predictions import index_prediction_files, sample_identifier

	DEFAULT_TEXT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
	DEFAULT_TEXT_MODEL_REVISION = "1110a243fdf4706b3f48f1d95db1a4f5529b4d41"



	def parse_args():
	parser = argparse.ArgumentParser()
	parser.add_argument("prediction_dirs", nargs="+", type=Path)
	parser.add_argument("--gt-json", type=Path, required=True)
	parser.add_argument("--ann-data", type=Path,
	help="Optional legacy image-size metadata; new predictions store image sizes.")
	parser.add_argument("--output", type=Path)
	parser.add_argument(
	"--subset-mode",
	choices=("available", "intersection"),
	default="available",
	help=(
	"Evaluate each directory's available files, or use the common file "
	"intersection for every directory."
	),
	)
	parser.add_argument("--with-text", action="store_true")
	parser.add_argument(
	"--with-meteor",
	action="store_true",
	help="Also run the paper's PTB-tokenized COCO METEOR computation.",
	)
	parser.add_argument("--text-model", default=DEFAULT_TEXT_MODEL)
	parser.add_argument(
	"--penalize-missing",
	action=argparse.BooleanOptionalAction,
	default=True,
	help=(
	"Count missing prediction files as empty outputs. The published API "
	"baseline rows used --no-penalize-missing (available outputs only)."
	),
	)
	return parser.parse_args()


	def decode_mask(segmentation, height: int, width: int):
	if isinstance(segmentation, list):
	rle = mask_utils.merge(mask_utils.frPyObjects(segmentation, height, width))
	elif isinstance(segmentation["counts"], list):
	rle = mask_utils.frPyObjects(segmentation, height, width)
	else:
	rle = segmentation
	return mask_utils.decode(rle).astype(bool)


	def binary_iou(left: np.ndarray, right: np.ndarray) -> float:
	# Counterfactual irrelevant-image replacement can put a distractor mask at
	# the original image slot with a different native resolution. Such masks
	# cannot describe the original target and are therefore exact non-matches.
	# Ordinary evaluation masks have matching shapes, so this leaves the paper
	# metric unchanged while making the replacement diagnostic well-defined.
	if left.shape != right.shape:
	return 0.0
	intersection = np.logical_and(left, right).sum()
	union = np.logical_or(left, right).sum()
	return float(intersection / union) if union else 0.0


	def greedy_mask_ious(pred_masks, gt_masks):
	matrix = np.zeros((len(pred_masks), len(gt_masks)), dtype=np.float64)
	for pred_index, pred_mask in enumerate(pred_masks):
	for gt_index, gt_mask in enumerate(gt_masks):
	matrix[pred_index, gt_index] = binary_iou(pred_mask, gt_mask)
	paired = []
	while matrix.size and matrix.max() > 0:
	pred_index, gt_index = np.unravel_index(matrix.argmax(), matrix.shape)
	paired.append(float(matrix[pred_index, gt_index]))
	matrix = np.delete(matrix, pred_index, axis=0)
	matrix = np.delete(matrix, gt_index, axis=1)
	return paired


	def greedy_mask_miou(pred_masks, gt_masks):
	paired = greedy_mask_ious(pred_masks, gt_masks)
	return float(np.mean(paired)) if paired else 0.0


	def extract_caption(saved) -> str:
	output = saved.get("output", [])
	raw = output if isinstance(output, str) else (output[0] if output else "")
	raw = "" if raw is None else str(raw)
	if "ASSISTANT:" in raw:
	raw = raw.split("ASSISTANT:", 1)[1]
	return " ".join(raw.strip().replace("</s>", "").split())


	def clean_caption(caption: str) -> str:
	for markup in (
	"<p> ",
	" </p> ",
	"[SEG]",
	"(IMAGE1)",
	"(IMAGE2)",
	"(IMAGE3)",
	):
	caption = caption.replace(markup, "")
	return re.sub(r"_\d+\b", "", caption).replace("_", " ").strip()


	def clean_reference_caption(caption: str) -> str:
	return re.sub(r"_\d+\b", "", caption).replace("_", " ").strip()


	def ptb_tokenize(captions):
	"""Run the same PTB tokenizer as COCOEvalCap with temp files in /tmp.

	The upstream package writes beside its JAR, which fails on read-only
	installations. Changing only the temporary-file location leaves the
	tokenizer command and post-processing identical.
	"""
	from pycocoevalcap.tokenizer import ptbtokenizer

	image_ids = [key for key, values in captions.items() for _ in values]
	sentences = "\n".join(
	value["caption"].replace("\n", " ")
	for values in captions.values()
	for value in values
	)
	jar_dir = Path(ptbtokenizer.__file__).resolve().parent
	jar_path = jar_dir / ptbtokenizer.STANFORD_CORENLP_3_4_1_JAR
	with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as handle:
	handle.write(sentences.encode())
	temporary_path = Path(handle.name)
	try:
	command = [
	"java",
	"-cp",
	str(jar_path),
	"edu.stanford.nlp.process.PTBTokenizer",
	"-preserveLines",
	"-lowerCase",
	str(temporary_path),
	]
	completed = subprocess.run(
	command,
	cwd=jar_dir,
	check=True,
	stdout=subprocess.PIPE,
	)
	finally:
	temporary_path.unlink(missing_ok=True)
	lines = completed.stdout.decode().split("\n")
	tokenized = {}
	for image_id, line in zip(image_ids, lines):
	words = [
	word
	for word in line.rstrip().split(" ")
	if word not in ptbtokenizer.PUNCTUATIONS
	]
	tokenized.setdefault(image_id, []).append(" ".join(words))
	return tokenized


	def coco_meteor(references, predictions) -> float:
	from pycocoevalcap.meteor.meteor import Meteor

	ground_truth = ptb_tokenize(references)
	results = ptb_tokenize(predictions)
	scorer = Meteor()
	score, _ = scorer.compute_score(ground_truth, results)
	return float(score)


	def convert_name(name: str) -> str:
	"""Normalize grounded labels using the paper's evaluation rules."""
	name = re.sub(r"\d+", "", name)
	name = name.replace("_", " ").strip()
	name = re.sub(r"_\d+\b", "", name)
	possessive = re.search(r"(?:'s\s\|’s\s)(.*)$", name, flags=re.IGNORECASE)
	if possessive:
	return possessive.group(1).strip()
	name = re.sub(r"\d+", "", name)
	return name.replace("_", "").strip()


	def extract_labels(caption: str, num_images: int):
	labels = [[] for _ in range(num_images)]
	pattern = re.compile(
	r"<p>\s(.?)\s</p>((?:\s\[SEG\]\s*\(IMAGE\d+\))+)",
	)
	for phrase, seg_block in pattern.findall(caption):
	phrase = convert_name(phrase.strip().lower())
	for image_number in re.findall(
	r"\[SEG\]\s*\(IMAGE(\d+)\)", seg_block
	):
	index = int(image_number) - 1
	if 0 <= index < num_images:
	labels[index].append(phrase)
	return labels


	def normalize_pred_masks(saved, num_images: int):
	raw = list(saved.get("pred_masks", []))
	# The model's saved structure has a sentinel at index zero.
	raw = raw[1 : num_images + 1]
	raw.extend([None] * (num_images - len(raw)))
	normalized = []
	for masks in raw:
	if masks is None or len(masks) == 0:
	normalized.append([])
	continue
	# Older cached runs wrap the per-image [N,H,W] tensor in a singleton
	# batch list, whereas newer evaluation saves that tensor directly.
	if isinstance(masks, (list, tuple)) and len(masks) == 1:
	masks = masks[0]
	if torch.is_tensor(masks):
	masks = masks.detach().cpu().numpy()
	else:
	masks = np.stack(
	[
	mask.detach().cpu().numpy() if torch.is_tensor(mask) else np.asarray(mask)
	for mask in masks
	]
	)
	normalized.append([(mask > 0) for mask in masks])
	return normalized


	def text_siou(left: str, right: str) -> float:
	left_words, right_words = set(left.split()), set(right.split())
	union = left_words \| right_words
	return len(left_words & right_words) / len(union) if union else 0.0


	class TextEmbedder:
	def __init__(self, name: str):
	from sentence_transformers import SentenceTransformer

	revision = DEFAULT_TEXT_MODEL_REVISION if name == DEFAULT_TEXT_MODEL else None
	self.model = SentenceTransformer(name, revision=revision)
	self.revision = getattr(self.model[0].auto_model.config, '_commit_hash', None)
	self.cache = {}

	def similarity(self, left: str, right: str) -> float:
	missing = [value for value in (left, right) if value not in self.cache]
	if missing:
	encoded = self.model.encode(missing, convert_to_numpy=True, normalize_embeddings=True)
	self.cache.update(zip(missing, encoded))
	return float(np.dot(self.cache[left], self.cache[right]))


	def grounded_matches(gt_masks, gt_labels, pred_masks, pred_labels, embedder):
	ious = np.zeros((len(gt_masks), len(pred_masks)), dtype=np.float64)
	for gt_index, gt_mask in enumerate(gt_masks):
	for pred_index, pred_mask in enumerate(pred_masks):
	ious[gt_index, pred_index] = binary_iou(gt_mask, pred_mask)
	similarities = np.zeros((len(gt_labels), len(pred_labels)), dtype=np.float64)
	sious = np.zeros_like(similarities)
	for gt_index, gt_label in enumerate(gt_labels):
	for pred_index, pred_label in enumerate(pred_labels):
	similarities[gt_index, pred_index] = embedder.similarity(gt_label, pred_label)
	sious[gt_index, pred_index] = text_siou(gt_label, pred_label)
	matches = 0
	similarity_sum = 0.0
	siou_sum = 0.0
	while ious.size:
	gt_index, pred_index = np.unravel_index(ious.argmax(), ious.shape)
	# The current evaluator treats a mask/label count mismatch as an
	# unmatched remainder instead of indexing outside the text matrix.
	if gt_index >= similarities.shape[0] or pred_index >= similarities.shape[1]:
	break
	if ious[gt_index, pred_index] < 0.5 or similarities[gt_index, pred_index] < 0.5:
	break
	matches += 1
	similarity_sum += similarities[gt_index, pred_index]
	siou_sum += sious[gt_index, pred_index]
	ious[gt_index, :] = 0
	ious[:, pred_index] = 0
	similarities[gt_index, :] = 0
	similarities[:, pred_index] = 0
	denom = max(matches, 1)
	return matches, similarity_sum / denom, siou_sum / denom


	def text_only_matches(gt_labels, pred_labels, embedder):
	similarities = np.zeros((len(gt_labels), len(pred_labels)), dtype=np.float64)
	sious = np.zeros_like(similarities)
	for gt_index, gt_label in enumerate(gt_labels):
	for pred_index, pred_label in enumerate(pred_labels):
	similarities[gt_index, pred_index] = embedder.similarity(gt_label, pred_label)
	sious[gt_index, pred_index] = text_siou(gt_label, pred_label)
	matches = 0
	similarity_sum = 0.0
	siou_sum = 0.0
	while similarities.size:
	gt_index, pred_index = np.unravel_index(similarities.argmax(), similarities.shape)
	if similarities[gt_index, pred_index] < 0.5:
	break
	matches += 1
	similarity_sum += similarities[gt_index, pred_index]
	siou_sum += sious[gt_index, pred_index]
	similarities[gt_index, :] = 0
	similarities[:, pred_index] = 0
	denom = max(matches, 1)
	return matches, similarity_sum / denom, siou_sum / denom


	def image_size(saved, image_index, groundings, image_info=None):
	"""Resolve native dimensions without requiring a separate annotation cache."""
	if image_info is not None:
	return int(image_info['height']), int(image_info['width'])
	sizes = saved.get('image_sizes', [])
	if image_index < len(sizes) and sizes[image_index] is not None:
	return tuple(map(int, sizes[image_index]))
	for grounding in groundings:
	segmentation = grounding['segmentation']
	if isinstance(segmentation, dict) and 'size' in segmentation:
	return tuple(map(int, segmentation['size']))
	paths = saved.get('image_paths', [])
	if image_index < len(paths) and Path(paths[image_index]).is_file():
	from PIL import Image
	with Image.open(paths[image_index]) as image:
	width, height = image.size
	return height, width
	return None


	def evaluate(
	prediction_dir: Path,
	prediction_files,
	records,
	selected_ids,
	ann_data=None,
	embedder=None,
	penalize_missing=False,
	with_meteor=False,
	):
	ann_data = ann_data or {}
	image_mious = []
	all_matched_pair_ious = []
	actual_positives = 0
	true_positives = 0
	image_matched_similarity = 0.0
	image_matched_siou = 0.0
	image_text_similarity = 0.0
	image_text_siou = 0.0
	samples = 0
	annotated_images = 0
	images_with_predictions = 0
	predicted_masks = 0
	load_errors = 0
	meteor_references = {}
	meteor_predictions = {}
	for record in records:
	pair_id = sample_identifier(record)
	if pair_id not in selected_ids:
	continue
	path = prediction_files.get(pair_id)
	if path is None and not penalize_missing:
	continue
	if path is None:
	saved = {"output": ["ASSISTANT: No prediction available. </s>"], "pred_masks": []}
	else:
	try:
	saved = torch.load(path, map_location="cpu")
	except (EOFError, OSError, RuntimeError):
	load_errors += 1
	if not penalize_missing:
	continue
	saved = {"output": ["ASSISTANT: No prediction available. </s>"], "pred_masks": []}
	num_images = len(record["img_ids"])
	caption = extract_caption(saved)
	labels_by_image = extract_labels(caption, num_images)
	masks_by_image = normalize_pred_masks(saved, num_images)
	samples += 1
	for image_index in range(num_images):
	groundings = record["groundings"].get(f"img_{image_index + 1}", [])
	if not groundings:
	continue
	pred_masks = masks_by_image[image_index]
	pred_labels = labels_by_image[image_index]
	gt_masks = []
	gt_labels = []
	dataset_name = str(record["dataset_name"])
	image_id = str(record["img_ids"][image_index])
	image_info = ann_data.get(dataset_name, {}).get(image_id)
	native_size = image_size(saved, image_index, groundings, image_info)
	if native_size is None:
	if pred_masks:
	raise ValueError(
	f'Native image size is missing for {pair_id}, image {image_index + 1}. '
	'Use predictions with image_sizes, restore the source images, or provide --ann-data.'
	)
	# An absent prediction scores zero without decoding polygon masks.
	image_mious.append(0.0)
	actual_positives += len(groundings)
	annotated_images += 1
	continue
	native_height, native_width = native_size
	for grounding in groundings:
	segmentation = grounding["segmentation"]
	gt_masks.append(decode_mask(segmentation, native_height, native_width))
	gt_label = str(grounding["name"]).lower()
	gt_labels.append(convert_name(gt_label))
	if len(pred_labels) < len(pred_masks):
	pred_labels.extend([""] * (len(pred_masks) - len(pred_labels)))
	paired_ious = greedy_mask_ious(pred_masks, gt_masks)
	image_mious.append(float(np.mean(paired_ious)) if paired_ious else 0.0)
	all_matched_pair_ious.extend(paired_ious)
	actual_positives += len(gt_labels)
	predicted_masks += len(pred_masks)
	annotated_images += 1
	images_with_predictions += int(bool(pred_masks))
	if with_meteor and pred_masks:
	caption_id = f"{pair_id}_{image_index + 1}"
	meteor_references[caption_id] = [
	{
	"caption": clean_reference_caption(
	str(record["a"])
	)
	}
	]
	meteor_predictions[caption_id] = [
	{"caption": clean_caption(caption)}
	]
	if embedder is not None and pred_masks and pred_labels:
	matched, similarity, siou = grounded_matches(
	gt_masks, gt_labels, pred_masks, pred_labels, embedder
	)
	true_positives += matched
	image_matched_similarity += similarity
	image_matched_siou += siou
	_, similarity, siou = text_only_matches(gt_labels, pred_labels, embedder)
	image_text_similarity += similarity
	image_text_siou += siou
	result = {
	"prediction_dir": str(prediction_dir),
	"samples": samples,
	"load_errors": load_errors,
	"annotated_images": annotated_images,
	"ground_truth_masks": actual_positives,
	"predicted_masks": predicted_masks,
	"miou": float(np.mean(image_mious)) if image_mious else 0.0,
	"matched_pairs": len(all_matched_pair_ious),
	"miou_over_all_matched_pairs": (
	float(np.mean(all_matched_pair_ious)) if all_matched_pair_ious else 0.0
	),
	"images_with_predictions": images_with_predictions,
	"miou_over_images_with_predictions": (
	float(np.sum(image_mious) / images_with_predictions)
	if images_with_predictions
	else 0.0
	),
	"images_with_positive_iou": sum(value > 0 for value in image_mious),
	"miou_over_images_with_positive_iou": (
	float(np.mean([value for value in image_mious if value > 0]))
	if any(value > 0 for value in image_mious)
	else 0.0
	),
	}
	if embedder is not None:
	result.update(
	{
	"recall": true_positives / max(actual_positives, 1),
	"iou_matched_semantic_similarity": image_matched_similarity
	/ max(annotated_images, 1),
	"iou_matched_text_siou": image_matched_siou / max(annotated_images, 1),
	"semantic_similarity": image_text_similarity / max(annotated_images, 1),
	"text_siou": image_text_siou / max(annotated_images, 1),
	}
	)
	if with_meteor:
	result["meteor_evaluated_images"] = len(meteor_predictions)
	result["METEOR"] = (
	coco_meteor(meteor_references, meteor_predictions)
	if meteor_predictions
	else 0.0
	)
	return result


	def score_predictions(prediction_dirs, records, ann_data=None, *,
	subset_mode="available", penalize_missing=True, with_text=False,
	with_meteor=False, text_model=DEFAULT_TEXT_MODEL, ground_truth=None):
	"""Score M4SEG predictions for the standalone CLI or post-generation workflow."""
	manifest_ids = {sample_identifier(record) for record in records}
	if len(manifest_ids) != len(records):
	raise ValueError('Evaluation annotations must have unique example identifiers.')
	prediction_files = []
	available = []
	for directory in prediction_dirs:
	files = index_prediction_files(directory)
	prediction_files.append(files)
	available.append(set(files) & manifest_ids)
	if subset_mode == "intersection":
	selected_sets = [set.intersection(available)] len(prediction_dirs)
	elif penalize_missing:
	selected_sets = [manifest_ids] * len(prediction_dirs)
	else:
	selected_sets = available
	embedder = TextEmbedder(text_model) if with_text else None
	report = {
	"ground_truth": str(ground_truth),
	"subset_mode": subset_mode,
	"text_model": text_model if with_text else None,
	"text_model_revision": getattr(embedder, 'revision', None),
	"penalize_missing": penalize_missing,
	"metric_definition": {
	"miou": "mean over annotated images of greedy matched-pair mean IoU",
	"missing_or_extra_masks": "unmatched masks are not added to the matched-pair denominator",
	"text_threshold": 0.5,
	"mask_iou_threshold": 0.5,
	},
	"results": [
	evaluate(
	directory,
	files,
	records,
	selected,
	ann_data,
	embedder,
	penalize_missing,
	with_meteor,
	)
	for directory, files, selected in zip(
	prediction_dirs, prediction_files, selected_sets
	)
	],
	}
	return report


	def main():
	args = parse_args()
	import gzip
	opener = gzip.open if str(args.gt_json).endswith('.gz') else open
	with opener(args.gt_json, 'rt', encoding='utf-8') as handle:
	records = json.load(handle)
	ann_data = torch.load(args.ann_data, map_location="cpu", weights_only=False) if args.ann_data else None
	report = score_predictions(
	args.prediction_dirs, records, ann_data,
	subset_mode=args.subset_mode, penalize_missing=args.penalize_missing,
	with_text=args.with_text, with_meteor=args.with_meteor,
	text_model=args.text_model, ground_truth=args.gt_json,
	)
	output = json.dumps(report, indent=2, sort_keys=True)
	print(output)
	if args.output:
	args.output.parent.mkdir(parents=True, exist_ok=True)
	args.output.write_text(output + "\n")


	if __name__ == "__main__":
	main()