Update app.py

app.py

@@ -0,0 +1,371 @@
+import gradio as gr
+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
+import tempfile
+import webbrowser
+import os
+from datetime import datetime
+
+@dataclass
+class PatientMetadata:
+    age: int
+    smoking_status: str
+    family_history: bool
+    menopause_status: str
+    previous_mammogram: bool
+    breast_density: str
+    hormone_therapy: bool
+
+@dataclass
+class AnalysisResult:
+    has_tumor: bool
+    tumor_size: str
+    metadata: PatientMetadata
+
+class BreastSinogramAnalyzer:
+    def __init__(self):
+        """Initialize the analyzer with required models."""
+        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:
+        """Initialize vision models for abnormality detection and size measurement."""
+        print("Loading detection models...")
+        self.tumor_detector = AutoModelForImageClassification.from_pretrained(
+            "SIATCN/vit_tumor_classifier"
+        ).to(self.device).eval()
+        self.tumor_processor = AutoImageProcessor.from_pretrained("SIATCN/vit_tumor_classifier")
+        
+        self.size_detector = AutoModelForImageClassification.from_pretrained(
+            "SIATCN/vit_tumor_radius_detection_finetuned"
+        ).to(self.device).eval()
+        self.size_processor = AutoImageProcessor.from_pretrained(
+            "SIATCN/vit_tumor_radius_detection_finetuned"
+        )
+
+    def _init_llm(self) -> None:
+        """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."""
+        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':
+            image = image.convert('RGB')
+        return image.resize((224, 224))
+
+    @torch.no_grad()
+    def _analyze_image(self, image: Image.Image) -> AnalysisResult:
+        """Perform abnormality detection and size measurement."""
+        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]
+
+        # 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, metadata)
+
+    def _generate_medical_report(self, analysis: AnalysisResult) -> str:
+        """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 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(', ')}
+
+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
+            )
+            
+            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}"""
+            
+            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 clear fallback report."""
+        if analysis.has_tumor:
+            return f"""FINDINGS AND RECOMMENDATIONS:
+
+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 """FINDINGS AND RECOMMENDATIONS:
+
+Finding: No abnormal areas were detected during this scan.
+
+Recommendation: Continue with routine screening as per standard guidelines."""
+
+    def _generate_print_preview(self, analysis_text: str, image: Image.Image) -> str:
+        """Generate an HTML print preview."""
+        temp_dir = tempfile.gettempdir()
+        temp_image_path = os.path.join(temp_dir, 'scan_image.png')
+        image.save(temp_image_path)
+
+        current_date = datetime.now().strftime("%B %d, %Y")
+        
+        html_content = f"""
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <title>Medical Imaging Report</title>
+            <style>
+                @media print {{
+                    body {{
+                        font-family: Arial, sans-serif;
+                        line-height: 1.6;
+                        padding: 20px;
+                        max-width: 800px;
+                        margin: 0 auto;
+                    }}
+                    .header {{
+                        text-align: center;
+                        margin-bottom: 30px;
+                        border-bottom: 2px solid #000;
+                        padding-bottom: 10px;
+                    }}
+                    .date {{
+                        text-align: right;
+                        margin-bottom: 20px;
+                    }}
+                    .content {{
+                        margin-bottom: 30px;
+                    }}
+                    .scan-image {{
+                        text-align: center;
+                        margin: 20px 0;
+                    }}
+                    .scan-image img {{
+                        max-width: 500px;
+                        height: auto;
+                    }}
+                    .footer {{
+                        margin-top: 50px;
+                        border-top: 1px solid #000;
+                        padding-top: 20px;
+                    }}
+                    @page {{
+                        size: A4;
+                        margin: 2cm;
+                    }}
+                    .no-print {{
+                        display: none;
+                    }}
+                }}
+                /* Screen-only styles */
+                body {{
+                    font-family: Arial, sans-serif;
+                    line-height: 1.6;
+                    padding: 20px;
+                    max-width: 800px;
+                    margin: 0 auto;
+                }}
+                .print-button {{
+                    background-color: #007bff;
+                    color: white;
+                    padding: 10px 20px;
+                    border: none;
+                    border-radius: 5px;
+                    cursor: pointer;
+                    margin-bottom: 20px;
+                }}
+                .print-button:hover {{
+                    background-color: #0056b3;
+                }}
+            </style>
+        </head>
+        <body>
+            <button onclick="window.print()" class="print-button no-print">Print Report</button>
+            
+            <div class="header">
+                <h1>Medical Imaging Report</h1>
+            </div>
+            
+            <div class="date">
+                Report Date: {current_date}
+            </div>
+            
+            <div class="scan-image">
+                <img src="file://{temp_image_path}" alt="Scan Image">
+            </div>
+            
+            <div class="content">
+                <pre style="white-space: pre-wrap; font-family: Arial, sans-serif;">{analysis_text}</pre>
+            </div>
+            
+            <div class="footer">
+                <p>This report is generated by an automated analysis system and should be reviewed by a qualified healthcare professional.</p>
+            </div>
+        </body>
+        </html>
+        """
+        
+        temp_html_path = os.path.join(temp_dir, 'report.html')
+        with open(temp_html_path, 'w', encoding='utf-8') as f:
+            f.write(html_content)
+        
+        return temp_html_path
+
+    def analyze(self, image: Image.Image) -> Tuple[str, str]:
+        """Main analysis pipeline."""
+        try:
+            processed_image = self._process_image(image)
+            analysis = self._analyze_image(processed_image)
+            report = self._generate_medical_report(analysis)
+            
+            analysis_text = 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 INFORMATION:
+• Age: {analysis.metadata.age} years
+• Risk Factors: {', '.join([
+    'family history of breast cancer' if analysis.metadata.family_history else '',
+    analysis.metadata.smoking_status.lower(),
+    'currently on hormone therapy' if analysis.metadata.hormone_therapy else '',
+    ]).strip(', ')}
+
+{report}"""
+
+            preview_path = self._generate_print_preview(analysis_text, image)
+            
+            return analysis_text, preview_path
+        except Exception as e:
+            return f"Error during analysis: {str(e)}", ""
+
+def open_print_preview(preview_path: str) -> None:
+    """Open the print preview in the default browser."""
+    if preview_path:
+        webbrowser.open(f'file://{preview_path}')
+    return None
+
+def create_interface() -> gr.Blocks:
+    """Create the Gradio interface."""
+    analyzer = BreastSinogramAnalyzer()
+    
+    with gr.Blocks() as interface:
+        gr.Markdown("# Breast Imaging Analysis System")
+        gr.Markdown("Upload a breast image for analysis and medical assessment.")
+        
+        with gr.Row():
+            input_image = gr.Image(type="pil", label="Upload Breast Image for Analysis")
+        
+        with gr.Row():
+            analyze_btn = gr.Button("Analyze Image", variant="primary")
+            print_btn = gr.Button("Open Print Preview")
+        
+        output_text = gr.Textbox(label="Analysis Results", lines=20)
+        preview_path = gr.Textbox(visible=False)
+        
+        analyze_btn.click(
+            fn=analyzer.analyze,
+            inputs=[input_image],
+            outputs=[output_text, preview_path]
+        )
+        
+        print_btn.click(
+            fn=open_print_preview,
+            inputs=[preview_path],
+            outputs=None
+        )
+    
+    return interface
+
+if __name__ == "__main__":
+    print("Starting application...")
+    interface = create_interface()
+    interface.launch(debug=True, share=True)