brainscore-vision 2.1__py3-none-any.whl

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

@@ -0,0 +1,186 @@
+import cv2
+from decord import VideoReader
+import numpy as np
+from PIL import Image as PILImage
+from typing import Tuple, Union
+from pathlib import Path
+
+from .base import Stimulus
+from .image import Image
+from brainscore_vision.model_helpers.activations.temporal.utils import batch_2d_resize
+
+
+EPS = 1e-9  
+
+def get_video_stats(video_path):
+    cap = cv2.VideoCapture(video_path)
+    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps    = cap.get(cv2.CAP_PROP_FPS)
+    size = (width, height)
+    duration = length / fps * 1000
+    cap.release()
+    return fps, duration, size
+
+
+def get_image_stats(image_path):
+    with PILImage.open(image_path) as img:
+        return img.size
+    return size
+
+
+class Video(Stimulus):
+    """Video object that represents a video clip."""
+
+    def __init__(
+            self, 
+            path: Union[str, Path], 
+            fps: float, 
+            start: float, 
+            end: float, 
+            size: Tuple[int, int]
+        ):
+        self._path = path
+        self._fps = fps
+        self._size = size
+        self._start = start
+        self._end = end
+        self._original_fps = None
+        self._original_duration = None
+        self._original_size = None
+
+    def __getattribute__(self, key):
+        if key.startswith("_original_"):
+            if super().__getattribute__(key) is None:
+                self._original_fps, self._original_duration, self._original_size = get_video_stats(self._path)
+        return super().__getattribute__(key)
+
+    def copy(self):
+        # return view
+        video = self.__class__(self._path, self._fps, self._start, self._end, self._size)
+        video._original_fps = self._original_fps
+        video._original_duration = self._original_duration
+        video._original_size = self._original_size
+        return video
+    
+    @property
+    def duration(self):
+        # in ms
+        return self._end - self._start
+
+    @property
+    def fps(self):
+        return self._fps
+
+    @property
+    def num_frames(self):
+        return int(self.duration * self.fps/1000 + EPS)
+    
+    @property
+    def original_num_frames(self):
+        return int(self._original_duration * self._original_fps/1000 + EPS)
+    
+    @property
+    def frame_size(self):
+        return self._size
+    
+    ### Transformations: return copy
+    
+    def set_fps(self, fps):
+        assert 1000/fps <= self.duration, f"fps {fps} is too low for duration {self.duration}ms."
+        video = self.copy()
+        video._fps = fps
+        return video
+    
+    def set_size(self, size):
+        # size: (width, height)
+        video = self.copy()
+        video._size = size
+        return video
+
+    def set_window(self, start, end, padding="repeat"):
+        # use ms as the time scale
+        if end < start:
+            raise ValueError("end time is earlier than start time")
+        
+        if padding != "repeat":
+            raise NotImplementedError()
+        
+        video = self.copy()
+        video._start = self._start + start
+        video._end = self._start + end
+        return video
+    
+    def get_frames(self, indices):
+        reader = VideoReader(self._path)
+        frames = reader.get_batch(indices).asnumpy()
+        del reader
+        return frames
+
+    ### I/O
+    def from_path(path):
+        fps, end, size = get_video_stats(path)
+        start = 0
+        return Video(path, fps, start, end, size)
+    
+    def from_img_path(img_path, duration, fps):
+        # duration in ms
+        size = get_image_stats(img_path)
+        return VideoFromImage(img_path, fps, 0, duration, size)
+    
+    def to_numpy(self):
+        # get the time stamps of frame samples
+        start_frame = self._start * self._original_fps / 1000
+        end_frame = self._end * self._original_fps / 1000
+        # avoid taking the last extra frame
+        samples = np.arange(start_frame, end_frame - EPS, self._original_fps/self.fps)
+        sample_indices = samples.astype(int)
+
+        # padding: repeat the first/last frame
+        original_num_frames = int(self._original_duration * self._original_fps/1000 - EPS)  # EPS to avoid last frame OOB error
+        sample_indices = np.clip(sample_indices, 0, original_num_frames-1)
+
+        # actual sampling
+        frames = self.get_frames(sample_indices)
+
+        # resizing
+        if self._size != (frames.shape[2], frames.shape[1]):
+            frames = batch_2d_resize(frames, self._size, "bilinear")
+
+        return frames
+    
+    def to_frames(self):
+        return [f for f in self.to_numpy()]
+
+    def to_pil_imgs(self):
+        return [PILImage.fromarray(frame) for frame in self.to_numpy()]
+
+    def to_path(self):
+        # use context manager ?
+        path = None  # make a temporal file
+        raise NotImplementedError()
+        return path
+
+    def store_to_path(self, path):
+        # pick format based on path filename
+        if path.endswith(".avi"):
+            fourcc = cv2.VideoWriter_fourcc(*'XVID')
+        elif path.endswith(".mp4"):
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        else:
+            raise ValueError("Unsupported video format.")
+
+        out = cv2.VideoWriter(path, fourcc, self._fps, self._size)
+        for frame in self.to_frames():
+            out.write(frame[...,::-1])  # to RGB
+        out.release()
+        return path
+    
+
+class VideoFromImage(Video):
+    def get_frames(self, indices):
+        data = Image.from_path(self._path).to_numpy()
+        N = len(indices)
+        ret = np.repeat(data[np.newaxis, ...], N, axis=0)
+        return ret

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

