update

app.py

@@ -4,6 +4,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import simple_slice_viewer as ssv
 import SimpleITK as sikt
 import gradio as gr
+import matplotlib.pyplot as plt
 
 device = torch.device('cpu')  # Set to 'cuda' if using a GPU
 dtype = torch.float32  # Data type for model processing
@@ -41,12 +42,30 @@ def process_image(image_path, question):
     generation = model.generate(image_pt, input_id, max_new_tokens=256, do_sample=True, top_p=0.9, temperature=1.0)
     generated_texts = tokenizer.batch_decode(generation, skip_special_tokens=True)
 
-    return generated_texts[0]
+    return generated_texts[0], image_np
 
 # Gradio Interface
 def gradio_interface(image, question):
-    response = process_image(image.name, question)
-    return response
+    response, image_np = process_image(image.name, question)
+    
+    # Extract slices from the image
+    slices = []
+    for i in range(image_np.shape[0]):  # Assuming the image is 3D
+        slices.append(image_np[i, :, :])  # Extract each slice
+    
+    # Plot the slices and save them as images
+    fig, axes = plt.subplots(1, len(slices), figsize=(15, 5))
+    if len(slices) == 1:
+        axes = [axes]
+    
+    for ax, slice_data in zip(axes, slices):
+        ax.imshow(slice_data, cmap='gray')
+        ax.axis('off')
+    
+    plt.tight_layout()
+    plt.savefig('slices.png')  # Save the slices as a PNG image
+
+    return response, 'slices.png'
 
 # Gradio App
 gr.Interface(
@@ -55,7 +74,10 @@ gr.Interface(
         gr.File(label="Upload .npy Image", type="filepath"),  # For uploading .npy image
         gr.Textbox(label="Enter your question", placeholder="Ask something about the image..."),
     ],
-    outputs=gr.Textbox(label="Model Response"),
+    outputs=[
+        gr.Textbox(label="Model Response"),
+        gr.Image(label="Image Slices", type="filepath", image_mode='L')
+    ],
     title="Medical Image Analysis",
-    description="Upload a .npy image and ask a question to analyze it using the model."
+    description="Upload a .npy image and ask a question to analyze it using the model. The image slices will be displayed."
 ).launch()