Update app.py

app.py

@@ -1,259 +1,156 @@
 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 transformers import AutoModelForCausalLM
+from janus.models import VLChatProcessor
 from PIL import Image
 import spaces
-from torchvision import transforms
 
-# Medical Imaging Configuration
+# Medical Image Analysis Configuration
 MEDICAL_CONFIG = {
-    "modality": "CT",  # Default imaging modality
-    "anatomical_region": "Chest",
-    "clinical_task": "analysis",
-    "report_style": "structured"
+    "echo_guidelines": "ASE 2023 Standards",
+    "histo_guidelines": "CAP Protocols 2024",
+    "cardiac_params": ["LVEF", "E/A Ratio", "Wall Motion"],
+    "histo_params": ["Nuclear Atypia", "Mitotic Count", "Stromal Invasion"]
 }
 
-# Load base model
+# Initialize Medical Imaging 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 adaptations
-vl_gpt = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    language_config=language_config,
-    trust_remote_code=True,
-    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):
+class MedicalImagingAdapter(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()
-
+        # Cardiac-specific projections
+        self.cardiac_proj = torch.nn.Linear(2048, 2048)
+        # Histopathology-specific projections
+        self.histo_proj = torch.nn.Linear(2048, 2048)
+        
     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)
+        return outputs
 
-vl_gpt.language_model = MedicalProjectionWrapper(vl_gpt.language_model)
+vl_gpt = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
+vl_gpt.language_model = MedicalImagingAdapter(vl_gpt.language_model)
 
 if torch.cuda.is_available():
-    vl_gpt = vl_gpt.cuda()
+    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
 
 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)
+# Medical Image Processing Pipelines
+def preprocess_echo(image):
+    """Process echocardiography images"""
+    img = Image.fromarray(image).convert('L')  # Grayscale
+    return np.array(img.resize((512, 512)))
+
+def preprocess_histo(image):
+    """Process histopathology slides"""
+    img = Image.fromarray(image)
+    return np.array(img.resize((1024, 1024)))
 
 @torch.inference_mode()
 @spaces.GPU(duration=120)
-def medical_image_analysis(image, clinical_query, seed=42):
-    torch.cuda.empty_cache()
-    torch.manual_seed(seed)
+def analyze_medical_case(image, clinical_context, modality):
+    # Preprocess based on modality
+    processed_img = preprocess_echo(image) if modality == "Echo" else preprocess_histo(image)
     
-    # Preprocess with medical transformations
-    medical_image = preprocess_medical_image(image)
+    # Create modality-specific prompt
+    system_prompt = f"""
+    Analyze this {modality} image following {MEDICAL_CONFIG['echo_guidelines' if modality=='Echo' else 'histo_guidelines']}.
+    Clinical Context: {clinical_context}
+    """
     
     conversation = [{
-        "role": "<|Radiologist|>",
-        "content": f"<medical_image>\nClinical Context: {clinical_query}",
-        "images": [medical_image],
+        "role": "<|Radiologist|>" if modality == "Echo" else "<|Pathologist|>",
+        "content": system_prompt,
+        "images": [processed_img],
     }, {"role": "<|AI_Assistant|>", "content": ""}]
-    
+
     inputs = vl_chat_processor(
         conversations=conversation,
-        images=[Image.fromarray(image)],
+        images=[Image.fromarray(processed_img)],
         force_batchify=True
-    ).to(cuda_device)
+    ).to(vl_gpt.device)
 
