Update app.py

app.py

@@ -4,201 +4,95 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 import matplotlib.pyplot as plt
-from datetime import datetime
-import json
-from PIL import Image
 
 # Model setup
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')  # Use 'cuda' if GPU is available
+dtype = torch.float32
 model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
 proj_out_num = 256
 
-# Create directory for saving chat histories and temp images
-os.makedirs('chat_histories', exist_ok=True)
-os.makedirs('temp_images', exist_ok=True)
-
 # Load model and tokenizer
-print("Loading model and tokenizer...")
 model = AutoModelForCausalLM.from_pretrained(
     model_name_or_path,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
+    torch_dtype=torch.float32,
+    device_map='cpu',
     trust_remote_code=True
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
     model_name_or_path,
-    model_max_length=4096,
+    model_max_length=512,
     padding_side="right",
     use_fast=False,
     trust_remote_code=True
 )
-print("Model loaded successfully!")
 
-# Session and chat history
+# Chat history storage
 chat_history = []
-current_image_path = None
-session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
-chat_metadata = {
-    "session_id": session_id,
-    "start_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-    "image_info": None
-}
-
-def save_chat_history():
-    """Save the chat history into a JSON file."""
-    if not chat_history:
-        return
-    filename = f"chat_histories/session_{session_id}.json"
-    data = {
-        "metadata": chat_metadata,
-        "conversation": [{"user": q, "assistant": a} for q, a in chat_history]
-    }
-    with open(filename, 'w', encoding='utf-8') as f:
-        json.dump(data, f, ensure_ascii=False, indent=2)
-    return filename
+current_image = None
 
 def extract_and_display_images(image_path):
-    """Extract slices from .npy medical file and create a JPEG preview."""
-    try:
-        npy_data = np.load(image_path)
-        
-        if npy_data.ndim == 4:
-            npy_data = npy_data[0]  # Take first batch
-        
-        if npy_data.shape[0] != 32:
-            return None, "Invalid .npy shape. Expected 32 slices."
-
-        # Normalize slices
-        npy_data = (npy_data - npy_data.min()) / (npy_data.max() - npy_data.min())
-        
-        # Create visualization grid
-        fig, axes = plt.subplots(4, 8, figsize=(16, 8))
-        for i, ax in enumerate(axes.flat):
-            ax.imshow(npy_data[i], cmap='gray')
-            ax.axis('off')
-            ax.set_title(f"Slice {i+1}", fontsize=8)
-        
-        plt.tight_layout()
-        temp_png = f"temp_images/preview_{session_id}.png"
-        plt.savefig(temp_png, dpi=150, bbox_inches='tight')
-        plt.close()
-
-        # Convert PNG to JPEG if needed
-        img = Image.open(temp_png).convert("RGB")
-        temp_jpeg = f"temp_images/preview_{session_id}.jpg"
-        img.save(temp_jpeg, format="JPEG", quality=95)
-        
-        # Update metadata
-        chat_metadata["image_info"] = {
-            "filename": os.path.basename(image_path),
-            "shape": npy_data.shape,
-            "processed_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        }
-        
-        return temp_jpeg, "Image processed successfully!"
+    npy_data = np.load(image_path)
+    if npy_data.ndim == 4 and npy_data.shape[1] == 32:
+        npy_data = npy_data[0]
+    elif npy_data.ndim != 3 or npy_data.shape[0] != 32:
+        return "Invalid .npy file format. Expected shape (1, 32, 256, 256) or (32, 256, 256)."
     
-    except Exception as e:
-        return None, f"Error: {str(e)}"
-
-def process_image_question(question):
-    """Process user question about uploaded medical image."""
-    if current_image_path is None:
-        return "Please upload a medical image first."
-    try:
-        # Create fake image patch tokens
-        image_tokens = "<im_patch>" * proj_out_num
-        input_prompt = image_tokens + question
-
-        # Tokenize input
-        input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids.to(device)
+    fig, axes = plt.subplots(4, 8, figsize=(12, 6))
+    for i, ax in enumerate(axes.flat):
+        ax.imshow(npy_data[i], cmap='gray')
+        ax.axis('off')
+    
+    image_output = "extracted_images.png"
+    plt.savefig(image_output, bbox_inches='tight')
+    plt.close()
+    return image_output
 
-        # Generate answer
-        output = model.generate(
-            input_ids=input_ids,
-            max_new_tokens=256,
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.7
-        )
 
-        answer = tokenizer.decode(output[0], skip_special_tokens=True)
-        if image_tokens in answer:
-            answer = answer.split(image_tokens)[-1]
+def process_image(question):
+    global current_image
+    if current_image is None:
+        return "Please upload an image first."
+    
+    image_np = np.load(current_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)
+    
+    image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)
+    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 answer.strip()
-    except Exception as e:
-        return f"Error answering question: {str(e)}"
 
 def chat_interface(question):