@@ -0,0 +1,2 @@
+from .base import ActivationWrapper
+from .pytorch import PytorchWrapper

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

@@ -0,0 +1,33 @@
+import numpy as np
+from ..inputs.base import Stimulus
+from ..core import ActivationsExtractor
+from ..core import TemporalInferencer, Inferencer
+
+from typing import List, Callable, Any, Dict
+
+
+class ActivationWrapper:
+    def __init__(
+            self, 
+            identifier : str, 
+            preprocessing : Callable[[List[Stimulus]], Any],
+            inferencer_cls : Inferencer = TemporalInferencer, 
+            **extractor_kwargs
+        ):
+        self.identifier = identifier
+        self.preprocessing = preprocessing
+        self.build_extractor(inferencer_cls, **extractor_kwargs)
+
+    # List[preprocessed_input] -> Dict[layer -> np.array]
+    def get_activations(self, inputs : List[Stimulus], layers : List[str]) -> Dict[str, np.array]:
+        raise NotImplementedError()
+
+    def __call__(self, *args, **kwargs):
+        return self._extractor(*args, **kwargs)
+    
+    def build_extractor(self, inferencer_cls, *args, **kwargs):
+        extractor = ActivationsExtractor(identifier=self.identifier, 
+                                         inferencer=inferencer_cls(get_activations=self.get_activations, 
+                                                        preprocessing=self.preprocessing, *args, **kwargs))
+        extractor.insert_attrs(self)
+        self._extractor = extractor

brainscore_vision/model_helpers/activations/temporal/model/pytorch.py

