fix prediction.py

prediction.py

@@ -136,8 +136,7 @@ def run_sequence_prediction(
 
 def run_image_prediction(
     sequence_input,
-    nucleus_image_path,
-    protein_image_path,
+    nucleus_image,
     model_ckpt_path,
     model_config_path,
     device
@@ -173,17 +172,6 @@ def run_image_prediction(
     # Convert SEQUENCE to sequence using dataset.tokenize_sequence()
     sequence = dataset.tokenize_sequence(sequence_input)
 
-    # Check if nucleus image path is provided and valid
-    if not os.path.exists(nucleus_image_path):
-        # Use default nucleus image from dataset and print warning
-        nucleus_image = dataset[0]["nucleus"]
-        print(
-            "Warning: No nucleus image provided. Using default nucleus image from dataset."
-        )
-    else:
-        # Load nucleus image from provided path
-        nucleus_image = process_image(nucleus_image_path)
-
     # 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:
@@ -209,59 +197,4 @@ def run_image_prediction(
     predicted_threshold = predicted_threshold.cpu()[0, 0]
     predicted_heatmap = predicted_heatmap.cpu()[0, 0]
 
-    # Create 3 or 4 panel plot depending on whether protein image path is provided
-    fig, axs = plt.subplots(1, 3 if protein_image_path is None else 4)
-    axs[0].imshow(nucleus_image)
-    axs[0].set_title("Nucleus Input")
-    axs[1].imshow(predicted_threshold)
-    axs[1].set_title("Predicted Threshold")
-    if protein_image_path is not None:
-        protein_image = process_image(protein_image_path)
-        axs[2].imshow(protein_image)
-        axs[2].set_title("Protein Image")
-    axs[-1].imshow(predicted_heatmap)
-    axs[-1].set_title("Predicted Heatmap")
-    plt.show()
-
-
-if __name__ == "__main__":
-    # Parse command line arguments for input parameters
-    parser = argparse.ArgumentParser(
-        description="Run Celle model with provided inputs."
-    )
-    parser.add_argument("--mode", type=str, default="", help="Sequence or Image")
-    parser.add_argument(
-        "--sequence", type=str, default="", help="Path to sequence file"
-    )
-    parser.add_argument(
-        "--nucleus_image_path",
-        type=str,
-        default="images/nucleus.jpg",
-        help="Path to nucleus image",
-    )
-    parser.add_argument(
-        "--protein_image_path",
-        type=str,
-        default=None,
-        help="Path to protein image (optional)",
-    )
-    parser.add_argument(
-        "--model_ckpt_path", type=str, required=True, help="Path to model checkpoint"
-    )
-    parser.add_argument(
-        "--model_config_path", type=str, required=True, help="Path to model config"
-    )
-    parser.add_argument(
-        "--device", type=str, default="cpu", required=True, help="device"
-    )
-    args = parser.parse_args()
-
-    run_model(
-        args.mode,
-        args.sequence,
-        args.nucleus_image_path,
-        args.protein_image_path,
-        args.model_ckpt_path,
-        args.model_config_path,
-        args.device
-    )
+    return predicted_threshold, predicted_heatmap