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