Update app.py

app.py

@@ -1,21 +1,100 @@
-def _generate_medical_report(self, has_tumor: bool, tumor_size: str, metadata: PatientMetadata) -> str:
-    """Generate a standardized medical report using Mistral."""
-    prompt = f"""<s>[INST] Generate a structured medical report for a breast imaging scan using the following format exactly.
+import gradio as gr
+from PIL import Image
+from dataclasses import dataclass
+import random
+from transformers import pipeline
+from huggingface_hub import InferenceClient, login
+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
+
+class MicrowaveBreastAnalyzer:
+    def __init__(self, hf_token: str):
+        """Initialize the analyzer with models."""
+        print("Initializing system...")
+        
+        # Login to Hugging Face
+        login(token=hf_token)
+        
+        # Initialize vision pipelines for tumor detection and size classification
+        self.tumor_classifier = pipeline(
+            "image-classification",
+            model="SIATCN/vit_tumor_classifier",
+            device="cpu"
+        )
+        
+        self.size_classifier = pipeline(
+            "image-classification",
+            model="SIATCN/vit_tumor_radius_detection_finetuned",
+            device="cpu"
+        )
+        
+        # Initialize Mistral client for report generation
+        self.report_generator = InferenceClient(
+            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
+            token=hf_token
+        )
+        
+        print("Initialization complete!")
+
+    def _generate_synthetic_metadata(self) -> PatientMetadata:
+        """Generate realistic patient metadata for 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))
+
+    def _generate_medical_report(self, has_tumor: bool, tumor_size: str, metadata: PatientMetadata) -> str:
+        """Generate a standardized report for microwave breast imaging."""
+        prompt = f"""<s>[INST] Generate a structured medical report for a microwave breast imaging scan using the following format exactly.
 Keep sections consistent and use proper medical terminology. Be concise yet thorough.
 
 EXAMINATION PERFORMED:
-- Diagnostic Digital Mammogram
+- Microwave Breast Imaging Scan
 - Date: {datetime.now().strftime('%B %d, %Y')}
 
-CLINICAL FINDINGS:
+IMAGING FINDINGS:
 Primary Finding: {'Abnormal area detected' if has_tumor else 'No abnormalities detected'}
-{f'Lesion Size: {tumor_size} cm' if has_tumor else ''}
+{f'Detected Mass Size: {tumor_size} cm' if has_tumor else ''}
 
 PATIENT HISTORY:
 - Age: {metadata.age} years
 - Menopausal Status: {metadata.menopause_status}
-- Previous Mammogram: {'Yes' if metadata.previous_mammogram else 'No'}
-- Breast Tissue Density: {metadata.breast_density}
+- Previous Screening: {'Yes' if metadata.previous_mammogram else 'No'}
+- Tissue Characteristics: {metadata.breast_density}
 
 RISK FACTORS:
 {f'• Family History: {"Present" if metadata.family_history else "None"}'}
@@ -25,32 +104,32 @@ RISK FACTORS:
 Please generate a report with these exact sections:
 
 1. DETAILED FINDINGS
-[Describe the mammographic findings in detail, including location, characteristics, and comparison with any previous studies if available]
+[Describe the microwave imaging findings in detail, including location and characteristics of any detected abnormalities]
 
-2. IMPRESSION
-[Provide a clear assessment using BI-RADS classification if applicable]
+2. INTERPRETATION
+[Provide a clear assessment of the microwave imaging results and their clinical significance]
 
 3. RECOMMENDATIONS
 [List specific follow-up actions and timeline]
 
-4. ADDITIONAL NOTES
-[Include any relevant screening considerations or risk-management suggestions]
-
-Format each section consistently and maintain professional medical terminology throughout. [/INST]</s>"""
-
-    # Generate response using Mistral
-    response = self.report_generator.text_generation(
-        prompt,
-        max_new_tokens=800,  # Increased for comprehensive report
-        temperature=0.3,     # Low temperature for consistency
-        top_p=0.9,
-        repetition_penalty=1.1,
-        do_sample=True,
-        seed=42
-    )
-        
-    # Post-process the response to ensure consistent formatting
-    formatted_response = f"""MAMMOGRAM REPORT
+4. TECHNICAL NOTES
+[Include any relevant information about the scan quality and any technical considerations]
+
+Format each section consistently and maintain professional medical terminology throughout. Note that this uses microwave imaging technology, not mammography. [/INST]</s>"""
+
+        # Generate response using Mistral
+        response = self.report_generator.text_generation(
+            prompt,
+            max_new_tokens=800,
+            temperature=0.3,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            do_sample=True,
+            seed=42
+        )
+            
+        # Post-process the response to ensure consistent formatting
+        formatted_response = f"""MICROWAVE BREAST IMAGING REPORT
 Date: {datetime.now().strftime('%B %d, %Y')}
 ----------------------------------------
 
@@ -58,40 +137,79 @@ Date: {datetime.now().strftime('%B %d, %Y')}
 
 ----------------------------------------
 NOTE: This report was generated using AI assistance and should be reviewed by a qualified healthcare professional.
-All findings and recommendations should be clinically verified."""
+This screening was performed using microwave imaging technology."""
 
-    return formatted_response
+        return formatted_response
 
-# Update the analyze method to ensure consistent output formatting
-def analyze(self, image: Image.Image) -> str:
-    """Main analysis pipeline with standardized output."""
-    try:
-        processed_image = self._process_image(image)
-        metadata = self._generate_synthetic_metadata()
+    def analyze(self, image: Image.Image) -> str:
+        """Main analysis pipeline with standardized output."""
+        try:
+            processed_image = self._process_image(image)
+            metadata = self._generate_synthetic_metadata()
 
-        # Detect tumor
-        tumor_result = self.tumor_classifier(processed_image)
-        has_tumor = tumor_result[0]['label'] == 'tumor'
-        tumor_confidence = tumor_result[0]['score']
+            # Detect tumor
+            tumor_result = self.tumor_classifier(processed_image)
+            has_tumor = tumor_result[0]['label'] == 'tumor'
+            tumor_confidence = tumor_result[0]['score']
 
-        # Measure size if tumor detected
-        size_result = self.size_classifier(processed_image)
-        tumor_size = size_result[0]['label'].replace('tumor-', '')
+            # Measure size if tumor detected
+            size_result = self.size_classifier(processed_image)
+            tumor_size = size_result[0]['label'].replace('tumor-', '')
 
-        # Generate report
-        report = self._generate_medical_report(has_tumor, tumor_size, metadata)
+            # Generate report
+            report = self._generate_medical_report(has_tumor, tumor_size, metadata)
 
-        return f"""BREAST IMAGING ANALYSIS REPORT
+            return f"""MICROWAVE BREAST IMAGING ANALYSIS
 ========================================
 
 INITIAL SCAN ASSESSMENT:
 {'⚠️ ABNORMAL FINDING DETECTED' if has_tumor else '✓ NO ABNORMALITIES DETECTED'}
 Detection Confidence: {tumor_confidence:.2%}
-{f'Estimated Size: {tumor_size} cm' if has_tumor else ''}
+{f'Estimated Mass Size: {tumor_size} cm' if has_tumor else ''}
 
 ----------------------------------------
 {report}"""
 
