Update app.py

app.py

@@ -1,28 +1,55 @@
 import gradio as gr
 import torch
+import numpy as np
 from transformers import AutoConfig, AutoModelForCausalLM
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 from janus.utils.io import load_pil_images
 from PIL import Image
-import numpy as np
-import os
-import time
 import spaces
+from torchvision import transforms
 
-# Load medical imaging-optimized model and processor
+# Medical Imaging Configuration
+MEDICAL_CONFIG = {
+    "modality": "CT",  # Default imaging modality
+    "anatomical_region": "Chest",
+    "clinical_task": "analysis",
+    "report_style": "structured"
+}
+
+# Load base model
 model_path = "deepseek-ai/Janus-Pro-1B"
 config = AutoConfig.from_pretrained(model_path)
 language_config = config.language_config
 language_config._attn_implementation = 'eager'
 
-# Initialize model with medical imaging parameters
+# Initialize model with medical adaptations
 vl_gpt = AutoModelForCausalLM.from_pretrained(
     model_path,
     language_config=language_config,
     trust_remote_code=True,
-    medical_head=True  # Assuming custom medical imaging head
+    hidden_dropout_prob=0.1,
+    attention_probs_dropout_prob=0.1,
+    output_attentions=True
 ).to(torch.bfloat16 if torch.cuda.is_available() else torch.float16)
 
+# Add medical projection layer
+class MedicalProjectionWrapper(torch.nn.Module):
+    def __init__(self, base_model):
+        super().__init__()
+        self.base_model = base_model
+        self.medical_proj = torch.nn.Linear(
+            base_model.config.hidden_size,
+            base_model.config.hidden_size * 2
+        )
+        self.activation = torch.nn.GELU()
+
+    def forward(self, *args, **kwargs):
+        outputs = self.base_model(*args, **kwargs)
+        medical_rep = self.activation(self.medical_proj(outputs.last_hidden_state))
+        return outputs.__class__(last_hidden_state=medical_rep)
+
+vl_gpt.language_model = MedicalProjectionWrapper(vl_gpt.language_model)
+
 if torch.cuda.is_available():
     vl_gpt = vl_gpt.cuda()
 
@@ -30,27 +57,41 @@ vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
 tokenizer = vl_chat_processor.tokenizer
 cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+# Medical image preprocessing
+def preprocess_medical_image(image):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    medical_transforms = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.Grayscale(num_output_channels=3),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
+    return medical_transforms(image).unsqueeze(0).to(cuda_device)
+
 @torch.inference_mode()
 @spaces.GPU(duration=120)
-def medical_image_analysis(medical_image, clinical_question, seed, top_p, temperature):
-    """Analyze medical images (CT, MRI, X-ray, histopathology) with clinical context."""
+def medical_image_analysis(image, clinical_query, seed=42):
     torch.cuda.empty_cache()
     torch.manual_seed(seed)
     
-    # Medical-specific conversation template
+    # Preprocess with medical transformations
+    medical_image = preprocess_medical_image(image)
+    
     conversation = [{
         "role": "<|Radiologist|>",
-        "content": f"<medical_image>\nClinical Context: {clinical_question}",
+        "content": f"<medical_image>\nClinical Context: {clinical_query}",
         "images": [medical_image],
     }, {"role": "<|AI_Assistant|>", "content": ""}]
     
-    processed_image = [Image.fromarray(medical_image)]
     inputs = vl_chat_processor(
-        conversations=conversation, 
-        images=processed_image,
+        conversations=conversation,
+        images=[Image.fromarray(image)],
         force_batchify=True
-    ).to(cuda_device, dtype=torch.bfloat16)
-    
+    ).to(cuda_device)
+
     inputs_embeds = vl_gpt.prepare_inputs_embeds(**inputs)
     
     # Medical-optimized generation parameters
@@ -58,115 +99,161 @@ def medical_image_analysis(medical_image, clinical_question, seed, top_p, temper
         inputs_embeds=inputs_embeds,
         attention_mask=inputs.attention_mask,
         max_new_tokens=512,
-        temperature=0.2,  # Lower for clinical precision
+        temperature=0.2,
         top_p=0.9,
-        repetition_penalty=1.2,  # Reduce hallucination
-        medical_mode=True
+        num_beams=5,
+        repetition_penalty=1.5,
+        early_stopping=True
     )
     
