app01.py

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:
- Microwave Breast Imaging Scan
- Date: {datetime.now().strftime('%B %d, %Y')}

IMAGING FINDINGS:
Primary Finding: {'Abnormal area detected' if has_tumor else 'No abnormalities detected'}
{f'Detected Mass Size: {tumor_size} cm' if has_tumor else ''}

PATIENT HISTORY:
- Age: {metadata.age} years
- Menopausal Status: {metadata.menopause_status}
- 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"}'}
• Smoking Status: {metadata.smoking_status}
• Hormone Therapy: {'Yes' if metadata.hormone_therapy else 'No'}

Please generate a report with these exact sections:

1. DETAILED FINDINGS
[Describe the microwave imaging findings in detail, including location and characteristics of any detected abnormalities]

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. 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')}
----------------------------------------

{response.strip()}

----------------------------------------
NOTE: This report was generated using AI assistance and should be reviewed by a qualified healthcare professional.
This screening was performed using microwave imaging technology."""

        return formatted_response

    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']

            # 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)

            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 Mass Size: {tumor_size} cm' if has_tumor else ''}

----------------------------------------
{report}"""

        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
    )