Update app.py

app.py

@@ -1,99 +1,296 @@
-import os  
-import numpy as np  
-import torch  
-from transformers import AutoTokenizer, AutoModelForCausalLM  
-import gradio as gr  
-import matplotlib.pyplot as plt  
+import os
+import numpy as np
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+import matplotlib.pyplot as plt
+from datetime import datetime
+import json
 
-# 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'  
-proj_out_num = 256  
+# Model setup
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+dtype = torch.float32
+model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
+proj_out_num = 256
 
-# Load model and tokenizer  
-model = AutoModelForCausalLM.from_pretrained(  
-    model_name_or_path,  
-    torch_dtype=torch.float32,  
-    device_map='cpu',  
-    trust_remote_code=True  
-)  
+# Create directory for saving chat histories
+os.makedirs('chat_histories', exist_ok=True)
 
-tokenizer = AutoTokenizer.from_pretrained(  
-    model_name_or_path,  
-    model_max_length=512,  
-    padding_side="right",  
-    use_fast=False,  
-    trust_remote_code=True  
-)  
+# Load model and tokenizer
+print("Loading model and tokenizer...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name_or_path,
+    torch_dtype=torch.float32,
+    device_map=device,
+    trust_remote_code=True
+)
 
-# Chat history storage  
-chat_history = []  
-current_image = None  
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name_or_path,
+    model_max_length=512,
+    padding_side="right",
+    use_fast=False,
+    trust_remote_code=True
+)
+print("Model loaded successfully!")
 
-def extract_and_display_images(image_path):  
-    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)."  
-    
-    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  
+# Chat and image storage
+chat_history = []
+current_image = 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 current chat history to 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
 
-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]  
+def extract_and_display_images(image_path):
+    """Process .npy file and create a visualization of the medical images"""
+    try:
+        npy_data = np.load(image_path)
+        
+        # Handle different possible shapes of the .npy file
+        if npy_data.ndim == 4 and npy_data.shape[1] == 32:
+            npy_data = npy_data[0]  # Extract first batch if batched
+        elif npy_data.ndim != 3 or npy_data.shape[0] != 32:
+            return None, "Invalid .npy file format. Expected shape (1, 32, 256, 256) or (32, 256, 256)."
+        
+        # Update metadata with image information
+        global chat_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")
+        }
+        
+        # Normalize for better visualization if needed
+        for i in range(npy_data.shape[0]):
+            slice_data = npy_data[i]
+            if slice_data.max() > 0:  # Avoid division by zero
+                npy_data[i] = (slice_data - slice_data.min()) / (slice_data.max() - slice_data.min())
+        
+        # Create grid visualization
+        rows, cols = 4, 8
+        fig, axes = plt.subplots(rows, cols, figsize=(16, 8))
+        for i, ax in enumerate(axes.flat):
+            if i < npy_data.shape[0]:
+                ax.imshow(npy_data[i], cmap='gray')
+                ax.set_title(f"Slice {i+1}", fontsize=8)
+            ax.axis('off')
+        
+        plt.tight_layout()
+        image_output = f"temp_images/extracted_{session_id}.png"
+        os.makedirs("temp_images", exist_ok=True)
+        plt.savefig(image_output, bbox_inches='tight', dpi=150)
+        plt.close()
+        
+        return image_output, "Image processed successfully!"
+    except Exception as e:
+        return None, f"Error processing image: {str(e)}"
 
+def process_image(question):
+    """Process a question about the current medical image using the AI model"""
+    global current_image
+    
+    if current_image is None:
+        return "Please upload a medical image (.npy file) first."
+    
+    try:
+        # Load the image data
+        image_np = np.load(current_image)
+        
+        # Prepare input for the model
+        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)
+        
+        # Convert image to tensor
+        image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)
+        
+        # Generate response from model
+        generation = model.generate(
+            image_pt, 
+            input_id, 
+            max_new_tokens=256, 
+            do_sample=True, 
+            top_p=0.9, 
+            temperature=0.8  # Slightly reduced for more consistent responses
+        )
+        
+        # Decode the generated text
+        generated_text = tokenizer.batch_decode(generation, skip_special_tokens=True)[0]
+        
+        # Remove the input prompt from the response if needed
+        if image_tokens in generated_text:
+            generated_text = generated_text.split(image_tokens)[-1]
+            
+        return generated_text
+        
+    except Exception as e:
+        return f"Error processing your question: {str(e)}"
 