-    """Handles chat conversation."""
     global chat_history
-    if not question.strip():
-        return chat_history
-    
-    response = process_image_question(question)
+    response = process_image(question)
     chat_history.append((question, response))
-    save_chat_history()
     return chat_history
 
-def upload_image(image):
-    """Handles image upload."""
-    global current_image_path
-    if image is None:
-        return "No file uploaded.", None
-
-    if not image.name.lower().endswith('.npy'):
-        return "Please upload a .npy file only.", None
-
-    current_image_path = image.name
-    extracted_image_path, status_message = extract_and_display_images(current_image_path)
-    if extracted_image_path is None:
-        return status_message, None
-    
-    return status_message, extracted_image_path
-
-def clear_conversation():
-    """Clears chat conversation."""
-    global chat_history
-    old_chat = chat_history.copy()
-    chat_history = []
-    return [], f"Conversation cleared. Saved to {save_chat_history()}."
-
-# Custom CSS
-custom_css = """
-.gradio-container {max-width: 1200px !important}
-#chat-history {height: 400px; overflow-y: auto;}
-"""
 
-# Build Gradio UI
-with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
+def upload_image(image):
+    global current_image
+    current_image = image.name
+    extracted_image_path = extract_and_display_images(current_image)
+    return "Image uploaded and processed successfully!", extracted_image_path
+
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as chat_ui:
+    gr.Markdown("ICliniq AI-Powered Medical Image Analysis Workspace")
     with gr.Row():
-        with gr.Column(scale=3):
-            gr.Markdown("# 🏥 ICliniq AI - Medical Image Analyzer")
-            gr.Markdown("""
-Upload a **.npy** medical scan file, view extracted slices, and ask clinical questions.
-""")
-            uploaded_image = gr.File(label="Upload Medical Image (.npy)", file_types=[".npy"], type="filepath")
-            upload_status = gr.Textbox(label="Upload Status", interactive=False)
-            extracted_image = gr.Image(label="Preview of Medical Image", elem_id="image-preview")
-        
+        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):
-            chat_list = gr.Chatbot(value=[], label="Conversation", elem_id="chat-history", height=500)
-            question_input = gr.Textbox(label="Ask your question", placeholder="e.g., Are there fractures visible?")
-            with gr.Row():
-                submit_button = gr.Button("Send Question", variant="primary")
-                clear_button = gr.Button("Clear Conversation", variant="secondary")
-            system_status = gr.Textbox(
-                value=f"Model loaded: {model_name_or_path}\nDevice: {device}\nSession ID: {session_id}",
-                interactive=False
-            )
-
-    uploaded_image.upload(upload_image, inputs=[uploaded_image], outputs=[upload_status, extracted_image])
-    submit_button.click(chat_interface, inputs=[question_input], outputs=[chat_list]).then(lambda: "", outputs=question_input)
-    question_input.submit(chat_interface, inputs=[question_input], outputs=[chat_list]).then(lambda: "", outputs=question_input)
-    clear_button.click(clear_conversation, inputs=[], outputs=[chat_list, system_status])
+            uploaded_image = gr.File(label="Upload .npy Image", type="filepath")
+            upload_status = gr.Textbox(label="Status", interactive=False)
+            extracted_image = gr.Image(label="Extracted Images")
+            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, extracted_image])
+    submit_button.click(chat_interface, question_input, chat_list)
+    question_input.submit(chat_interface, question_input, chat_list)
 
-# Run
-if __name__ == "__main__":
-    print("Launching Medical Image Analyzer...")
-    demo.launch(share=True)
+chat_ui.launch()