app.py

import gradio as gr
import spaces
import torch
from PIL import Image
from pathlib import Path
from main import load_and_preprocess_image, genomic_plip_predictions, classify_tiles

@spaces.GPU
def run_load_and_preprocess_image(image_path, clip_processor_path):
 image_tensor = load_and_preprocess_image(image_path, clip_processor_path)
 if image_tensor is not None:
 image_tensor = image_tensor.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
 return image_tensor

@spaces.GPU
def run_genomic_plip_predictions(image_tensor, model_path):
 if image_tensor is None:
 return None # Handle the case where image_tensor is None
 pred_data = genomic_plip_predictions(image_tensor, model_path)
 return pred_data

@spaces.GPU
def run_classify_tiles(pred_data, model_path):
 if pred_data is None:
 return "Error: Prediction data is None." # Handle the case where pred_data is None
 output = classify_tiles(pred_data, model_path)
 return output


example_files = list(Path("sample_tiles").glob("*.jpeg"))

with gr.Blocks() as demo:
 gr.Markdown("## Cancer Risk Prediction from Tissue Slide")

 image_file = gr.Image(type="filepath", label="Upload Image File")
 clip_processor_path = gr.Textbox(label="CLIP Processor Path", value="./genomic_plip_model")
 genomic_plip_model_path = gr.Textbox(label="Genomic PLIP Model Path", value="./genomic_plip_model")
 classifier_model_path = gr.Textbox(label="Classifier Model Path", value="./models/classifier.pth")

 image_tensor_output = gr.State()
 pred_data_output = gr.State()
 result_output = gr.Textbox(label="Result")

 preprocess_button = gr.Button("Preprocess Image")
 predict_button = gr.Button("Features with Genomic PLIP")
 classify_button = gr.Button("Identify Risk")

 preprocess_button.click(run_load_and_preprocess_image, inputs=[image_file, clip_processor_path], outputs=image_tensor_output)
 predict_button.click(run_genomic_plip_predictions, inputs=[image_tensor_output, genomic_plip_model_path], outputs=pred_data_output)
 classify_button.click(run_classify_tiles, inputs=[pred_data_output, classifier_model_path], outputs=result_output)

 gr.Markdown("## Step by Step Workflow")
 with gr.Row():
 preprocess_status = gr.Checkbox(label="Preprocessed Image")
 predict_status = gr.Checkbox(label="Features with Genomic PLIP")
 classify_status = gr.Checkbox(label="Identify Risk")

 def update_status(status, result):
 return status, result is not None

 preprocess_button.click(update_status, inputs=[preprocess_status, image_tensor_output], outputs=preprocess_status)
 predict_button.click(update_status, inputs=[predict_status, pred_data_output], outputs=predict_status)
 classify_button.click(update_status, inputs=[classify_status, result_output], outputs=classify_status)

 gr.Markdown("## Example Images")
 gr.Examples(example_files, inputs=image_file)

demo.launch()