brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/models/temporal_model_S3D_text_video/__init__.py

@@ -0,0 +1,14 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry["s3d-HowTo100M"] = lambda: commit_model("s3d-HowTo100M")

brainscore_vision/models/temporal_model_S3D_text_video/model.py

@@ -0,0 +1,65 @@
+import torch
+import numpy as np
+from torchvision import transforms
+from s3dg_howto100m import S3D
+
+from brainscore_vision.model_helpers.activations.temporal.model.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.s3 import load_weight_file
+
+
+img_transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+])
+
+def transform_video(video):
+    frames = video.to_numpy() / 255.
+    frames = torch.Tensor(frames)
+    frames = frames.permute(0, 3, 1, 2)
+    frames = img_transform(frames)
+    return frames.permute(1, 0, 2, 3)
+
+
+def get_model(identifier="s3d-HowTo100M"):
+    inferencer_kwargs = {
+        "fps": 24,  # common YouTube frame rate
+        "layer_activation_format": 
+        {
+            "conv1": "CTHW",
+            "conv_2c": "CTHW",
+            "mixed_3c": "CTHW",
+            "mixed_4b": "CTHW",
+            "mixed_4d": "CTHW",
+            "mixed_4f": "CTHW",
+            "mixed_5c": "CTHW",
+            "fc": "C"
+        },
+    }
+    process_output = None
+
+    model_name = identifier
+
+    model_pth = load_weight_file(
+        bucket="brainscore-vision",
+        relative_path="temporal_model_S3D_text_video/s3d_howto100m.pth",
+        version_id="hRp6I8bpwreIMUVL0H.zCdK0hqRggL7n",
+        sha1="31e99d2a1cd48f2259ca75e719ac82c8b751ea75"
+    )
+
+    dict_pth = load_weight_file(
+        bucket="brainscore-vision",
+        relative_path="temporal_model_S3D_text_video/s3d_dict.npy",
+        version_id="4NxVLe8DSL6Uue0F7e2rz8HZuOk.tkBI",
+        sha1="d368ff7d397ec8240f1f963b5efe8ff245bac35f"
+    )
+
+    # Instantiate the model
+    model = S3D(dict_pth, 512)
+
+    # Load the model weights
+    model.load_state_dict(torch.load(model_pth))
+
+    wrapper = PytorchWrapper(identifier, model, transform_video, 
+                             process_output=process_output,
+                             **inferencer_kwargs)
+    
+    return wrapper

brainscore_vision/models/temporal_model_S3D_text_video/requirements.txt

@@ -0,0 +1 @@
+S3D_HowTo100M @ git+https://github.com/YingtianDt/S3D_HowTo100M

brainscore_vision/models/temporal_model_S3D_text_video/test.py

@@ -0,0 +1,15 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "s3d-HowTo100M",
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None

brainscore_vision/models/temporal_model_SeLaVi/__init__.py

@@ -0,0 +1,17 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry["SeLaVi-Kinetics400"] = lambda: commit_model("SeLaVi-Kinetics400")
+model_registry["SeLaVi-Kinetics-Sound"] = lambda: commit_model("SeLaVi-Kinetics-Sound")
+model_registry["SeLaVi-VGG-Sound"] = lambda: commit_model("SeLaVi-VGG-Sound")
+model_registry["SeLaVi-AVE"] = lambda: commit_model("SeLaVi-AVE")

brainscore_vision/models/temporal_model_SeLaVi/model.py

