prediction.py

import os
os.chdir('..')
base_dir = os.getcwd()
from dataloader import CellLoader
from celle_main import instantiate_from_config
from omegaconf import OmegaConf

def run_sequence_prediction(
 sequence_input,
 nucleus_image,
 protein_image,
 model_ckpt_path,
 model_config_path,
 device
):
 """
 Run Celle model with provided inputs and display results.

 :param sequence: Path to sequence file
 :param nucleus_image_path: Path to nucleus image
 :param protein_image_path: Path to protein image (optional)
 :param model_ckpt_path: Path to model checkpoint
 :param model_config_path: Path to model config
 """
 
 # Instantiate dataset object
 dataset = CellLoader(
 sequence_mode="embedding",
 vocab="esm2",
 split_key="val",
 crop_method="center",
 resize=600,
 crop_size=256,
 text_seq_len=1000,
 pad_mode="end",
 threshold="median",
 )

 # Check if sequence is provided and valid
 if len(sequence_input) == 0:
 raise ValueError("Sequence must be provided.")

 if "<mask>" not in sequence_input:
 print("Warning: Sequence does not contain any masked positions to predict.")

 # Convert SEQUENCE to sequence using dataset.tokenize_sequence()
 sequence = dataset.tokenize_sequence(sequence_input)

 # Load model config and set ckpt_path if not provided in config
 config = OmegaConf.load(model_config_path)
 if config["model"]["params"]["ckpt_path"] is None:
 config["model"]["params"]["ckpt_path"] = model_ckpt_path

 # Set condition_model_path and vqgan_model_path to None
 config["model"]["params"]["condition_model_path"] = None
 config["model"]["params"]["vqgan_model_path"] = None
 
 os.chdir(os.path.dirname(model_ckpt_path))

 # Instantiate model from config and move to device
 model = instantiate_from_config(config.model).to(device)

 # Sample from model using provided sequence and nucleus image
 _, predicted_sequence, _ = model.celle.sample_text(
 text=sequence.to(device),
 condition=nucleus_image.to(device),
 image=protein_image.to(device),
 force_aas=True,
 temperature=1,
 progress=False,
 )
 
 os.chdir(base_dir)

 return predicted_sequence