Captions (pyavs.captions)

Loading transcribed participant captions and official MS-COCO captions, and computing multilingual caption embeddings – see Semantic Captioning Task.

Note

Like pyavs.scenes.embeddings, this module is intentionally not re-exported from top-level pyavs – it’s imported directly by the scripts that use it (e.g. captions/analyze_caption_similarity.py), not part of the core top-level API surface.

Loading Captions

Caption loading functions for pyAVS.

This module provides functions to load transcribed and MSCOCO captions from explog files.

pyavs.captions.load.parse_mscoco_captions(caption_string)[source]

Parse MSCOCO captions from string format to list of individual captions.

The captions are stored as a string representation of a list: “[‘caption1’, ‘caption2’, ‘caption3’, ‘caption4’, ‘caption5’]”

But often they appear concatenated without proper separators, so we need to split them using sentence patterns.

Parameters:

caption_string (str or list) – MSCOCO captions in string or list format

Returns:

List of individual caption strings (up to 5)

Return type:

list

pyavs.captions.load.load_coco_captions_for_scenes(scene_ids: List[int], coco_annotations_paths: str | List[str]) → Dict[int, List[str]][source]

Load COCO captions directly from annotations files for specific scene IDs.

This function can load from multiple annotation files (train + val) since AVS scenes are sampled from both COCO train and validation sets.

Parameters:

  • scene_ids (list of int) – List of scene IDs (which are COCO image IDs)

  • coco_annotations_paths (str or list of str) – Path(s) to COCO annotations JSON file(s). Can be a single file or list of files.

Returns:

Dictionary mapping scene_id to list of captions

Return type:

dict

pyavs.captions.load.find_coco_annotations(search_root: str) → List[str][source]

Try to find MSCOCO annotation files under search_root.

Since AVS scenes come from both COCO train and val sets, we need to find both. This function searches for and returns all available annotation files.

Note that the AVS release does not ship the raw MSCOCO annotation archives — only per-scene transformed annotations under stimuli/annotations/. Download annotations_trainval2017.zip from https://cocodataset.org/#download and point search_root (or load_captions(coco_annotations_path=...)) at it to use the COCO API.

Parameters:

search_root (str) – Directory to search in.

Returns:

List of paths to annotations files found

Return type:

list of str

pyavs.captions.load.load_captions(subjects: int | List[int], sessions: int | List[int], data_path: str | None = None, coco_annotations_path: str | List[str] | None = None, use_coco: bool = True) → DataFrame[source]

Load transcribed and MSCOCO captions from explog files.

Parameters:

  • subjects (int or list of int) – Subject ID(s) to load

  • sessions (int or list of int) – Session number(s) to load

  • data_path (str, optional) – Path to data directory (default: None, uses configured path)

  • coco_annotations_path (str or list of str, optional) – Path(s) to COCO annotations file(s) (default: None, auto-search if use_coco=True)

  • use_coco (bool, default True) – Whether to try loading COCO captions via API (falls back to parsing if fails)

Returns:

DataFrame with columns: subject, session, trial, block, scene_ID, transcribed_caption, mscoco_captions, caption_task

Return type:

pd.DataFrame

pyavs.captions.load.load_captions_for_scenes(scene_ids: List[int], subjects: int | List[int], sessions: int | List[int], data_path: str | None = None) → DataFrame[source]

Load captions for specific scene IDs.

Parameters:

  • scene_ids (list of int) – Scene IDs to load captions for

  • subjects (int or list of int) – Subject ID(s) to search

  • sessions (int or list of int) – Session number(s) to search

  • data_path (str, optional) – Path to data directory

Returns:

Filtered DataFrame containing only the specified scenes

Return type:

pd.DataFrame

pyavs.captions.load.inspect_explog_columns(subject: int, session: int, data_path: str | None = None) → List[str][source]

Inspect available columns in an explog file.

Parameters:

  • subject (int) – Subject ID

  • session (int) – Session number

  • data_path (str, optional) – Path to data directory

Returns:

Column names in the explog file

Return type:

list of str

MS-COCO Caption Loading

COCO caption loading functions for pyAVS.

This module provides functions to load captions directly from COCO annotations using the official pycocotools API, which is more reliable than parsing strings.

pyavs.captions.coco_loader.load_coco_captions_from_annotations(coco_annotations_path: str, scene_ids: List[int]) → Dict[int, List[str]][source]

Load COCO captions directly from annotations file.

Parameters:

  • coco_annotations_path (str) – Path to COCO annotations JSON file (e.g., instances_val2014.json)

  • scene_ids (list of int) – List of scene IDs (which are COCO image IDs)

Returns:

Dictionary mapping scene_id to list of captions

Return type:

dict

pyavs.captions.coco_loader.load_captions_with_coco(subjects: int | List[int], sessions: int | List[int], data_path: str | None = None, coco_annotations_path: str | None = None, fallback_to_parsing: bool = True) → DataFrame[source]