@@ -0,0 +1,68 @@
+import torch
+
+from selavi.model import load_model
+from selavi.video_transforms import clip_augmentation
+
+from brainscore_vision.model_helpers.activations.temporal.model import PytorchWrapper
+from brainscore_vision.model_helpers.activations.temporal.utils import download_weight_file
+
+
+def transform_video(video):
+    arr = video.to_numpy()
+    arr = torch.as_tensor(arr)
+    return clip_augmentation(arr)
+
+
+def get_model(identifier):
+
+    assert identifier.startswith("SeLaVi-")
+    dataset = "-".join(identifier.split("-")[1:])
+
+    if dataset == "Kinetics400":
+        model_name = "selavi_kinetics.pth"
+        num_classes = 400
+    elif dataset == "Kinetics-Sound":
+        model_name = "selavi_kinetics_sound.pth"
+        num_classes = 32
+    elif dataset == "VGG-Sound":
+        model_name = "selavi_vgg_sound.pth"
+        num_classes = 309
+    elif dataset == "AVE":
+        model_name = "selavi_ave.pth"
+        num_classes = 28
+    else:
+        raise ValueError(f"Unknown dataset: {dataset}")
+    
+    url = f"https://dl.fbaipublicfiles.com/selavi/{model_name}"
+    pth = download_weight_file(url, folder="temporal_model_SeLaVi")
+
+    # Load model
+    model = load_model(
+        vid_base_arch="r2plus1d_18",
+        aud_base_arch="resnet9",
+        use_mlp=True,
+        num_classes=num_classes,
+        pretrained=False,
+        norm_feat=False,
+        use_max_pool=False,
+        headcount=10,
+    )
+
+    model = model.video_network  # Remove audio network
+
+    # Load weights
+    state_dict_ = torch.load(pth, map_location="cpu")['model']
+    state_dict = {}
+    for k, v in list(state_dict_.items()):
+        if k.startswith("module.video_network."):
+            k = k[len("module.video_network."):]
+            state_dict[k] = v
+    model.load_state_dict(state_dict)
+
+    layer_activation_format = {
+        "base.stem": "CTHW",
+        **{f"base.layer{i}": "CTHW" for i in range(1, 5)},
+        # "base.fc": "C",  # no fc
+    }
+
+    return PytorchWrapper(identifier, model, transform_video, fps=30, layer_activation_format=layer_activation_format)

brainscore_vision/models/temporal_model_SeLaVi/requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+selavi @ git+https://github.com/YingtianDt/selavi.git

brainscore_vision/models/temporal_model_SeLaVi/test.py

@@ -0,0 +1,18 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "SeLaVi-Kinetics400",
+    "SeLaVi-Kinetics-Sound",
+    "SeLaVi-VGG-Sound",
+    "SeLaVi-AVE"
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None

brainscore_vision/models/temporal_model_VideoMAE/__init__.py

@@ -0,0 +1,15 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry["VideoMAE-V1-B"] = lambda: commit_model("VideoMAE-V1-B")
+model_registry["VideoMAE-V1-L"] = lambda: commit_model("VideoMAE-V1-L")

brainscore_vision/models/temporal_model_VideoMAE/model.py

