Update app.py

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
-from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
 import torch
 from typing import Tuple, Optional, Dict, Any
 from dataclasses import dataclass
 import random
-from datetime import datetime, timedelta
 
 @dataclass
 class PatientMetadata:
@@ -21,7 +20,6 @@ class PatientMetadata:
 class AnalysisResult:
     has_tumor: bool
     tumor_size: str
-    confidence: float
     metadata: PatientMetadata
 
 class BreastSinogramAnalyzer:
@@ -51,18 +49,15 @@ class BreastSinogramAnalyzer:
         )
 
     def _init_llm(self) -> None:
-        """Initialize the language model for report generation."""
-        print("Loading language model pipeline...")
-        self.pipe = pipeline(
-            "text-generation",
-            model="Qwen/QwQ-32B-Preview",
-            torch_dtype=torch.float16,
-            device_map="auto",
-            model_kwargs={
-                "load_in_4bit": False,
-                "bnb_4bit_compute_dtype": torch.float16,
-            }
+        """Initialize the Qwen language model for report generation."""
+        print("Loading Qwen language model...")
+        self.model_name = "Qwen/QwQ-32B-Preview"
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            torch_dtype="auto",
+            device_map="auto"
         )
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
 
     def _generate_synthetic_metadata(self) -> PatientMetadata:
         """Generate realistic patient metadata for breast cancer screening."""
@@ -96,7 +91,6 @@ 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
@@ -104,61 +98,73 @@ class BreastSinogramAnalyzer:
         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
+        # Measure size if tumor detected
         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()
         sizes = ["no-tumor", "0.5", "1.0", "1.5"]
         tumor_size = sizes[size_pred.argmax().item()]
 
-        return AnalysisResult(has_tumor, tumor_size, confidence, metadata)
+        return AnalysisResult(has_tumor, tumor_size, metadata)
 
     def _generate_medical_report(self, analysis: AnalysisResult) -> str:
-        """Generate a simplified medical report."""
-        prompt = f"""<|system|>You are a radiologist providing clear and concise medical reports.</s>
-<|user|>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}
+        """Generate a clear medical report using Qwen."""
+        try:
+            messages = [
+                {
+                    "role": "system",
+                    "content": "You are a radiologist providing clear and straightforward medical reports. Focus on clarity and actionable recommendations."
+                },
+                {
+                    "role": "user",
+                    "content": f"""Generate a clear medical report for this breast imaging scan:
+
+Scan Results:
+- Finding: {'Abnormal area detected' if analysis.has_tumor else 'No abnormalities detected'}
+{f'- Size of abnormal area: {analysis.tumor_size} cm' if analysis.has_tumor else ''}
+
+Patient Information:
+- Age: {analysis.metadata.age} years
 - Risk factors: {', '.join([
-    'family history' if analysis.metadata.family_history else '',
-    analysis.metadata.smoking_status.lower(),
-    'hormone therapy' if analysis.metadata.hormone_therapy else ''
+    'family history of breast cancer' if analysis.metadata.family_history else '',
+    f'{analysis.metadata.smoking_status.lower()}',
+    'currently on hormone therapy' if analysis.metadata.hormone_therapy else ''
     ]).strip(', ')}
 
-Provide:
-1. One sentence interpreting the findings
-2. One clear management recommendation</s>
-<|assistant|>"""
-
-        try:
-            response = self.pipe(
-                prompt,
+Please provide:
+1. A clear interpretation of the findings
+2. A specific recommendation for next steps"""
+                }
+            ]
+
+            text = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+            
+            model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
+            
+            generated_ids = self.model.generate(
+                **model_inputs,
                 max_new_tokens=128,
                 temperature=0.3,
                 top_p=0.9,
                 repetition_penalty=1.1,
-                do_sample=True,
-                num_return_sequences=1
-            )[0]["generated_text"]
-
-            # Extract assistant's response
-            if "<|assistant|>" in response:
-                report = response.split("<|assistant|>")[-1].strip()
-            else:
-                report = response[len(prompt):].strip()
-
-            # Simple validation
-            if len(report.split()) >= 10:
-                return f"""INTERPRETATION AND RECOMMENDATION:
-{report}"""
+                do_sample=True
+            )
+            
+            generated_ids = [
+                output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+            ]
+            
+            response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+            if len(response.split()) >= 10:
+                return f"""FINDINGS AND RECOMMENDATIONS:
+{response}"""
             
-            print("Report too short, using fallback")
             return self._generate_fallback_report(analysis)
 
         except Exception as e:
@@ -166,17 +172,19 @@ Provide:
             return self._generate_fallback_report(analysis)
 
     def _generate_fallback_report(self, analysis: AnalysisResult) -> str:
-        """Generate a simple fallback report."""
+        """Generate a clear 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.
+            return f"""FINDINGS AND RECOMMENDATIONS:
 
-{'Immediate conventional imaging and clinical correlation recommended.' if analysis.tumor_size in ['1.0', '1.5'] else 'Follow-up imaging recommended in 6 months.'}"""
+Finding: An abnormal area measuring {analysis.tumor_size} cm was detected during the scan.
+
+Recommendation: {'An immediate follow-up with conventional mammogram and ultrasound is required.' if analysis.tumor_size in ['1.0', '1.5'] else 'A follow-up scan is recommended in 6 months.'}"""
         else:
-            return f"""INTERPRETATION AND RECOMMENDATION:
-Microwave imaging shows normal dielectric properties with {analysis.confidence:.1%} confidence level.
+            return """FINDINGS AND RECOMMENDATIONS:
+
+Finding: No abnormal areas were detected during this scan.
 
-Routine screening recommended per standard protocol."""
+Recommendation: Continue with routine screening as per standard guidelines."""
 
     def analyze(self, image: Image.Image) -> str:
         """Main analysis pipeline."""
@@ -185,20 +193,18 @@ Routine screening recommended per standard protocol."""
             analysis = self._analyze_image(processed_image)
             report = self._generate_medical_report(analysis)
             
-            return f"""MICROWAVE IMAGING ANALYSIS:
-• Detection: {'Positive' if analysis.has_tumor else 'Negative'}
-• Size: {analysis.tumor_size} cm
-
+            return f"""SCAN RESULTS:
+{'⚠️ Abnormal area detected' if analysis.has_tumor else '✓ No abnormalities detected'}
+{f'Size of abnormal area: {analysis.tumor_size} cm' if analysis.has_tumor else ''}
 
-PATIENT INFO:
+PATIENT INFORMATION:
 • Age: {analysis.metadata.age} years
 • Risk Factors: {', '.join([
-    'family history' if analysis.metadata.family_history else '',
+    'family history of breast cancer' if analysis.metadata.family_history else '',
     analysis.metadata.smoking_status.lower(),
-    'hormone therapy' if analysis.metadata.hormone_therapy else '',
+    'currently on hormone therapy' if analysis.metadata.hormone_therapy else '',
     ]).strip(', ')}
 
-REPORT:
 {report}"""
         except Exception as e:
             return f"Error during analysis: {str(e)}"
@@ -210,13 +216,13 @@ def create_interface() -> gr.Interface:
     interface = gr.Interface(
         fn=analyzer.analyze,
         inputs=[
-            gr.Image(type="pil", label="Upload Breast Microwave Image")
+            gr.Image(type="pil", label="Upload Breast Image for Analysis")
         ],
         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.",
+        title="Breast Imaging Analysis System",
+        description="Upload a breast image for analysis and medical assessment.",
     )
     
     return interface