brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/model_helpers/activations/temporal/core/inferencer/video/base.py

@@ -0,0 +1,134 @@
+import numpy as np
+from typing import Union, Tuple, Callable, Hashable, List, Dict
+from pathlib import Path
+
+from brainscore_vision.model_helpers.activations.temporal.inputs import Video, Stimulus
+from brainscore_vision.model_helpers.activations.temporal.utils import assembly_align_to_fps, stack_with_nan_padding
+from brainio.assemblies import NeuroidAssembly
+
+from ..base import Inferencer
+from . import time_aligner as time_aligners 
+
+
+class TemporalInferencer(Inferencer):
+    """Inferencer for video stimuli. The model takes video stimuli as input and generate the activations over time.
+    Then, the activations will be aligned to video time by the time_aligner specified in the constructor. The aligned
+    activations will be again unified to the fps specified within the constructor (self.fps). Finally, the activations
+    will be packaged into a NeuroidAssembly.
+    
+    NOTE: for all the time_alignment method, the inference of time bins will only be done with the longest video, but ignore all other input videos.
+
+    Example:
+        temporal_inferencer = TemporalInferenver(..., fps=10)
+        model_assembly = temporal_inferencer(video_paths[1000ms], layers)
+        model_assembly.time_bins -> [(0, 100), (100, 200), ..., (900, 1000)]  # 1000ms, 10fps
+
+    Parameters
+    ----------
+    fps: float
+        frame rate of the model sampling.
+
+    num_frames: int, or (int, int)
+        - If None, the model accepts videos of any length.
+        - If a single int is passed, specify how many frames the model takes. 
+        - If a tuple of two ints is passed, specify the range of the number of frames the model takes (inclusive). If you need to specify infinite, use np.inf.
+    
+    duration: float, or (float, float)
+        - If None, the model accepts videos of any length.
+        - If a single float is passed, specify the duration of the model takes, in ms.
+        - If a tuple of two floats is passed, specify the range of the duration the model takes (inclusive). If you need to specify infinite, use np.inf.
+    
+    time_alignment: str
+        specify the method to align the activations in time.
+        The options and specifications are in the time_aligners module. The current options are:
+        - evenly_spaced: align the activations to have evenly spaced time bins across the whole video time span.
+        - ignore_time: ignore the time information and make a single time bin of the entire video.
+        - estimate_layer_fps: estimate the fps of the layer based on the video fps.
+        - per_frame_aligned: align the activations to the video frames.
+
+    convert_img_to_video: bool
+        whether to convert the input images to videos.
+    img_duration: float
+        specify the duration of the images, in ms. This will work only if convert_img_to_video is True.
+    batch_size: int
+        number of stimuli to process in each batch.
+    batch_grouper: function
+        function that takes a stimulus and return the property based on which the stimuli can be grouped.
+    """
+    def __init__(
+            self,
+            *args,
+            fps : float,
+            num_frames : Union[int, Tuple[int, int]] = None,
+            duration : Union[float, Tuple[float, float]] = None,
+            time_alignment : str = "evenly_spaced",
+            convert_img_to_video : bool = True,
+            img_duration : float = 1000.0,
+            batch_size : int = 32,
+            batch_grouper : Callable[[Video], Hashable] = lambda video: (round(video.duration, 6), video.fps),  # not including video.frame_size because most preprocessors will change the frame size to be the same
+            **kwargs,
+    ):
+        super().__init__(*args, stimulus_type=Video, batch_size=batch_size, 
+                         batch_grouper=batch_grouper, **kwargs)
+        self.fps = fps
+        self.num_frames = self._make_range(num_frames, type="num_frames")
+        self.duration = self._make_range(duration, type="duration")
+        assert hasattr(time_aligners, time_alignment), f"Unknown time alignment method: {time_alignment}"
+        self.time_aligner = getattr(time_aligners, time_alignment)
+
+        if convert_img_to_video:
+            assert img_duration is not None, "img_duration should be specified if convert_img_to_video is True"
+        self.img_duration = img_duration
+        self.convert_to_video = convert_img_to_video
+
+    @property
+    def identifier(self) -> str:
+        id = f"{super().identifier}.{self.time_aligner.__name__}.fps={float(self.fps)}"
+        if self.convert_to_video:
+            id += f".img_dur={float(self.img_duration)}"
+        return id
+
+    def load_stimulus(self, path: Union[str, Path]) -> Video:
+        if self.convert_to_video and Stimulus.is_image_path(path):
+            video = Video.from_img_path(path, self.img_duration, self.fps)
+        else:
+            video = Video.from_path(path)
+        video = video.set_fps(self.fps)
+        self._check_video(video)
+        return video
+    
+    def package_layer(
+            self, 
+            layer_activations : List[np.array],
+            layer_spec : str, 
+            stimuli : List[Stimulus]
+        ):
+        assert len(layer_activations) == len(stimuli)
+        longest_stimulus = stimuli[np.argmax(np.array([stimulus.duration for stimulus in stimuli]))]
+        ignore_time = self.time_aligner is time_aligners.ignore_time
+        channels = self._map_dims(layer_spec)
+        layer_activations = stack_with_nan_padding(layer_activations)
+        assembly = self._package(layer_activations, ["stimulus_path"] + channels)
+        # align to the longest stimulus
+        assembly = self.time_aligner(assembly, longest_stimulus)
+        if "channel_temporal" in channels and not ignore_time: 
+            channels.remove("channel_temporal")
+        assembly = self._stack_neuroid(assembly, channels)
+        if not ignore_time:
+            assembly = assembly_align_to_fps(assembly, self.fps)
+        return assembly
+
+    def _make_range(self, num, type="num_frames"):
+        if num is None:
+            return (1 if type=='num_frames' else 0, np.inf)
+        if isinstance(num, (tuple, list)):
+            return num
+        else:
+            return (num, num)
+
+    def _check_video(self, video: Video):
+        if self.num_frames is not None:
+            estimated_num_frames = int(self.fps * video.duration / 1000)
+            assert self.num_frames[0] <= estimated_num_frames <= self.num_frames[1], f"The number of frames must be within {self.num_frames}, but got {estimated_num_frames}"
+        if self.duration is not None:
+            assert self.duration[0] <= video.duration <= self.duration[1], f"The duration must be within {self.duration}, but got {video.duration}"