@@ -0,0 +1,100 @@
+from timm import create_model
+from videomae import *
+import torch as th
+
+from brainscore_vision.model_helpers.activations.temporal.model import PytorchWrapper
+from brainscore_vision.model_helpers.s3 import load_weight_file
+from torchvision import transforms
+
+
+LAYER_SELECTION_STEP = 2
+
+class VideoMAEv1Wrapper(PytorchWrapper):
+    def forward(self, inputs):
+        tensor = th.stack(inputs)
+        tensor = tensor.to(self._device)
+        return self._model.forward_encoder(tensor, mask=None)  # encoder only
+    
+
+input_mean = [0.485, 0.456, 0.406] # IMAGENET_DEFAULT_MEAN
+input_std = [0.229, 0.224, 0.225] # IMAGENET_DEFAULT_STD
+transform_img = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.Normalize(mean=input_mean, std=input_std),
+])
+
+def transform_video(video):
+    import torch
+    frames = torch.Tensor(video.to_numpy() / 255.0).permute(0, 3, 1, 2)
+    frames = transform_img(frames)
+    return frames.permute(1, 0, 2, 3)
+
+
+def get_model(identifier, num_frames=16):
+    assert identifier.startswith("VideoMAE-V1")
+    pretrain_only = True
+
+    if identifier == "VideoMAE-V1-B":
+        model_name = "pretrain_videomae_base_patch16_224"
+        pth = load_weight_file(
+            bucket="brainscore-vision",
+            relative_path='temporal_model_VideoMAE/vit_b_k400_pt_1200e.pth',
+            version_id="Oi3VboRZujNyZAcwf5q7XZ2M8q1cPO6o",
+            sha1="8faf42df13f619a8970d653695e98f0643760b54"
+        )
+        num_blocks = 12
+    elif identifier == "VideoMAE-V1-L":
+        model_name = "pretrain_videomae_large_patch16_224"
+        pth = load_weight_file(
+            bucket="brainscore-vision",
+            relative_path='temporal_model_VideoMAE/vit_l_k400_pt_1200e.pth',
+            version_id="MiPfczDaVponDGuUBrEPqmT.BiVvh_I1",
+            sha1="7ff6acbba221f85d7148223ec932ad7c57f2f40c"
+        )
+        num_blocks = 24
+
+    # Instantiate the model
+
+    net = create_model(
+        model_name,
+        pretrained=False,
+    )
+
+    # Load the model weights
+    if pretrain_only:
+        st = th.load(pth, map_location='cpu')['model']
+        msg = net.load_state_dict(st, strict=False)  # encoder only
+    else:
+        pth = weight_registry['VideoMAE/videomae-v1-vit-b-k400-finetune.pth']
+        st_ = th.load(pth, map_location='cpu')['module']
+        st = {}
+        for k, v in st_.items():
+            st['encoder.'+k] = v
+        msg = net.load_state_dict(st, strict=False) 
+    for layer in msg.missing_keys:
+        assert layer.startswith("decoder.")
+
+    feature_map_size = 14
+
+    inferencer_kwargs = {
+        "fps": 6.25,
+        "layer_activation_format": {
+            "encoder.patch_embed": "THWC",
+            **{f"encoder.blocks.{i}": "THWC" for i in range(0, num_blocks, LAYER_SELECTION_STEP)},
+        },
+        "num_frames": num_frames,
+    }
+
+    for layer in inferencer_kwargs["layer_activation_format"].keys():
+        assert "decoder" not in layer, "Decoder layers are not supported."
+
+    def process_activation(layer, layer_name, inputs, output):
+        B = output.shape[0]
+        C = output.shape[-1]
+        output = output.reshape(B, -1, feature_map_size, feature_map_size, C)
+        return output
+
+    wrapper = VideoMAEv1Wrapper(identifier, net, transform_video, 
+                                process_output=process_activation,
+                                **inferencer_kwargs)
+    return wrapper

brainscore_vision/models/temporal_model_VideoMAE/requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+timm
+videomae @ git+https://github.com/YingtianDt/VideoMAE.git
+iopath
+simplejson

brainscore_vision/models/temporal_model_VideoMAE/test.py

@@ -0,0 +1,16 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "VideoMAE-V1-B",
+    "VideoMAE-V1-L",
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None

brainscore_vision/models/temporal_model_VideoMAEv2/__init__.py

@@ -0,0 +1,14 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+model_registry["VideoMAE-V2-B"] = lambda: commit_model("VideoMAE-V2-B")
+model_registry["VideoMAE-V2-G"] = lambda: commit_model("VideoMAE-V2-G")

brainscore_vision/models/temporal_model_VideoMAEv2/model.py

