{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "5082cdc1-b052-433d-b018-29ba6fa3d313", "metadata": {}, "outputs": [], "source": [ "# %% [Block 1]\n", "from glob import glob\n", "from tqdm.auto import tqdm\n", "import os\n", "from PIL import Image\n", "import numpy as np\n", "import torch\n", "from scipy.io import savemat\n", "from bdpy.dl.torch import FeatureExtractor\n", "from torchvision import models, transforms\n", "import yaml\n", "import h5py" ] }, { "cell_type": "code", "execution_count": 2, "id": "a31be173-61ff-4adf-a558-63a3efe74cff", "metadata": {}, "outputs": [], "source": [ "# %% [Block 2]\n", "# Custom transform class to convert RGB to BGR\n", "class ConvertRGBtoBGR:\n", " def __call__(self, image):\n", " return image[[2, 1, 0], :, :] \n", " \n", "# Custom transform class to convert Image format to tensor while keeping pixel ranges\n", "class ToTensorWithoutScaling:\n", " def __call__(self, image):\n", " image = np.array(image).astype(np.float32) # Convert PIL Image to NumPy array\n", " tensor = torch.from_numpy(image).permute(2, 0, 1).float() # Convert NumPy array to tensor and rearrange dimensions\n", " return tensor\n", " \n", "class Convert32to16:\n", " def __call__(self, image):\n", " return image.type(torch.float16)\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "969abfe4-e4b6-4fa2-bd09-17f77ddf2292", "metadata": {}, "outputs": [], "source": [ "# %% [Block 3]\n", "def load_model(network):\n", " \n", " if network == \"VGG19_ILSVRC_19_layers\":\n", " from bdpy.dl.torch.models import layer_map, model_factory\n", " model = model_factory('vgg19')\n", " encoder_param_file = '/weka/proj-fmri/ckadirt/spurious_reconstruction/analysis/VGG_ILSVRC_19_layers/VGG_ILSVRC_19_layers.pt'\n", " model.load_state_dict(torch.load(encoder_param_file))\n", " model.eval()\n", " layer_mapping = layer_map('vgg19')\n", " # mean_image = [104., 117., 123.] # BGR\n", " preprocess = transforms.Compose([\n", " transforms.Resize((256, 256), interpolation=Image.BICUBIC),\n", " # ToTensorWithoutScaling(), # Custom transform added\n", " # ConvertRGBtoBGR(), # Custom transform added\n", " # transforms.Normalize(mean=mean_image, std=[1.,1.,1.]),\n", " ])\n", " def preprocess_and_check(img):\n", " print(\"Original shape:\", img.shape)\n", " img = transforms.Resize((224, 224), interpolation=Image.BICUBIC)(img)\n", " print(\"After ToTensor:\", img.shape)\n", " print(\"After Resize:\", img.shape)\n", " img = transforms.Normalize(mean=mean_image, std=[1.,1.,1.])(img)\n", " return img\n", " \n", " elif network == \"vgg19_torchvision\":\n", " model = models.vgg19(pretrained=True)\n", " model.eval()\n", " \n", " preprocess = transforms.Compose([\n", " transforms.Resize((224, 224)), \n", " transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n", " ]) \n", " def preprocess_and_check(img):\n", " print(\"Original shape:\", img.shape)\n", " img = transforms.ToTensor()(img)\n", " print(\"After ToTensor:\", img.shape)\n", " img = transforms.Resize(224)(img)\n", " print(\"After Resize:\", img.shape)\n", " img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)\n", " return img\n", " else:\n", " raise ValueError(\"Network not supported. Please choose from: VGG19_ILSVRC_19_layers\"\n", " +\"or define by yourself. The output should be model and preprocessing functions from RGB image to model input.\")\n", " \n", " return model, preprocess\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "8020d232-72cd-42ed-bb8e-183c8c0ee9b6", "metadata": {}, "outputs": [], "source": [ "# %% [Block 4]\n", "def load_config(config_path):\n", " '''Load configuration file.'''\n", " with open(config_path, 'r') as file:\n", " config = yaml.safe_load(file)\n", " return config\n" ] }, { "cell_type": "code", "execution_count": 5, "id": "a766c128-7e43-419c-9d7e-6f317e44753e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "/weka/proj-medarc/shared/mindeyev2_dataset//wds/subj01/train/{0..39}.tar\n", "/weka/proj-medarc/shared/mindeyev2_dataset//wds/subj01/new_test/0.tar\n", "Loaded test dl for subj1!\n", "\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "f28e9699465f4df1af719bcc78a898e2", "version_major": 2, "version_minor": 0 }, "text/plain": [ "0it [00:00, ?it/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "542a05f56b9a40ab82d55751c552b7bb", "version_major": 2, "version_minor": 0 }, "text/plain": [ "0it [00:00, ?it/s]" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "import webdataset as wds\n", "import random\n", "\n", "def my_split_by_node(urls): return urls\n", "data_path = '/weka/proj-medarc/shared/mindeyev2_dataset/'\n", "subj_list = [1]\n", "num_sessions = 40\n", "multi_subject = False\n", "batch_size = 1\n", "data_type = torch.float16\n", "subj = 1\n", "\n", "train_data = {}\n", "train_dl = {}\n", "num_voxels = {}\n", "voxels = {}\n", "\n", "train_url = f\"{data_path}/wds/subj0{subj}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n", "print(train_url)\n", "\n", "train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\\\n", " .decode(\"torch\")\\\n", " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", "train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n", "\n", "# f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n", "# betas = f['betas'][:]\n", "# betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n", "# num_voxels_list.append(betas[0].shape[-1])\n", "# num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n", "# voxels[f'subj0{s}'] = betas\n", "# print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n", "\n", "\n", "num_test=3000\n", "test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n", "print(test_url)\n", "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n", " .decode(\"torch\")\\\n", " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", "test_dl = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False, drop_last=True, pin_memory=True)\n", "print(f\"Loaded test dl for subj{subj}!\\n\")\n", "\n", "all_indexes_train = []\n", "\n", "for behav0, past_behav0, future_behav0, old_behav0 in tqdm(train_dl):\n", " img_idx = behav0[:,0,0].cpu().long().numpy()[0]\n", " all_indexes_train.append(img_idx)\n", "\n", "all_indexes_test = []\n", "\n", "for behav0, past_behav0, future_behav0, old_behav0 in tqdm(test_dl):\n", " img_idx = behav0[:,0,0].cpu().long().numpy()[0]\n", " all_indexes_test.append(img_idx)" ] }, { "cell_type": "code", "execution_count": 6, "id": "17a4377a-67dc-49fc-b89c-7d641ef218ce", "metadata": {}, "outputs": [], "source": [ "def extract_features(config, device='cuda'):\n", " '''Extract features based on the configuration.'''\n", " print(\"Extracting features using network:\", config['network'])\n", " \n", " f = h5py.File(f'{config[\"image path\"]}/coco_images_224_float16.hdf5', 'r')\n", " images = f['images']\n", " print(images.shape)\n", " model, preprocess = load_model(config['network'])\n", " layers = config[\"features\"]\n", " feature_extractor = FeatureExtractor(model, layers, device=device, detach=True)\n", " \n", " output_dir = os.path.join(config[\"output base dir\"], \"pytorch\", config['network'])\n", "\n", " all_indexes_to_compute = set(all_indexes_train + list(set(all_indexes_test)))\n", " \n", " # First, ensure all layer directories exist\n", " # for layer in layers:\n", " # layer_output_dir = os.path.join(output_dir, layer)\n", " # os.makedirs(layer_output_dir, exist_ok=True)\n", " # img = torch.Tensor(images[0])\n", "\n", " # x = preprocess(img).unsqueeze(0).to(device)\n", " # # Extract features\n", " # features = feature_extractor.run(x)\n", " \n", " # # Get the feature for the current layer\n", " # f = features.get(layer)\n", " # print(f.shape)\n", " \n", "\n", " \n", " for layer in tqdm(layers, desc=\"Processing layers\"): # Iterate over each layer first\n", " layer_output_dir = os.path.join(output_dir, layer)\n", "\n", " with h5py.File(f'{layer_output_dir}/{subj}.h5', 'a') as hdf5_file:\n", "\n", " # initial_shape = (0, 100) # Start with 0 rows and 100 columns (or your array shape)\n", " # max_shape = (None, 100) # Allow unlimited rows\n", " \n", " # if 'dataset' not in hdf5_file:\n", " # dataset = hdf5_file.create_dataset(\n", " # 'dataset', \n", " # shape=initial_shape, \n", " # maxshape=max_shape, \n", " # dtype='float64'\n", " # )\n", " # else:\n", " # dataset = hdf5_file['dataset']\n", " for i, image_index in tqdm(enumerate(all_indexes_to_compute), total = len(all_indexes_to_compute)):\n", " output_file = os.path.join(layer_output_dir, f\"{image_index}.npy\")\n", " \n", " if os.path.exists(output_file):\n", " continue # Skip if the feature for this image and layer already exists\n", " \n", " # Load and preprocess the image\n", " img = torch.Tensor(images[image_index])\n", " x = preprocess(img).unsqueeze(0).to(device)\n", " \n", " # Extract features\n", " features = feature_extractor.run(x)\n", " # Get the feature for the current layer\n", " f = features.get(layer)\n", " if f is None:\n", " print(f\"Warning: Layer '{layer}' not found in the extracted features.\")\n", " continue\n", "\n", " if i == 0:\n", " initial_shape = f.shape\n", " max_shape = (None,) + initial_shape[1:]\n", "\n", " if 'dataset' not in hdf5_file:\n", " dataset = hdf5_file.create_dataset(\n", " 'dataset', \n", " shape=initial_shape, \n", " maxshape=max_shape, \n", " dtype='float16'\n", " )\n", " else:\n", " dataset = hdf5_file['dataset']\n", "\n", " dataset.resize(dataset.shape[0] + f.shape[0], axis=0)\n", " \n", " # Append the data\n", " dataset[-f.shape[0]:] = f.astype(np.float16)\n", " \n", " # # Save the feature\n", " # np.save(output_file, f.astype(np.float16))\n", " \n", " # Optional: Break early for testing\n", " # if layer == some_condition:\n", " # break\n", "\n", " print('All done')" ] }, { "cell_type": "code", "execution_count": 7, "id": "5dc66f30-9406-4e7a-adca-81161df3611a", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Extracting features using network: VGG19_ILSVRC_19_layers\n", "(73000, 3, 224, 224)\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "8d90c3331b2149948797a4b66d9c9e32", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Processing layers: 0%| | 0/2 [00:00 1\u001b[0m \u001b[43muuuu\u001b[49m\n", "\u001b[0;31mNameError\u001b[0m: name 'uuuu' is not defined" ] } ], "source": [ "uuuu" ] }, { "cell_type": "code", "execution_count": null, "id": "35e4f307-1652-4af7-ad82-3ee4c3f142fd", "metadata": {}, "outputs": [], "source": [ "device = 'cuda'\n", "'''Extract features based on the configuration.'''\n", "print(\"Extracting features using network:\", config['network'])\n", "\n", "f = h5py.File(f'{config[\"image path\"]}/coco_images_224_float16.hdf5', 'r')\n", "images = f['images']\n", "print(images.shape)\n", "model, preprocess = load_model(config['network'])\n", "layers = config[\"features\"]\n", "feature_extractor = FeatureExtractor(model, layers, device=device, detach=True)\n", "\n", "output_dir = os.path.join(config[\"output base dir\"], \"pytorch\", config['network'])\n", "\n", "all_indexes_to_compute = set(all_indexes_train + list(set(all_indexes_test)))\n", "\n", "# First, ensure all layer directories exist\n", "# for layer in layers:\n", "# layer_output_dir = os.path.join(output_dir, layer)\n", "# os.makedirs(layer_output_dir, exist_ok=True)\n", "# img = torch.Tensor(images[0])\n", "\n", "# x = preprocess(img).unsqueeze(0).to(device)\n", "# # Extract features\n", "# features = feature_extractor.run(x)\n", " \n", "# # Get the feature for the current layer\n", "# f = features.get(layer)\n", "# print(f.shape)\n", " \n", "\n", "\n", "for layer in tqdm(layers, desc=\"Processing layers\"): # Iterate over each layer first\n", " layer_output_dir = os.path.join(output_dir, layer)\n", "\n", " # with h5py.File(f'{layer_output_dir}/{subj}.h5', 'a') as hdf5_file:\n", "\n", " # initial_shape = (0, 100) # Start with 0 rows and 100 columns (or your array shape)\n", " # max_shape = (None, 100) # Allow unlimited rows\n", " \n", " # if 'dataset' not in hdf5_file:\n", " # dataset = hdf5_file.create_dataset(\n", " # 'dataset', \n", " # shape=initial_shape, \n", " # maxshape=max_shape, \n", " # dtype='float64'\n", " # )\n", " # else:\n", " # dataset = hdf5_file['dataset']\n", " for i, image_index in tqdm(enumerate(all_indexes_to_compute), total = len(all_indexes_to_compute)):\n", " output_file = os.path.join(layer_output_dir, f\"{image_index}.npy\")\n", " \n", " if os.path.exists(output_file):\n", " continue # Skip if the feature for this image and layer already exists\n", " \n", " # Load and preprocess the image\n", " img = torch.Tensor(images[image_index])\n", " x = preprocess(img).unsqueeze(0).to(device)\n", " \n", " # Extract features\n", " features = feature_extractor.run(x)\n", " ththth\n", " # Get the feature for the current layer\n", " f = features.get(layer)\n", " if f is None:\n", " print(f\"Warning: Layer '{layer}' not found in the extracted features.\")\n", " continue\n", "\n", " if i == 0:\n", " initial_shape = f.shape\n", " max_shape = (None,) + initial_shape[1:]\n", "\n", " if 'dataset' not in hdf5_file:\n", " dataset = hdf5_file.create_dataset(\n", " 'dataset', \n", " shape=initial_shape, \n", " maxshape=max_shape, \n", " dtype='float16'\n", " )\n", " else:\n", " dataset = hdf5_file['dataset']\n", "\n", " dataset.resize(dataset.shape[0] + f.shape[0], axis=0)\n", "\n", " # Append the data\n", " dataset[-f.shape[0]:] = f.astype(np.float16)\n", " \n", " # # Save the feature\n", " # np.save(output_file, f.astype(np.float16))\n", "\n", "# Optional: Break early for testing\n", "# if layer == some_condition:\n", "# break\n", "\n", "print('All done')" ] }, { "cell_type": "code", "execution_count": null, "id": "638ad922-e988-4331-9e56-4287648abe4e", "metadata": {}, "outputs": [], "source": [ "model" ] }, { "cell_type": "code", "execution_count": null, "id": "05a07c1f-9047-49e4-8a82-9133a55f469f", "metadata": {}, "outputs": [], "source": [ "from torchvision.transforms import ToPILImage\n", "class EncoderFeatureExtractor(torch.nn.Module):\n", " def __init__(self, encoder, target_layer=0):\n", " super(EncoderFeatureExtractor, self).__init__()\n", " self.encoder = encoder\n", " self.target_layer = target_layer\n", " self.features = torch.nn.Sequential(*list(self.encoder.features.children())[:self.target_layer + 1])\n", "\n", " def forward(self, x):\n", " return self.features(x)\n", "\n", "feature_extractor = EncoderFeatureExtractor(model, target_layer=0).to(device)\n", "\n", "\n", "from PIL import Image\n", "img\n", "\n", "# Preprocess the image\n", "preprocess = transforms.Compose([\n", " transforms.Resize((256, 256)),\n", " # transforms.ToTensor(),\n", " # transforms.Normalize(mean=[0.485, 0.456, 0.406], # Example normalization (ImageNet)\n", " # std=[0.229, 0.224, 0.225]),\n", "])\n", "\n", "# Preprocess the image\n", "image_tensor = preprocess(img).unsqueeze(0).to(device)\n", "\n", "# show processed image\n", "to_pil = ToPILImage()\n", "pil_image = to_pil(image_tensor.squeeze())\n", "pil_image.show()\n", "\n", "# Generate features from the image\n", "with torch.no_grad():\n", " image_features = feature_extractor(image_tensor)\n", "\n", "# Define target features\n", "target_features = image_features.to(device).float()" ] }, { "cell_type": "code", "execution_count": null, "id": "53d2ed5f-f166-4b56-ac5d-f995db7ffbd7", "metadata": {}, "outputs": [], "source": [ "target_features[0,0,0,10:30].cpu()" ] }, { "cell_type": "code", "execution_count": null, "id": "63d92bf0-79fe-4584-8f76-53e3d6fa216d", "metadata": {}, "outputs": [], "source": [ "features['features[0]'].shape" ] }, { "cell_type": "code", "execution_count": null, "id": "ebb6206b-084b-4204-b5d0-c8f920dea6c0", "metadata": {}, "outputs": [], "source": [ "features['features[0]'][0,0,0,10:30]" ] }, { "cell_type": "code", "execution_count": null, "id": "276faec3-0d46-4218-88fa-5c3fa103c81e", "metadata": {}, "outputs": [], "source": [ "torch.save(features['features[0]'], './ff.pt')" ] }, { "cell_type": "code", "execution_count": null, "id": "7e2049a1-0ccc-43d3-8fcb-4ff1531b1e9e", "metadata": {}, "outputs": [], "source": [ "# from PIL import Image\n", "# import numpy as np\n", "\n", "# # Create a random image tensor with dimensions (3, 224, 224) \n", "# # and values between 0 and 255 to simulate an RGB image\n", "\n", "# image = Image.fromarray((img.numpy().transpose(1, 2, 0) * 255).astype(np.uint8))\n", "\n", "# # Display the image\n", "# image.show()\n" ] }, { "cell_type": "code", "execution_count": null, "id": "a189b39e-8493-4dff-806d-ec7768acada6", "metadata": {}, "outputs": [], "source": [ "# features['features[0]'].min()" ] }, { "cell_type": "code", "execution_count": null, "id": "dd3fc063-68bd-4c31-b794-4baaf30d629b", "metadata": {}, "outputs": [], "source": [ "# features['features[0]']" ] }, { "cell_type": "code", "execution_count": null, "id": "2d6c4809-0aab-4257-8994-4e892e23de9f", "metadata": {}, "outputs": [], "source": [ "# with h5py.File(f'/weka/proj-fmri/ckadirt/spurious_reconstruction/analysis/0_preprocessing/features/pytorch/vgg19_torchvision/features[0]/1.h5', 'a') as hdf5_file:\n", "# dataset = hdf5_file['dataset']\n", "# print(len(dataset))\n", "# print(dataset[10].shape, dataset[10].dtype)\n", "# print(dataset.shape)" ] }, { "cell_type": "code", "execution_count": null, "id": "fcf6d112-e24c-4fd5-ac59-f21e3611ebbc", "metadata": {}, "outputs": [], "source": [ "# dataset" ] } ], "metadata": { "kernelspec": { "display_name": "mindeye", "language": "python", "name": "mindeye" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.10" } }, "nbformat": 4, "nbformat_minor": 5 }