Update app.py

app.py

@@ -8,18 +8,25 @@ from main import load_and_preprocess_image, genomic_plip_predictions, classify_t
 @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:
@@ -35,8 +42,8 @@ with gr.Blocks() as demo:
  result_output = gr.Textbox(label="Result")
 
  preprocess_button = gr.Button("Preprocess Image")
- predict_button = gr.Button("Predict with Genomic PLIP")
- classify_button = gr.Button("Classify Tiles")
+ 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)
@@ -45,8 +52,8 @@ with gr.Blocks() as demo:
  gr.Markdown("## Step by Step Workflow")
  with gr.Row():
  preprocess_status = gr.Checkbox(label="Preprocessed Image")
- predict_status = gr.Checkbox(label="Predicted with Genomic PLIP")
- classify_status = gr.Checkbox(label="Classified Tiles")
+ 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