@@ -0,0 +1,107 @@
+from collections import OrderedDict
+import typing
+from typing import Callable, List, Any, Dict
+
+import logging
+
+from brainscore_vision.model_helpers.utils import fullname
+from .base import ActivationWrapper
+from ..inputs import Stimulus
+
+
+SUBMODULE_SEPARATOR = '.'
+
+
+def default_process_activation(layer, layer_name, inputs, output):
+    # (torch.nn.Module, str, torch.Tensor, torch.Tensor) -> torch.Tensor
+    return output
+
+class PytorchWrapper(ActivationWrapper):
+    def __init__(
+            self, 
+            identifier : str, 
+            model, 
+            preprocessing : Callable[[List[Stimulus]], Any], 
+            process_output = None, 
+            *args, 
+            **kwargs
+        ):
+        import torch
+        logger = logging.getLogger(fullname(self))
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.debug(f"Using device {self._device}")
+        self._model = model.to(self._device)
+        # preprocessing: input.Stimulus -> actual model inputs
+        self._preprocess = preprocessing
+        self._process_activation = default_process_activation if process_output is None else process_output
+        super().__init__(identifier, preprocessing, *args, **kwargs)
+
+    def forward(self, inputs : List[Any]) -> Any: 
+        # this function gets a list of preprocessed inputs and does the forward pass
+        import torch
+        tensor = torch.stack(inputs)
+        tensor = tensor.to(self._device)
+        return self._model(tensor)
+
+    def get_activations(self, inputs : List[Any], layer_names : List[str]) -> typing.OrderedDict[str, Any]:
+        import torch
+        self._model.eval()
+        layer_results = OrderedDict()
+        hooks = []
+
+        for layer_name in layer_names:
+            layer = self.get_layer(layer_name)
+            hook = self._register_hook(layer, layer_name, target_dict=layer_results)
+            hooks.append(hook)
+
+        with torch.no_grad():
+            self.forward(inputs)
+
+        for hook in hooks:
+            hook.remove()
+        return layer_results
+    
+    def get_layer(self, layer_name : str):
+        # the layer_name is a string that represents the module hierarchy up to the target layer,
+        # seperated by ".", e.g., "module1.submodule2.relu".
+        if layer_name == 'logits':
+            return self._output_layer()
+        module = self._model
+        for part in layer_name.split(SUBMODULE_SEPARATOR):
+            module = module._modules.get(part)
+            assert module is not None, f"No submodule found for layer {layer_name}, at part {part}"
+        return module
+
+    def _output_layer(self):
+        module = self._model
+        while module._modules:
+            module = module._modules[next(reversed(module._modules))]
+        return module
+
+    @classmethod
+    def _tensor_to_numpy(cls, output):
+        return output.cpu().data.numpy()
+
+    def _register_hook(self, layer, layer_name, target_dict):
+        def hook_function(_layer, _input, output, name=layer_name, target_dict=target_dict):
+            output = self._process_activation(_layer, name, _input, output)
+            target_dict[name] = PytorchWrapper._tensor_to_numpy(output)
+
+        hook = layer.register_forward_hook(hook_function)
+        return hook
+
+    def __repr__(self):
+        return repr(self._model)
+
+    def layers(self):
+        for name, module in self._model.named_modules():
+            if len(list(module.children())) > 0:  # this module only holds other modules
+                continue
+            yield name, module
+
+    def graph(self):
+        import networkx as nx
+        g = nx.DiGraph()
+        for layer_name, layer in self.layers():
+            g.add_node(layer_name, object=layer, type=type(layer))
+        return g

brainscore_vision/model_helpers/activations/temporal/utils.py