Load captions using COCO annotations API with fallback to string parsing.

Parameters:

  • subjects (int or list of int) – Subject ID(s) to load

  • sessions (int or list of int) – Session number(s) to load

  • data_path (str, optional) – Path to data directory

  • coco_annotations_path (str, optional) – Path to COCO annotations file. If None, will try to find it automatically.

  • fallback_to_parsing (bool, default True) – Whether to fall back to string parsing if COCO loading fails

Returns:

DataFrame with COCO captions loaded properly

Return type:

pd.DataFrame

pyavs.captions.coco_loader.find_coco_annotations(data_path: str | None = None) → str | None[source]

Try to find COCO annotations file in common locations.

Parameters:

data_path (str, optional) – Base data path to search in

Returns:

Path to annotations file if found

Return type:

str or None

pyavs.captions.coco_loader.get_coco_info(coco_annotations_path: str) → Dict[source]

Get information about the COCO dataset.

Parameters:

coco_annotations_path (str) – Path to COCO annotations file

Returns:

Information about the dataset

Return type:

dict

Caption Embeddings

Caption embedding functions for pyAVS.

This module provides functions to encode captions into embeddings using various language models. Default model is multilingual BERT for cross-language support.

pyavs.captions.embedding.encode_captions(captions: List[str] | Series, model_name: str = 'distiluse-base-multilingual-cased', model_type: str = 'sentence-transformers', batch_size: int = 32, max_length: int = 512, device: str | None = None, return_tensors: bool = False) → ndarray[source]

Encode captions into embeddings using specified language model.

Parameters:

  • captions (list of str or pd.Series) – Captions to encode

  • model_name (str, default 'distiluse-base-multilingual-cased') – Model name/path. Options: - ‘distiluse-base-multilingual-cased’ (default, fast multilingual) - ‘sentence-transformers/all-MiniLM-L12-v2’ (English) - ‘sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2’ (multilingual) - Any sentence-transformers or HuggingFace model name

  • model_type (str, default 'sentence-transformers') – Type of model loading: ‘sentence-transformers’ or ‘transformers’

  • batch_size (int, default 32) – Batch size for encoding

  • max_length (int, default 512) – Maximum sequence length

  • device (str, optional) – Device to use (‘cuda’, ‘cpu’, ‘mps’). Auto-detected if None.

  • return_tensors (bool, default False) – Return torch tensors instead of numpy arrays

Returns:

Embeddings array of shape (n_captions, embedding_dim)

Return type:

np.ndarray or torch.Tensor

pyavs.captions.embedding.encode_caption_dataframe(df: DataFrame, caption_columns: List[str] = ['transcribed_caption'], model_name: str = 'distiluse-base-multilingual-cased', model_type: str = 'sentence-transformers', batch_size: int = 32, max_length: int = 512, device: str | None = None, suffix: str = '_embedding') → DataFrame[source]

Encode caption columns in a DataFrame and add embedding columns.

Parameters:

  • df (pd.DataFrame) – DataFrame containing caption columns

  • caption_columns (list of str, default ['transcribed_caption']) – Column names containing captions to encode

  • model_name (str, default 'bert-base-multilingual-cased') – Model name for encoding

  • model_type (str, default 'transformers') – Type of model loading

  • batch_size (int, default 32) – Batch size for encoding

  • max_length (int, default 512) – Maximum sequence length

  • device (str, optional) – Device to use

  • suffix (str, default '_embedding') – Suffix to add to embedding column names

Returns:

DataFrame with additional embedding columns

Return type:

pd.DataFrame

pyavs.captions.embedding.encode_mscoco_captions(df: DataFrame, mscoco_column: str = 'mscoco_captions', model_name: str = 'distiluse-base-multilingual-cased', model_type: str = 'sentence-transformers', aggregation: str = 'mean', batch_size: int = 32, max_length: int = 512, device: str | None = None) → DataFrame[source]

Encode MSCOCO caption lists and aggregate them.

Parameters:

  • df (pd.DataFrame) – DataFrame containing MSCOCO caption lists

  • mscoco_column (str, default 'mscoco_captions') – Column name containing lists of MSCOCO captions

  • model_name (str, default 'bert-base-multilingual-cased') – Model name for encoding

  • model_type (str, default 'transformers') – Type of model loading

  • aggregation (str, default 'mean') – How to aggregate multiple captions: ‘mean’, ‘max’, ‘concat’, ‘individual’

  • batch_size (int, default 32) – Batch size for encoding

  • max_length (int, default 512) – Maximum sequence length

  • device (str, optional) – Device to use

Returns:

DataFrame with MSCOCO embedding column(s)

Return type:

pd.DataFrame

pyavs.captions.embedding.get_available_models() → Dict[str, List[str]][source]

Get list of recommended models for different use cases.

Returns:

Dictionary of model categories and recommended models

Return type:

dict