@@ -0,0 +1,109 @@
+
+import numpy as np
+import torch
+from timm.models import create_model
+from torchvision import transforms
+
+# NOTE: Do not comment `import models`, it is used to register models
+from videomae_v2 import *  # noqa: F401
+
+from brainscore_vision.model_helpers.activations.temporal.model import PytorchWrapper
+from brainscore_vision.model_helpers.s3 import load_weight_file
+
+
+LAYER_SELECT_STEP = 2
+
+def to_normalized_float_tensor(vid):
+    vid = torch.Tensor(vid.to_numpy())
+    return vid.permute(3, 0, 1, 2).to(torch.float32) / 255
+
+def resize(vid, size, interpolation='bilinear'):
+    scale = None
+    if isinstance(size, int):
+        scale = float(size) / min(vid.shape[-2:])
+        size = None
+    return torch.nn.functional.interpolate(
+        vid,
+        size=size,
+        scale_factor=scale,
+        mode=interpolation,
+        align_corners=False)
+
+class ToFloatTensorInZeroOne(object):
+    def __call__(self, vid):
+        return to_normalized_float_tensor(vid)
+
+
+class Resize(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, vid):
+        return resize(vid, self.size)
+    
+
+transform_video = transforms.Compose(
+    [ToFloatTensorInZeroOne(),
+        Resize((224, 224))])
+
+def get_model(identifier):
+
+    if identifier == "VideoMAE-V2-G":
+        model_name = "vit_giant_patch14_224"
+        pth = load_weight_file(
+            bucket="brainscore-vision",
+            relative_path="temporal_model_VideoMAEv2/vit_g_hybrid_pt_1200e.pth",
+            version_id="TxtkfbeMV105dzpzTwi0Kn5glnvQvIrq",
+            sha1="9048f2bc0b0c7ba4d0e5228f3a7c0bef4dbaca69"
+        )
+        num_blocks = 40
+        feature_map_size = 16
+    elif identifier == "VideoMAE-V2-B":
+        model_name = "vit_base_patch16_224"
+        pth = load_weight_file(
+            bucket="brainscore-vision",
+            relative_path="temporal_model_VideoMAEv2/vit_b_hybrid_pt_800e.pth",
+            version_id="rRjpYq21dAQ5KaCLbEHK.YaLZ_fbMPKw",
+            sha1="1e3602691964b1eb6f7c33529119243a5b235635"
+        )
+        num_blocks = 12
+        feature_map_size = 14
+        
+    num_frames = 16
+
+    model = create_model(model_name)
+    
+    ckpt = torch.load(pth, map_location='cpu')
+    for model_key in ['model', 'module']:
+        if model_key in ckpt:
+            ckpt = ckpt[model_key]
+            break
+
+    encoder_ckpt = {}
+    for k, v in ckpt.items():
+        if k.startswith("encoder."):
+            encoder_ckpt[k[8:]] = v
+
+    msg = model.load_state_dict(encoder_ckpt, strict=False)
+    print(msg)
+
+    inferencer_kwargs = {
+        "fps": 6.25,
+        "layer_activation_format": {
+            "patch_embed": "THWC",
+            **{f"blocks.{i}": "THWC" for i in range(0, num_blocks, LAYER_SELECT_STEP)},
+            # "head": "THWC"  # weight not available
+        },
+        "num_frames": num_frames,
+    }
+
+    def process_activation(layer, layer_name, inputs, output):
+        B = output.shape[0]
+        C = output.shape[-1]
+        output = output.reshape(B, -1, feature_map_size, feature_map_size, C)
+        return output
+
+    wrapper = PytorchWrapper(identifier, model, transform_video, 
+                                process_output=process_activation,
+                                **inferencer_kwargs)
+    return wrapper

brainscore_vision/models/temporal_model_VideoMAEv2/requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+timm
+videomae_v2 @ git+https://github.com/YingtianDt/VideoMAEv2.git

brainscore_vision/models/temporal_model_VideoMAEv2/test.py

@@ -0,0 +1,16 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "VideoMAE-V2-B",
+    "VideoMAE-V2-G",
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None

brainscore_vision/models/temporal_model_mae_st/__init__.py

@@ -0,0 +1,15 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry["MAE-ST-L"] = lambda: commit_model("MAE-ST-L")
+model_registry["MAE-ST-G"] = lambda: commit_model("MAE-ST-G")

brainscore_vision/models/temporal_model_mae_st/model.py

