commit

app.py

@@ -2,14 +2,11 @@ import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-import matplotlib.pyplot as plt
-from PIL import Image
-import io
 
 # Model setup
-device = torch.device("cpu")  # Use 'cuda' if GPU is available
-dtype = torch.float32
-model_name_or_path = "GoodBaiBai88/M3D-LaMed-Phi-3-4B"
+device = torch.device('cpu')  # Use 'cuda' if GPU is available
+dtype = torch.float32  # Data type for model processing
+model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
 proj_out_num = 256
 
 # Load model and tokenizer
@@ -30,39 +27,31 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 # Chat history storage
 chat_history = []
-current_image = None
-
-# Convert .npy to JPEG
-def npy_to_jpeg(npy_file):
-    image_np = np.load(npy_file)
-    image = Image.fromarray((image_np * 255).astype(np.uint8))  # Normalize and convert to uint8
-    img_bytes = io.BytesIO()
-    image.save(img_bytes, format='JPEG')
-    return img_bytes.getvalue()
+current_image = None  # To store the uploaded image
 
 def process_image(question):
     global current_image
     if current_image is None:
-        return "Please upload an image first.", None
+        return "Please upload an image first."
     
-    image_bytes = npy_to_jpeg(current_image)  # Convert image to JPEG
+    image_np = np.load(current_image)  # Load the stored .npy 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)
+    input_id = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device)
+    
+    # Prepare image for model
+    image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)
     
     # Generate response
-    generation = model.generate(input_id, max_new_tokens=256, do_sample=True, top_p=0.9, temperature=1.0)
+    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_bytes
+    return generated_texts[0]
 
 # Function to update chat
 def chat_interface(question):
     global chat_history
-    response, image_bytes = process_image(question)
-    if image_bytes:
-        chat_history.append((question, response, image_bytes))
-    else:
-        chat_history.append((question, response, None))
+    response = process_image(question)
+    chat_history.append((question, response))
     return chat_history
 
 # Function to handle image upload
@@ -74,14 +63,17 @@ def upload_image(image):
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as chat_ui:
     gr.Markdown("# 🏥 Medical Image Analysis Chatbot")
-    chat_list = gr.Chatbot(label="Chat History", elem_id="chat-history")
-    uploaded_image = gr.File(label="Upload .npy Image", type="filepath")
-    upload_status = gr.Textbox(label="Status", interactive=False)
-    question_input = gr.Textbox(label="Ask a question", placeholder="Ask something about the image...")
-    submit_button = gr.Button("Send")
+    with gr.Row():
+        with gr.Column(scale=1, min_width=200):
+            chat_list = gr.Chatbot(value=[], label="Chat History", elem_id="chat-history")
+        with gr.Column(scale=4):
+            uploaded_image = gr.File(label="Upload .npy Image", type="filepath")
+            upload_status = gr.Textbox(label="Status", interactive=False)
+            question_input = gr.Textbox(label="Ask a question", placeholder="Ask something about the image...")
+            submit_button = gr.Button("Send")
     
     uploaded_image.upload(upload_image, uploaded_image, upload_status)
     submit_button.click(chat_interface, question_input, chat_list)
     question_input.submit(chat_interface, question_input, chat_list)
 
-chat_ui.launch()
+chat_ui.launch()