Update app.py

app.py

@@ -6,10 +6,8 @@ import gradio as gr
 import matplotlib.pyplot as plt
 
 # Model setup
-device = torch.device('cpu')  # Use 'cuda' if GPU is available
-dtype = torch.float32
+device = torch.device('cpu')  # Use 'cuda' if available
 model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
-proj_out_num = 256
 
 # Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
@@ -18,7 +16,6 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map='cpu',
     trust_remote_code=True
 )
-
 tokenizer = AutoTokenizer.from_pretrained(
     model_name_or_path,
     model_max_length=512,
@@ -27,7 +24,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     trust_remote_code=True
 )
 
-# Chat history storage
+# Storage
 chat_history = []
 current_image = None
 
@@ -36,86 +33,115 @@ def extract_and_display_images(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)."
+        return "Invalid .npy format. Expected (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')
+    output_path = "converted_image_preview.png"
+    plt.savefig(output_path, bbox_inches='tight')
     plt.close()
-    return image_output
+    return output_path
 
-def process_image(question):
+def process_question(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
+    image_tokens = "<im_patch>" * 256
     input_txt = image_tokens + question
-    input_id = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device)
+    input_ids = 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)
+    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_interface(question):
-    global chat_history
-    response = process_image(question)
-    chat_history.append((question, response))
-    return chat_history
-
 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
+    preview_path = extract_and_display_images(current_image)
+    return "Image uploaded successfully!", preview_path
 
-# Gradio UI with Animation and Styling
+def chat_with_model(user_message):
+    global chat_history
+    response = process_question(user_message)
+    chat_history.append((user_message, response))
+    return chat_history
+
+# UI Design
 with gr.Blocks(css="""
-    .gr-chatbot-container {
-        transition: all 0.3s ease-in-out;
-        opacity: 0;
-    }
-    .gr-chatbot-container.show {
-        opacity: 1;
-    }
-    .gr-chatbot-message {
-        transition: all 0.3s ease-in-out;
-        transform: translateX(-10px);
-    }
-    .gr-chatbot-message.show {
-        transform: translateX(0);
-    }
-    .gr-image-container img {
-        transition: transform 0.5s ease;
-    }
-    .gr-image-container img:hover {
-        transform: scale(1.1);
-    }
-    .gr-button:hover {
-        background-color: #007bff;
-        transform: scale(1.1);
-        transition: all 0.3s ease-in-out;
-    }
-""") as chat_ui:
-    gr.Markdown("ICliniq AI-Powered Medical Image Analysis Workspace")
+body {
+    background: linear-gradient(135deg, #00b4db, #0083b0);
+    font-family: 'Poppins', sans-serif;
+    color: white;
+}
+
+.gr-box {
+    border-radius: 16px;
+    background: rgba(255,255,255,0.1);
+    padding: 20px;
+    backdrop-filter: blur(10px);
+    box-shadow: 0 8px 32px 0 rgba( 31, 38, 135, 0.37 );
+}
+
+h1 {
+    text-align: center;
+    font-size: 2.5em;
+    margin-bottom: 20px;
+    color: #ffffff;
+}
+
+.gr-chatbot-container {
+    overflow-y: auto;
+    max-height: 500px;
+}
+
+.gr-chatbot-message {
+    margin-bottom: 15px;
+    padding: 10px;
+    border-radius: 10px;
+    background: rgba(0,0,0,0.3);
+    transition: 0.3s;
+}
+
+.gr-chatbot-message:hover {
+    transform: scale(1.02);
+    background: rgba(255,255,255,0.1);
+}
+
+.gr-button {
+    background-color: #ff7e5f;
+    border: none;
+    padding: 10px 20px;
+    border-radius: 20px;
+    color: white;
+    font-weight: bold;
+    transition: 0.3s;
+}
+
+.gr-button:hover {
+    background-color: #feb47b;
+    transform: scale(1.05);
+}
+""") as app:
+    gr.Markdown("# 🚀 AI Powered Medical Image Analysis System 🚀")
+
     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")
+        with gr.Column(scale=1, min_width=250):
+            chat_history_box = gr.Chatbot(value=[], label="🗂 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")
-            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)
+            preview_image = gr.Image(label="🖼 Image Preview")
+            user_input = gr.Textbox(label="💬 Ask a question about the image...")
+            send_button = gr.Button("📨 Send")
+
+    uploaded_image.upload(upload_image, uploaded_image, [upload_status, preview_image])
+    send_button.click(chat_with_model, user_input, chat_history_box)
+    user_input.submit(chat_with_model, user_input, chat_history_box)
 
-chat_ui.launch()
+app.launch()