@@ -0,0 +1,120 @@
+import torch
+from iopath.common.file_io import g_pathmgr as pathmgr
+from mae_st import models_vit
+from mae_st.models_vit import VisionTransformer, nn, partial
+from mae_st.util import misc 
+from mae_st.util.pos_embed import interpolate_pos_embed
+
+from brainscore_vision.model_helpers.activations.temporal.model import PytorchWrapper
+from brainscore_vision.model_helpers.s3 import load_weight_file
+
+
+def vit_huge_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16,
+        embed_dim=1280,
+        depth=32,
+        num_heads=16,
+        mlp_ratio=4,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs,
+    )
+    return model
+
+models_vit.__dict__["vit_huge_patch16"] = vit_huge_patch16
+
+
+LAYER_SELECT_STEP = 2
+mean = (0.45, 0.45, 0.45)
+std = (0.225, 0.225, 0.225)
+
+from torchvision import transforms
+
+transform_img = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.Normalize(mean, std),
+])
+
+
+def transform_video(video):
+    import torch
+    frames = torch.Tensor(video.to_numpy() / 255.0).permute(0, 3, 1, 2)
+    frames = transform_img(frames)
+    return frames.permute(1, 0, 2, 3)
+
+
+
+def get_model(identifier):
+
+    if identifier == "MAE-ST-L":
+        model_name = "vit_large_patch16"
+        num_blocks = 24
+        feature_map_size = 14
+        load_path = load_weight_file(
+            bucket="brainscore-vision", 
+            relative_path="temporal_model_mae_st/mae_pretrain_vit_large_k400.pth", 
+            version_id="cPcP4AzpG95CimQ5Pn.CHKnGUJlLXM3m",
+            sha1="c7fb91864a4ddf8b99309440121a3abe66b846bb"
+        )
+
+    elif identifier == "MAE-ST-G":
+        model_name = "vit_huge_patch16"
+        num_blocks = 32
+        feature_map_size = 14
+        load_path = load_weight_file(
+            bucket="brainscore-vision", 
+            relative_path="temporal_model_mae_st/mae_pretrain_vit_huge_k400.pth", 
+            version_id="IYKa8QiocgBzo3EhsBouS62HboK6iqYT",
+            sha1="177e48577142ca01949c08254834ffa1198b9eb4"
+        )
+
+    num_frames = 16
+    t_patch_size = 2
+
+    model = models_vit.__dict__[model_name](
+        num_frames=num_frames,
+        t_patch_size=t_patch_size,
+        cls_embed=True,
+        sep_pos_embed=True
+    )
+
+    with pathmgr.open(load_path, "rb") as f:
+        checkpoint = torch.load(f, map_location="cpu")
+
+    print("Load pre-trained checkpoint from: %s" % load_path)
+    if "model" in checkpoint.keys():
+        checkpoint_model = checkpoint["model"]
+    else:
+        checkpoint_model = checkpoint["model_state"]
+    # interpolate position embedding
+    interpolate_pos_embed(model, checkpoint_model)
+
+    checkpoint_model = misc.convert_checkpoint(checkpoint_model)
+
+    # load pre-trained model
+    msg = model.load_state_dict(checkpoint_model, strict=False)
+    print(msg)
+
+    inferencer_kwargs = {
+        "fps": 6.25,
+        "layer_activation_format": {
+            "patch_embed": "THWC",
+            **{f"blocks.{i}": "THWC" for i in range(0, num_blocks, LAYER_SELECT_STEP)},
+            # "head": "THWC"  # weight not available
+        },
+        "num_frames": num_frames,
+    }
+
+    def process_activation(layer, layer_name, inputs, output):
+        B = output.shape[0]
+        C = output.shape[-1]
+        if layer_name.startswith("blocks"):
+            output = output[:, 1:]  # remove cls token
+        output = output.reshape(B, -1, feature_map_size, feature_map_size, C)
+        return output
+
+    wrapper = PytorchWrapper(identifier, model, transform_video, 
+                                process_output=process_activation,
+                                **inferencer_kwargs)
+
+    return wrapper

brainscore_vision/models/temporal_model_mae_st/requirements.txt

@@ -0,0 +1,3 @@
+mae_st @ git+https://github.com/YingtianDt/mae_st.git
+torch
+torchvision

brainscore_vision/models/temporal_model_mae_st/test.py

@@ -0,0 +1,16 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "MAE-ST-L",
+    "MAE-ST-G",
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None