| 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), |
| ) |
|
|
| |
|
|
| 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() |
|
|
