Update app.py

app.py

@@ -1,13 +1,14 @@
 import gradio as gr
-from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import torch
-from typing import Tuple, Optional, Dict, Any
+from typing import Tuple, Optional, Dict, Any, List
 from dataclasses import dataclass
 import random
 from datetime import datetime, timedelta
 import os
 from qwen_agent.agents import Assistant
+from qwen_agent.gui.web_ui import WebUI
 
 @dataclass
 class PatientMetadata:
@@ -26,15 +27,30 @@ class AnalysisResult:
     confidence: float
     metadata: PatientMetadata
 
-class BreastSinogramAnalyzer:
+class BreastCancerAgent(Assistant):
     def __init__(self):
-        """Initialize the analyzer with required models."""
+        super().__init__(
+            llm={
+                'model': os.environ.get("MODELNAME", "qwen-vl-chat"),
+                'generate_cfg': {
+                    'max_input_tokens': 32768,
+                    'max_retries': 10,
+                    'temperature': float(os.environ.get("T", 0.001)),
+                    'repetition_penalty': float(os.environ.get("R", 1.0)),
+                    "top_k": int(os.environ.get("K", 20)),
+                    "top_p": float(os.environ.get("P", 0.8)),
+                }
+            },
+            name='Breast Cancer Analyzer',
+            description='Medical imaging analysis system specializing in breast cancer detection and reporting.',
+            system_message='You are an expert medical imaging system analyzing breast cancer scans. Provide clear, accurate, and professional analysis.'
+        )
+        
         print("Initializing system...")
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {self.device}")
         
         self._init_vision_models()
-        self._init_llm()
         print("Initialization complete!")
 
     def _init_vision_models(self) -> None:
@@ -52,50 +68,6 @@ class BreastSinogramAnalyzer:
             "SIATCN/vit_tumor_radius_detection_finetuned"
         )
 
-    def _init_llm(self) -> None:
-        """Initialize the Qwen model for report generation."""
-        print("Loading language model...")
-        self.agent = Assistant(
-            llm={
-                'model': os.environ.get("MODELNAME"),
-                'generate_cfg': {
-                    'max_input_tokens': 32768,
-                    'max_retries': 10,
-                    'temperature': float(os.environ.get("T", 0.001)),
-                    'repetition_penalty': float(os.environ.get("R", 1.0)),
-                    "top_k": int(os.environ.get("K", 20)),
-                    "top_p": float(os.environ.get("P", 0.8)),
-                }
-            },
-            name='QwQ-32B-preview',
-            description='Medical report generation model based on QwQ-32B-Preview',
-            system_message='You are an experienced radiologist providing clear and concise medical reports. You should think step-by-step and be precise in your analysis.',
-            rag_cfg={'max_ref_token': 32768, 'rag_searchers': []},
-        )
-
-    def _generate_synthetic_metadata(self) -> PatientMetadata:
-        """Generate realistic patient metadata for breast cancer screening."""
-        age = random.randint(40, 75)
-        smoking_status = random.choice(["Never Smoker", "Former Smoker", "Current Smoker"])
-        family_history = random.choice([True, False])
-        menopause_status = "Post-menopausal" if age > 50 else "Pre-menopausal"
-        previous_mammogram = random.choice([True, False])
-        breast_density = random.choice(["A: Almost entirely fatty", 
-                                      "B: Scattered fibroglandular",
-                                      "C: Heterogeneously dense",
-                                      "D: Extremely dense"])
-        hormone_therapy = random.choice([True, False])
-
-        return PatientMetadata(
-            age=age,
-            smoking_status=smoking_status,
-            family_history=family_history,
-            menopause_status=menopause_status,
-            previous_mammogram=previous_mammogram,
-            breast_density=breast_density,
-            hormone_therapy=hormone_therapy
-        )
-
     def _process_image(self, image: Image.Image) -> Image.Image:
         """Process input image for model consumption."""
         if image.mode != 'RGB':
@@ -105,17 +77,14 @@ class BreastSinogramAnalyzer:
     @torch.no_grad()
     def _analyze_image(self, image: Image.Image) -> AnalysisResult:
         """Perform abnormality detection and size measurement."""
-        # Generate metadata
         metadata = self._generate_synthetic_metadata()
-
-        # Detect abnormality
+        
         tumor_inputs = self.tumor_processor(image, return_tensors="pt").to(self.device)
         tumor_outputs = self.tumor_detector(**tumor_inputs)
         tumor_probs = tumor_outputs.logits.softmax(dim=-1)[0].cpu()
         has_tumor = tumor_probs[1] > tumor_probs[0]
         confidence = float(tumor_probs[1] if has_tumor else tumor_probs[0])
 
-        # Measure size
         size_inputs = self.size_processor(image, return_tensors="pt").to(self.device)
         size_outputs = self.size_detector(**size_inputs)
         size_pred = size_outputs.logits.softmax(dim=-1)[0].cpu()
@@ -124,63 +93,40 @@ class BreastSinogramAnalyzer:
 
         return AnalysisResult(has_tumor, tumor_size, confidence, metadata)
 