brainscore_vision/model_helpers/activations/temporal/core/inferencer/video/temporal_context/__init__.py

@@ -0,0 +1,2 @@
+from .causal import CausalInferencer
+from .block import BlockInferencer

brainscore_vision/model_helpers/activations/temporal/core/inferencer/video/temporal_context/base.py

@@ -0,0 +1,99 @@
+from typing import Union, List
+from pathlib import Path
+
+from brainscore_vision.model_helpers.activations.temporal.inputs import Video, Stimulus
+from ..base import TemporalInferencer
+
+
+MS_ROUNDING_DIGITS = 3
+
+class TemporalContextInferencerBase(TemporalInferencer):
+    """Inferencer base that computes the temporal context for concrete context-based inferencers,
+    like CausalInferencer and BlockInferencer. 
+
+    The context range is first determined by the num_frames and duration. Then, temporal_context_strategy
+    is used to determine the lower bound of temporal context and the expected temporal context.
+    
+    Parameters
+    ----------
+    temporal_context_strategy: str
+        specify how the length of temporal context for causal inference is determined.
+        Options:
+        - "greedy": the length of the temporal context is determined by the maximum of num_frames and duration.
+        - "conservative": the length of the temporal context is determined by the minimum of num_frames and duration.
+        - "fix": the length of the temporal context is determined by the specified "fixed_temporal_context".
+    
+    fixed_temporal_context: float
+        specify the fixed length of the temporal context, in ms. It will be used only if temporal_context_strategy is "fix".
+    
+    out_of_bound_strategy: str
+        specify how to handle the out-of-bound temporal context.
+        Options:
+        - "repeat": the out-of-bound temporal context will be repeated.
+        - TODO: "black": the out-of-bound temporal context will be zero-padded.
+        - TODO: "gray": the out-of-bound temporal context will be 128-padded.
+    """
+    def __init__(
+            self, 
+            *args,
+            temporal_context_strategy : str = "greedy",
+            fixed_temporal_context : float = None,
+            out_of_bound_strategy : str = "repeat",
+            **kwargs
+        ):
+        self.temporal_context_strategy = temporal_context_strategy
+        self.fixed_temporal_context = fixed_temporal_context
+        self.out_of_bound_strategy = out_of_bound_strategy
+        if self.temporal_context_strategy == "fix" and self.fixed_temporal_context is None:
+            raise ValueError("fixed_temporal_context must be specified if temporal_context_strategy is 'fix'.")
+        super().__init__(*args, **kwargs)
+
+    @property
+    def identifier(self) -> str:
+        lower, context = self._compute_temporal_context()
+        lower = round(lower, MS_ROUNDING_DIGITS)
+        context = round(context, MS_ROUNDING_DIGITS)
+        to_add = f".context={context}>{lower}.oob={self.out_of_bound_strategy}"
+        return f"{super().identifier}{to_add}"
+        
+    def load_stimulus(self, path: Union[str, Path]) -> Video:
+        if self.convert_to_video and Stimulus.is_image_path(path):
+            video = Video.from_img_path(path, self.img_duration, self.fps)
+        else:
+            video = Video.from_path(path)
+        video = video.set_fps(self.fps)
+        # does no check here
+        return video
+
+    def _overlapped_range(self, start1, end1, start2, end2):
+        # compute the overlapped range of two ranges (start1, end1) and (start2, end2)
+        lower, upper = max(start1, start2), min(end1, end2)
+        if lower > upper:
+            raise ValueError(f"Ranges [{start1}, {end1}] and [{start2}, {end2}] do not overlap.")
+        return lower, upper
+        
+    def _compute_temporal_context(self):
+        duration = self.duration
+        num_frames = self.num_frames
+        strategy = self.temporal_context_strategy
+
+        interval = 1000 / self.fps
+        num_frames_implied_ran = (num_frames[0] * interval, num_frames[1] * interval)
+        ran = self._overlapped_range(*num_frames_implied_ran, *duration)
+        lower = ran[0]
+
+        if strategy in ["greedy", "conservative"]:
+            if strategy == "greedy":
+                return lower, ran[1]
+            elif strategy == "conservative":
+                return lower, ran[0]
+
+        elif strategy == "fix":
+            context = self.fixed_temporal_context
+            assert ran[0] <= context <= ran[1], f"Fixed temporal context {context} is not within the range {ran}"
+
+        else:
+            raise ValueError(f"Unknown temporal context strategy: {strategy}")
+
+        return lower, context
+    