-    except Exception as e:
-        import traceback
-        return f"ERROR IN ANALYSIS:\n{str(e)}\n\nDebug Information:\n{traceback.format_exc()}"
+        except Exception as e:
+            import traceback
+            return f"Error during analysis: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+
+def create_interface(hf_token: str) -> gr.Interface:
+    """Create the Gradio interface."""
+    analyzer = MicrowaveBreastAnalyzer(hf_token)
+    
+    interface = gr.Interface(
+        fn=analyzer.analyze,
+        inputs=[
+            gr.Image(type="pil", label="Upload Microwave Breast Image for Analysis")
+        ],
+        outputs=[
+            gr.Textbox(label="Analysis Results", lines=20)
+        ],
+        title="Microwave Breast Imaging Analysis System",
+        description="""Upload a microwave breast image for comprehensive analysis. The system will:
+        1. Detect the presence of tumors using microwave imaging technology
+        2. Classify tumor size if present
+        3. Generate a detailed medical report with recommendations
+        
+        Note: This system uses microwave imaging technology for breast screening, which offers a safe, 
+        radiation-free alternative to traditional mammography.""",
+    )
+    
+    return interface
+
+if __name__ == "__main__":
+    print("Starting microwave breast imaging analysis system...")
+    # Load HuggingFace token from secrets
+    HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
+    if not HF_TOKEN:
+        raise ValueError("Please set HUGGINGFACE_TOKEN environment variable")
+        
+    interface = create_interface(HF_TOKEN)
+    # Modified launch parameters for Spaces
+    interface.launch(
+        debug=True,
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )