import os
import numpy as np
from pathlib import Path
from typing import Optional, List, Dict
import imageio
from pynwb import NWBFile
from pynwb.image import IndexSeries
from allensdk.brain_observatory.behavior.data_files import BehaviorStimulusFile
from allensdk.core import DataObject
from allensdk.core import \
NwbReadableInterface
from allensdk.brain_observatory.behavior.data_files.stimulus_file import \
StimulusFileReadableInterface
from allensdk.core import \
NwbWritableInterface
from allensdk.brain_observatory.behavior.data_objects.stimuli.presentations \
import \
Presentations
from allensdk.brain_observatory.behavior.stimulus_processing import \
get_stimulus_templates
from allensdk.brain_observatory.behavior.data_objects.stimuli \
.stimulus_templates import \
StimulusTemplate, StimulusTemplateFactory, StimulusMovieTemplateFactory
from allensdk.brain_observatory.behavior.write_nwb.extensions\
.stimulus_template.ndx_stimulus_template import \
StimulusTemplateExtension
from allensdk.internal.core.lims_utilities import safe_system_path
[docs]
class Templates(DataObject, StimulusFileReadableInterface,
NwbReadableInterface, NwbWritableInterface):
def __init__(self, templates: Dict[str, StimulusTemplate]):
super().__init__(name='stimulus_templates', value=templates)
# Grab the keys from the input dictionary. The "images" key is assumed
# to be the key in the dictionary that does not have "movie" in its key
# name. For VBO and VBN releases, there should only be at most 2
# keys in the dictionary.
image_template_keys = [
key for key in templates.keys()
if 'movie' not in key.lower()]
self._image_template_key = None
error_message = ""
if len(image_template_keys) == 1:
self._image_template_key = image_template_keys[0]
elif len(image_template_keys) > 1:
error_message += (
"Found multiple image StimulusTemplates "
f"{image_template_keys}. ")
movie_template_keys = [
key for key in templates.keys()
if 'movie' in key.lower()]
self._fingerprint_movie_template_key = None
if len(movie_template_keys) == 1:
self._fingerprint_movie_template_key = movie_template_keys[0]
elif len(movie_template_keys) > 1:
error_message += (
"Found multiple fingerprint movie StimulusTemplates "
f"{movie_template_keys}. ")
if len(error_message) > 0:
error_message += (
"This is not currently supported. "
"Please limit input to one image template and/or one "
"fingerprint movie template.")
raise NotImplementedError(error_message)
[docs]
@classmethod
def from_stimulus_file(
cls, stimulus_file: BehaviorStimulusFile,
limit_to_images: Optional[List] = None,
load_stimulus_movie: bool = False) -> "Templates":
"""Get stimulus templates (movies, scenes) for behavior session."""
# TODO: Eventually the `grating_images_dict` should be provided by the
# BehaviorLimsExtractor/BehaviorJsonExtractor classes.
# - NJM 2021/2/23
gratings_dir = "/allen/programs/braintv/production/visualbehavior"
gratings_dir = os.path.join(gratings_dir,
"prod5/project_VisualBehavior")
grating_images_dict = {
"gratings_0.0": {
"warped": np.asarray(imageio.imread(
safe_system_path(os.path.join(gratings_dir,
"warped_grating_0.png")))),
"unwarped": np.asarray(imageio.imread(
safe_system_path(os.path.join(
gratings_dir, "masked_unwarped_grating_0.png"))))
},
"gratings_90.0": {
"warped": np.asarray(imageio.imread(
safe_system_path(os.path.join(gratings_dir,
"warped_grating_90.png")))),
"unwarped": np.asarray(imageio.imread(
safe_system_path(os.path.join(
gratings_dir, "masked_unwarped_grating_90.png"))))
},
"gratings_180.0": {
"warped": np.asarray(imageio.imread(
safe_system_path(os.path.join(gratings_dir,
"warped_grating_180.png")))),
"unwarped": np.asarray(imageio.imread(
safe_system_path(os.path.join(
gratings_dir, "masked_unwarped_grating_180.png"))))
},
"gratings_270.0": {
"warped": np.asarray(imageio.imread(
safe_system_path(os.path.join(gratings_dir,
"warped_grating_270.png")))),
"unwarped": np.asarray(imageio.imread(
safe_system_path(os.path.join(
gratings_dir, "masked_unwarped_grating_270.png"))))
}
}
pkl = stimulus_file.data
stim_template = get_stimulus_templates(
pkl=pkl,
grating_images_dict=grating_images_dict,
limit_to_images=limit_to_images)
t = {stim_template.image_set_name: stim_template}
has_fingerprint_stimulus = (
"fingerprint" in pkl["items"]["behavior"]["items"]
)
if has_fingerprint_stimulus and load_stimulus_movie:
movie_data = np.load(
Path(pkl['items']['behavior']['items'][
'fingerprint']['static_stimulus']['movie_path'])
)
movie_template = StimulusMovieTemplateFactory.from_unprocessed(
movie_name="natural_movie_one",
movie_frames=movie_data,
)
t[movie_template.image_set_name] = movie_template
return Templates(templates=t)
[docs]
@classmethod
def from_nwb(cls, nwbfile: NWBFile) -> "Templates":
templates = {}
for image_set_name, image_data in nwbfile.stimulus_template.items():
image_attributes = [
{'image_name': image_name}
for image_name in image_data.control_description
]
templates[image_set_name] = StimulusTemplateFactory.from_processed(
image_set_name=image_set_name,
image_attributes=image_attributes,
warped=image_data.data[:],
unwarped=image_data.unwarped[:]
)
return Templates(templates=templates)
[docs]
def to_nwb(self, nwbfile: NWBFile,
stimulus_presentations: Presentations) -> NWBFile:
for key, stimulus_templates in self.value.items():
unwarped_images = []
warped_images = []
image_names = []
for image_name, image_data in stimulus_templates.items():
image_names.append(image_name)
unwarped_images.append(image_data.unwarped)
warped_images.append(image_data.warped)
image_index = np.zeros(len(image_names))
image_index[:] = np.nan
visual_stimulus_image_series = \
StimulusTemplateExtension(
name=stimulus_templates.image_set_name,
data=warped_images,
unwarped=unwarped_images,
control=list(range(len(image_names))),
control_description=image_names,
unit='NA',
format='raw',
timestamps=image_index)
nwbfile.add_stimulus_template(visual_stimulus_image_series)
if 'image_index' in stimulus_presentations.value \
and self._image_template_key is not None:
nwbfile = self._add_image_index_to_nwb(
nwbfile=nwbfile, presentations=stimulus_presentations)
return nwbfile
def _add_image_index_to_nwb(
self, nwbfile: NWBFile, presentations: Presentations):
"""Adds the image index and start_time for all stimulus templates
to NWB"""
stimulus_templates = self.value[self._image_template_key]
presentations = presentations.value
nwb_template = nwbfile.stimulus_template[
stimulus_templates.image_set_name]
stimulus_name = 'image_set' \
if 'image_set' in presentations else 'stimulus_name'
stimulus_index = presentations[
presentations[stimulus_name] == nwb_template.name]
image_index = IndexSeries(
name=nwb_template.name,
data=stimulus_index['image_index'].values,
unit='N/A',
indexed_timeseries=nwb_template,
timestamps=stimulus_index['start_time'].values)
nwbfile.add_stimulus(image_index)
return nwbfile
@property
def image_template_key(self) -> str:
"""
Name of the image template in template dictionary.
"""
return self._image_template_key
@property
def fingerprint_movie_template_key(self) -> str:
"""
Name of the fingerprint movie template in template dictionary.
"""
return self._fingerprint_movie_template_key