""" Caption Child Agent — scene captioning powered by ArcisVLM. Responsibilities: - Process frames with task token prefix - Generate rich scene descriptions and summaries - Parse model output into structured captioning results - Support varying levels of detail (brief, standard, detailed) """ from __future__ import annotations import re import logging from pathlib import Path from typing import Any import torch from torchvision import transforms from agents.base import BaseAgent, Task, Result, EXPERT_CAPTION from agents.postprocess import enrich_result from model.vlm import VLJEPAModel from model.tokenizer import BPETokenizer logger = logging.getLogger(__name__) _DEFAULT_TRANSFORM = transforms.Compose([ transforms.Resize((448, 448)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # Scene attributes vocabulary for parsing captions _SCENE_TYPES = [ "indoor", "outdoor", "street", "road", "parking", "office", "warehouse", "lobby", "hallway", "entrance", "garden", "park", "intersection", "building", "store", "shop", "restaurant", "residential", ] _LIGHTING_TERMS = [ "bright", "dark", "dim", "daylight", "nighttime", "overcast", "sunny", "cloudy", "shadow", "illuminated", "fluorescent", ] _WEATHER_TERMS = [ "rain", "snow", "fog", "mist", "clear", "cloudy", "windy", "storm", ] class CaptionAgent(BaseAgent): """ Scene captioning agent. Uses the VL-JEPA + MoE model with the task token to produce descriptive scene summaries. Supports different verbosity levels through query phrasing. Args: agent_id: Unique identifier. model: Pre-loaded VLJEPAModel. tokenizer: BPETokenizer instance. device: Torch device. max_new_tokens: Maximum tokens to generate per query. temperature: Sampling temperature. transform: Image preprocessing transform. """ def __init__( self, agent_id: str, model: VLJEPAModel, tokenizer: BPETokenizer, device: str = "cpu", max_new_tokens: int = 256, temperature: float = 0.0, transform: transforms.Compose | None = None, ) -> None: super().__init__(agent_id=agent_id, expert_type=EXPERT_CAPTION) self.model = model self.tokenizer = tokenizer self.device = torch.device(device) self.max_new_tokens = max_new_tokens self.temperature = temperature self.transform = transform or _DEFAULT_TRANSFORM # (task tokens removed — training uses plain BOS/EOS from tokenizer) # -- Lifecycle ----------------------------------------------------------- def on_start(self) -> None: self.model.to(self.device) self.model.eval() logger.info("CaptionAgent %s: model on %s", self.agent_id, self.device) def on_stop(self) -> None: self.model.cpu() # -- Image loading ------------------------------------------------------- def _load_image(self, image_ref: str) -> torch.Tensor: path = Path(image_ref) if path.exists() and path.is_file(): from PIL import Image img = Image.open(path).convert("RGB") tensor = self.transform(img) return tensor.unsqueeze(0).to(self.device) raise FileNotFoundError(f"Image not found: {image_ref}") # -- Query preparation --------------------------------------------------- def _prepare_query(self, query: str) -> tuple[torch.Tensor, torch.Tensor]: """Tokenize query with BOS/EOS and pad to max_q=64 to match training format.""" max_q = 64 # Must match training ids = self.tokenizer.encode(query, add_special=True) # BOS + tokens + EOS pad_id = self.tokenizer.pad_id if len(ids) > max_q: ids = ids[:max_q] else: ids = ids + [pad_id] * (max_q - len(ids)) query_ids = torch.tensor([ids], dtype=torch.long, device=self.device) padding_mask = (query_ids != pad_id).long() return query_ids, padding_mask # -- Parsing model output ------------------------------------------------ @staticmethod def _parse_caption(answer: str) -> dict[str, Any]: """ Parse the model's free-form caption into structured scene attributes. Extracts: - scene_type: detected scene category (indoor, outdoor, street, etc.) - lighting: lighting conditions mentioned - weather: weather conditions if mentioned - objects_mentioned: list of notable objects referenced in the caption - sentence_count: number of sentences in the caption """ answer_lower = answer.lower() result: dict[str, Any] = { "scene_type": "unknown", "lighting": [], "weather": [], "objects_mentioned": [], "sentence_count": 0, } # Detect scene type for scene in _SCENE_TYPES: if re.search(rf"\b{re.escape(scene)}\b", answer_lower): result["scene_type"] = scene break # Detect lighting conditions for term in _LIGHTING_TERMS: if re.search(rf"\b{re.escape(term)}\b", answer_lower): result["lighting"].append(term) # Detect weather for term in _WEATHER_TERMS: if re.search(rf"\b{re.escape(term)}\b", answer_lower): result["weather"].append(term) # Extract mentioned objects object_pattern = re.compile( r"\b(person|people|car|vehicle|truck|bus|bicycle|motorcycle|" r"dog|cat|bird|chair|table|tree|building|sign|pedestrian|" r"traffic light|bench|fence|wall|door|window)\b", re.IGNORECASE, ) for match in object_pattern.finditer(answer_lower): obj = match.group(1).lower() if obj == "people": obj = "person" if obj not in result["objects_mentioned"]: result["objects_mentioned"].append(obj) # Count sentences sentences = re.split(r"[.!?]+", answer.strip()) result["sentence_count"] = len([s for s in sentences if s.strip()]) return result # -- Core processing ----------------------------------------------------- def process(self, task: Task) -> Result: """Run captioning inference for a single task.""" image_ref = task.image_ref query = task.query or "Describe this scene in detail." if not image_ref: return Result(answer="", error="No image_ref provided", expert_used=EXPERT_CAPTION) # Load image payload_tensor = task.payload.get("image_tensor") if payload_tensor is not None and isinstance(payload_tensor, torch.Tensor): if payload_tensor.dim() == 3: payload_tensor = payload_tensor.unsqueeze(0) images = payload_tensor.to(self.device) else: images = self._load_image(image_ref) # Prepare query with token query_ids, query_mask = self._prepare_query(query) # Generate with torch.no_grad(): generated_ids = self.model.generate( images=images, query_ids=query_ids, query_padding_mask=query_mask, max_new_tokens=self.max_new_tokens, temperature=self.temperature, ) answer = self.tokenizer.decode(generated_ids[0].tolist()) # Estimate confidence with torch.no_grad(): embed = self.model.get_embedding(images, query_ids, query_mask) norm = embed.norm(dim=-1).mean().item() confidence = min(1.0, max(0.0, norm / (norm + 1.0))) # Parse caption attributes caption_info = self._parse_caption(answer) result = Result( answer=answer, confidence=round(confidence, 4), expert_used=EXPERT_CAPTION, metadata={ "image_ref": image_ref, "query": query, "scene_type": caption_info["scene_type"], "lighting": caption_info["lighting"], "weather": caption_info["weather"], "objects_mentioned": caption_info["objects_mentioned"], "sentence_count": caption_info["sentence_count"], "tokens_generated": generated_ids.shape[1], }, ) return enrich_result(result, image_path=image_ref) # -- Health check -------------------------------------------------------- def health_check(self) -> bool: try: return self._started except Exception: return False