-    def _generate_medical_report(self, analysis: AnalysisResult) -> str:
-        """Generate a medical report using Qwen model."""
-        prompt = f"""Generate a brief medical report for this microwave breast imaging scan:
-
-Findings:
-- {'Abnormal' if analysis.has_tumor else 'Normal'} dielectric properties
-- Size: {analysis.tumor_size} cm
-- Confidence: {analysis.confidence:.2%}
-- Patient age: {analysis.metadata.age}
-- Risk factors: {', '.join([
-    'family history' if analysis.metadata.family_history else '',
-    analysis.metadata.smoking_status.lower(),
-    'hormone therapy' if analysis.metadata.hormone_therapy else ''
-    ]).strip(', ')}
-
-Provide:
-1. One sentence interpreting the findings
-2. One clear management recommendation"""
-
-        try:
-            response = self.agent.chat(prompt)
-            if len(response.split()) >= 10:
-                return f"""INTERPRETATION AND RECOMMENDATION:
-{response}"""
-            
-            print("Report too short, using fallback")
-            return self._generate_fallback_report(analysis)
-
-        except Exception as e:
-            print(f"Error in report generation: {str(e)}")
-            return self._generate_fallback_report(analysis)
-
-    def _generate_fallback_report(self, analysis: AnalysisResult) -> str:
-        """Generate a simple fallback report."""
-        if analysis.has_tumor:
-            return f"""INTERPRETATION AND RECOMMENDATION:
-Microwave imaging reveals abnormal dielectric properties measuring {analysis.tumor_size} cm with {analysis.confidence:.1%} confidence level.
-
-{'Immediate conventional imaging and clinical correlation recommended.' if analysis.tumor_size in ['1.0', '1.5'] else 'Follow-up imaging recommended in 6 months.'}"""
-        else:
-            return f"""INTERPRETATION AND RECOMMENDATION:
-Microwave imaging shows normal dielectric properties with {analysis.confidence:.1%} confidence level.
+    def _generate_synthetic_metadata(self) -> PatientMetadata:
+        """Generate realistic patient metadata for breast cancer screening."""
+        age = random.randint(40, 75)
+        smoking_status = random.choice(["Never Smoker", "Former Smoker", "Current Smoker"])
+        family_history = random.choice([True, False])
+        menopause_status = "Post-menopausal" if age > 50 else "Pre-menopausal"
+        previous_mammogram = random.choice([True, False])
+        breast_density = random.choice(["A: Almost entirely fatty", 
+                                      "B: Scattered fibroglandular",
+                                      "C: Heterogeneously dense",
+                                      "D: Extremely dense"])
+        hormone_therapy = random.choice([True, False])
 
-Routine screening recommended per standard protocol."""
+        return PatientMetadata(
+            age=age,
+            smoking_status=smoking_status,
+            family_history=family_history,
+            menopause_status=menopause_status,
+            previous_mammogram=previous_mammogram,
+            breast_density=breast_density,
+            hormone_therapy=hormone_therapy
+        )
 
-    def analyze(self, image: Image.Image) -> str:
-        """Main analysis pipeline."""
+    def run(self, image_path: str) -> str:
+        """Run analysis on an image."""
         try:
+            image = Image.open(image_path)
             processed_image = self._process_image(image)
             analysis = self._analyze_image(processed_image)
-            report = self._generate_medical_report(analysis)
             
-            return f"""MICROWAVE IMAGING ANALYSIS:
+            report = f"""MICROWAVE IMAGING ANALYSIS:
 • Detection: {'Positive' if analysis.has_tumor else 'Negative'}
 • Size: {analysis.tumor_size} cm
 
-
 PATIENT INFO:
 • Age: {analysis.metadata.age} years
 • Risk Factors: {', '.join([
@@ -190,29 +136,33 @@ PATIENT INFO:
     ]).strip(', ')}
 
 REPORT:
-{report}"""
+{'Abnormal scan showing potential mass.' if analysis.has_tumor else 'Normal scan with no significant findings.'} 
+Confidence level: {analysis.confidence:.1%}
+
+RECOMMENDATION:
+{('Immediate follow-up imaging recommended.' if analysis.tumor_size in ['1.0', '1.5'] else 'Follow-up imaging in 6 months recommended.') if analysis.has_tumor else 'Continue routine screening per protocol.'}"""
+            
+            return report
         except Exception as e:
             return f"Error during analysis: {str(e)}"
 
-def create_interface() -> gr.Interface:
-    """Create the Gradio interface."""
-    analyzer = BreastSinogramAnalyzer()
+def run_interface():
+    """Create and run the WebUI interface."""
+    agent = BreastCancerAgent()
     
-    interface = gr.Interface(
-        fn=analyzer.analyze,
-        inputs=[
-            gr.Image(type="pil", label="Upload Breast Microwave Image")
-        ],
-        outputs=[
-            gr.Textbox(label="Analysis Results", lines=20)
-        ],
-        title="Breast Cancer Microwave Imaging Analysis System",
-        description="Upload a breast microwave image for comprehensive analysis and medical assessment.",
-    )
+    chatbot_config = {
+        'user.name': 'Medical Staff',
+        'input.placeholder': 'Upload a breast microwave image for analysis...',
+        'prompt.suggestions': [
+            {'text': 'Can you analyze this mammogram?'},
+            {'text': 'What should I look for in the results?'},
+            {'text': 'How reliable is the detection?'}
+        ]
+    }
     
-    return interface
+    app = WebUI(agent, chatbot_config=chatbot_config)
+    app.run(share=True, concurrency_limit=80)
 
 if __name__ == "__main__":
     print("Starting application...")
-    interface = create_interface()
-    interface.launch(debug=True, share=True)
+    run_interface()