brainscore_vision/model_helpers/activations/temporal/core/inferencer/video/temporal_context/block.py

@@ -0,0 +1,77 @@
+import numpy as np
+from collections import OrderedDict
+from tqdm import tqdm
+
+from .base import TemporalContextInferencerBase
+from brainscore_vision.model_helpers.activations.temporal.utils import stack_with_nan_padding
+
+
+class BlockInferencer(TemporalContextInferencerBase):
+    """Inferencer that divides the original video into smaller blocks and does inference on the blocks separately.
+    Finally, the activations are joint along the temporal dimension for the final activations. 
+
+    Specifically, suppose the video lasts for 1000ms and the block size is 400ms.
+    Then, the video is segmented into [0~400ms], [400~800ms], [800~1200ms] (1000~1200ms padded).
+    The activations for each segment will be stacked together.
+
+    The block size is determined by the temporal parameters (num_frames & duration) and temporal_context_strategy.
+    If num_frames or duration is given, the model's temporal context will be set to match the two.
+    """
+    
+    def inference(self, stimuli, layers):
+        _, context = self._compute_temporal_context()
+        self._time_ends = []
+
+        if np.isinf(context):
+            # if the context is inf, pass the whole video directively
+            self._time_ends = [inp.duration for inp in stimuli]
+            layer_activations = super().inference(stimuli, layers)
+            layer_activations = OrderedDict([(layer, [[a] for a in activations]) 
+                                             for layer, activations in layer_activations.items()]) 
+                                             # make one-clip video activations
+        else:
+            num_clips = []
+            for stimulus in stimuli:
+                duration = stimulus.duration
+                videos = []
+                # for each stimulus, divide it into block clips with the specified context
+                for time_start in np.arange(0, duration, context):
+                    time_end = time_start + context
+                    clip = stimulus.set_window(time_start, time_end, padding=self.out_of_bound_strategy)
+                    videos.append(clip)
+                # record the actual time_end (which could be larger than the duration of the original video)
+                # so that we can align the time correctly when packaging the activations
+                self._time_ends.append(time_end)
+                self._executor.add_stimuli(videos)
+                num_clips.append(len(videos))  # record the number of clips for each video
+
+            activations = self._executor.execute(layers)
+            layer_activations = OrderedDict()
+            for layer in layers:
+                clip_start = 0
+                for num_clip in num_clips:
+                    # retrieve clips from a video by num_clip
+                    clips = activations[layer][clip_start:clip_start+num_clip]  # clips for this video
+                    layer_activations.setdefault(layer, []).append(clips)
+                    clip_start += num_clip
+    
+        # concat the activations from the clips of the same video
+        ret = OrderedDict()
+        for layer in layers:
+            activation_dims = self.layer_activation_format[layer]
+            for clips in layer_activations[layer]:
+                # make T the first dimension, as [T, ...] for easy concatenation
+                # the package_layer will change accordingly
+                if 'T' in activation_dims:
+                    time_index = activation_dims.index('T')
+                    clips = [np.moveaxis(a, time_index, 0) for a in clips]
+                    ret.setdefault(layer, []).append(np.concatenate(clips, axis=0))
+                else:
+                    ret.setdefault(layer, []).append(np.stack(clips, axis=0))
+
+        return ret
+    
+    def package_layer(self, activations, layer_spec, stimuli):
+        layer_spec = "T" + layer_spec.replace('T', '')  # T has been moved to the first dimension
+        stimuli = [stimulus.set_window(0, time_end) for stimulus, time_end in zip(stimuli, self._time_ends)]
+        return super().package_layer(activations, layer_spec, stimuli) 

