Spaces:
Running
Running
| import numpy as np | |
| import pandas as pd | |
| import streamlit as st | |
| from PIL import Image | |
| import sys | |
| import io | |
| import os | |
| import glob | |
| import json | |
| import zipfile | |
| from tqdm import tqdm | |
| from itertools import chain | |
| import torch | |
| from torch.utils.data import DataLoader | |
| sys.path.insert(0, os.path.abspath("src/")) | |
| from clip.clip import _transform | |
| from training.datasets import CellPainting | |
| from clip.model import convert_weights, CLIPGeneral | |
| from rdkit import Chem | |
| from rdkit.Chem import Draw | |
| from rdkit.Chem import AllChem | |
| from rdkit.Chem import DataStructs | |
| st.set_page_config(layout="wide") | |
| basepath = os.path.dirname(__file__) | |
| datapath = os.path.join(basepath, "data") | |
| CLOOME_PATH = "/home/ana/gitrepos/hti-cloob" | |
| MODEL_PATH = os.path.join(datapath, "epoch_55.pt") | |
| npzs = os.path.join(datapath, "npzs") | |
| molecule_features = os.path.join(datapath, "all_molecule_cellpainting_features.pkl") | |
| mol_index_file = os.path.join(datapath, "cellpainting-unique-molecule.csv") | |
| image_features = os.path.join(datapath, "subset_image_cellpainting_features.pkl") | |
| images_arr = os.path.join(datapath, "subset_npzs_dict_.npz") | |
| img_index_file = os.path.join(datapath, "cellpainting-all-imgpermol.csv") | |
| imgname = "I1" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model_type = "RN50" | |
| ######### CLOOME FUNCTIONS ######### | |
| def convert_models_to_fp32(model): | |
| for p in model.parameters(): | |
| p.data = p.data.float() | |
| if p.grad: | |
| p.grad.data = p.grad.data.float() | |
| def load(model_path, device, model, image_resolution): | |
| state_dict = torch.load(model_path, map_location=device) | |
| state_dict = state_dict["state_dict"] | |
| model_config_file = f"{model.replace('/', '-')}.json" | |
| print('Loading model from', model_config_file) | |
| assert os.path.exists(model_config_file) | |
| with open(model_config_file, 'r') as f: | |
| model_info = json.load(f) | |
| model = CLIPGeneral(**model_info) | |
| convert_weights(model) | |
| convert_models_to_fp32(model) | |
| if str(device) == "cpu": | |
| model.float() | |
| print(device) | |
| new_state_dict = {k[len('module.'):]: v for k,v in state_dict.items()} | |
| model.load_state_dict(new_state_dict) | |
| model.to(device) | |
| model.eval() | |
| return model | |
| def get_features(dataset, model, device): | |
| all_image_features = [] | |
| all_text_features = [] | |
| all_ids = [] | |
| print(f"get_features {device}") | |
| print(len(dataset)) | |
| with torch.no_grad(): | |
| for batch in tqdm(DataLoader(dataset, num_workers=1, batch_size=64)): | |
| if type(batch) is dict: | |
| imgs = batch | |
| text_features = None | |
| mols = None | |
| elif type(batch) is torch.Tensor: | |
| mols = batch | |
| imgs = None | |
| else: | |
| imgs, mols = batch | |
| if mols is not None: | |
| text_features = model.encode_text(mols.to(device)) | |
| text_features = text_features / text_features.norm(dim=-1, keepdim=True) | |
| all_text_features.append(text_features) | |
| molecules_exist = True | |
| if imgs is not None: | |
| images = imgs["input"] | |
| ids = imgs["ID"] | |
| img_features = model.encode_image(images.to(device)) | |
| img_features = img_features / img_features.norm(dim=-1, keepdim=True) | |
| all_image_features.append(img_features) | |
| all_ids.append(ids) | |
| all_ids = list(chain.from_iterable(all_ids)) | |
| if imgs is not None and mols is not None: | |
| return torch.cat(all_image_features), torch.cat(all_text_features), all_ids | |
| elif imgs is not None: | |
| return torch.cat(all_image_features), all_ids | |
| elif mols is not None: | |
| return torch.cat(all_text_features), all_ids | |
| return | |
| def read_array(file): | |
| t = torch.load(file) | |
| features = t["mol_features"] | |
| ids = t["mol_ids"] | |
| return features, ids | |
| def main(df, model_path, model, img_path=None, mol_path=None, image_resolution=None): | |
| # Load the model | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(torch.cuda.device_count()) | |
| model = load(model_path, device, model, image_resolution) | |
| preprocess_val = _transform(image_resolution, image_resolution, is_train=False, normalize="dataset", preprocess="downsize") | |
| # Load the dataset | |
| val = CellPainting(df, | |
| img_path, | |
| mol_path, | |
| transforms = preprocess_val) | |
| # Calculate the image features | |
| print("getting_features") | |
| result = get_features(val, model, device) | |
| if len(result) > 2: | |
| val_img_features, val_text_features, val_ids = result | |
| return val_img_features, val_text_features, val_ids | |
| else: | |
| val_img_features, val_ids = result | |
| return val_img_features, val_ids | |
| def img_to_numpy(file): | |
| img = Image.open(file) | |
| arr = np.array(img) | |
| return arr | |
| def illumination_threshold(arr, perc=0.0028): | |
| """ Return threshold value to not display a percentage of highest pixels""" | |
| perc = perc/100 | |
| h = arr.shape[0] | |
| w = arr.shape[1] | |
| # find n pixels to delete | |
| total_pixels = h * w | |
| n_pixels = total_pixels * perc | |
| n_pixels = int(np.around(n_pixels)) | |
| # find indexes of highest pixels | |
| flat_inds = np.argpartition(arr, -n_pixels, axis=None)[-n_pixels:] | |
| inds = np.array(np.unravel_index(flat_inds, arr.shape)).T | |
| max_values = [arr[i, j] for i, j in inds] | |
| threshold = min(max_values) | |
| return threshold | |
| def process_image(arr): | |
| threshold = illumination_threshold(arr) | |
| scaled_img = sixteen_to_eight_bit(arr, threshold) | |
| return scaled_img | |
| def sixteen_to_eight_bit(arr, display_max, display_min=0): | |
| threshold_image = ((arr.astype(float) - display_min) * (arr > display_min)) | |
| scaled_image = (threshold_image * (256. / (display_max - display_min))) | |
| scaled_image[scaled_image > 255] = 255 | |
| scaled_image = scaled_image.astype(np.uint8) | |
| return scaled_image | |
| def process_image(arr): | |
| threshold = illumination_threshold(arr) | |
| scaled_img = sixteen_to_eight_bit(arr, threshold) | |
| return scaled_img | |
| def process_sample(imglst, channels, filenames, outdir, outfile): | |
| sample = np.zeros((520, 696, 5), dtype=np.uint8) | |
| filenames_dict, channels_dict = {}, {} | |
| for i, (img, channel, fname) in enumerate(zip(imglst, channels, filenames)): | |
| print(channel) | |
| arr = img_to_numpy(img) | |
| arr = process_image(arr) | |
| sample[:,:,i] = arr | |
| channels_dict[i] = channel | |
| filenames_dict[channel] = fname | |
| sample_dict = dict(sample=sample, | |
| channels=channels_dict, | |
| filenames=filenames_dict) | |
| outfile = outfile + ".npz" | |
| outpath = os.path.join(outdir, outfile) | |
| np.savez(outpath, sample=sample, channels=channels, filenames=filenames) | |
| return sample_dict, outpath | |
| def display_cellpainting(sample): | |
| arr = sample["sample"] | |
| r = arr[:, :, 0].astype(np.float32) | |
| g = arr[:, :, 3].astype(np.float32) | |
| b = arr[:, :, 4].astype(np.float32) | |
| rgb_arr = np.dstack((r, g, b)) | |
| im = Image.fromarray(rgb_arr.astype("uint8")) | |
| im_rgb = im.convert("RGB") | |
| return im_rgb | |
| def morgan_from_smiles(smiles, radius=3, nbits=1024, chiral=True): | |
| mol = Chem.MolFromSmiles(smiles) | |
| fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=3, nBits=nbits, useChirality=chiral) | |
| arr = np.zeros((0,), dtype=np.int8) | |
| DataStructs.ConvertToNumpyArray(fp,arr) | |
| return arr | |
| def save_hdf(fps, index, outfile_hdf): | |
| ids = [i for i in range(len(fps))] | |
| columns = [str(i) for i in range(fps[0].shape[0])] | |
| df = pd.DataFrame(fps, index=ids, columns=columns) | |
| df.to_hdf(outfile_hdf, key="df", mode="w") | |
| return outfile_hdf | |
| def create_index(outdir, ids, filename): | |
| filepath = os.path.join(outdir, filename) | |
| if type(ids) is str: | |
| values = [ids] | |
| else: | |
| values = ids | |
| data = {"SAMPLE_KEY": values} | |
| print(data) | |
| df = pd.DataFrame(data) | |
| df.to_csv(filepath) | |
| return filepath | |
| def draw_molecules(smiles_lst): | |
| mols = [Chem.MolFromSmiles(s) for s in smiles_lst] | |
| mol_imgs = [Chem.Draw.MolToImage(m) for m in mols] | |
| return mol_imgs | |
| def reshape_image(arr): | |
| c, h, w = arr.shape | |
| reshaped_image = np.empty((h, w, c)) | |
| reshaped_image[:,:,0] = arr[0] | |
| reshaped_image[:,:,1] = arr[1] | |
| reshaped_image[:,:,2] = arr[2] | |
| reshaped_pil = Image.fromarray(reshaped_image.astype("uint8")) | |
| return reshaped_pil | |
| # missing functions: save morgan to to_hdf, create index, load features, calculate similarities | |
| ##### STREAMLIT FUNCTIONS ###### | |
| st.title('CLOOME. Bioimage database retrieval from chemical structures (and viceversa)') | |
| def main_page(top_n, model_path): | |
| st.markdown( | |
| """ | |
| Contrastive learning for self-supervised representation learning has brought a strong improvement to many application areas, such as computer vision and natural language processing. | |
| With the availability of large collections of unlabeled data in vision and language, contrastive learning of language and image representations has shown impressive results. | |
| The contrastive learning methods CLIP and CLOOB have demonstrated that the learned representations are highly transferable to a large set of diverse tasks when trained on multi-modal data from two different domains. | |
| In life sciences, similar large, multi-modal datasets comprising both cell-based microscopy images and chemical structures of molecules are available. | |
| However, contrastive learning has not yet been used for this type of multi-modal data, although this would allow to design cross-modal retrieval systems for bioimaing and chemical databases. | |
| In this work, we present a such a contrastive learning method, the retrieval systems, and the transferability of the learned representations. Our method, Contrastive Learning and leave-One-Out-boost for Molecule Encoders (CLOOME), is based on both CLOOB and CLIP and comprises an encoder for microscopy data, an encoder for chemical structures and a contrastive learning objective, which produce rich embeddings of bioimages and chemical structures. | |
| On the benchmark dataset ”Cell Painting”, we demonstrate that the embddings can be used to form a retrieval system for bioimaging and chemical databases. We also show that CLOOME learns transferable representations by performing linear probing for activity prediction tasks. Furthermore, the image embeddings can identify new cell phenotypes, as we show in a zero-shot classification task. | |
| """ | |
| ) | |
| def molecules_from_image(top_n, model_path): | |
| ## TODO: Check if expander can be automatically collapsed | |
| exp = st.expander("Upload a microscopy image") | |
| with exp: | |
| channels = ['Mito', 'ERSyto', 'ERSytoBleed', 'Ph_golgi', 'Hoechst'] | |
| imglst, filenames = [], [] | |
| for c in channels: | |
| file_obj = st.file_uploader(f'Choose a TIF image for {c}:', ".tif") | |
| if file_obj is not None: | |
| imglst.append(file_obj) | |
| filenames.append(file_obj.name) | |
| if imglst: | |
| if not os.path.isdir(npzs): | |
| os.mkdir(npzs) | |
| sample_dict, imgpath = process_sample(imglst, channels, filenames, npzs, imgname) | |
| print(imglst) | |
| i = display_cellpainting(sample_dict) | |
| st.image(i) | |
| uploaded_file = st.file_uploader("Choose a molecule file to retrieve from (optional)") | |
| if imglst: | |
| if uploaded_file is not None: | |
| molecule_df = pd.read_csv(uploaded_file) | |
| smiles = molecule_df["SMILES"].tolist() | |
| morgan = [morgan_from_smiles(s) for s in smiles] | |
| molnames = [f"M{i}" for i in range(len(morgan))] | |
| mol_index_fname = "mol_index.csv" | |
| mol_index = create_index(datapath, molnames, mol_index_fname) | |
| molpath = os.path.join(datapath, "mols.hdf") | |
| fps_fname = save_hdf(morgan, molnames, molpath) | |
| mol_imgs = draw_molecules(smiles) | |
| mol_features, mol_ids = main(mol_index, model_path, model_type, mol_path=molpath, image_resolution=image_resolution) | |
| predefined_features = False | |
| else: | |
| mol_index = pd.read_csv(mol_index_file) | |
| mol_features_torch = torch.load(molecule_features, map_location=device) | |
| mol_features = mol_features_torch["mol_features"] | |
| mol_ids = mol_features_torch["mol_ids"] | |
| print(len(mol_ids)) | |
| predefined_features = True | |
| img_index_fname = "img_index.csv" | |
| img_index = create_index(datapath, imgname, img_index_fname) | |
| img_features, img_ids = main(img_index, MODEL_PATH, model_type, img_path=npzs, image_resolution=image_resolution) | |
| print(img_features.shape) | |
| print(mol_features.shape) | |
| top_n = int(top_n) | |
| logits = img_features @ mol_features.T | |
| mol_probs = (30.0 * logits).softmax(dim=-1) | |
| top_probs, top_labels = mol_probs.cpu().topk(top_n, dim=-1) | |
| # Delete this if want to allow retrieval for multiple images | |
| top_probs = torch.flatten(top_probs) | |
| top_labels = torch.flatten(top_labels) | |
| print(top_probs.shape) | |
| print(top_labels.shape) | |
| if predefined_features: | |
| mol_index.set_index(["SAMPLE_KEY"], inplace=True) | |
| top_ids = [mol_ids[i] for i in top_labels] | |
| smiles = mol_index.loc[top_ids]["SMILES"].tolist() | |
| mol_imgs = draw_molecules(smiles) | |
| with st.container(): | |
| if uploaded_file is not None: | |
| st.write("Retrieved molecules") | |
| else: | |
| st.write("Retrieved molecules from the Cell Painting database") | |
| columns = st.columns(len(top_probs)) | |
| for i, col in enumerate(columns): | |
| if predefined_features: | |
| image_id = i | |
| else: | |
| image_id = top_labels[i] | |
| index = i+1 | |
| col.image(mol_imgs[image_id], width=140, caption=index) | |
| print(mol_probs.sum(dim=-1)) | |
| print((top_probs, top_labels)) | |
| def images_from_molecule(top_n, model_path): | |
| #st.markdown("Enter a query molecule in [SMILES](https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system) format",) | |
| smiles = st.text_input("Enter a query molecule in SMILES format", value="CC(=O)OC1=CC=CC=C1C(=O)O", placeholder="CC(=O)OC1=CC=CC=C1C(=O)O") | |
| if smiles: | |
| smiles = [smiles] | |
| morgan = [morgan_from_smiles(s) for s in smiles] | |
| molnames = [f"M{i}" for i in range(len(morgan))] | |
| mol_index_fname = "mol_index.csv" | |
| mol_index = create_index(datapath, molnames, mol_index_fname) | |
| molpath = os.path.join(datapath, "mols.hdf") | |
| fps_fname = save_hdf(morgan, molnames, molpath) | |
| mol_imgs = draw_molecules(smiles) | |
| mol_features, mol_ids = main(mol_index, model_path, model_type, mol_path=molpath, image_resolution=image_resolution) | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| st.write("") | |
| with col2: | |
| st.image(mol_imgs, width = 140) | |
| with col3: | |
| st.write("") | |
| st.markdown('##') | |
| top_n = int(top_n) | |
| img_features_torch = torch.load(image_features, map_location=device) | |
| img_features = img_features_torch["img_features"] | |
| img_ids = img_features_torch["img_ids"] | |
| logits = mol_features @ img_features.T | |
| img_probs = (30.0 * logits).softmax(dim=-1) | |
| top_probs, top_labels = img_probs.cpu().topk(top_n, dim=-1) | |
| top_probs = torch.flatten(top_probs) | |
| top_labels = torch.flatten(top_labels) | |
| img_index = pd.read_csv(img_index_file) | |
| img_index.set_index(["SAMPLE_KEY"], inplace=True) | |
| top_ids = [img_ids[i] for i in top_labels] | |
| images_dict = np.load(images_arr, allow_pickle = True) | |
| st.write("Retrieved images from the Cell Painting database") | |
| with st.container(): | |
| n_columns = 5 | |
| columns = st.columns(n_columns) | |
| for i, col in enumerate(columns): | |
| l = list(range(len(top_ids))) | |
| col_ids = [m for m in l if m % 5 == i] | |
| for n in col_ids: | |
| id = top_ids[n] | |
| id = f"{id}.npz" | |
| image = images_dict[id] | |
| ## TODO: generalize and functionalize | |
| im = reshape_image(image) | |
| index = n+1 | |
| score = float(top_probs[n]) * 100 | |
| col.image(im, caption=f"Top {index}. Score: {round(score, 2)}", width=200) | |
| page_names_to_funcs = { | |
| "Microscopy images from a molecule": images_from_molecule, | |
| "Molecules from a microscopy image": molecules_from_image, | |
| "About CLOOME": main_page, | |
| } | |
| selected_page = st.sidebar.selectbox("What would you like to retrieve?", page_names_to_funcs.keys()) | |
| st.sidebar.markdown('') | |
| n_objects = st.sidebar.selectbox( | |
| "How many objects would you like to retrieve?", | |
| ("5", "10", "20")) | |
| selected_model = st.sidebar.selectbox( | |
| "Select a CLOOME model to load", | |
| ("CLOOME (default)", "CLOOME (CLIP imgres 320)", "CLOOME (fullres)", "CLOOME (imgres 320)")) | |
| model_dict = { | |
| "CLOOME (default)" : "cloome_default.pt", | |
| "CLOOME (CLIP imgres 320)" : "cloome_cli_imres320.pt", | |
| "CLOOME (fullres)" : "cloome_fullres.pt", | |
| "CLOOME (imgres 320)" : "cloome_imres320.pt" | |
| } | |
| model_file = model_dict[selected_model] | |
| model_path = os.path.join(datapath, model_file) | |
| if model_path.endswith("320).pt"): | |
| image_resolution = 320 | |
| else: | |
| image_resolution = 520 | |
| page_names_to_funcs[selected_page](n_objects, model_path) | |