-def chat_interface(question):  
-    global chat_history  
-    response = process_image(question)  
-    chat_history.append((question, response))  # Save dynamic chat history  
-    return chat_history  # Return updated chat history  
+def chat_interface(question):
+    """Handle the chat interface and maintain conversation history"""
+    global chat_history
+    
+    if not question.strip():
+        return chat_history
+    
+    # Process the question
+    response = process_image(question)
+    
+    # Add to chat history
+    chat_history.append((question, response))
+    
+    # Save chat history periodically
+    save_chat_history()
+    
+    # Return the updated chat history for display
+    return chat_history
 
+def upload_image(image):
+    """Handle image upload and processing"""
+    global current_image
+    
+    if image is None:
+        return "No file uploaded.", None
+    
+    # Check if file exists and is .npy
+    if not os.path.exists(image.name) or not image.name.lower().endswith('.npy'):
+        return "Please upload a valid .npy file.", None
+    
+    # Set as current image
+    current_image = image.name
+    
+    # Process and extract images
+    extracted_image_path, status_message = extract_and_display_images(current_image)
+    
+    if extracted_image_path is None:
+        return status_message, None
+    
+    return status_message, extracted_image_path
 
-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  
+def clear_conversation():
+    """Clear the current conversation history"""
+    global chat_history
+    old_history = chat_history.copy()
+    chat_history = []
+    return [], f"Conversation cleared. Previous conversation saved to {save_chat_history()}"
 
-# Gradio UI  
-with gr.Blocks(theme="soft") as chat_ui:  
-    with gr.Row():  
-        # Dynamic Chat Area  
-        with gr.Column(scale=4):  
-            gr.Markdown("### ICliniq AI-Powered Medical Image Analysis")  
-            chat_list = gr.Chatbot(value=[], label="Chat History", type='messages', elem_id="chat-history")  # Dynamic chat with messages type  
-            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")  
+# CSS for better UI
+custom_css = """
+.gradio-container {max-width: 1200px !important}
+#chat-history {height: 400px; overflow-y: auto;}
+.image-preview {border-radius: 10px; border: 1px solid #ddd;}
+"""
 
-            # Upload and Processing Interactions  
-            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)  
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as chat_ui:
+    with gr.Row():
+        with gr.Column(scale=3):
+            gr.Markdown("# ICliniq AI-Powered Medical Image Analysis")
+            gr.Markdown("""
+            This system analyzes medical images in .npy format and answers your questions.
+            
+            ## How to use:
+            1. Upload your medical image (.npy format)
+            2. Wait for the image to be processed
+            3. Ask questions about the image
+            """)
+            
+            with gr.Row():
+                with gr.Column(scale=1):
+                    uploaded_image = gr.File(
+                        label="Upload Medical Image (.npy format)",
+                        file_types=[".npy"],
+                        type="file"
+                    )
+                
+                with gr.Column(scale=1):
+                    upload_status = gr.Textbox(
+                        label="Upload Status",
+                        interactive=False
+                    )
+            
+            extracted_image = gr.Image(
+                label="Processed Image Preview",
+                elem_id="image-preview"
+            )
+            
+        with gr.Column(scale=4):
+            chat_list = gr.Chatbot(
+                value=[],
+                label="Conversation",
+                elem_id="chat-history",
+                height=500
+            )
+            
+            with gr.Row():
+                question_input = gr.Textbox(
+                    label="Ask about the medical image",
+                    placeholder="What abnormalities do you see in this scan?",
+                    lines=2
+                )
+                
+            with gr.Row():
+                clear_button = gr.Button("Clear Conversation", variant="secondary")
+                submit_button = gr.Button("Send Question", variant="primary")
+            
+            gr.Markdown("### System Status")
+            system_status = gr.Textbox(
+                label="",
+                value=f"Model loaded: {model_name_or_path}\nDevice: {device}\nSession ID: {session_id}",
+                interactive=False
+            )
+    
+    # Set up event handlers
+    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: "",  # Clear input after sending
+        outputs=question_input
+    )
+    
+    question_input.submit(
+        chat_interface,
+        inputs=[question_input],
+        outputs=[chat_list]
+    ).then(
+        lambda: "",  # Clear input after sending
+        outputs=question_input
+    )
+    
+    clear_button.click(
+        clear_conversation,
+        inputs=[],
+        outputs=[chat_list, system_status]
+    )
 
-chat_ui.launch()  
+# Launch the interface
+if __name__ == "__main__":
+    print("Starting ICliniq Medical Image Analysis System...")
+    chat_ui.launch(share=True)