brainscore_vision/model_helpers/activations/temporal/core/inferencer/video/temporal_context/causal.py

@@ -0,0 +1,86 @@
+import numpy as np
+from collections import OrderedDict
+
+from .base import TemporalContextInferencerBase
+from brainscore_vision.model_helpers.activations.temporal.inputs.video import Video
+from brainscore_vision.model_helpers.activations.temporal.utils import stack_with_nan_padding
+
+
+class CausalInferencer(TemporalContextInferencerBase):
+    """Inferencer that ensures the activations are causal by feeding in the video in a causal manner. 
+
+    Specifically, suppose the video lasts for 1000ms and the model samples every 100ms.
+    Then, the activations from the last time step of the activations from the following videos:
+    [0 ~ 100ms], [0 ~ 200ms], ..., [0 ~ 1000ms].
+    will be stacked together to form the final activations for the video. In this way, the 
+    activations are always "causal", which means that the temporal contexts of the model are 
+    always from the past, not the future.
+
+    If num_frames or duration is given, the model's temporal context will be set to match the two.
+    """
+    def __init__(
+            self, 
+            *args,
+            **kwargs
+        ):
+        if "time_alignment" in kwargs:
+            if kwargs["time_alignment"] != "per_frame_aligned":
+                raise ValueError("CausalInferencer enforces time_alignment='per_frame_aligned'.")
+            else:
+                del kwargs["time_alignment"]
+        super().__init__(*args, **kwargs, time_alignment="per_frame_aligned")
+
+    def inference(self, stimuli, layers):
+        interval = 1000 / self.fps
+        lower, context = self._compute_temporal_context()
+        num_clips = []
+        latest_time_end = 0
+        for inp in stimuli:
+            duration = inp.duration
+            videos = []
+            # here we ensure that the covered time range at least include the whole duration
+            for time_end in np.arange(interval, duration+interval, interval):
+                # see if the model only receive limited context
+                time_start = self._get_time_start(time_end, context, lower)
+                clip = inp.set_window(time_start, time_end, padding=self.out_of_bound_strategy)
+                latest_time_end = max(time_end, latest_time_end)
+                videos.append(clip)
+
+            self._executor.add_stimuli(videos)
+            num_clips.append(len(videos))
+
+        activations = self._executor.execute(layers)
+        layer_activations = OrderedDict()
+        for layer in layers:
+            activation_dims = self.layer_activation_format[layer]
+            clip_start = 0
+            for num_clip in num_clips:
+                video_activations = activations[layer][clip_start:clip_start+num_clip]  # clips for this video
+                # make T the first dimension, as [T, ...]
+                if 'T' in activation_dims:
+                    time_index = activation_dims.index('T')
+                    video_activations = [a.take(-1, axis=time_index) for a in video_activations]
+                layer_activations.setdefault(layer, []).append(np.stack(video_activations, axis=0))
+                clip_start += num_clip
+
+        return layer_activations
+    
+    def package_layer(self, activations, layer_spec, stimuli):
+        layer_spec = "T" + layer_spec.replace('T', '')  # T has been moved to the first dimension
+        return super().package_layer(activations, layer_spec, stimuli) 
+    
+    def _get_time_start(self, time_end, context, lower):
+        assert context >= lower, f"Temporal context {context} is not within the range {lower}"
+        if self.temporal_context_strategy == "fix":
+            return time_end - context
+        elif self.temporal_context_strategy == "greedy":
+            proposed_time_start = time_end - context
+            if proposed_time_start >= 0:
+                return proposed_time_start
+            else:
+                if time_end < lower:
+                    return time_end - lower
+                else:
+                    return 0 
+        elif self.temporal_context_strategy == "conservative":
+            return time_end - context

brainscore_vision/model_helpers/activations/temporal/core/inferencer/video/time_aligner.py

@@ -0,0 +1,73 @@
+import numpy as np
+from brainio.assemblies import DataAssembly
+from brainscore_vision.model_helpers.activations.temporal.inputs.video import Video
+
+
+"""This module includes different time alignment strategies for the activations of a temporal neural network.
+
+    A time alignment strategy is a function that takes the DataAssembly (with channel_temporal) and the video stimuli,
+    and aligns the activations to the video time. The channel_temporal dimension will be changed into time_bin dimension,
+    and the time_bin_start and time_bin_end will be added as coordinates of it.
+"""
+
+
+def _convert(assembly, time_bin_starts, time_bin_ends):
+    # this function converts the "channel_temporal" dimension to "time_bin" dimension
+    # if "channel_temporal" is not present, it adds a "time_bin" dimension
+    asm_type = assembly.__class__
+    if "channel_temporal" in assembly.dims:
+        assembly = assembly.drop_vars("channel_temporal")
+        assembly = assembly.rename({"channel_temporal": "time_bin"})
+    else:
+        assembly = assembly.expand_dims("time_bin")
+    assembly = assembly.assign_coords({
+        "time_bin_start": ("time_bin", time_bin_starts),
+        "time_bin_end": ("time_bin", time_bin_ends)
+    })
+    return asm_type(assembly)
+
+def estimate_layer_fps(assembly : DataAssembly, video : Video) -> DataAssembly:
+    # in a temporal neural net, different layers may have different temporal resolutions
+    # this function estimates the temporal resolution of a layer, based on the video fps
+    fps = video.fps
+    num_t = assembly.sizes['channel_temporal'] if "channel_temporal" in assembly.dims else 1
+    duration = video.duration
+    model_frame_interval = 1000 / fps
+    estimated_frame_interval = duration / num_t
+    estimated_multiplier = estimated_frame_interval / model_frame_interval
+    estimated_multiplier = int(round(estimated_multiplier))  # round to nearest integer
+    estimated_interval = model_frame_interval * estimated_multiplier
+    time_bin_starts = np.arange(0, num_t) * estimated_interval
+    time_bin_ends = time_bin_starts + estimated_interval
+    return _convert(assembly, time_bin_starts, time_bin_ends)
+
+def evenly_spaced(assembly : DataAssembly, video : Video) -> DataAssembly:
+    # this function assumes that the activation of different time steps is evenly spaced
+    num_t = assembly.sizes['channel_temporal'] if "channel_temporal" in assembly.dims else 1
+    interval = video.duration / num_t
+    time_bin_starts = np.linspace(0, video.duration, num_t+1)[:-1]
+    time_bin_ends = time_bin_starts + interval
+    time_bin_ends[-1] = video.duration
+    return _convert(assembly, time_bin_starts, time_bin_ends)
+
+def per_frame_aligned(assembly : DataAssembly, video : Video) -> DataAssembly:
+    # this function assumes that the activation of different time steps is aligned with the video frames
+    num_t = assembly.sizes['channel_temporal'] if "channel_temporal" in assembly.dims else 1
+    assert video.num_frames <= num_t
+    interval = 1000 / video.fps
+    time_bin_starts = np.arange(0, num_t) * interval
+    time_bin_ends = time_bin_starts + interval
+    return _convert(assembly, time_bin_starts, time_bin_ends)
+
+def ignore_time(assembly : DataAssembly, video : Video) -> DataAssembly:
+    # this function treats the activations from the entire video as from a single time bin,
+    # and treat the "channel_temporal" as a regular channel dimension and does no conversion 
+    asm_type = assembly.__class__
+    assembly = assembly.expand_dims("time_bin")
+    time_bin_starts = [0]
+    time_bin_ends = [video.duration]
+    assembly = assembly.assign_coords({
+        "time_bin_start": ("time_bin", time_bin_starts),
+        "time_bin_end": ("time_bin", time_bin_ends)
+    })
+    return asm_type(assembly)

brainscore_vision/model_helpers/activations/temporal/inputs/__init__.py

@@ -0,0 +1,3 @@
+from .video import Video
+from .image import Image
+from .base import Stimulus

brainscore_vision/model_helpers/activations/temporal/inputs/base.py

@@ -0,0 +1,17 @@
+from typing import Union
+from pathlib import Path
+
+
+class Stimulus:
+    def from_path(self, path):
+        raise NotImplementedError("Choose a concrete Stimulus type to use.")
+    
+    @staticmethod
+    def is_video_path(path: Union[str, Path]) -> bool:
+        extension = path.split('.')[-1].lower()
+        return extension in ['mp4', 'avi', 'mov', 'flv', 'wmv', 'webm', 'mkv', 'gif']
+    
+    @staticmethod
+    def is_image_path(path: Union[str, Path]) -> bool:
+        extension = path.split('.')[-1].lower()
+        return extension in ['jpg', 'jpeg', 'png', 'bmp', 'tiff']

brainscore_vision/model_helpers/activations/temporal/inputs/image.py

@@ -0,0 +1,50 @@
+import numpy as np
+from PIL import Image as PILImage
+
+from .base import Stimulus
+from brainscore_vision.model_helpers.activations.temporal.utils import batch_2d_resize
+
+
+class Image(Stimulus):
+    def __init__(self, path: str, size: int):
+        self._path = path
+        self._size = size
+
+    def copy(self):
+        return Image(self._path, self._size)
+    
+    @property
+    def size(self):
+        return self._size
+    
+    def set_size(self, size):
+        img = self.copy()
+        img._size = size
+        return img
+    
+    def from_path(path):
+        return Image(path, get_image_size(path))
+    
+    def to_pil_img(self):
+        return PILImage.fromarray(self.to_numpy())
+
+    def get_frame(self):
+        return np.array(PILImage.open(self._path).convert('RGB'))
+    
+    # return (H, W, C[RGB])
+    def to_numpy(self):
+        arr = self.get_frame()
+
+        if arr.shape[:2][::-1] != self._size:
+            arr = batch_2d_resize(arr[None,:], self._size, "bilinear")[0]
+
+        return arr
+
+    def store_to_path(self, path):
+        self.to_img().save(path)
+        return path
+
+def get_image_size(path):
+    with PILImage.open(path) as img:
+        size = img.size
+    return size