-    inputs_embeds = vl_gpt.prepare_inputs_embeds(**inputs)
-    
-    # Medical-optimized generation parameters
-    outputs = vl_gpt.language_model.generate(
-        inputs_embeds=inputs_embeds,
+    outputs = vl_gpt.generate(
+        inputs_embeds=vl_gpt.prepare_inputs_embeds(**inputs),
         attention_mask=inputs.attention_mask,
         max_new_tokens=512,
-        temperature=0.2,
+        temperature=0.1,
         top_p=0.9,
-        num_beams=5,
-        repetition_penalty=1.5,
-        early_stopping=True
+        repetition_penalty=1.5
     )
     
-    report = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-    return format_medical_report(report)
+    report = vl_chat_processor.tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
+    return format_medical_report(report, modality)
 
-def format_medical_report(raw_text):
+def format_medical_report(text, modality):
+    # Structure report based on modality
     sections = {
-        "Findings": "",
-        "Impression": "",
-        "Recommendations": ""
+        "Echo": [
+            ("Chamber Dimensions", "LVEDD", "LVESD"),
+            ("Valvular Function", "Aortic Valve", "Mitral Valve"),
+            ("Hemodynamics", "E/A Ratio", "LVEF")
+        ],
+        "Histo": [
+            ("Architecture", "Gland Formation", "Stromal Pattern"),
+            ("Cellular Features", "Nuclear Atypia", "Mitotic Count"),
+            ("Diagnostic Impression", "Tumor Grade", "Margin Status")
+        ]
     }
     
-    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=12345, guidance=7, temperature=0.6):
-    torch.cuda.empty_cache()
-    if seed is not None:
-        torch.manual_seed(seed)
-
-    medical_prompt = f"{prompt} [Modality: {MEDICAL_CONFIG['modality']}, Anatomy: {MEDICAL_CONFIG['anatomical_region']}]"
-    
-    messages = [{
-        'role': '<|Clinician|>', 
-        'content': medical_prompt
-    }]
-    
-    text = vl_chat_processor.apply_chat_template(
-        messages,
-        system_prompt='Generate educational medical imaging data'
-    )
-    
-    input_ids = torch.LongTensor(tokenizer.encode(text)).to(cuda_device)
-    
-    # Medical image generation parameters
-    generated_tokens, patches = vl_gpt.generate(
-        input_ids,
-        width=512,
-        height=512,
-        cfg_weight=guidance,
-        temperature=temperature,
-        parallel_size=3,
-        image_token_num_per_image=576,
-        patch_size=16
+    formatted = f"**{modality} Analysis Report**\n\n"
+    for section in sections[modality]:
+        header = section[0]
+        formatted += f"### {header}\n"
+        for sub in section[1:]:
+            if sub in text:
+                start = text.find(sub)
+                end = text.find("\n\n", start)
+                formatted += f"- **{sub}:** {text[start+len(sub)+1:end].strip()}\n"
+    return formatted
+
+# Medical Imaging Interface
+with gr.Blocks(title="Cardiac & Histopathology AI", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    ## Medical Imaging Analysis Platform
+    *Analyzes echocardiograms and histopathology slides - Research Use Only*
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Upload Medical Image")
+            modality_select = gr.Radio(
+                ["Echo", "Histo"], 
+                label="Image Modality",
+                info="Select 'Echo' for cardiac ultrasound, 'Histo' for biopsy slides"
+            )
+            clinical_input = gr.Textbox(
+                label="Clinical Context",
+                placeholder="e.g., 'Assess LV function' or 'Evaluate for malignancy'"
+            )
+            analyze_btn = gr.Button("Analyze Case", variant="primary")
+        
+        with gr.Column():
+            report_output = gr.Markdown(label="AI Clinical Report")
+    
+    # Preloaded examples from space files
+    gr.Examples(
+        examples=[
+            ["Evaluate LV systolic function", "case1.png", "Echo"],
+            ["Assess mitral valve function", "case2.jpg", "Echo"],
+            ["Analyze for malignant features", "case3.png", "Histo"],
+            ["Evaluate tumor margins", "case4.png", "Histo"]
+        ],
+        inputs=[clinical_input, image_input, modality_select],
+        label="Example Medical Cases"
     )
-    
-    synthetic_images = postprocess_medical_images(patches)
-    return [Image.fromarray(img).resize((512, 512)) for img in synthetic_images]
 
-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():
-            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")
-
-        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():
-            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)
-
-    # Event handlers
-    analysis_btn.click(
-        medical_image_analysis,
-        [med_image, clinical_input],
-        report_output
-    )
-    
-    generate_btn.click(
-        generate_medical_image,
-        [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
-        )
+analyze_btn.click(
+    analyze_medical_case,
+    [image_input, clinical_input, modality_select],
+    report_output
+)
 
-demo.launch(share=True, server_port=7860)
+demo.launch(share=True)