Update app.py

app.py

@@ -1,222 +1,101 @@
-import os
-import numpy as np
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import gradio as gr
-import matplotlib.pyplot as plt
-
-device = torch.device('cpu')
-model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name_or_path,
-    torch_dtype=torch.float32,
-    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
-)
-
-chat_history = []
-current_image = None
-
-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 format. Expected (1, 32, 256, 256) or (32, 256, 256)."
+import os  
+import numpy as np  
+import torch  
+from transformers import AutoTokenizer, AutoModelForCausalLM  
+import gradio as gr  
+import matplotlib.pyplot as plt  
+
+# 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  
+
+# Load model and tokenizer  
+model = AutoModelForCausalLM.from_pretrained(  
+    model_name_or_path,  
+    torch_dtype=torch.float32,  
+    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  
+)  
+
+# Chat history storage  
+chat_history = []  
+current_image = None  
+
+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')
+    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')  
     
-    output_path = "converted_image_preview.png"
-    plt.savefig(output_path, bbox_inches='tight')
-    plt.close()
-    return output_path
+    image_output = "extracted_images.png"  
+    plt.savefig(image_output, bbox_inches='tight')  
+    plt.close()  
+    return image_output  
 
-def upload_image(image):
-    global current_image
-    if image is None:
-        return "", None
-    current_image = image.name
-    preview_path = extract_and_display_images(current_image)
-    return "Image uploaded successfully!", preview_path
 
-def process_question(question):
-    global current_image
-    if current_image is None:
-        return "Please upload an image first."
+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>" * 256
-    input_txt = image_tokens + question
-    input_ids = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device)
+    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=torch.float32, device=device)
-    generation = model.generate(image_pt, input_ids, 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 chat_with_model(user_message):
-    global chat_history
-    if not user_message.strip():
-        return chat_history
-    response = process_question(user_message)
-    chat_history.append((user_message, response))
-    return chat_history
-
-# Function to export chat history to a text file
-def export_chat_history():
-    history_text = ""
-    for user_msg, model_reply in chat_history:
-        history_text += f"User: {user_msg}\nAI: {model_reply}\n\n"
-    with open("chat_history.txt", "w") as f:
-        f.write(history_text)
-    return "Chat history exported as chat_history.txt"
-
-# UI
-with gr.Blocks(css="""
-body {
-    background: #f5f5f5;
-    font-family: 'Inter', sans-serif;
-    color: #333333;
-}
-
-h1 {
-    text-align: center;
-    font-size: 2em;
-    margin-bottom: 20px;
-    color: #222;
-}
-
-.gr-box {
-    background: #ffffff;
-    padding: 20px;
-    border-radius: 10px;
-    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.1);
-}
-
-.gr-chatbot-container {
-    overflow-y: auto;
-    max-height: 500px;
-    scroll-behavior: smooth;
-}
-
-.gr-chatbot-message {
-    margin-bottom: 10px;
-    padding: 8px;
-    border-radius: 8px;
-    background: #f5f5f5;
-    animation: fadeIn 0.5s ease-out;
-}
-
-.gr-button {
-    background-color: #4CAF50;
-    color: white;
-    border: none;
-    padding: 8px 16px;
-    border-radius: 6px;
-    cursor: pointer;
-}
-
-.gr-button:hover {
-    background-color: #45a049;
-}
+    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]  
 
-#message-box {
-    display: flex;
-    align-items: center;
-    position: relative;
-    transition: all 0.3s ease;
-}
 
-#message-input {
-    width: 100%;
-    padding: 10px;
-    border-radius: 6px;
-    border: 1px solid #ddd;
-    font-size: 14px;
-    margin-right: 40px; /* To give space for the icon */
-}
+def chat_interface(question):  
+    global chat_history  
+    response = process_image(question)  
+    chat_history.append((question, response))  
+    return chat_history  
 
-#upload-icon {
-    position: absolute;
-    right: 10px;
-    cursor: pointer;
-    font-size: 24px;
-    color: #4CAF50;
-    animation: bounce 0.6s infinite alternate;
-}
 
-#upload-icon:hover {
-    color: #45a049;
-    transform: scale(1.1);
-}
+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  
 
-#loading-spinner {
-    display: none;
-    text-align: center;
-}
-
-#loading-spinner img {
-    width: 50px;
-    height: 50px;
-}
-
-@keyframes bounce {
-    0% { transform: translateY(0); }
-    100% { transform: translateY(-10px); }
-}
-
-@keyframes fadeIn {
-    0% { opacity: 0; }
-    100% { opacity: 1; }
-}
-""") as app:
-    gr.Markdown("# AI Powered Medical Image Analysis System")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            chatbot_ui = gr.Chatbot(value=[], label="Chat History")
-        with gr.Column(scale=2):
-            # Message input area with the '+' icon inside it
-            with gr.Box(elem_id="message-box"):
-                message_input = gr.Textbox(placeholder="Type your question here...", label="Your Message", elem_id="message-input")
-                upload_button = gr.HTML('<span id="upload-icon">+</span>')  # The + icon inside the message box
-
-            upload_section = gr.File(label="Upload NPY Image", type="filepath", visible=False)
-            upload_status = gr.Textbox(label="Status", interactive=False)
-            preview_img = gr.Image(label="Image Preview", interactive=False)
-            send_button = gr.Button("Send")
-            export_button = gr.Button("Export Chat History")
-            loading_spinner = gr.HTML('<div id="loading-spinner"><img src="https://i.imgur.com/llf5Jjs.gif" alt="Loading..."></div>')
-
-    # Handle click event for the '+' icon
-    upload_button.click(lambda: upload_section.update(visible=True), None, upload_section)
-
-    # Handle file upload
-    upload_section.upload(lambda *args: loading_spinner.update("<div id='loading-spinner'><img src='https://i.imgur.com/llf5Jjs.gif' alt='Loading...'></div>"), upload_section, None)
-    upload_section.upload(upload_image, upload_section, [upload_status, preview_img])
-
-    # Display loading spinner while processing question
-    send_button.click(lambda *args: loading_spinner.update("<div id='loading-spinner'><img src='https://i.imgur.com/llf5Jjs.gif' alt='Loading...'></div>"), None, None)
-    send_button.click(chat_with_model, message_input, chatbot_ui)
-    send_button.click(lambda *args: loading_spinner.update(''), None, None)
-    message_input.submit(chat_with_model, message_input, chatbot_ui)
-
-    # Export chat history functionality
-    export_button.click(export_chat_history)
-
-    # Auto-focus typing box and scroll to bottom after message sent
-    message_input.submit(lambda: gr.update(focus=True), None, message_input)
-    send_button.click(lambda: gr.update(focus=True), None, message_input)
+# 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=1):  
+            chat_list = gr.Chatbot(value=[], label="Chat History", elem_id="chat-history")  
+        
+        with gr.Column(scale=2):  
+            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", type="filepath")  
+            question_input = gr.Textbox(label="Ask a question", placeholder="Ask something about the image...", lines=2)  
+            submit_button = gr.Button("Send")  
+    
+    # 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)  
 
-app.launch()
+chat_ui.launch()