-    findings = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-    return f"Clinical Findings:\n{findings}"
+    report = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
+    return format_medical_report(report)
+
+def format_medical_report(raw_text):
+    sections = {
+        "Findings": "",
+        "Impression": "",
+        "Recommendations": ""
+    }
+    
+    current_section = None
+    for line in raw_text.split('\n'):
+        if "FINDINGS:" in line:
+            current_section = "Findings"
+        elif "IMPRESSION:" in line:
+            current_section = "Impression"
+        elif "RECOMMENDATIONS:" in line:
+            current_section = "Recommendations"
+        elif current_section:
+            sections[current_section] += line.strip() + '\n'
+    
+    return f"""**Clinical Report**
+    
+**Findings:**
+{sections['Findings'] or 'No significant findings'}
+
+**Impression:**
+{sections['Impression'] or 'No conclusive diagnosis'}
+
+**Recommendations:**
+{sections['Recommendations'] or 'Follow-up as clinically indicated'}"""
 
+# Medical image generation components
 @torch.inference_mode()
 @spaces.GPU(duration=120)
-def generate_medical_image(prompt, seed=None, guidance=5, t2i_temperature=0.5):
-    """Generate synthetic medical images for educational/research purposes."""
+def generate_medical_image(prompt, seed=12345, guidance=7, temperature=0.6):
     torch.cuda.empty_cache()
     if seed is not None:
         torch.manual_seed(seed)
-    
-    # Medical image generation parameters
-    medical_config = {
-        'width': 512,
-        'height': 512,
-        'parallel_size': 3,
-        'modality': 'mri',  # Can specify CT, X-ray, etc.
-        'anatomy': 'brain'  # Target anatomy
-    }
+
+    medical_prompt = f"{prompt} [Modality: {MEDICAL_CONFIG['modality']}, Anatomy: {MEDICAL_CONFIG['anatomical_region']}]"
     
     messages = [{
         'role': '<|Clinician|>', 
-        'content': f"{prompt} [Modality: {medical_config['modality']}, Anatomy: {medical_config['anatomy']}]"
+        'content': medical_prompt
     }]
     
