update

app.py

@@ -1,61 +1,70 @@
 import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import simple_slice_viewer as ssv
-import SimpleITK as sikt
 import gradio as gr
 
-device = torch.device('cpu')  # Set to 'cuda' if using a GPU
-dtype = torch.float32  # Data type for model processing
-
+# Set device & model details
+device = torch.device('cpu')
+dtype = torch.float32
 model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
-proj_out_num = 256  # Number of projection outputs required for the image
+proj_out_num = 256  # Number of projection outputs required
 
-# Load model and tokenizer
+# Load model & tokenizer
 model = AutoModelForCausalLM.from_pretrained(
-    model_name_or_path,
-    torch_dtype=torch.float32,
-    device_map='cpu',
-    trust_remote_code=True
+    model_name_or_path, torch_dtype=dtype, device_map='cpu', trust_remote_code=True
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
-    model_name_or_path,
-    model_max_length=512,
-    padding_side="right",
-    use_fast=False,
-    trust_remote_code=True
+    model_name_or_path, model_max_length=512, padding_side="right", use_fast=False, trust_remote_code=True
 )
 
-def process_image(image_path, question):
-    # Load the image
-    image_np = np.load(image_path)  # Load the .npy image
+# Image placeholder (to maintain session context)
+uploaded_image = None
+
+def process_image(question, history):
+    global uploaded_image
+    if uploaded_image is None:
+        return "⚠️ Please upload an image first!"
+
+    # Load the .npy image
+    image_np = np.load(uploaded_image)
     image_tokens = "<im_patch>" * proj_out_num
     input_txt = image_tokens + question
     input_id = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device)
 
-    # Prepare image for model
+    # Convert image to tensor
     image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)
 
-    # Generate model response
+    # Generate response
     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]
 
-# Gradio Interface
-def gradio_interface(image, question):
-    response = process_image(image.name, question)
-    return response
-
-# Gradio App
-gr.Interface(
-    fn=gradio_interface,
-    inputs=[
-        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"),
-    title="Medical Image Analysis",
-    description="Upload a .npy image and ask a question to analyze it using the model."
-).launch()
+def upload_image(image):
+    """ Stores the uploaded image path to be used in chat """
+    global uploaded_image
+    uploaded_image = image.name
+    return f"✅ Image uploaded successfully: {image.name}"
+
+# Chat Interface with File Upload
+with gr.Blocks(theme="soft") as chat_ui:
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("# 🏥 Medical Image Chatbot")
+            uploaded_file = gr.File(label="Upload .npy Image", type="filepath")
+            upload_button = gr.Button("Upload")
+            status = gr.Markdown("")
+            chat = gr.Chatbot(height=400)
+
+        with gr.Column(scale=3):
+            input_box = gr.Textbox(placeholder="Ask something about the image...")
+            send_button = gr.Button("Send ✉️")
+
+    # Handle image upload
+    upload_button.click(upload_image, inputs=[uploaded_file], outputs=[status])
+
+    # Handle chat interaction
+    send_button.click(process_image, inputs=[input_box, chat], outputs=[chat])
+
+chat_ui.launch()