@@ -0,0 +1,228 @@
+import os
+import numpy as np
+
+from brainscore_vision.model_helpers.brain_transformation.temporal import assembly_time_align
+
+
+def concat_with_nan_padding(arr_list, axis=0, dtype=np.float16):
+    # Get shapes of all arrays
+    shapes = [np.array(arr.shape) for arr in arr_list]
+    max_shape = np.max(shapes, axis=0)
+    len_axis = sum([arr.shape[axis] for arr in arr_list])
+
+    # Allocate concatenated array with NaN padding
+    result_shape = [s if i != axis else len_axis for i, s in enumerate(max_shape)]
+    result = np.full(result_shape, np.nan, dtype=dtype)
+
+    # Fill in individual arrays
+    offset = 0
+    for arr in arr_list:
+        slices = [slice(0, s) if i != axis else slice(offset, offset+s) 
+                  for i, s in enumerate(arr.shape)]
+        offset += arr.shape[axis]
+        result[slices] = arr
+
+    return result
+
+
+def stack_with_nan_padding_(arr_list, axis=0, dtype=np.float16):
+    # Get shapes of all arrays
+    shapes = [np.array(arr.shape) for arr in arr_list]
+    max_shape = np.max(shapes, axis=0)
+
+    # Allocate concatenated array with NaN padding
+    result = np.full(np.concatenate(([len(arr_list)], max_shape)), np.nan, dtype=dtype)
+
+    # Fill in individual arrays
+    for i, arr in enumerate(arr_list):
+        slices = tuple(slice(0, s) for s in arr.shape)
+        result[i][slices] = arr
+
+    result = np.swapaxes(result, 0, axis)
+
+    return result
+
+
+def stack_with_nan_padding(arr_list, axis=0, dtype=None):
+    # Get shapes of all arrays
+    shapes = [np.array(arr.shape) for arr in arr_list]
+    max_shape = np.max(shapes, axis=0)
+
+    results = []
+
+    # Fill in individual arrays
+    for arr in arr_list:
+        result = np.pad(arr, [(0, max_shape[j]-s) for j, s in enumerate(arr.shape)], 
+                        mode='constant', constant_values=np.nan)
+        results.append(result)
+
+    result = np.stack(results, axis=axis)
+    result = np.swapaxes(result, 0, axis)
+    if dtype is not None and result.dtype != dtype:
+        result = result.astype(dtype)
+
+    return result
+
+
+def batch_2d_resize(arr, size, mode):
+    import cv2
+    # arr [N, H, W, C]
+    N, H, W, C = arr.shape
+    arr = arr.transpose(1, 2, 3, 0).reshape(H, W, C*N)
+    if mode == "bilinear":
+        mode = cv2.INTER_LINEAR
+        ret = cv2_resize(arr, size, mode)
+    elif mode == "pool":
+        # ret = proportional_average_pooling(arr, size)
+        mode = cv2.INTER_AREA
+        ret = cv2_resize(arr, size, mode)
+    else:
+        raise ValueError(f"Unknown mode {mode}")
+    ret = ret.reshape(size[1], size[0], C, N).transpose(3, 0, 1, 2)
+    return ret
+
+
+def proportional_average_pooling(arr, size):
+    import cv2
+    mode = cv2.INTER_AREA
+    ret = cv2_resize(arr, size, mode)
+    return ret
+
+
+def proportional_average_pooling_(arr, size):
+    H, W, C = arr.shape
+    w, h = size
+
+    ret = np.zeros((h, w, C))
+    for r in range(h):
+        for c in range(w):
+            y0 = r * H / h
+            x0 = c * W / w
+            y1 = (r + 1) * H / h
+            x1 = (c + 1) * W / w
+            r_start = int(y0)
+            r_end = np.ceil(y1).astype(int)
+            c_start = int(x0)
+            c_end = np.ceil(x1).astype(int)
+            val = arr[r_start:r_end, c_start:c_end]
+            y_overlap_start = np.maximum(y0, np.arange(r_start, r_end))
+            y_overlap_end = np.minimum(y1, np.arange(r_start, r_end)+1)
+            x_overlap_start = np.maximum(x0, np.arange(c_start, c_end))
+            x_overlap_end = np.minimum(x1, np.arange(c_start, c_end)+1)
+            y_overlap = y_overlap_end - y_overlap_start
+            x_overlap = x_overlap_end - x_overlap_start
+            areas = np.outer(y_overlap, x_overlap)
+            areas = areas / np.sum(areas)
+            val = (val * areas[...,None]).reshape(-1, C).sum(0)
+            ret[r, c] = val
+    return ret.astype(arr.dtype)
+
+
+# cv2 has the wierd bug of cannot handling too large channel size
+def cv2_resize(arr, size, mode, batch_size=4):
+    # arr [H, W, C]
+    import cv2
+    ori_dtype = arr.dtype
+    arr = arr.astype(float)
+    C = arr.shape[-1]
+    ret = []
+    for i in range(0, C, batch_size):
+        val = cv2.resize(arr[..., i:i+batch_size], size, interpolation=mode)
+        if len(val.shape)<3: val = val[...,None]
+        ret.append(val)
+    return np.concatenate(ret, axis=-1).astype(ori_dtype)
+
+
+def parallelize(func, iterable, n_jobs=1, verbose=0, mode="multiprocess"):
+    if mode == "multiprocess":
+        from joblib import Parallel, delayed
+        return Parallel(n_jobs=n_jobs, verbose=verbose)(delayed(func)(item) for item in iterable)
+    elif mode == "threading":
+        from concurrent.futures import ThreadPoolExecutor
+        with ThreadPoolExecutor(max_workers=n_jobs) as executor:
+            return list(executor.map(func, iterable))
+    else:
+        raise ValueError(f"Unknown mode {mode}")
+
+
+def assembly_align_to_fps(output_assembly, fps, mode="portion"):
+    EPS = 1e-9  # prevent the duration from being slightly larger than the last time bin
+    interval = 1000 / fps
+    duration = output_assembly["time_bin_end"].values[-1]
+    target_time_bin_starts = np.arange(0, duration-EPS, interval)
+    target_time_bin_ends = target_time_bin_starts + interval
+    target_time_bin_ends[-1] = duration  # use this to avoid numerical error
+    target_time_bins = [(start, end) for start, end in zip(target_time_bin_starts, target_time_bin_ends)]
+    return assembly_time_align(output_assembly, target_time_bins, mode=mode)
+
+
+## model utils
+
+# get the inferencer from any model
+def get_inferencer(any_model):
+    from brainscore_vision.model_helpers.activations.temporal.core import Inferencer, ActivationsExtractor
+    from brainscore_vision.model_helpers.activations.temporal.model.base import ActivationWrapper
+    from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+
+    if isinstance(any_model, Inferencer): return any_model
+    if isinstance(any_model, ActivationWrapper): return any_model._extractor.inferencer
+    if isinstance(any_model, ModelCommitment): return any_model.activations_model._extractor.inferencer
+    if isinstance(any_model, ActivationsExtractor): return any_model.inferencer
+    raise ValueError(f"Cannot find inferencer from the model {any_model}")
+
+def get_base_model(any_model):
+    from brainscore_vision.model_helpers.activations.temporal.model.base import ActivationWrapper
+    from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+
+    if isinstance(any_model, ActivationWrapper): return any_model
+    if isinstance(any_model, ModelCommitment): return any_model.activations_model
+    raise ValueError(f"Cannot find inferencer from the model {any_model}")
+
+
+# get the layers from the layer_activation_format
+def get_specified_layers(any_model):
+    inferencer = get_inferencer(any_model)
+    return list(inferencer.layer_activation_format.keys())
+
+# switch the inferencer at any level
+# specify key='same' to retrive the same parameter from the original inferencer
+def switch_inferencer(any_model, new_inferencer_cls, **kwargs):
+    inferencer = get_inferencer(any_model)
+    base_model = get_base_model(any_model)
+    for k, v in kwargs.items():
+        if v == 'same': kwargs[k] = getattr(inferencer, k)
+    base_model.build_extractor(new_inferencer_cls, **kwargs)
+
+
+## Others
+    
+def download_weight_file(url, folder=None):
+    import requests
+    from tqdm import tqdm
+    weight_fname = os.path.basename(url)
+    brainscore_home = os.getenv("BRAINSCORE_HOME", os.path.expanduser("~/.brain-score"))
+    model_cache = os.path.join(brainscore_home, "models")
+    if folder: model_cache = os.path.join(model_cache, folder)
+    os.makedirs(model_cache, exist_ok=True)
+    weight_path = os.path.join(model_cache, weight_fname)
+
+    if os.path.exists(weight_path):
+        return weight_path
+
+    # Streaming, so we can iterate over the response.
+    response = requests.get(url, stream=True)
+
+    # Sizes in bytes.
+    total_size = int(response.headers.get("content-length", 0))
+    block_size = 1024
+
+    with tqdm(total=total_size, unit="B", unit_scale=True, desc=f"{weight_fname}") as progress_bar:
+        with open(weight_path, "wb") as file:
+            for data in response.iter_content(block_size):
+                progress_bar.update(len(data))
+                file.write(data)
+
+    if total_size != 0 and progress_bar.n != total_size:
+        raise RuntimeError("Could not download file")
+
+    return weight_path

