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