src/context_analysis.py

from __future__ import annotations

from dataclasses import dataclass

from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity

from .ijepa_localization import IJepaPatchLocalizer
from .obstacle_dataset import YoloBox


@dataclass
class ObjectContextResult:
    object_index: int
    class_name: str
    object_context_similarity: float
    scene_context_similarity: float
    context_pattern: str
    context_strength: str
    object_box_xyxy: tuple[int, int, int, int]
    context_box_xyxy: tuple[int, int, int, int]


def analyze_object_contexts(
    image: Image.Image,
    yolo_boxes: list[YoloBox],
    localizer: IJepaPatchLocalizer,
    margin_ratio: float = 0.75,
) -> list[ObjectContextResult]:
    if not yolo_boxes:
        return []

    scene_embedding = localizer.embed_image(image)
    results = []
    for index, box in enumerate(yolo_boxes, start=1):
        object_box = box.to_xyxy(*image.size)
        context_box = expand_box(object_box, image.size, margin_ratio=margin_ratio)
        object_crop = image.crop(object_box)
        context_crop = image.crop(context_box)
        if object_crop.width < 2 or object_crop.height < 2 or context_crop.width < 2 or context_crop.height < 2:
            continue
        object_embedding = localizer.embed_image(object_crop)
        context_embedding = localizer.embed_image(context_crop)

        object_context_similarity = vector_similarity(object_embedding, context_embedding)
        scene_context_similarity = vector_similarity(scene_embedding, context_embedding)
        results.append(
            ObjectContextResult(
                object_index=index,
                class_name=box.class_name,
                object_context_similarity=object_context_similarity,
                scene_context_similarity=scene_context_similarity,
                context_pattern=describe_context_pattern(
                    object_context_similarity,
                    scene_context_similarity,
                ),
                context_strength=describe_context_strength(
                    object_context_similarity,
                    scene_context_similarity,
                ),
                object_box_xyxy=object_box,
                context_box_xyxy=context_box,
            )
        )
    return results


def expand_box(
    box: tuple[int, int, int, int],
    image_size: tuple[int, int],
    margin_ratio: float,
) -> tuple[int, int, int, int]:
    image_width, image_height = image_size
    x1, y1, x2, y2 = box
    width = x2 - x1
    height = y2 - y1
    margin_x = width * margin_ratio
    margin_y = height * margin_ratio
    return (
        int(max(0, x1 - margin_x)),
        int(max(0, y1 - margin_y)),
        int(min(image_width - 1, x2 + margin_x)),
        int(min(image_height - 1, y2 + margin_y)),
    )


def vector_similarity(left, right) -> float:
    return float(cosine_similarity([left], [right])[0, 0])


def describe_context_pattern(object_context_similarity: float, scene_context_similarity: float) -> str:
    if object_context_similarity >= 0.92 and scene_context_similarity >= 0.92:
        return "near other objects / scene-embedded"
    if object_context_similarity >= 0.9:
        return "isolated / object-dominant"
    if scene_context_similarity >= 0.9:
        return "group / crowd context"
    return "distinct object/context"


def describe_context_strength(object_context_similarity: float, scene_context_similarity: float) -> str:
    strength = max(object_context_similarity, scene_context_similarity)
    if strength >= 0.93:
        return "high"
    if strength >= 0.86:
        return "medium"
    return "low"