import os
import numpy as np
from PIL import Image, ImageSequence
import json
import pandas as pd
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import torchvision.transforms.functional as TF
from celle.utils import replace_outliers
def simple_conversion(seq):
"""Create 26-dim embedding"""
chars = [
"-",
"M",
"R",
"H",
"K",
"D",
"E",
"S",
"T",
"N",
"Q",
"C",
"U",
"G",
"P",
"A",
"V",
"I",
"F",
"Y",
"W",
"L",
"O",
"X",
"Z",
"B",
"J",
]
nums = range(len(chars))
seqs_x = np.zeros(len(seq))
for idx, char in enumerate(seq):
lui = chars.index(char)
seqs_x[idx] = nums[lui]
return torch.tensor([seqs_x]).long()
class CellLoader(Dataset):
"""imports mined opencell images with protein sequence"""
def __init__(
self,
data_csv=None,
dataset=None,
split_key=None,
resize=600,
crop_size=600,
crop_method="random",
sequence_mode="simple",
vocab="bert",
threshold="median",
text_seq_len=0,
pad_mode="random",
):
self.data_csv = data_csv
self.dataset = dataset
self.image_folders = []
self.crop_method = crop_method
self.resize = resize
self.crop_size = crop_size
self.sequence_mode = sequence_mode
self.threshold = threshold
self.text_seq_len = int(text_seq_len)
self.vocab = vocab
self.pad_mode = pad_mode
if self.sequence_mode == "embedding" or self.sequence_mode == "onehot":
if self.vocab == "esm1b" or self.vocab == "esm2":
from esm import Alphabet
self.tokenizer = Alphabet.from_architecture(
"ESM-1b"
).get_batch_converter()
self.text_seq_len += 2
if data_csv:
data = pd.read_csv(data_csv)
self.parent_path = os.path.dirname(data_csv).split(data_csv)[0]
if split_key == "train":
self.data = data[data["split"] == "train"]
elif split_key == "val":
self.data = data[data["split"] == "val"]
else:
self.data = data
self.data = self.data.reset_index(drop=True)
def __len__(self):
return len(self.data)
def __getitem__(
self,
idx,
get_sequence=True,
get_images=True,
):
if get_sequence and self.text_seq_len > 0:
protein_vector = self.get_protein_vector(idx)
else:
protein_vector = torch.zeros((1, 1))
if get_images:
nucleus, target, threshold = self.get_images(idx, self.dataset)
else:
nucleus, target, threshold = torch.zeros((3, 1))
data_dict = {
"nucleus": nucleus.float(),
"target": target.float(),
"threshold": threshold.float(),
"sequence": protein_vector.long(),
}
return data_dict
def get_protein_vector(self, idx):
if "protein_sequence" not in self.data.columns:
metadata = self.retrieve_metadata(idx)
protein_sequence = metadata["sequence"]
else:
protein_sequence = self.data.iloc[idx]["protein_sequence"]
protein_vector = self.tokenize_sequence(protein_sequence)
return protein_vector
def get_images(self, idx, dataset):
if dataset == "HPA":
nucleus = Image.open(
os.path.join(
self.parent_path, self.data.iloc[idx]["nucleus_image_path"]
)
)
target = Image.open(
os.path.join(self.parent_path, self.data.iloc[idx]["target_image_path"])
)
nucleus = TF.to_tensor(nucleus)[0]
target = TF.to_tensor(target)[0]
image = torch.stack([nucleus, target], axis=0)
normalize = (0.0655, 0.0650), (0.1732, 0.1208)
elif dataset == "OpenCell":
image = Image.open(
os.path.join(self.parent_path, self.data.iloc[idx]["image_path"])
)
nucleus, target = [page.copy() for page in ImageSequence.Iterator(image)]
nucleus = replace_outliers(torch.divide(TF.to_tensor(nucleus), 65536))[0]
target = replace_outliers(torch.divide(TF.to_tensor(target), 65536))[0]
image = torch.stack([nucleus, target], axis=0)
normalize = (
(0.0272, 0.0244),
(0.0486, 0.0671),
)
# # from https://discuss.pytorch.org/t/how-to-apply-same-transform-on-a-pair-of-picture/14914
t_forms = [transforms.Resize(self.resize, antialias=None)]
if self.crop_method == "random":
t_forms.append(transforms.RandomCrop(self.crop_size))
t_forms.append(transforms.RandomHorizontalFlip(p=0.5))
t_forms.append(transforms.RandomVerticalFlip(p=0.5))
elif self.crop_method == "center":
t_forms.append(transforms.CenterCrop(self.crop_size))
t_forms.append(transforms.Normalize(normalize[0], normalize[1]))
image = transforms.Compose(t_forms)(image)
nucleus, target = image
nucleus /= torch.abs(nucleus).max()
target -= target.min()
target /= target.max()
nucleus = nucleus.unsqueeze(0)
target = target.unsqueeze(0)
threshold = target
if self.threshold == "mean":
threshold = 1.0 * (threshold > (torch.mean(threshold)))
elif self.threshold == "median":
threshold = 1.0 * (threshold > (torch.median(threshold)))
elif self.threshold == "1090_IQR":
p10 = torch.quantile(threshold, 0.1, None)
p90 = torch.quantile(threshold, 0.9, None)
threshold = torch.clip(threshold, p10, p90)
nucleus = torch.nan_to_num(nucleus, 0.0, 1.0, 0.0)
target = torch.nan_to_num(target, 0.0, 1.0, 0.0)
threshold = torch.nan_to_num(threshold, 0.0, 1.0, 0.0)
return nucleus, target, threshold
def retrieve_metadata(self, idx):
with open(
os.path.join(self.parent_path, self.data.iloc[idx]["metadata_path"])
) as f:
metadata = json.load(f)
return metadata
def tokenize_sequence(self, protein_sequence):
pad_token = 0
if self.sequence_mode == "simple":
protein_vector = simple_conversion(protein_sequence)
elif self.sequence_mode == "center":
protein_sequence = protein_sequence.center(self.text_seq_length, "-")
protein_vector = simple_conversion(protein_sequence)
elif self.sequence_mode == "alternating":
protein_sequence = protein_sequence.center(self.text_seq_length, "-")
protein_sequence = protein_sequence[::18]
protein_sequence = protein_sequence.center(
int(self.text_seq_length / 18) + 1, "-"
)
protein_vector = simple_conversion(protein_sequence)
elif self.sequence_mode == "embedding":
if self.vocab == "esm1b" or self.vocab == "esm2":
pad_token = 1
protein_vector = self.tokenizer([("", protein_sequence)])[-1]
if protein_vector.shape[-1] < self.text_seq_len:
diff = self.text_seq_len - protein_vector.shape[-1]
if self.pad_mode == "end":
protein_vector = torch.nn.functional.pad(
protein_vector, (0, diff), "constant", pad_token
)
elif self.pad_mode == "random":
split = diff - np.random.randint(0, diff + 1)
protein_vector = torch.cat(
[torch.ones(1, split) * 0, protein_vector], dim=1
)
protein_vector = torch.nn.functional.pad(
protein_vector, (0, diff - split), "constant", pad_token
)
elif protein_vector.shape[-1] > self.text_seq_len:
start_int = np.random.randint(
0, protein_vector.shape[-1] - self.text_seq_len
)
protein_vector = protein_vector[
:, start_int : start_int + self.text_seq_len
]
return protein_vector.long()