-    text = vl_chat_processor.apply_medical_template(
+    text = vl_chat_processor.apply_chat_template(
         messages,
-        system_prompt='Generate education-quality medical imaging data'
+        system_prompt='Generate educational medical imaging data'
     )
     
     input_ids = torch.LongTensor(tokenizer.encode(text)).to(cuda_device)
-    generated_tokens, patches = vl_gpt.generate_medical_image(
+    
+    # Medical image generation parameters
+    generated_tokens, patches = vl_gpt.generate(
         input_ids,
-        **medical_config,
+        width=512,
+        height=512,
         cfg_weight=guidance,
-        temperature=t2i_temperature
+        temperature=temperature,
+        parallel_size=3,
+        image_token_num_per_image=576,
+        patch_size=16
     )
     
-    # Post-processing for medical imaging standards
-    synthetic_images = postprocess_medical_images(patches, **medical_config)
+    synthetic_images = postprocess_medical_images(patches)
     return [Image.fromarray(img).resize((512, 512)) for img in synthetic_images]
 
-# Medical-optimized Gradio interface
-with gr.Blocks(title="Medical Imaging AI Suite") as demo:
-    gr.Markdown("""## Medical Image Analysis Suite v2.1
-    *For research use only - not for clinical diagnosis*""")
+def postprocess_medical_images(patches):
+    patches = patches.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
+    patches = np.clip((patches + 1) / 2 * 255, 0, 255).astype(np.uint8)
+    return [patches[i] for i in range(patches.shape[0])]
+
+# Medical-optimized interface
+with gr.Blocks(title="Medical Imaging AI", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""## Medical Imaging Analysis Suite v3.2
+    *Research use only - Not for clinical decision-making*""")
     
     with gr.Tab("Clinical Image Analysis"):
+        gr.Markdown("### Upload medical scan and clinical context")
         with gr.Row():
-            medical_image_input = gr.Image(label="Upload Medical Scan")
-            clinical_question = gr.Textbox(label="Clinical Query", 
-                placeholder="E.g.: 'Assess tumor progression in this MRI series'")
-            
-        with gr.Accordion("Advanced Parameters", open=False):
-            und_seed = gr.Number(42, label="Reproducibility Seed")
-            analysis_top_p = gr.Slider(0.8, 1.0, 0.95, label="Diagnostic Certainty")
-            analysis_temp = gr.Slider(0.1, 0.5, 0.2, label="Analysis Precision")
-            
-        analysis_btn = gr.Button("Analyze Scan", variant="primary")
-        clinical_report = gr.Textbox(label="AI Analysis Report", interactive=False)
-        
-        gr.Examples(
-            examples=[
-                ["Identify pulmonary nodules in this CT scan", "ct_chest.png"],
-                ["Assess MRI for multiple sclerosis lesions", "brain_mri.jpg"],
-                ["Histopathology analysis: tumor grading", "biopsy_slide.png"]
-            ],
-            inputs=[clinical_question, medical_image_input]
-        )
+            with gr.Column(scale=1):
+                med_image = gr.Image(label="Medical Imaging Study", type="numpy")
+                med_upload_btns = gr.Row([
+                    gr.Button("CT Scan"),
+                    gr.Button("MRI"),
+                    gr.Button("X-ray")
+                ])
+                
+            with gr.Column(scale=2):
+                clinical_input = gr.Textbox(label="Clinical Context", lines=3,
+                                          placeholder="Patient history and clinical question...")
+                analysis_btn = gr.Button("Analyze Study", variant="primary")
+                report_output = gr.Markdown(label="AI Analysis Report")
 
-    with gr.Tab("Medical Imaging Synthesis"):
-        gr.Markdown("**Educational Image Generation**")
-        synth_prompt = gr.Textbox(label="Synthesis Prompt",
-            placeholder="E.g.: 'Synthetic brain MRI showing glioblastoma multiforme'")
-        
+        gr.Examples([
+            ["Evaluate lung nodules in this CT scan", "ct_chest.png"],
+            ["Assess brain MRI for metastatic lesions", "brain_mri.jpg"],
+            ["Analyze bone structure in this wrist X-ray", "wrist_xray.png"]
+        ], [clinical_input, med_image])
+
+    with gr.Tab("Educational Image Synthesis"):
+        gr.Markdown("### Generate synthetic medical images for training")
         with gr.Row():
-            synth_guidance = gr.Slider(3, 7, 5, label="Anatomical Accuracy")
-            synth_temp = gr.Slider(0.3, 1.0, 0.6, label="Synthesis Variability")
-            
-        synth_btn = gr.Button("Generate Educational Images", variant="secondary")
-        synthetic_gallery = gr.Gallery(label="Synthetic Medical Images", 
-                                     columns=3, object_fit="contain")
-        
-        gr.Examples(
-            examples=[
-                "High-resolution CT of healthy lung parenchyma",
-                "T2-weighted MRI of lumbar spine with herniated disc",
-                "Histopathology slide of benign breast tissue"
-            ],
-            inputs=synth_prompt
-        )
+            with gr.Column():
+                synth_prompt = gr.Textbox(label="Synthesis Prompt", lines=2,
+                                        placeholder="Describe the desired medical image...")
+                gr.Markdown("**Modality Options**")
+                modality_btns = gr.Row([
+                    gr.Button("CT"),
+                    gr.Button("MRI"),
+                    gr.Button("X-ray")
+                ])
+                
+            with gr.Column():
+                synth_params = gr.Accordion("Advanced Parameters", open=False)
+                with synth_params:
+                    gr.Row([
+                        gr.Slider(3, 7, 5, label="Anatomical Accuracy"),
+                        gr.Slider(0.3, 1.0, 0.6, label="Synthesis Variability")
+                    ])
+                generate_btn = gr.Button("Generate Educational Images", variant="secondary")
+                
+        synth_gallery = gr.Gallery(label="Synthetic Images", columns=3, height=400)
 
-    # Connect functionality
+    # Event handlers
     analysis_btn.click(
         medical_image_analysis,
-        inputs=[medical_image_input, clinical_question, und_seed, analysis_top_p, analysis_temp],
-        outputs=clinical_report
+        [med_image, clinical_input],
+        report_output
     )
     
-    synth_btn.click(
+    generate_btn.click(
         generate_medical_image,
-        inputs=[synth_prompt, und_seed, synth_guidance, synth_temp],
-        outputs=synthetic_gallery
+        [synth_prompt, synth_params],
+        synth_gallery
     )
+    
+    for btn in [*med_upload_btns.children, *modality_btns.children]:
+        btn.click(
+            lambda m: MEDICAL_CONFIG.update(modality=m),
+            [btn],
+            None
+        ).then(
+            lambda: gr.Info(f"Modality set to {MEDICAL_CONFIG['modality']}"),
+            None,
+            None
+        )
 
 demo.launch(share=True, server_port=7860)