brainscore_vision/model_helpers/brain_transformation/__init__.py

@@ -0,0 +1,97 @@
+from brainscore_vision import load_benchmark
+from brainscore_vision.model_helpers.brain_transformation.temporal import TemporalAligned
+from brainscore_vision.model_interface import BrainModel
+from brainscore_vision.utils import LazyLoad
+from .behavior import BehaviorArbiter, LabelBehavior, ProbabilitiesMapping, OddOneOut
+from .neural import LayerMappedModel, LayerSelection, LayerScores
+
+STANDARD_REGION_BENCHMARKS = {
+    'V1': LazyLoad(lambda: load_benchmark('FreemanZiemba2013public.V1-pls')),
+    'V2': LazyLoad(lambda: load_benchmark('FreemanZiemba2013public.V2-pls')),
+    'V4': LazyLoad(lambda: load_benchmark('MajajHong2015public.V4-pls')),
+    'IT': LazyLoad(lambda: load_benchmark('MajajHong2015public.IT-pls')),
+}
+
+
+class ModelCommitment(BrainModel):
+    """
+    Standard process to convert a BaseModel (e.g. standard Machine Learning/Computer Vision model)
+    into a BrainModel (e.g. commitments on layer-to-region, pixel-to-degrees, ...).
+    """
+
+    def __init__(self, identifier,
+                 activations_model, layers, behavioral_readout_layer=None, region_layer_map=None,
+                 visual_degrees=8):
+        self.layers = layers
+        self.activations_model = activations_model
+        # We set the visual degrees of the ActivationsExtractorHelper here to avoid changing its signature.
+        #  The ideal solution would be to not expose the _extractor of the activations_model here, but to change
+        #  the signature of the ActivationsExtractorHelper. See https://github.com/brain-score/vision/issues/554
+        self.activations_model._extractor.set_visual_degrees(visual_degrees)  # for microsaccades
+        self._visual_degrees = visual_degrees
+        # region-layer mapping
+        if region_layer_map is None:
+            layer_selection = LayerSelection(model_identifier=identifier,
+                                             activations_model=activations_model, layers=layers,
+                                             visual_degrees=visual_degrees)
+            region_layer_map = RegionLayerMap(layer_selection=layer_selection,
+                                              region_benchmarks=STANDARD_REGION_BENCHMARKS)
+        # neural
+        layer_model = LayerMappedModel(identifier=identifier, activations_model=activations_model,
+                                       region_layer_map=region_layer_map)
+        self.layer_model = TemporalAligned(layer_model)
+        logits_behavior = LabelBehavior(identifier=identifier, activations_model=activations_model)
+        behavioral_readout_layer = behavioral_readout_layer or layers[-1]
+        probabilities_behavior = ProbabilitiesMapping(identifier=identifier, activations_model=activations_model,
+                                                      layer=behavioral_readout_layer)
+        odd_one_out = OddOneOut(identifier=identifier, activations_model=activations_model,
+                                layer=behavioral_readout_layer)
+        self.behavior_model = BehaviorArbiter({BrainModel.Task.label: logits_behavior,
+                                               BrainModel.Task.probabilities: probabilities_behavior,
+                                               BrainModel.Task.odd_one_out: odd_one_out,
+                                               })
+        self.do_behavior = False
+
+    def visual_degrees(self) -> int:
+        return self._visual_degrees
+
+    def start_task(self, task: BrainModel.Task, *args, **kwargs):
+        if task != BrainModel.Task.passive:
+            self.behavior_model.start_task(task, *args, **kwargs)
+            self.do_behavior = True
+        else:
+            self.do_behavior = False
+
+    def look_at(self, stimuli, number_of_trials: int = 1, require_variance: bool = False):
+        if self.do_behavior:
+            return self.behavior_model.look_at(stimuli, number_of_trials=number_of_trials, require_variance=require_variance)
+        else:
+            return self.layer_model.look_at(stimuli, number_of_trials=number_of_trials)
+
+    def start_recording(self, recording_target, time_bins):
+        return self.layer_model.start_recording(recording_target, time_bins)
+
+    @property
+    def identifier(self):
+        return self.layer_model.identifier
+
+
+class RegionLayerMap(dict):
+    """
+    mapping of regions to layers that only evaluates lazily and only once.
+    """
+
+    def __init__(self, layer_selection, region_benchmarks):
+        super(RegionLayerMap, self).__init__()
+        self.layer_selection = layer_selection
+        self.region_benchmarks = region_benchmarks
+
+    def __getitem__(self, region):
+        if region not in self:  # not yet committed
+            self.commit_region(region)
+        return super(RegionLayerMap, self).__getitem__(region)
+
+    def commit_region(self, region):
+        benchmark = self.region_benchmarks[region]
+        best_layer = self.layer_selection(selection_identifier=region, benchmark=benchmark)
+        self[region] = best_layer