"""
Object detection and mapping for pyAVS package.
This module provides memory-efficient functions for mapping eye tracking fixations
to MSCOCO objects in scene images used in the Active Visual Semantics experiment.
Key features:
- Compressed mask storage using RLE encoding (90-95% space reduction)
- Spatial indexing for faster coordinate lookups
- On-demand mask computation and loading
- Memory usage scales with active objects, not total objects
"""
import os
import json
import numpy as np
import pandas as pd
from typing import List, Optional, Tuple, Dict, Union
from PIL import Image
import pycocotools.mask
from pycocotools.coco import COCO
from scipy.spatial.distance import euclidean
import pickle
from dataclasses import dataclass, asdict
from pathlib import Path
from ..layout import get_layout
from .cocostuff_classes import (
COCOSTUFF_CLASSES,
get_class_name,
is_thing_class,
is_stuff_class
)
[docs]
class CocoObjectMasker:
"""
Memory-efficient MSCOCO object masker using compressed storage.
Instead of storing full boolean masks, this class:
1. Stores RLE-compressed masks from COCO annotations directly
2. Creates spatial indices for fast coordinate lookups
3. Only decompresses masks when needed
This provides 90-95% reduction in storage space compared to full mask storage.
"""
[docs]
def __init__(self, annotation_dir: str, output_dir: str, mscoco_image_dir: str):
"""
Initialize the object masker.
Parameters
----------
annotation_dir : str
Path to MSCOCO annotation directory
output_dir : str
Path to output directory for compressed masks
mscoco_image_dir : str
Path to MSCOCO images directory
"""
self.annotation_dir = annotation_dir
self.output_dir = output_dir
self.mscoco_image_dir = mscoco_image_dir
os.makedirs(output_dir, exist_ok=True)
# Use JSON for metadata and separate files for compressed masks
self.metadata_file = os.path.join(output_dir, 'object_masks_metadata.json')
self.compressed_masks_dir = os.path.join(output_dir, 'compressed_masks')
os.makedirs(self.compressed_masks_dir, exist_ok=True)
self.metadata = self._load_metadata()
self._init_annotation_database()
def _load_metadata(self) -> Dict[str, List[ObjectMaskMetadata]]:
"""Load existing metadata or create empty structure."""
if os.path.exists(self.metadata_file):
with open(self.metadata_file, 'r') as f:
data = json.load(f)
# Convert back to ObjectMaskMetadata objects
metadata = {}
for scene_id, mask_list in data.items():
metadata[scene_id] = [ObjectMaskMetadata(**mask_data) for mask_data in mask_list]
return metadata
return {}
def _save_metadata(self):
"""Save metadata to disk."""
# Convert ObjectMaskMetadata objects to dicts for JSON serialization
serializable_data = {}
for scene_id, mask_list in self.metadata.items():
serializable_data[scene_id] = [asdict(mask_data) for mask_data in mask_list]
with open(self.metadata_file, 'w') as f:
json.dump(serializable_data, f, indent=2)
def _init_annotation_database(self):
"""Initialize the MSCOCO annotation database."""
self.coco = {}
for dataset_name in ['train2017', 'val2017']:
annotation_fname = os.path.join(self.annotation_dir, f'instances_{dataset_name}.json')
if os.path.exists(annotation_fname):
print(f'Loading annotation file: {annotation_fname}')
self.coco[dataset_name] = COCO(annotation_fname)
else:
print(f'Warning: Annotation file not found: {annotation_fname}')
[docs]
def compute_masks_for_image(self, coco_id: int):
"""
Compute and store compressed object masks for a single image.
Parameters
----------
coco_id : int
COCO image ID
"""
coco_id_str = str(coco_id)
# Skip if already processed
if coco_id_str in self.metadata:
return
# Find image dataset
dataset_name = None
for dataset in ['train2017', 'val2017']:
image_fname = f"{coco_id_str.zfill(12)}.jpg"
full_image_path = os.path.join(self.mscoco_image_dir, dataset, image_fname)
if os.path.exists(full_image_path):
dataset_name = dataset
break
if dataset_name is None or dataset_name not in self.coco:
print(f"Warning: Cannot process image {coco_id}")
return
# Get image dimensions
img_info = self.coco[dataset_name].loadImgs(coco_id)[0]
img_height, img_width = img_info['height'], img_info['width']
# Get annotations
ann_ids = self.coco[dataset_name].getAnnIds(imgIds=coco_id, iscrowd=None)
annotations = self.coco[dataset_name].loadAnns(ann_ids)
# Process annotations and create compressed masks
scene_metadata = []
# Group annotations by category to merge overlapping segments
category_annotations = {}
for ann in annotations:
cat_id = ann['category_id']
if cat_id not in category_annotations:
category_annotations[cat_id] = []
category_annotations[cat_id].append(ann)
for category_id, cat_annotations in category_annotations.items():
# Merge all masks for this category
merged_mask = np.zeros((img_height, img_width), dtype=bool)
for ann in cat_annotations:
# Convert annotation to mask
if 'segmentation' in ann:
mask = self.coco[dataset_name].annToMask(ann)
merged_mask = np.logical_or(merged_mask, mask.astype(bool))
if not merged_mask.any():
continue
# Calculate bounding box of merged mask
rows, cols = np.where(merged_mask)
if len(rows) == 0:
continue
bbox = (int(np.min(cols)), int(np.min(rows)),
int(np.max(cols) - np.min(cols) + 1),
int(np.max(rows) - np.min(rows) + 1))
# Create compressed mask key
mask_key = f"{coco_id}_{category_id}"
# Save compressed mask using RLE
rle = pycocotools.mask.encode(np.asfortranarray(merged_mask.astype(np.uint8)))
compressed_file = os.path.join(self.compressed_masks_dir, f"{mask_key}.rle")
# Save RLE data
with open(compressed_file, 'wb') as f:
pickle.dump(rle, f)
# Get category name
category_name = self.coco[dataset_name].loadCats(ids=category_id)[0]['name']
# Create metadata entry
mask_metadata = ObjectMaskMetadata(
scene_id=coco_id,
category_id=category_id,
category_name=category_name,
bbox=bbox,
area=int(np.sum(merged_mask)),
compressed_mask_key=mask_key
)
scene_metadata.append(mask_metadata)
# Store metadata
self.metadata[coco_id_str] = scene_metadata
self._save_metadata()
[docs]
def compute_masks(self, coco_ids: Union[int, List[int]]) -> str:
"""
Compute compressed object masks for multiple images.
Parameters
----------
coco_ids : int or list of int
COCO image ID(s)
Returns
-------
str
Path to metadata file
"""
if not isinstance(coco_ids, (list, tuple)):
coco_ids = [coco_ids]
for coco_id in coco_ids:
self.compute_masks_for_image(int(coco_id))
return self.metadata_file
[docs]
def load_mask_for_category(self, coco_id: int, category_id: int) -> Optional[np.ndarray]:
"""
Load and decompress mask for a specific category.
Parameters
----------
coco_id : int
COCO image ID
category_id : int
Object category ID
Returns
-------
np.ndarray or None
Decompressed boolean mask or None if not found
"""
coco_id_str = str(coco_id)
if coco_id_str not in self.metadata:
return None
# Find the mask metadata for this category
mask_meta = None
for meta in self.metadata[coco_id_str]:
if meta.category_id == category_id:
mask_meta = meta
break
if mask_meta is None:
return None
# Load and decompress the mask
compressed_file = os.path.join(self.compressed_masks_dir, f"{mask_meta.compressed_mask_key}.rle")
if not os.path.exists(compressed_file):
return None
try:
with open(compressed_file, 'rb') as f:
rle = pickle.load(f)
# Decode RLE to mask
mask = pycocotools.mask.decode(rle).astype(bool)
return mask
except Exception as e:
print(f"Error loading mask for {coco_id}/{category_id}: {e}")
return None
[docs]
def close(self):
"""Close any open resources."""
pass
[docs]
class FixationObjectChecker:
"""
Fixation object checker using transformed AVS scene annotations.
This class works with pre-transformed annotations that match the processed
scene format used in the AVS experiment. Annotations are loaded from
JSON files created by the AVS scene annotation transformer.
"""
[docs]
def __init__(self, transformed_annotations_dir: str, use_cocostuff: bool = False):
"""
Initialize the fixation object checker.
Parameters
----------
transformed_annotations_dir : str
Path to directory containing transformed annotation JSON files
use_cocostuff : bool, default=False
If True, use COCO-Stuff annotations (183 classes: 1 unlabeled + 80 things + 91 stuff + 11 missing).
If False, use standard COCO annotations (80 thing classes only).
"""
self.annotations_dir = Path(transformed_annotations_dir)
self.annotation_cache = {} # Cache for loaded annotations
self.use_cocostuff = use_cocostuff
self.num_classes = 183 if use_cocostuff else 91 # COCO-Stuff: 0-182, COCO: 1-90
# Get scene dimensions from config
from ..config.config import PyAVSConfig
config = PyAVSConfig()
self.scene_width = int(config.screen_size_pixels[0] * config.screen_usage)
self.scene_height = int(config.screen_size_pixels[1] * config.screen_usage)
def _load_scene_annotations(self, coco_id: int) -> Dict:
"""Load transformed annotations for a scene."""
if coco_id in self.annotation_cache:
return self.annotation_cache[coco_id]
annotation_file = self.annotations_dir / f"{coco_id}_transformed.json"
if not annotation_file.exists():
self.annotation_cache[coco_id] = {}
return {}
try:
with open(annotation_file, 'r') as f:
annotations = json.load(f)
# Validate annotation format
if not self._validate_annotation_format(annotations):
print(f"Warning: Invalid annotation format for scene {coco_id}")
self.annotation_cache[coco_id] = annotations
return annotations
except Exception as e:
print(f"Error loading annotations for scene {coco_id}: {e}")
self.annotation_cache[coco_id] = {}
return {}
def _validate_annotation_format(self, annotations: Dict) -> bool:
"""
Validate that category IDs in annotations are in expected range.
Parameters
----------
annotations : Dict
Loaded annotation dictionary with 'categories' key
Returns
-------
bool
True if all category IDs are valid, False otherwise
"""
if not annotations or 'categories' not in annotations:
return True # Empty annotations are valid
categories = annotations['categories']
for category_id in categories.keys():
cat_id_int = int(category_id)
if self.use_cocostuff:
# COCO-Stuff: valid range 0-182
if not (0 <= cat_id_int < 183):
print(f"Warning: Category ID {cat_id_int} out of COCO-Stuff range [0, 182]")
return False
else:
# Standard COCO: valid range 1-90
if not (1 <= cat_id_int <= 90):
print(f"Warning: Category ID {cat_id_int} out of COCO range [1, 90]")
return False
return True
def _get_category_name(self, category_id: int, fallback_name: str = None) -> str:
"""
Get category name from ID.
Parameters
----------
category_id : int
COCO or COCO-Stuff category ID
fallback_name : str, optional
Fallback name if lookup fails
Returns
-------
str
Category name
"""
if self.use_cocostuff:
# Use COCO-Stuff class list
name = get_class_name(category_id)
if name != 'unknown':
return name
# Fallback: use provided name or generate from ID
if fallback_name:
return fallback_name
return f"category_{category_id}"
def _decode_rle_mask(self, rle_data: Dict) -> np.ndarray:
"""Decode RLE mask data."""
# Reconstruct RLE format for pycocotools
rle = {
'size': rle_data['size'],
'counts': rle_data['counts'].encode('utf-8') if isinstance(rle_data['counts'], str) else rle_data['counts']
}
return pycocotools.mask.decode(rle).astype(bool)
def _compute_distance_to_nearest_object_pixel(self, mask: np.ndarray,
fix_x: int, fix_y: int,
search_radius: int) -> float:
"""
Compute Euclidean distance from fixation to nearest object pixel within search area.
Parameters
----------
mask : np.ndarray
Boolean object mask
fix_x, fix_y : int
Fixation coordinates in image space
search_radius : int
Search radius in pixels
Returns
-------
float
Distance to nearest object pixel, or np.inf if no object found
"""
from scipy.spatial.distance import euclidean
# Define search area bounds
y_start = max(0, fix_y - search_radius)
y_end = min(mask.shape[0], fix_y + search_radius + 1)
x_start = max(0, fix_x - search_radius)
x_end = min(mask.shape[1], fix_x + search_radius + 1)
# Extract search area from mask
search_area = mask[y_start:y_end, x_start:x_end]
if not search_area.any():
return np.inf
# Find all object pixels in search area
obj_pixels_local = np.where(search_area)
if len(obj_pixels_local[0]) == 0:
return np.inf
# Convert to absolute coordinates
obj_y = obj_pixels_local[0] + y_start
obj_x = obj_pixels_local[1] + x_start
# Calculate distances to all object pixels
distances = np.sqrt((obj_x - fix_x)**2 + (obj_y - fix_y)**2)
return float(np.min(distances))
[docs]
def get_fixated_objects(self, coco_id: int, x_pos: Union[float, np.ndarray],
y_pos: Union[float, np.ndarray],
error_margin_pixels: int = 10) -> Tuple[List[int], List[str]]:
"""
Check which objects are fixated at given coordinates with error margin tolerance.
This method first checks for direct hits at the exact fixation coordinates. If no
object is found, it searches within an error margin around the fixation to account
for eye tracker noise and calibration drift.
Parameters
----------
coco_id : int
COCO image ID
x_pos : float or array
Screen-centered x coordinates (pixels)
y_pos : float or array
Screen-centered y coordinates (pixels)
error_margin_pixels : int, optional
Search radius in pixels around fixation for nearest object (default: 10)
This accounts for eye tracker noise and calibration drift.
Returns
-------
tuple
(object_category_ids, object_category_names)
"""
# Load annotations for this scene
annotations = self._load_scene_annotations(coco_id)
if not annotations:
# No annotations available
x_pos = np.atleast_1d(x_pos)
return [-1] * len(x_pos), ['None'] * len(x_pos)
# Convert to arrays
x_pos = np.atleast_1d(x_pos)
y_pos = np.atleast_1d(y_pos)
if x_pos.shape != y_pos.shape:
raise ValueError("x_pos and y_pos must have the same shape")
object_cats = []
category_names = []
# Process each fixation position
for i in range(len(x_pos)):
# Convert screen-centered coordinates to image coordinates
x_img = int(x_pos[i] + self.scene_width / 2)
y_img = int(abs(y_pos[i] - self.scene_height / 2))
# Check if fixation is outside scene boundaries
if (x_img < 0 or x_img >= self.scene_width or
y_img < 0 or y_img >= self.scene_height):
object_cats.append(-2)
category_names.append('outside')
continue
# Step 1: Check for direct hits at exact fixation coordinates
direct_hit_objects = []
for cat_id_str, obj_data in annotations.items():
# Check bounding box first (fast pre-filter)
bbox_x, bbox_y, bbox_w, bbox_h = obj_data['bbox']
if (bbox_x <= x_img <= bbox_x + bbox_w and
bbox_y <= y_img <= bbox_y + bbox_h):
# Decode and check actual mask for exact hit
try:
mask = self._decode_rle_mask(obj_data['rle'])
if (y_img < mask.shape[0] and x_img < mask.shape[1] and
mask[y_img, x_img]):
direct_hit_objects.append(obj_data)
except Exception as e:
print(f"Error decoding mask for category {cat_id_str}: {e}")
continue
# Handle direct hits
if len(direct_hit_objects) == 1:
# Single direct hit
obj = direct_hit_objects[0]
object_cats.append(obj['category_id'])
category_names.append(obj['category_name'])
elif len(direct_hit_objects) > 1:
# Multiple direct hits: choose smallest area
min_area_obj = min(direct_hit_objects, key=lambda x: x['area'])
object_cats.append(min_area_obj['category_id'])
category_names.append(min_area_obj['category_name'])
else:
# Step 2: No direct hit - search within error margin
candidate_objects = []
for cat_id_str, obj_data in annotations.items():
# Expand bounding box check to include error margin
bbox_x, bbox_y, bbox_w, bbox_h = obj_data['bbox']
# Check if fixation is within error margin of bounding box
if (bbox_x - error_margin_pixels <= x_img <= bbox_x + bbox_w + error_margin_pixels and
bbox_y - error_margin_pixels <= y_img <= bbox_y + bbox_h + error_margin_pixels):
try:
mask = self._decode_rle_mask(obj_data['rle'])
# Compute distance to nearest object pixel within search radius
distance = self._compute_distance_to_nearest_object_pixel(
mask, x_img, y_img, error_margin_pixels
)
# Only consider objects within the error margin
if distance <= error_margin_pixels:
candidate_objects.append((obj_data, distance))
except Exception as e:
print(f"Error processing object {cat_id_str} for error margin search: {e}")
continue
# Select closest object within error margin
if candidate_objects:
# Sort by distance and choose closest
closest_obj, closest_distance = min(candidate_objects, key=lambda x: x[1])
object_cats.append(closest_obj['category_id'])
category_names.append(closest_obj['category_name'])
else:
# No objects found even with error margin
object_cats.append(-1)
category_names.append('None')
return object_cats, category_names
[docs]
def clear_cache(self):
"""Clear the annotation cache."""
self.annotation_cache.clear()
[docs]
def get_fixated_objects(events_df: pd.DataFrame,
transformed_annotations_dir: str,
verbose: bool = False,
error_margin_pixels: int = 10,
use_cocostuff: bool = True) -> pd.DataFrame:
"""
Add object labels to fixation events using transformed AVS scene annotations.
This function uses pre-transformed annotations that match the processed scene
format used in the AVS experiment, providing more accurate object detection.
Includes error margin tolerance to account for eye tracker noise and calibration drift.
Parameters
----------
events_df : pd.DataFrame
Eye tracking events dataframe
transformed_annotations_dir : str
Path to directory containing transformed annotation JSON files
verbose : bool, optional
Whether to print progress information (default: False)
error_margin_pixels : int, optional
Search radius in pixels around fixation for nearest object (default: 10)
This accounts for eye tracker noise and calibration drift.
use_cocostuff : bool, optional
If True, use COCO-Stuff annotations (172 classes: 80 things + 91 stuff + 1 unlabeled).
If False, use standard COCO annotations (80 thing classes only).
Default is True (COCO-Stuff mode).
Returns
-------
pd.DataFrame
Events dataframe with object_label and object_id columns added
Notes
-----
COCO-Stuff mode (use_cocostuff=True) provides better coverage by including
amorphous background regions like sky, grass, walls, water, etc. This typically
increases fixation labeling coverage by 20-40% compared to COCO-only mode.
"""
# Initialize fixation object checker
checker = FixationObjectChecker(transformed_annotations_dir, use_cocostuff=use_cocostuff)
# Add object label columns
events_df = events_df.copy()
events_df['object_label'] = pd.Series(dtype=str)
events_df['object_id'] = pd.Series(dtype=float)
def center_pixel_coords(pix_coords, screen_size_pix):
"""Center pixel coordinates around zero."""
return pix_coords - screen_size_pix / 2
# Get screen size from config
from ..config.config import PyAVSConfig
config = PyAVSConfig()
screen_x_pix, screen_y_pix = config.screen_size_pixels
# Process each subject and trial
subjects = events_df.subject.unique()
total_processed = 0
for subject in subjects:
subject_mask = events_df.subject == subject
for trial in events_df[subject_mask].trial.unique():
if pd.isna(trial):
continue
# Get scene ID for this trial
trial_mask = (events_df.subject == subject) & (events_df.trial == trial)
scene_ids = events_df[trial_mask].sceneID.dropna().unique()
if len(scene_ids) == 0:
continue
scene_id = int(scene_ids[0])
# Process fixations and saccades separately
for et_type in ['fixation', 'saccade']:
type_mask = trial_mask & (events_df['type'] == et_type)
if not type_mask.any():
continue
trial_events = events_df[type_mask].reset_index(drop=True)
# Get appropriate coordinates
coord_type = 'mean' if et_type == 'fixation' else 'end'
x_coords = center_pixel_coords(trial_events[f"{coord_type}_gx"], screen_x_pix)
y_coords = center_pixel_coords(trial_events[f"{coord_type}_gy"], screen_y_pix)
# Get object labels using transformed annotations with error margin
try:
object_cat_ids, object_cat_labels = checker.get_fixated_objects(
coco_id=scene_id,
x_pos=x_coords,
y_pos=y_coords,
error_margin_pixels=error_margin_pixels
)
# Add labels to dataframe
events_df.loc[type_mask, 'object_label'] = object_cat_labels
events_df.loc[type_mask, 'object_id'] = object_cat_ids
total_processed += len(object_cat_ids)
except Exception as e:
if verbose:
print(f'Error processing scene {scene_id}: {e}')
continue
if verbose:
print(f"Processed {total_processed} events with object labels")
# Convert object_id to integer
events_df['object_id'] = events_df['object_id'].astype('Int64')
return events_df
_MASKS_NOT_IN_RELEASE = (
"Precomputed RLE object masks are not part of the public AVS release. The release "
"ships transformed scene annotations at stimuli/annotations/{coco_objects,cocostuff} "
"instead — use get_fixated_objects() or FixationObjectChecker, which read those "
"directly. To use RLE masks, build them yourself from a raw MSCOCO annotation "
"download with CocoObjectMasker and pass the output directory as masks_dir."
)
[docs]
def load_object_masks(scene_ids: Union[int, List[int]],
masks_dir: Optional[str] = None) -> Dict[int, Dict[str, np.ndarray]]:
"""
Load precomputed RLE object masks for specified scene IDs.
Parameters
----------
scene_ids : int or list of int
Scene ID(s) to load masks for
masks_dir : str
Directory holding ``object_masks_metadata.json`` and
``compressed_masks/``, as produced by :class:`CocoObjectMasker`.
Returns
-------
dict
Dictionary mapping scene IDs to object masks
Raises
------
FileNotFoundError
If ``masks_dir`` is not given, or does not contain the metadata file.
These masks are **not** part of the public release — see
:func:`get_fixated_objects` for the annotation-based equivalent.
"""
if masks_dir is None:
raise FileNotFoundError(_MASKS_NOT_IN_RELEASE)
metadata_file = os.path.join(masks_dir, 'object_masks_metadata.json')
mask_files_dir = os.path.join(masks_dir, 'compressed_masks')
if not os.path.exists(metadata_file):
raise FileNotFoundError(
f"Object mask metadata not found: {metadata_file}. {_MASKS_NOT_IN_RELEASE}"
)
if isinstance(scene_ids, int):
scene_ids = [scene_ids]
# Load metadata
with open(metadata_file, 'r') as f:
data = json.load(f)
metadata = {}
for scene_id, mask_list in data.items():
metadata[scene_id] = [ObjectMaskMetadata(**mask_data) for mask_data in mask_list]
masks = {}
for scene_id in scene_ids:
scene_id_str = str(scene_id)
if scene_id_str in metadata:
masks[scene_id] = {}
for mask_meta in metadata[scene_id_str]:
# Load compressed mask
compressed_file = os.path.join(mask_files_dir, f"{mask_meta.compressed_mask_key}.rle")
if os.path.exists(compressed_file):
try:
with open(compressed_file, 'rb') as f:
rle = pickle.load(f)
mask = pycocotools.mask.decode(rle).astype(bool)
masks[scene_id][str(mask_meta.category_id)] = mask
except Exception as e:
print(f"Error loading mask for scene {scene_id}, category {mask_meta.category_id}: {e}")
else:
print(f"Warning: No masks found for scene {scene_id}")
return masks
[docs]
def map_fixations_to_objects(fixations_df: pd.DataFrame,
scene_id: int,
x_col: str = 'mean_gx',
y_col: str = 'mean_gy',
data_path: Optional[str] = None,
use_cocostuff: bool = False,
transformed_annotations_dir: Optional[str] = None) -> pd.DataFrame:
"""
Map fixations to objects for a single scene.
Parameters
----------
fixations_df : pd.DataFrame
Dataframe containing fixation data
scene_id : int
COCO scene ID
x_col : str, optional
Column name for x coordinates (default: 'mean_gx')
y_col : str, optional
Column name for y coordinates (default: 'mean_gy')
data_path : str, optional
``avs-public`` root. If None, uses the configured data path.
use_cocostuff : bool, optional
Use the COCO-Stuff annotations (183 classes) rather than the 80 COCO
thing classes (default: False).
transformed_annotations_dir : str, optional
Explicit annotation directory, overriding ``data_path``.
Returns
-------
pd.DataFrame
Fixations dataframe with ``object_id`` and ``object_label`` columns added.
"""
if transformed_annotations_dir is None:
kind = 'cocostuff' if use_cocostuff else 'coco_objects'
transformed_annotations_dir = get_layout(data_path).annotations_dir(kind)
fix_checker = FixationObjectChecker(str(transformed_annotations_dir),
use_cocostuff=use_cocostuff)
object_ids, object_labels = fix_checker.get_fixated_objects(
coco_id=scene_id,
x_pos=fixations_df[x_col].values,
y_pos=fixations_df[y_col].values,
look_up_closest=True
)
result_df = fixations_df.copy()
result_df['object_id'] = object_ids
result_df['object_label'] = object_labels
return result_df
# Original 80 COCO object classes (backward compatibility)
# For COCO-Stuff (172 classes: 80 things + 91 stuff + 1 unlabeled), use cocostuff_classes.COCOSTUFF_CLASSES
MSCOCO_CLASSES = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush'
]
# RSA Analysis Category Dictionary
RSA_CATEGORIES = {
# ANIMATE
'animate': {
'human': ['person'],
'mammal_large': ['bear', 'cow', 'elephant', 'giraffe', 'horse', 'zebra'],
'mammal_small': ['cat', 'dog', 'sheep'],
'bird': ['bird']
},
# INANIMATE
'inanimate': {
# Transportation
'vehicle_air': ['airplane'],
'vehicle_ground': ['bicycle', 'bus', 'car', 'motorcycle', 'skateboard', 'train', 'truck'],
'vehicle_water': ['boat', 'surfboard'],
# Food & Kitchen
'food_natural': ['apple', 'banana', 'broccoli', 'carrot', 'orange'],
'food_prepared': ['cake', 'donut', 'hot dog', 'pizza', 'sandwich'],
'kitchenware': ['bottle', 'bowl', 'cup', 'fork', 'knife', 'spoon', 'wine glass'],
'appliances': ['microwave', 'oven', 'refrigerator', 'toaster'],
# Furniture & Indoor
'furniture': ['bed', 'bench', 'chair', 'couch', 'dining table'],
'household_items': ['clock', 'potted plant', 'sink', 'toilet', 'tv', 'vase'],
# Technology & Electronics
'electronics': ['cell phone', 'keyboard', 'laptop', 'mouse', 'remote'],
'tools_appliances': ['hair drier', 'scissors', 'toothbrush'],
# Personal Items & Accessories
'bags_accessories': ['backpack', 'handbag', 'suitcase', 'tie', 'umbrella'],
# Sports & Recreation
'sports_equipment': ['baseball bat', 'baseball glove', 'frisbee', 'kite', 'skis',
'snowboard', 'sports ball', 'tennis racket'],
# Books & Media
'media': ['book'],
# Urban/Public Objects
'urban_infrastructure': ['fire hydrant', 'parking meter', 'stop sign', 'traffic light'],
# Toys
'toys': ['teddy bear']
}
}
# Create flattened lookup dictionary
_CATEGORY_LOOKUP = {}
for main_category, subcategories in RSA_CATEGORIES.items():
for subcategory, items in subcategories.items():
for item in items:
_CATEGORY_LOOKUP[item] = {
'main_category': main_category,
'subcategory': subcategory,
'hierarchical': f"{main_category}_{subcategory}"
}
[docs]
def categorize_objects(object_names: List[str], level: str = 'subcategory') -> List[str]:
"""
Categorize object names into broader categories for RSA analysis.
Parameters
----------
object_names : list of str
List of object names to categorize
level : str, optional
Categorization level: 'main_category', 'subcategory', or 'hierarchical'
Default: 'subcategory'
Returns
-------
list of str
List of category names for each object
Examples
--------
>>> categorize_objects(['person', 'car', 'dog'], level='main_category')
['animate', 'inanimate', 'animate']
>>> categorize_objects(['person', 'car', 'dog'], level='subcategory')
['human', 'vehicle_ground', 'mammal_small']
"""
categories = []
for obj_name in object_names:
if obj_name in _CATEGORY_LOOKUP:
categories.append(_CATEGORY_LOOKUP[obj_name][level])
else:
categories.append('unknown')
return categories
[docs]
def sort_objects_by_category(object_names: List[str], level: str = 'subcategory') -> Tuple[List[str], List[int]]:
"""
Sort objects by their categories and return sorted objects with indices.
Parameters
----------
object_names : list of str
List of object names to sort
level : str, optional
Categorization level for sorting: 'main_category', 'subcategory', or 'hierarchical'
Default: 'subcategory'
Returns
-------
tuple
(sorted_objects, sort_indices) where sort_indices maps new positions to original positions
Examples
--------
>>> objects = ['car', 'person', 'dog']
>>> sorted_objs, indices = sort_objects_by_category(objects)
>>> print(sorted_objs) # ['person', 'dog', 'car'] (animate first, then inanimate)
>>> print(indices) # [1, 2, 0] (person was at index 1, dog at 2, car at 0)
"""
# Get categories for all objects
categories = categorize_objects(object_names, level=level)
# Create tuples of (category, object_name, original_index)
obj_tuples = [(cat, obj, i) for i, (cat, obj) in enumerate(zip(categories, object_names))]
# Sort by category name (this will group similar categories together)
obj_tuples.sort(key=lambda x: x[0])
# Extract sorted objects and indices
sorted_objects = [obj for _, obj, _ in obj_tuples]
sort_indices = [orig_idx for _, _, orig_idx in obj_tuples]
return sorted_objects, sort_indices
# =============================
# COCO SUPERCATEGORY MAPPING
# =============================
# Official COCO-Stuff supercategory mapping (covers all 171 RSA object labels)
COCO_SUPERCATEGORY_MAP = {
# person
'person': 'person',
# animal
'bird': 'animal', 'cat': 'animal', 'dog': 'animal', 'horse': 'animal',
'sheep': 'animal', 'cow': 'animal', 'elephant': 'animal', 'bear': 'animal',
'zebra': 'animal', 'giraffe': 'animal',
# vehicle
'bicycle': 'vehicle', 'car': 'vehicle', 'motorcycle': 'vehicle',
'airplane': 'vehicle', 'bus': 'vehicle', 'train': 'vehicle',
'truck': 'vehicle', 'boat': 'vehicle',
# outdoor (things)
'traffic light': 'outdoor', 'fire hydrant': 'outdoor',
'stop sign': 'outdoor', 'parking meter': 'outdoor',
'bench': 'outdoor', 'street sign': 'outdoor',
# sports
'frisbee': 'sports', 'skis': 'sports', 'snowboard': 'sports',
'sports ball': 'sports', 'kite': 'sports',
'baseball bat': 'sports', 'baseball glove': 'sports',
'skateboard': 'sports', 'surfboard': 'sports', 'tennis racket': 'sports',
# accessory
'backpack': 'accessory', 'umbrella': 'accessory',
'handbag': 'accessory', 'tie': 'accessory',
'suitcase': 'accessory', 'shoe': 'accessory',
'eye glasses': 'accessory', 'hat': 'accessory',
# appliance
'microwave': 'appliance', 'oven': 'appliance',
'toaster': 'appliance', 'sink': 'appliance',
'refrigerator': 'appliance', 'blender': 'appliance',
# electronic
'tv': 'electronic', 'laptop': 'electronic',
'mouse': 'electronic', 'remote': 'electronic',
'keyboard': 'electronic', 'cell phone': 'electronic',
# furniture (things + stuff combined)
'chair': 'furniture', 'couch': 'furniture', 'potted plant': 'furniture',
'bed': 'furniture', 'dining table': 'furniture', 'toilet': 'furniture',
'window': 'furniture', 'desk': 'furniture', 'mirror': 'furniture',
'door': 'furniture',
'furniture-other': 'furniture', 'stairs': 'furniture', 'light': 'furniture',
'counter': 'furniture', 'cupboard': 'furniture', 'cabinet': 'furniture',
'shelf': 'furniture', 'table': 'furniture',
# food (things + stuff combined)
'banana': 'food', 'apple': 'food', 'sandwich': 'food',
'orange': 'food', 'broccoli': 'food', 'carrot': 'food',
'hot dog': 'food', 'pizza': 'food', 'donut': 'food', 'cake': 'food',
'food-other': 'food', 'vegetable': 'food', 'salad': 'food', 'fruit': 'food',
# kitchen
'bottle': 'kitchen', 'wine glass': 'kitchen', 'cup': 'kitchen',
'fork': 'kitchen', 'knife': 'kitchen', 'spoon': 'kitchen',
'bowl': 'kitchen', 'plate': 'kitchen',
# indoor misc (things)
'hair brush': 'indoor', 'toothbrush': 'indoor', 'hair dryer': 'indoor',
'teddy bear': 'indoor', 'scissors': 'indoor', 'vase': 'indoor',
'clock': 'indoor', 'book': 'indoor',
# water
'water-other': 'water', 'waterdrops': 'water', 'sea': 'water',
'river': 'water', 'fog': 'water',
# ground
'ground-other': 'ground', 'playingfield': 'ground', 'platform': 'ground',
'railroad': 'ground', 'pavement': 'ground', 'road': 'ground',
'gravel': 'ground', 'dirt': 'ground', 'snow': 'ground', 'sand': 'ground',
# solid
'solid-other': 'solid', 'hill': 'solid', 'mountain': 'solid',
'stone': 'solid', 'rock': 'solid', 'wood': 'solid',
# sky
'sky-other': 'sky', 'clouds': 'sky',
# plant
'plant-other': 'plant', 'straw': 'plant', 'moss': 'plant',
'branch': 'plant', 'flower': 'plant', 'leaves': 'plant',
'bush': 'plant', 'tree': 'plant', 'grass': 'plant',
# structural
'structural-other': 'structural', 'railing': 'structural',
'net': 'structural', 'cage': 'structural', 'fence': 'structural',
# building
'building-other': 'building', 'roof': 'building', 'tent': 'building',
'brick': 'building', 'skyscraper': 'building', 'house': 'building',
# textile
'textile-other': 'textile', 'banner': 'textile', 'pillow': 'textile',
'blanket': 'textile', 'curtain': 'textile', 'cloth': 'textile',
'clothes': 'textile', 'napkin': 'textile', 'towel': 'textile',
'mat': 'textile', 'rug': 'textile',
# window
'window-other': 'window', 'window-blind': 'window',
# floor
'floor-other': 'floor', 'floor-stone': 'floor', 'floor-marble': 'floor',
'floor-wood': 'floor', 'floor-tile': 'floor', 'floor-carpet': 'floor',
# ceiling
'ceiling-other': 'ceiling', 'ceiling-tile': 'ceiling',
# wall
'wall-other': 'wall', 'wall-concrete': 'wall', 'wall-stone': 'wall',
'wall-brick': 'wall', 'wall-wood': 'wall', 'wall-panel': 'wall',
'wall-tile': 'wall',
# raw material
'metal': 'raw material', 'plastic': 'raw material',
'paper': 'raw material', 'cardboard': 'raw material',
}
# Preferred display order for supercategories (animate → inanimate things → stuff)
SUPERCATEGORY_ORDER = [
'person', 'animal', # animate
'vehicle', 'outdoor', 'sports', 'accessory', # inanimate things
'appliance', 'electronic', 'kitchen', 'food',
'furniture', 'indoor',
'water', 'sky', 'plant', 'ground', 'solid', # stuff: outdoor
'building', 'structural', 'wall', 'ceiling', # stuff: built environment
'floor', 'window', 'textile', 'raw material', # stuff: surfaces & materials
]
_SUPERCATEGORY_PALETTE = None # lazily initialised
[docs]
def get_supercategory_palette():
"""Return {supercategory: (r,g,b)} using husl palette, cached."""
global _SUPERCATEGORY_PALETTE
if _SUPERCATEGORY_PALETTE is None:
import seaborn as sns
colors = sns.color_palette('husl', len(SUPERCATEGORY_ORDER))
_SUPERCATEGORY_PALETTE = dict(zip(SUPERCATEGORY_ORDER, colors))
_SUPERCATEGORY_PALETTE['unknown'] = (0.6, 0.6, 0.6)
return _SUPERCATEGORY_PALETTE