Update app.py

app.py

@@ -6,37 +6,38 @@ 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')
-dtype = torch.float32
 model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
 proj_out_num = 256
 
-# Create directory for saving chat histories
+# 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.float32,
-    device_map=device,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",
     trust_remote_code=True
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
     model_name_or_path,
-    model_max_length=512,
+    model_max_length=4096,
     padding_side="right",
     use_fast=False,
     trust_remote_code=True
 )
 print("Model loaded successfully!")
 
-# Chat and image storage
+# Session and chat history
 chat_history = []
-current_image = None
+current_image_path = None
 session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
 chat_metadata = {
     "session_id": session_id,
@@ -45,252 +46,159 @@ chat_metadata = {
 }
 
 def save_chat_history():
-    """Save the current chat history to a JSON file"""
+    """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
 
 def extract_and_display_images(image_path):
-    """Process .npy file and create a visualization of the medical images"""
+    """Extract slices from .npy medical file and create a JPEG preview."""
     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")
-        }
+        if npy_data.ndim == 4:
+            npy_data = npy_data[0]  # Take first batch
         
-        # 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())
+        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 grid visualization
-        rows, cols = 4, 8
-        fig, axes = plt.subplots(rows, cols, figsize=(16, 8))
+        # Create visualization grid
+        fig, axes = plt.subplots(4, 8, 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.imshow(npy_data[i], cmap='gray')
             ax.axis('off')
+            ax.set_title(f"Slice {i+1}", fontsize=8)
         
         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)
+        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)
         
-        return image_output, "Image processed successfully!"
+        # 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!"
+    
     except Exception as e:
-        return None, f"Error processing image: {str(e)}"
+        return None, f"Error: {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."
-    
+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:
-        # Load the image data
-        image_np = np.load(current_image)
-        
-        # Prepare input for the model
+        # Create fake image patch tokens
         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
+        input_prompt = image_tokens + question
+
+        # Tokenize input
+        input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids.to(device)
+
+        # Generate answer
+        output = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=256,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.7
         )
-        
-        # 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
-        
+
+        answer = tokenizer.decode(output[0], skip_special_tokens=True)
+        if image_tokens in answer:
+            answer = answer.split(image_tokens)[-1]
+
+        return answer.strip()
     except Exception as e:
-        return f"Error processing your question: {str(e)}"
+        return f"Error answering question: {str(e)}"
 
 def chat_interface(question):
-    """Handle the chat interface and maintain conversation history"""
+    """Handles chat conversation."""
     global chat_history
-    
     if not question.strip():
         return chat_history
     
-    # Process the question
-    response = process_image(question)
-    
-    # Add to chat history
+    response = process_image_question(question)
     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
-    
+    """Handles image upload."""
+    global current_image_path
     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 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():
-    """Clear the current conversation history"""
+    """Clears chat conversation."""
     global chat_history
-    old_history = chat_history.copy()
+    old_chat = chat_history.copy()
     chat_history = []
-    return [], f"Conversation cleared. Previous conversation saved to {save_chat_history()}"
+    return [], f"Conversation cleared. Saved to {save_chat_history()}."
 
-# CSS for better UI
+# Custom CSS
 custom_css = """
 .gradio-container {max-width: 1200px !important}
 #chat-history {height: 400px; overflow-y: auto;}
-.image-preview {border-radius: 10px; border: 1px solid #ddd;}
 """
 
-# Gradio UI
-with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as chat_ui:
+# Build Gradio UI
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
     with gr.Row():
         with gr.Column(scale=3):
-            gr.Markdown("# ICliniq AI-Powered Medical Image Analysis")
+            gr.Markdown("# 🏥 ICliniq AI - Medical Image Analyzer")
             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="filepath"
-                    )
-                
-                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"
-            )
-            
+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=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
-                )
-                
+            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():
-                clear_button = gr.Button("Clear Conversation", variant="secondary")
                 submit_button = gr.Button("Send Question", variant="primary")
-            
-            gr.Markdown("### System Status")
+                clear_button = gr.Button("Clear Conversation", variant="secondary")
             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]
-    )
 
-# Launch the interface
+    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])
+
+# Run
 if __name__ == "__main__":
-    print("Starting ICliniq Medical Image Analysis System...")
-    chat_ui.launch(share=True)
+    print("Launching Medical Image Analyzer...")
+    demo.launch(share=True)