u

app.py

@@ -4,7 +4,6 @@ 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
@@ -42,30 +41,12 @@ 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], image_np
+    return generated_texts[0]
 
 # Gradio Interface
 def gradio_interface(image, question):
-    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'
+    response = process_image(image.name, question)
+    return response
 
 # Gradio App
 gr.Interface(
@@ -74,10 +55,7 @@ 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"),
-        gr.Image(label="Image Slices", type="filepath", image_mode='L')
-    ],
+    outputs=gr.Textbox(label="Model Response"),
     title="Medical Image Analysis",
-    description="Upload a .npy image and ask a question to analyze it using the model. The image slices will be displayed."
+    description="Upload a .npy image and ask a question to analyze it using the model."
 ).launch()