Update prediction.py

prediction.py

@@ -1,16 +1,11 @@
 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(
+
+def run_image_prediction(
  sequence_input,
  nucleus_image,
- protein_image,
- model_ckpt_path,
- model_config_path,
+ model,
  device
 ):
  """
@@ -22,7 +17,6 @@ def run_sequence_prediction(
  :param model_ckpt_path: Path to model checkpoint
  :param model_config_path: Path to model config
  """
- 
  # Instantiate dataset object
  dataset = CellLoader(
  sequence_mode="embedding",
@@ -36,40 +30,20 @@ def run_sequence_prediction(
  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(
+ _, _, _, predicted_threshold, predicted_heatmap = model.celle.sample(
  text=sequence.to(device),
  condition=nucleus_image.to(device),
- image=protein_image.to(device),
- force_aas=True,
+ timesteps=1,
  temperature=1,
  progress=False,
  )
- 
- os.chdir(base_dir)
 
- return predicted_sequence
+ # Move predicted_threshold and predicted_heatmap to CPU and select first element of batch
+ predicted_threshold = predicted_threshold.cpu()[0, 0]
+ predicted_heatmap = predicted_heatmap.cpu()[0, 0]
+
+ return predicted_threshold, predicted_heatmap