Restore all files to state before comment removal - fix broken code

mediSync/__init__.py

@@ -1 +1,17 @@
-__version__ = "0.1.0"
+"""
+U-HRES: Unified Health Record Exchange System AI Layer
+=======================================================
+
+A healthcare solution that combines X-ray image analysis with patient report text processing
+to provide comprehensive medical insights.
+
+This package contains the following modules:
+- models: Image and text analysis models, along with multimodal fusion
+- utils: Utility functions for preprocessing and visualization
+- app: Main application with Gradio interface
+
+Author: AI Development Team
+License: MIT
+"""
+
+__version__ = "0.1.0"

mediSync/app.py

@@ -8,11 +8,11 @@ import gradio as gr
 import matplotlib.pyplot as plt
 from PIL import Image
 
-                              
+# Add parent directory to path
 parent_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(parent_dir)
 
-                    
+# Import our modules
 from models.multimodal_fusion import MultimodalFusion
 from utils.preprocessing import enhance_xray_image, normalize_report_text
 from utils.visualization import (
@@ -21,7 +21,7 @@ from utils.visualization import (
     plot_report_entities,
 )
 
-                
+# Set up logging
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
@@ -29,7 +29,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-                                                                
+# Create temporary directory for sample data if it doesn't exist
 os.makedirs(os.path.join(parent_dir, "data", "sample"), exist_ok=True)
 
 
@@ -43,7 +43,7 @@ class UHRESApp:
         self.logger = logging.getLogger(__name__)
         self.logger.info("Initializing U-HRES application")
 
-                                                      
+        # Initialize models with None for lazy loading
         self.fusion_model = None
         self.image_model = None
         self.text_model = None
@@ -79,38 +79,38 @@ class UHRESApp:
             tuple: (image, image_results_html, plot_as_html)
         """
         try:
-                                      
+            # Ensure models are loaded
             if not self.load_models() or self.image_model is None:
                 return image, "Error: Models not loaded properly.", None
 
-                                                     
+            # Save uploaded image to a temporary file
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "upload.png")
 
             if isinstance(image, str):
-                                              
+                # Copy the file if it's a path
                 from shutil import copyfile
 
                 copyfile(image, temp_path)
             else:
-                                                          
+                # Save if it's a Gradio UploadButton image
                 image.save(temp_path)
 
-                                
+            # Run image analysis
             self.logger.info(f"Analyzing image: {temp_path}")
             results = self.image_model.analyze(temp_path)
 
-                                  
+            # Create visualization
             fig = plot_image_prediction(
                 image,
                 results.get("predictions", []),
                 f"Primary Finding: {results.get('primary_finding', 'Unknown')}",
             )
 
-                                         
+            # Convert to HTML for display
             plot_html = self.fig_to_html(fig)
 
-                                    
+            # Format results as HTML
             html_result = f"""
             <h2>X-ray Analysis Results</h2>
             <p><strong>Primary Finding:</strong> {results.get("primary_finding", "Unknown")}</p>
@@ -121,13 +121,13 @@ class UHRESApp:
             <ul>
             """
 
-                                   
+            # Add top 5 predictions
             for label, prob in results.get("predictions", [])[:5]:
                 html_result += f"<li>{label}: {prob:.1%}</li>"
 
             html_result += "</ul>"
 
-                             
+            # Add explanation
             explanation = self.image_model.get_explanation(results)
             html_result += f"<h3>Analysis Explanation:</h3><p>{explanation}</p>"
 
@@ -148,11 +148,11 @@ class UHRESApp:
             tuple: (text, text_results_html, entities_plot_html)
         """
         try:
-                                      
+            # Ensure models are loaded
             if not self.load_models() or self.text_model is None:
                 return text, "Error: Models not loaded properly.", None
 
-                                  
+            # Check for empty text
             if not text or len(text.strip()) < 10:
                 return (
                     text,
@@ -160,21 +160,21 @@ class UHRESApp:
                     None,
                 )
 
-                            
+            # Normalize text
             normalized_text = normalize_report_text(text)
 
-                               
+            # Run text analysis
             self.logger.info("Analyzing medical report text")
             results = self.text_model.analyze(normalized_text)
 
-                                                   
+            # Get entities and create visualization
             entities = results.get("entities", {})
             fig = plot_report_entities(normalized_text, entities)
 
-                                         
+            # Convert to HTML for display
             entities_plot_html = self.fig_to_html(fig)
 
-                                    
+            # Format results as HTML
             html_result = f"""
             <h2>Medical Report Analysis Results</h2>
             <p><strong>Severity Level:</strong> {results.get("severity", {}).get("level", "Unknown")}</p>
@@ -185,7 +185,7 @@ class UHRESApp:
             <ul>
             """
 
-                          
+            # Add findings
             findings = results.get("findings", [])
             if findings:
                 for finding in findings:
@@ -195,14 +195,14 @@ class UHRESApp:
 
             html_result += "</ul>"
 
-                          
+            # Add entities
             html_result += "<h3>Extracted Medical Entities:</h3>"
 
             for category, items in entities.items():
                 if items:
                     html_result += f"<p><strong>{category.capitalize()}:</strong> {', '.join(items)}</p>"
 
-                                           
+            # Add follow-up recommendations
             html_result += "<h3>Follow-up Recommendations:</h3><ul>"
             followups = results.get("followup_recommendations", [])
 
@@ -232,11 +232,11 @@ class UHRESApp:
             tuple: (results_html, multimodal_plot_html)
         """
         try:
-                                      
+            # Ensure models are loaded
             if not self.load_models() or self.fusion_model is None:
                 return "Error: Models not loaded properly.", None
 
-                                    
+            # Check for empty inputs
             if image is None:
                 return "Error: Please upload an X-ray image for analysis.", None
 
@@ -246,36 +246,36 @@ class UHRESApp:
                     None,
                 )
 
-                                                     
+            # Save uploaded image to a temporary file
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "upload.png")
 
             if isinstance(image, str):
-                                              
+                # Copy the file if it's a path
                 from shutil import copyfile
 
                 copyfile(image, temp_path)
             else:
-                                                          
+                # Save if it's a Gradio UploadButton image
                 image.save(temp_path)
 
-                            
+            # Normalize text
             normalized_text = normalize_report_text(text)
 
-                                     
+            # Run multimodal analysis
             self.logger.info("Performing multimodal analysis")
             results = self.fusion_model.analyze(temp_path, normalized_text)
 
-                                  
+            # Create visualization
             fig = plot_multimodal_results(results, image, text)
 
-                                         
+            # Convert to HTML for display
             plot_html = self.fig_to_html(fig)
 
-                                  
+            # Generate explanation
             explanation = self.fusion_model.get_explanation(results)
 
-                                    
+            # Format results as HTML
             html_result = f"""
             <h2>Multimodal Medical Analysis Results</h2>
             
@@ -289,7 +289,7 @@ class UHRESApp:
             <ul>
             """
 
-                          
+            # Add findings
             findings = results.get("findings", [])
             if findings:
                 for finding in findings:
@@ -299,7 +299,7 @@ class UHRESApp:
 
             html_result += "</ul>"
 
-                                           
+            # Add follow-up recommendations
             html_result += "<h3>Recommended Follow-up</h3><ul>"
             followups = results.get("followup_recommendations", [])
 
@@ -313,7 +313,7 @@ class UHRESApp:
 
             html_result += "</ul>"
 
-                                 
+            # Add confidence note
             confidence = results.get("severity", {}).get("confidence", 0)
             html_result += f"""
             <p><em>Note: This analysis has a confidence level of {confidence:.0%}. 
@@ -340,31 +340,31 @@ class UHRESApp:
             if image is None:
                 return None
 
-                                                     
+            # Save uploaded image to a temporary file
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "upload.png")
 
             if isinstance(image, str):
-                                              
+                # Copy the file if it's a path
                 from shutil import copyfile
 
                 copyfile(image, temp_path)
             else:
-                                                          
+                # Save if it's a Gradio UploadButton image
                 image.save(temp_path)
 
-                           
+            # Enhance image
             self.logger.info(f"Enhancing image: {temp_path}")
             output_path = os.path.join(temp_dir, "enhanced.png")
             enhance_xray_image(temp_path, output_path)
 
-                                 
+            # Load enhanced image
             enhanced = Image.open(output_path)
             return enhanced
 
         except Exception as e:
             self.logger.error(f"Error enhancing image: {e}")
-            return image                                  
+            return image  # Return original image on error
 
     def fig_to_html(self, fig):
         """Convert matplotlib figure to HTML for display in Gradio."""
@@ -390,7 +390,7 @@ def create_interface():
 
     app = UHRESApp()
 
-                                     
+    # Example medical report for demo
     example_report = """
     CHEST X-RAY EXAMINATION
     
@@ -409,7 +409,7 @@ def create_interface():
     RECOMMENDATIONS: Follow-up chest CT to further characterize the nodular opacity in the right lower lobe.
     """
 
-                                        
+    # Get sample image path if available
     sample_images_dir = Path(parent_dir) / "data" / "sample"
     sample_images = list(sample_images_dir.glob("*.png")) + list(
         sample_images_dir.glob("*.jpg")
@@ -419,7 +419,7 @@ def create_interface():
     if sample_images:
         sample_image_path = str(sample_images[0])
 
-                      
+    # Define interface
     with gr.Blocks(
         title="U-HRES: Unified Health Record Exchange System AI Layer", theme=gr.themes.Soft()
     ) as interface:
@@ -456,7 +456,7 @@ def create_interface():
                     multi_results = gr.HTML(label="Analysis Results")
                     multi_plot = gr.HTML(label="Visualization")
 
-                                                    
+            # Set up examples if sample image exists
             if sample_image_path:
                 gr.Examples(
                     examples=[[sample_image_path, example_report]],
@@ -476,7 +476,7 @@ def create_interface():
                     img_results = gr.HTML(label="Analysis Results")
                     img_plot = gr.HTML(label="Visualization")
 
-                                                   
+            # Set up example if sample image exists
             if sample_image_path:
                 gr.Examples(
                     examples=[[sample_image_path]],
@@ -500,7 +500,7 @@ def create_interface():
                     text_results = gr.HTML(label="Analysis Results")
                     text_plot = gr.HTML(label="Entity Visualization")
 
-                            
+            # Set up example
             gr.Examples(
                 examples=[[example_report]],
                 inputs=[text_input],
@@ -529,7 +529,7 @@ def create_interface():
             This tool is for educational and research purposes only. It is not intended to provide medical advice or replace professional healthcare. Always consult with qualified healthcare providers for medical decisions.
             """)
 
-                               
+        # Set up event handlers
         multi_img_enhance.click(
             app.enhance_image, inputs=multi_img_input, outputs=multi_img_input
         )
@@ -552,7 +552,7 @@ def create_interface():
             outputs=[text_output, text_results, text_plot],
         )
 
-                       
+    # Run the interface
     interface.launch(show_api=False)
 
 

mediSync/models/__init__.py

@@ -1,5 +1,16 @@
-from .image_analyzer import XRayImageAnalyzer
-from .multimodal_fusion import MultimodalFusion
-from .text_analyzer import MedicalReportAnalyzer
-
-__all__ = ["XRayImageAnalyzer", "MedicalReportAnalyzer", "MultimodalFusion"]
+"""
+U-HRES: Models Module
+=====================
+
+This module contains the core machine learning models for the U-HRES system:
+
+1. XRayImageAnalyzer: Analyzes X-ray images using pre-trained vision models
+2. MedicalReportAnalyzer: Extracts information from medical reports using NLP models
+3. MultimodalFusion: Combines insights from both image and text analysis
+"""
+
+from .image_analyzer import XRayImageAnalyzer
+from .multimodal_fusion import MultimodalFusion
+from .text_analyzer import MedicalReportAnalyzer
+
+__all__ = ["XRayImageAnalyzer", "MedicalReportAnalyzer", "MultimodalFusion"]

mediSync/models/image_analyzer.py

@@ -5,13 +5,28 @@ import torch
 from PIL import Image
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 
+
 class XRayImageAnalyzer:
+    """
+    A class for analyzing medical X-ray images using pre-trained models from Hugging Face.
+
+    This analyzer uses the DeiT (Data-efficient image Transformers) model fine-tuned
+    on chest X-ray images to detect abnormalities.
+    """
 
     def __init__(
         self, model_name="codewithdark/vit-chest-xray", device=None
     ):
+        """
+        Initialize the X-ray image analyzer with a specific pre-trained model.
+
+        Args:
+            model_name (str): The Hugging Face model name to use
+            device (str, optional): Device to run the model on ('cuda' or 'cpu')
+        """
         self.logger = logging.getLogger(__name__)
 
+        # Determine device (CPU or GPU)
         if device is None:
             self.device = "cuda" if torch.cuda.is_available() else "cpu"
         else:
@@ -19,6 +34,8 @@ class XRayImageAnalyzer:
 
         self.logger.info(f"Using device: {self.device}")
 
+        # Load model and feature extractor
+        # Use HF_TOKEN from environment if available (for Hugging Face Spaces)
         hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
         try:
             self.feature_extractor = AutoFeatureExtractor.from_pretrained(
@@ -28,9 +45,10 @@ class XRayImageAnalyzer:
                 model_name, token=hf_token
             )
             self.model.to(self.device)
-            self.model.eval()
+            self.model.eval()  # Set to evaluation mode
             self.logger.info(f"Successfully loaded model: {model_name}")
 
+            # Map labels to more informative descriptions
             self.labels = self.model.config.id2label
 
         except Exception as e:
@@ -38,14 +56,26 @@ class XRayImageAnalyzer:
             raise
 
     def preprocess_image(self, image_path):
+        """
+        Preprocess an X-ray image for model input.
+
+        Args:
+            image_path (str or PIL.Image): Path to image or PIL Image object
+
+        Returns:
+            dict: Processed inputs ready for the model
+        """
         try:
+            # Load image if path is provided
             if isinstance(image_path, str):
                 if not os.path.exists(image_path):
                     raise FileNotFoundError(f"Image file not found: {image_path}")
                 image = Image.open(image_path).convert("RGB")
             else:
+                # Assume it's already a PIL Image
                 image = image_path.convert("RGB")
 
+            # Apply feature extraction
             inputs = self.feature_extractor(images=image, return_tensors="pt")
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
@@ -56,22 +86,42 @@ class XRayImageAnalyzer:
             raise
 
     def analyze(self, image_path, threshold=0.5):
+        """
+        Analyze an X-ray image and detect abnormalities.
+
+        Args:
+            image_path (str or PIL.Image): Path to the X-ray image or PIL Image object
+            threshold (float): Classification threshold for positive findings
+
+        Returns:
+            dict: Analysis results including:
+                - predictions: List of (label, probability) tuples
+                - primary_finding: The most likely abnormality
+                - has_abnormality: Boolean indicating if abnormalities were detected
+                - confidence: Confidence score for the primary finding
+        """
         try:
+            # Preprocess the image
             inputs, original_image = self.preprocess_image(image_path)
 
+            # Run inference
             with torch.no_grad():
                 outputs = self.model(**inputs)
 
+            # Process predictions
             probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
             probabilities = probabilities.cpu().numpy()
 
+            # Get predictions sorted by probability
             predictions = []
             for i, p in enumerate(probabilities):
                 label = self.labels[i]
                 predictions.append((label, float(p)))
 
+            # Sort by probability (descending)
             predictions.sort(key=lambda x: x[1], reverse=True)
 
+            # Determine if there's an abnormality and the primary finding
             normal_idx = [
                 i
                 for i, (label, _) in enumerate(predictions)
@@ -99,6 +149,15 @@ class XRayImageAnalyzer:
             raise
 
     def get_explanation(self, results):
+        """
+        Generate a human-readable explanation of the analysis results.
+
+        Args:
+            results (dict): The results returned by the analyze method
+
+        Returns:
+            str: A text explanation of the findings
+        """
         if not results["has_abnormality"]:
             explanation = (
                 f"The X-ray appears normal with {results['confidence']:.1%} confidence."
@@ -110,17 +169,23 @@ class XRayImageAnalyzer:
                 f"Other potential findings include:\n"
             )
 
+            # Add top 3 other findings (skipping the first one which is primary)
             for label, prob in results["predictions"][1:4]:
-                if prob > 0.05:
+                if prob > 0.05:  # Only include if probability > 5%
                     explanation += f"- {label}: {prob:.1%}\n"
 
         return explanation
 
+
+# Example usage
 if __name__ == "__main__":
+    # Set up logging
     logging.basicConfig(level=logging.INFO)
 
+    # Test on a sample image if available
     analyzer = XRayImageAnalyzer()
 
+    # Check if sample data directory exists
     sample_dir = "../data/sample"
     if os.path.exists(sample_dir) and os.listdir(sample_dir):
         sample_image = os.path.join(sample_dir, os.listdir(sample_dir)[0])

mediSync/models/multimodal_fusion.py

@@ -3,11 +3,30 @@ import logging
 from .image_analyzer import XRayImageAnalyzer
 from .text_analyzer import MedicalReportAnalyzer
 
+
 class MultimodalFusion:
+    """
+    A class for fusing insights from image analysis and text analysis of medical data.
+
+    This fusion approach combines the strengths of both modalities:
+    - Images provide visual evidence of abnormalities
+    - Text reports provide context, history and radiologist interpretations
+
+    The combined analysis provides a more comprehensive understanding than either modality alone.
+    """
 
     def __init__(self, image_model=None, text_model=None, device=None):
+        """
+        Initialize the multimodal fusion module with image and text analyzers.
+
+        Args:
+            image_model (str, optional): Model to use for image analysis
+            text_model (str, optional): Model to use for text analysis
+            device (str, optional): Device to run models on ('cuda' or 'cpu')
+        """
         self.logger = logging.getLogger(__name__)
 
+        # Determine device
         if device is None:
             import torch
 
@@ -17,6 +36,7 @@ class MultimodalFusion:
 
         self.logger.info(f"Using device: {self.device}")
 
+        # Initialize image analyzer
         try:
             self.image_analyzer = XRayImageAnalyzer(
                 model_name=image_model
@@ -29,6 +49,7 @@ class MultimodalFusion:
             self.logger.error(f"Failed to initialize image analyzer: {e}")
             self.image_analyzer = None
 
+        # Initialize text analyzer
         try:
             self.text_analyzer = MedicalReportAnalyzer(
                 classifier_model=text_model if text_model else "medicalai/ClinicalBERT",
@@ -40,6 +61,15 @@ class MultimodalFusion:
             self.text_analyzer = None
 
     def analyze_image(self, image_path):
+        """
+        Analyze a medical image.
+
+        Args:
+            image_path (str): Path to the medical image
+
+        Returns:
+            dict: Image analysis results
+        """
         if not self.image_analyzer:
             self.logger.warning("Image analyzer not available")
             return {"error": "Image analyzer not available"}
@@ -51,6 +81,15 @@ class MultimodalFusion:
             return {"error": str(e)}
 
     def analyze_text(self, text):
+        """
+        Analyze medical report text.
+
+        Args:
+            text (str): Medical report text
+
+        Returns:
+            dict: Text analysis results
+        """
         if not self.text_analyzer:
             self.logger.warning("Text analyzer not available")
             return {"error": "Text analyzer not available"}
@@ -62,24 +101,41 @@ class MultimodalFusion:
             return {"error": str(e)}
 
     def _calculate_agreement_score(self, image_results, text_results):
+        """
+        Calculate agreement score between image and text analyses.
+
+        Args:
+            image_results (dict): Results from image analysis
+            text_results (dict): Results from text analysis
+
+        Returns:
+            float: Agreement score (0-1, where 1 is perfect agreement)
+        """
         try:
+            # Default to neutral agreement
             agreement = 0.5
 
+            # Check if image detected abnormality
             image_abnormal = image_results.get("has_abnormality", False)
 
+            # Check text severity
             text_severity = text_results.get("severity", {}).get("level", "Unknown")
             text_abnormal = text_severity not in ["Normal", "Unknown"]
 
+            # Basic agreement check
             if image_abnormal == text_abnormal:
                 agreement += 0.25
             else:
                 agreement -= 0.25
 
+            # Check if specific findings match
             image_finding = image_results.get("primary_finding", "").lower()
 
+            # Extract problem entities from text
             problems = text_results.get("entities", {}).get("problem", [])
             problem_text = " ".join(problems).lower()
 
+            # Check for common keywords in both
             common_conditions = [
                 "pneumonia",
                 "effusion",
@@ -104,35 +160,51 @@ class MultimodalFusion:
 
                 if in_image and in_text:
                     matching_conditions += 1
-                    agreement += 0.05
+                    agreement += 0.05  # Boost agreement for each matching condition
 
+            # Calculate condition match ratio if any conditions were mentioned
             if total_mentioned > 0:
                 match_ratio = matching_conditions / total_mentioned
                 agreement += match_ratio * 0.2
 
+            # Normalize agreement to 0-1 range
             agreement = max(0, min(1, agreement))
 
             return agreement
 
         except Exception as e:
             self.logger.error(f"Error calculating agreement score: {e}")
-            return 0.5
+            return 0.5  # Return neutral agreement on error
 
     def _get_confidence_weighted_finding(self, image_results, text_results, agreement):
+        """
+        Get the most confident finding weighted by modality confidence.
+
+        Args:
+            image_results (dict): Results from image analysis
+            text_results (dict): Results from text analysis
+            agreement (float): Agreement score between modalities
+
+        Returns:
+            str: Most confident finding
+        """
         try:
             image_finding = image_results.get("primary_finding", "")
             image_confidence = image_results.get("confidence", 0.5)
 
+            # For text, use the most severe problem as primary finding
             problems = text_results.get("entities", {}).get("problem", [])
 
             text_confidence = text_results.get("severity", {}).get("confidence", 0.5)
 
             if not problems:
+                # No problems identified in text
                 if image_confidence > 0.7:
                     return image_finding
                 else:
                     return "No significant findings"
 
+            # Simple confidence-weighted selection
             if image_confidence > text_confidence + 0.2:
                 return image_finding
             elif problems and text_confidence > image_confidence + 0.2:
@@ -142,11 +214,14 @@ class MultimodalFusion:
                     else "Unknown finding"
                 )
             else:
+                # Similar confidence, check agreement
                 if agreement > 0.7:
+                    # High agreement, try to find the specific condition mentioned in both
                     for problem in problems:
                         if problem.lower() in image_finding.lower():
                             return problem
 
+                    # Default to image finding if high confidence
                     if image_confidence > 0.6:
                         return image_finding
                     elif problems:
@@ -154,6 +229,7 @@ class MultimodalFusion:
                     else:
                         return image_finding
                 else:
+                    # Low agreement, include both perspectives
                     if image_finding and problems:
                         return f"{image_finding} (image) / {problems[0]} (report)"
                     elif image_finding:
@@ -168,9 +244,21 @@ class MultimodalFusion:
             return "Unable to determine primary finding"
 
     def _merge_followup_recommendations(self, image_results, text_results):
+        """
+        Merge follow-up recommendations from both modalities.
+
+        Args:
+            image_results (dict): Results from image analysis
+            text_results (dict): Results from text analysis
+
+        Returns:
+            list: Combined follow-up recommendations
+        """
         try:
+            # Get text-based recommendations
             text_recommendations = text_results.get("followup_recommendations", [])
 
+            # Create image-based recommendations based on findings
             image_recommendations = []
 
             if image_results.get("has_abnormality", False):
@@ -198,8 +286,10 @@ class MultimodalFusion:
                         "Consider clinical correlation and potential follow-up."
                     )
 
+            # Combine recommendations, removing duplicates
             all_recommendations = text_recommendations + image_recommendations
 
+            # Remove near-duplicates (similar recommendations)
             unique_recommendations = []
             for rec in all_recommendations:
                 if not any(
@@ -215,34 +305,55 @@ class MultimodalFusion:
             return ["Follow-up recommended based on findings."]
 
     def _is_similar_recommendation(self, rec1, rec2):
+        """Check if two recommendations are semantically similar."""
+        # Convert to lowercase for comparison
         rec1_lower = rec1.lower()
         rec2_lower = rec2.lower()
 
+        # Check for significant overlap
         words1 = set(rec1_lower.split())
         words2 = set(rec2_lower.split())
 
+        # Calculate Jaccard similarity
         intersection = words1.intersection(words2)
         union = words1.union(words2)
 
         similarity = len(intersection) / len(union) if union else 0
 
+        # Consider similar if more than 60% overlap
         return similarity > 0.6
 
     def _get_final_severity(self, image_results, text_results, agreement):
+        """
+        Determine final severity based on both modalities.
+
+        Args:
+            image_results (dict): Results from image analysis
+            text_results (dict): Results from text analysis
+            agreement (float): Agreement score between modalities
+
+        Returns:
+            dict: Final severity assessment
+        """
         try:
+            # Get text-based severity
             text_severity = text_results.get("severity", {})
             text_level = text_severity.get("level", "Unknown")
             text_score = text_severity.get("score", 0)
             text_confidence = text_severity.get("confidence", 0.5)
 
+            # Convert image findings to severity
             image_abnormal = image_results.get("has_abnormality", False)
             image_confidence = image_results.get("confidence", 0.5)
 
+            # Default severity mapping from image
             image_severity = "Normal" if not image_abnormal else "Moderate"
             image_score = 0 if not image_abnormal else 2.0
 
+            # Adjust image severity based on specific findings
             primary_finding = image_results.get("primary_finding", "").lower()
 
+            # Map certain conditions to severity levels
             severity_mapping = {
                 "pneumonia": ("Moderate", 2.5),
                 "pneumothorax": ("Severe", 3.0),
@@ -256,15 +367,19 @@ class MultimodalFusion:
                 "consolidation": ("Moderate", 2.0),
             }
 
+            # Check if any key terms are in the primary finding
             for key, (severity, score) in severity_mapping.items():
                 if key in primary_finding:
                     image_severity = severity
                     image_score = score
                     break
 
+            # Weight based on confidence and agreement
             if agreement > 0.7:
+                # High agreement - weight equally
                 final_score = (image_score + text_score) / 2
             else:
+                # Lower agreement - weight by confidence
                 total_confidence = image_confidence + text_confidence
                 if total_confidence > 0:
                     image_weight = image_confidence / total_confidence
@@ -275,6 +390,7 @@ class MultimodalFusion:
                 else:
                     final_score = (image_score + text_score) / 2
 
+            # Map score to severity level
             severity_levels = {
                 0: "Normal",
                 1: "Mild",
@@ -283,6 +399,7 @@ class MultimodalFusion:
                 4: "Critical",
             }
 
+            # Round to nearest level
             level_index = round(min(4, max(0, final_score)))
             final_level = severity_levels[level_index]
 
@@ -297,30 +414,48 @@ class MultimodalFusion:
             return {"level": "Unknown", "score": 0, "confidence": 0}
 
     def fuse_analyses(self, image_results, text_results):
+        """
+        Fuse the results from image and text analyses.
+
+        Args:
+            image_results (dict): Results from image analysis
+            text_results (dict): Results from text analysis
+
+        Returns:
+            dict: Fused analysis results
+        """
         try:
+            # Calculate agreement between modalities
             agreement = self._calculate_agreement_score(image_results, text_results)
             self.logger.info(f"Agreement score between modalities: {agreement:.2f}")
 
+            # Get confidence-weighted primary finding
             primary_finding = self._get_confidence_weighted_finding(
                 image_results, text_results, agreement
             )
 
+            # Merge follow-up recommendations
             followup = self._merge_followup_recommendations(image_results, text_results)
 
+            # Get final severity assessment
             severity = self._get_final_severity(image_results, text_results, agreement)
 
+            # Create comprehensive findings list
             findings = []
 
+            # Add text-extracted findings
             text_findings = text_results.get("findings", [])
             if text_findings:
                 findings.extend(text_findings)
 
+            # Add primary image finding if not already included
             image_finding = image_results.get("primary_finding", "")
             if image_finding and not any(
                 image_finding.lower() in f.lower() for f in findings
             ):
                 findings.append(f"Image finding: {image_finding}")
 
+            # Create fused result
             fused_result = {
                 "agreement_score": round(agreement, 2),
                 "primary_finding": primary_finding,
@@ -340,11 +475,24 @@ class MultimodalFusion:
             }
 
     def analyze(self, image_path, report_text):
+        """
+        Perform multimodal analysis of medical image and report.
+
+        Args:
+            image_path (str): Path to the medical image
+            report_text (str): Medical report text
+
+        Returns:
+            dict: Fused analysis results
+        """
         try:
+            # Analyze image
             image_results = self.analyze_image(image_path)
 
+            # Analyze text
             text_results = self.analyze_text(report_text)
 
+            # Fuse the analyses
             return self.fuse_analyses(image_results, text_results)
 
         except Exception as e:
@@ -352,9 +500,19 @@ class MultimodalFusion:
             return {"error": str(e)}
 
     def get_explanation(self, fused_results):
+        """
+        Generate a human-readable explanation of the fused analysis.
+
+        Args:
+            fused_results (dict): Results from the fused analysis
+
+        Returns:
+            str: A text explanation of the fused analysis
+        """
         try:
             explanation = []
 
+            # Add overview section
             primary_finding = fused_results.get("primary_finding", "Unknown")
             severity = fused_results.get("severity", {}).get("level", "Unknown")
 
@@ -363,6 +521,7 @@ class MultimodalFusion:
             explanation.append(f"Primary finding: **{primary_finding}**\n")
             explanation.append(f"Severity level: **{severity}**\n")
 
+            # Add agreement information
             agreement = fused_results.get("agreement_score", 0)
             agreement_text = (
                 "High" if agreement > 0.7 else "Moderate" if agreement > 0.4 else "Low"
@@ -372,6 +531,7 @@ class MultimodalFusion:
                 f"Image and text analysis agreement: **{agreement_text}** ({agreement:.0%})\n"
             )
 
+            # Add findings section
             explanation.append("\n## Detailed Findings\n")
             findings = fused_results.get("findings", [])
 
@@ -381,6 +541,7 @@ class MultimodalFusion:
             else:
                 explanation.append("No specific findings detailed.\n")
 
+            # Add follow-up section
             explanation.append("\n## Recommended Follow-up\n")
             followups = fused_results.get("followup_recommendations", [])
 
@@ -390,6 +551,7 @@ class MultimodalFusion:
             else:
                 explanation.append("No specific follow-up recommendations provided.\n")
 
+            # Add confidence note
             confidence = fused_results.get("severity", {}).get("confidence", 0)
             explanation.append(
                 f"\n*Note: This analysis has a confidence level of {confidence:.0%}. "
@@ -402,26 +564,68 @@ class MultimodalFusion:
             self.logger.error(f"Error generating explanation: {e}")
             return "Error generating analysis explanation."
 
+
+# Example usage
 if __name__ == "__main__":
+    # Set up logging
     logging.basicConfig(level=logging.INFO)
 
+    # Test on sample data if available
     import os
 
     fusion = MultimodalFusion()
 
+    # Sample text report
     sample_report = """
     CHEST X-RAY EXAMINATION
-
+    
     CLINICAL HISTORY: 55-year-old male with cough and fever.
-
+    
     FINDINGS: The heart size is at the upper limits of normal. The lungs are clear without focal consolidation, 
     effusion, or pneumothorax. There is mild prominence of the pulmonary vasculature. No pleural effusion is seen. 
     There is a small nodular opacity noted in the right lower lobe measuring approximately 8mm, which is suspicious 
     and warrants further investigation. The mediastinum is unremarkable. The visualized bony structures show no acute abnormalities.
-
+    
     IMPRESSION:
     1. Mild cardiomegaly.
     2. 8mm nodular opacity in the right lower lobe, recommend follow-up CT for further evaluation.
     3. No acute pulmonary parenchymal abnormality.
-
+    
     RECOMMENDATIONS: Follow-up chest CT to further characterize the nodular opacity in the right lower lobe.
+    """
+
+    # Check if sample data directory exists and contains images
+    sample_dir = "../data/sample"
+    if os.path.exists(sample_dir) and os.listdir(sample_dir):
+        sample_image = os.path.join(sample_dir, os.listdir(sample_dir)[0])
+        print(f"Analyzing sample image: {sample_image}")
+
+        # Perform multimodal analysis
+        fused_results = fusion.analyze(sample_image, sample_report)
+        explanation = fusion.get_explanation(fused_results)
+
+        print("\nFused Analysis Results:")
+        print(explanation)
+    else:
+        print("No sample images found. Only analyzing text report.")
+
+        # Analyze just the text
+        text_results = fusion.analyze_text(sample_report)
+
+        print("\nText Analysis Results:")
+        print(
+            f"Severity: {text_results['severity']['level']} (Score: {text_results['severity']['score']})"
+        )
+
+        print("\nKey Findings:")
+        for finding in text_results["findings"]:
+            print(f"- {finding}")
+
+        print("\nEntities:")
+        for category, items in text_results["entities"].items():
+            if items:
+                print(f"- {category.capitalize()}: {', '.join(items)}")
+
+        print("\nFollow-up Recommendations:")
+        for rec in text_results["followup_recommendations"]:
+            print(f"- {rec}")

mediSync/models/text_analyzer.py

@@ -1,357 +1,483 @@
-import logging
-import os
-import re
-
-import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
-
-class MedicalReportAnalyzer:
-
-    def __init__(
-        self,
-        ner_model="samrawal/bert-base-uncased_medical-ner",
-        classifier_model="medicalai/ClinicalBERT",
-        device=None,
-    ):
-        self.logger = logging.getLogger(__name__)
-
-        if device is None:
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        else:
-            self.device = device
-
-        self.logger.info(f"Using device: {self.device}")
-
-        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-
-        try:
-            self.ner_pipeline = pipeline(
-                "token-classification",
-                model=ner_model,
-                aggregation_strategy="simple",
-                device=0 if self.device == "cuda" else -1,
-                token=hf_token,
-            )
-            self.logger.info(f"Successfully loaded NER model: {ner_model}")
-        except Exception as e:
-            self.logger.error(f"Failed to load NER model: {e}")
-            self.ner_pipeline = None
-
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                classifier_model, token=hf_token
-            )
-            self.classifier = AutoModelForSequenceClassification.from_pretrained(
-                classifier_model, token=hf_token
-            )
-            self.classifier.to(self.device)
-            self.classifier.eval()
-            self.logger.info(
-                f"Successfully loaded classifier model: {classifier_model}"
-            )
-        except Exception as e:
-            self.logger.error(f"Failed to load classifier model: {e}")
-            self.classifier = None
-
-        self.severity_levels = {
-            0: "Normal",
-            1: "Mild",
-            2: "Moderate",
-            3: "Severe",
-            4: "Critical",
-        }
-
-        self.finding_severity = {
-            "pneumonia": 3,
-            "fracture": 3,
-            "tumor": 4,
-            "nodule": 2,
-            "mass": 3,
-            "edema": 2,
-            "effusion": 2,
-            "hemorrhage": 3,
-            "opacity": 1,
-            "atelectasis": 2,
-            "pneumothorax": 3,
-            "consolidation": 2,
-            "cardiomegaly": 2,
-        }
-
-    def extract_entities(self, text):
-        if not self.ner_pipeline:
-            self.logger.warning("NER model not available")
-            return {}
-
-        try:
-            entities = self.ner_pipeline(text)
-
-            grouped_entities = {
-                "problem": [],
-                "test": [],
-                "treatment": [],
-                "anatomy": [],
-            }
-
-            for entity in entities:
-                entity_type = entity.get("entity_group", "").lower()
-
-                if entity_type in ["problem", "disease", "condition", "diagnosis"]:
-                    category = "problem"
-                elif entity_type in ["test", "procedure", "examination"]:
-                    category = "test"
-                elif entity_type in ["treatment", "medication", "drug"]:
-                    category = "treatment"
-                elif entity_type in ["body_part", "anatomy", "organ"]:
-                    category = "anatomy"
-                else:
-                    continue
-
-                word = entity.get("word", "")
-                score = entity.get("score", 0)
-
-                if score > 0.7 and word not in grouped_entities[category]:
-                    grouped_entities[category].append(word)
-
-            return grouped_entities
-
-        except Exception as e:
-            self.logger.error(f"Error extracting entities: {e}")
-            return {}
-
-    def assess_severity(self, text):
-        if not self.classifier:
-            self.logger.warning("Classifier model not available")
-            return {"level": "Unknown", "score": 0.0}
-
-        try:
-            severity_score = 0
-            confidence = 0.5
-
-            severe_keywords = [
-                "severe",
-                "critical",
-                "urgent",
-                "emergency",
-                "immediate attention",
-            ]
-            moderate_keywords = ["moderate", "concerning", "follow-up", "monitor"]
-            mild_keywords = ["mild", "minimal", "slight", "minor"]
-            normal_keywords = [
-                "normal",
-                "unremarkable",
-                "no abnormalities",
-                "within normal limits",
-            ]
-
-            text_lower = text.lower()
-            severe_count = sum(text_lower.count(word) for word in severe_keywords)
-            moderate_count = sum(text_lower.count(word) for word in moderate_keywords)
-            mild_count = sum(text_lower.count(word) for word in mild_keywords)
-            normal_count = sum(text_lower.count(word) for word in normal_keywords)
-
-            if severe_count > 0:
-                severity_score += min(severe_count, 2) * 1.5
-                confidence += 0.1
-            if moderate_count > 0:
-                severity_score += min(moderate_count, 3) * 0.75
-                confidence += 0.05
-            if mild_count > 0:
-                severity_score += min(mild_count, 3) * 0.25
-                confidence += 0.05
-            if normal_count > 0:
-                severity_score -= min(normal_count, 3) * 0.75
-                confidence += 0.1
-
-            for finding, level in self.finding_severity.items():
-                if finding in text_lower:
-                    severity_score += level * 0.5
-                    confidence += 0.05
-
-            severity_score = max(0, min(4, severity_score))
-            severity_level = int(round(severity_score))
-
-            severity = self.severity_levels.get(severity_level, "Moderate")
-
-            confidence = min(0.95, confidence)
-
-            return {
-                "level": severity,
-                "score": round(severity_score, 1),
-                "confidence": round(confidence, 2),
-            }
-
-        except Exception as e:
-            self.logger.error(f"Error assessing severity: {e}")
-            return {"level": "Unknown", "score": 0.0, "confidence": 0.0}
-
-    def extract_findings(self, text):
-        try:
-            sentences = re.split(r"[.!?]\s+", text)
-            findings = []
-
-            finding_markers = [
-                "finding",
-                "observed",
-                "noted",
-                "shows",
-                "reveals",
-                "demonstrates",
-                "indicates",
-                "evident",
-                "apparent",
-                "consistent with",
-                "suggestive of",
-            ]
-
-            negation_markers = ["no", "not", "none", "negative", "without", "denies"]
-
-            for sentence in sentences:
-                if len(sentence.split()) < 3:
-                    continue
-
-                sentence = sentence.strip()
-
-                contains_finding_marker = any(
-                    marker in sentence.lower() for marker in finding_markers
-                )
-
-                contains_negation = any(
-                    marker in sentence.lower().split() for marker in negation_markers
-                )
-
-                if contains_finding_marker or (
-                    contains_negation
-                    and any(
-                        term in sentence.lower()
-                        for term in self.finding_severity.keys()
-                    )
-                ):
-                    findings.append(sentence)
-
-            return findings
-
-        except Exception as e:
-            self.logger.error(f"Error extracting findings: {e}")
-            return []
-
-    def suggest_followup(self, text, entities, severity):
-        try:
-            followups = []
-
-            severity_level = severity.get("level", "Unknown")
-            severity_score = severity.get("score", 0)
-
-            problems = entities.get("problem", [])
-
-            followup_mentioned = any(
-                phrase in text.lower()
-                for phrase in [
-                    "follow up",
-                    "follow-up",
-                    "followup",
-                    "return",
-                    "refer",
-                    "consult",
-                ]
-            )
-
-            if severity_level == "Critical":
-                followups.append("Immediate specialist consultation recommended.")
-
-            elif severity_level == "Severe":
-                followups.append("Prompt follow-up with specialist is recommended.")
-
-                for problem in problems:
-                    if "pneumonia" in problem.lower():
-                        followups.append(
-                            "Consider antibiotic therapy and close monitoring."
-                        )
-                    elif "fracture" in problem.lower():
-                        followups.append(
-                            "Orthopedic consultation for treatment planning."
-                        )
-                    elif "mass" in problem.lower() or "tumor" in problem.lower():
-                        followups.append(
-                            "Further imaging and possible biopsy recommended."
-                        )
-
-            elif severity_level == "Moderate":
-                followups.append("Follow-up with primary care physician recommended.")
-                if not followup_mentioned and problems:
-                    followups.append(
-                        "Consider additional imaging or tests for further evaluation."
-                    )
-
-            elif severity_level == "Mild":
-                if problems:
-                    followups.append(
-                        "Routine follow-up with primary care physician as needed."
-                    )
-                else:
-                    followups.append("No immediate follow-up required.")
-
-            else:
-                followups.append(
-                    "No specific follow-up indicated based on this report."
-                )
-
-            for critical_term in ["mass", "tumor", "nodule", "opacity"]:
-                if (
-                    critical_term in text.lower()
-                    and "follow-up" not in " ".join(followups).lower()
-                ):
-                    followups.append(
-                        f"Follow-up imaging recommended to monitor {critical_term}."
-                    )
-                    break
-
-            return followups
-
-        except Exception as e:
-            self.logger.error(f"Error suggesting follow-up: {e}")
-            return ["Unable to generate follow-up recommendations."]
-
-    def analyze(self, text):
-        try:
-            entities = self.extract_entities(text)
-
-            severity = self.assess_severity(text)
-
-            findings = self.extract_findings(text)
-
-            followups = self.suggest_followup(text, entities, severity)
-
-            report = {
-                "entities": entities,
-                "severity": severity,
-                "findings": findings,
-                "followup_recommendations": followups,
-            }
-
-            return report
-
-        except Exception as e:
-            self.logger.error(f"Error analyzing report: {e}")
-            return {"error": str(e)}
-
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
-
-    analyzer = MedicalReportAnalyzer()
-
-    sample_report = """
-    CHEST X-RAY EXAMINATION
-
-    CLINICAL HISTORY: 55-year-old male with cough and fever.
-
-    FINDINGS: The heart size is at the upper limits of normal. The lungs are clear without focal consolidation, 
-    effusion, or pneumothorax. There is mild prominence of the pulmonary vasculature. No pleural effusion is seen. 
-    There is a small nodular opacity noted in the right lower lobe measuring approximately 8mm, which is suspicious 
-    and warrants further investigation. The mediastinum is unremarkable. The visualized bony structures show no acute abnormalities.
-
-    IMPRESSION:
-    1. Mild cardiomegaly.
-    2. 8mm nodular opacity in the right lower lobe, recommend follow-up CT for further evaluation.
-    3. No acute pulmonary parenchymal abnormality.
-
-    RECOMMENDATIONS: Follow-up chest CT to further characterize the nodular opacity in the right lower lobe.
+import logging
+import os
+import re
+
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+
+
+class MedicalReportAnalyzer:
+    """
+    A class for analyzing medical text reports using pre-trained NLP models from Hugging Face.
+
+    This analyzer can:
+    1. Extract medical entities (conditions, treatments, tests)
+    2. Classify report severity
+    3. Extract key findings
+    4. Identify suggested follow-up actions
+    """
+
+    def __init__(
+        self,
+        ner_model="samrawal/bert-base-uncased_medical-ner",
+        classifier_model="medicalai/ClinicalBERT",
+        device=None,
+    ):
+        """
+        Initialize the text analyzer with specific pre-trained models.
+
+        Args:
+            ner_model (str): Model for named entity recognition
+            classifier_model (str): Model for text classification
+            device (str, optional): Device to run models on ('cuda' or 'cpu')
+        """
+        self.logger = logging.getLogger(__name__)
+
+        # Determine device
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+
+        self.logger.info(f"Using device: {self.device}")
+
+        # Use HF_TOKEN from environment if available (for Hugging Face Spaces)
+        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+
+        # Load NER model for entity extraction
+        try:
+            self.ner_pipeline = pipeline(
+                "token-classification",
+                model=ner_model,
+                aggregation_strategy="simple",
+                device=0 if self.device == "cuda" else -1,
+                token=hf_token,
+            )
+            self.logger.info(f"Successfully loaded NER model: {ner_model}")
+        except Exception as e:
+            self.logger.error(f"Failed to load NER model: {e}")
+            self.ner_pipeline = None
+
+        # Load classifier model for severity assessment
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                classifier_model, token=hf_token
+            )
+            self.classifier = AutoModelForSequenceClassification.from_pretrained(
+                classifier_model, token=hf_token
+            )
+            self.classifier.to(self.device)
+            self.classifier.eval()
+            self.logger.info(
+                f"Successfully loaded classifier model: {classifier_model}"
+            )
+        except Exception as e:
+            self.logger.error(f"Failed to load classifier model: {e}")
+            self.classifier = None
+
+        # Severity levels mapping
+        self.severity_levels = {
+            0: "Normal",
+            1: "Mild",
+            2: "Moderate",
+            3: "Severe",
+            4: "Critical",
+        }
+
+        # Common medical findings and their severity levels
+        self.finding_severity = {
+            "pneumonia": 3,
+            "fracture": 3,
+            "tumor": 4,
+            "nodule": 2,
+            "mass": 3,
+            "edema": 2,
+            "effusion": 2,
+            "hemorrhage": 3,
+            "opacity": 1,
+            "atelectasis": 2,
+            "pneumothorax": 3,
+            "consolidation": 2,
+            "cardiomegaly": 2,
+        }
+
+    def extract_entities(self, text):
+        """
+        Extract medical entities from the report text.
+
+        Args:
+            text (str): Medical report text
+
+        Returns:
+            dict: Dictionary of entity lists by category
+        """
+        if not self.ner_pipeline:
+            self.logger.warning("NER model not available")
+            return {}
+
+        try:
+            # Run NER
+            entities = self.ner_pipeline(text)
+
+            # Group entities by type
+            grouped_entities = {
+                "problem": [],  # Medical conditions
+                "test": [],  # Tests/procedures
+                "treatment": [],  # Treatments/medications
+                "anatomy": [],  # Anatomical locations
+            }
+
+            for entity in entities:
+                entity_type = entity.get("entity_group", "").lower()
+
+                # Map entity types to our categories
+                if entity_type in ["problem", "disease", "condition", "diagnosis"]:
+                    category = "problem"
+                elif entity_type in ["test", "procedure", "examination"]:
+                    category = "test"
+                elif entity_type in ["treatment", "medication", "drug"]:
+                    category = "treatment"
+                elif entity_type in ["body_part", "anatomy", "organ"]:
+                    category = "anatomy"
+                else:
+                    continue  # Skip other entity types
+
+                word = entity.get("word", "")
+                score = entity.get("score", 0)
+
+                # Only include if confidence is reasonable
+                if score > 0.7 and word not in grouped_entities[category]:
+                    grouped_entities[category].append(word)
+
+            return grouped_entities
+
+        except Exception as e:
+            self.logger.error(f"Error extracting entities: {e}")
+            return {}
+
+    def assess_severity(self, text):
+        """
+        Assess the severity level of the medical report.
+
+        Args:
+            text (str): Medical report text
+
+        Returns:
+            dict: Severity assessment including level and confidence
+        """
+        if not self.classifier:
+            self.logger.warning("Classifier model not available")
+            return {"level": "Unknown", "score": 0.0}
+
+        try:
+            # Use rule-based approach along with model
+            severity_score = 0
+            confidence = 0.5  # Start with neutral confidence
+
+            # Check for severe keywords
+            severe_keywords = [
+                "severe",
+                "critical",
+                "urgent",
+                "emergency",
+                "immediate attention",
+            ]
+            moderate_keywords = ["moderate", "concerning", "follow-up", "monitor"]
+            mild_keywords = ["mild", "minimal", "slight", "minor"]
+            normal_keywords = [
+                "normal",
+                "unremarkable",
+                "no abnormalities",
+                "within normal limits",
+            ]
+
+            # Count keyword occurrences
+            text_lower = text.lower()
+            severe_count = sum(text_lower.count(word) for word in severe_keywords)
+            moderate_count = sum(text_lower.count(word) for word in moderate_keywords)
+            mild_count = sum(text_lower.count(word) for word in mild_keywords)
+            normal_count = sum(text_lower.count(word) for word in normal_keywords)
+
+            # Adjust severity based on keyword counts
+            if severe_count > 0:
+                severity_score += min(severe_count, 2) * 1.5
+                confidence += 0.1
+            if moderate_count > 0:
+                severity_score += min(moderate_count, 3) * 0.75
+                confidence += 0.05
+            if mild_count > 0:
+                severity_score += min(mild_count, 3) * 0.25
+                confidence += 0.05
+            if normal_count > 0:
+                severity_score -= min(normal_count, 3) * 0.75
+                confidence += 0.1
+
+            # Check for specific medical findings
+            for finding, level in self.finding_severity.items():
+                if finding in text_lower:
+                    severity_score += level * 0.5
+                    confidence += 0.05
+
+            # Normalize severity score to 0-4 range
+            severity_score = max(0, min(4, severity_score))
+            severity_level = int(round(severity_score))
+
+            # Map to severity level
+            severity = self.severity_levels.get(severity_level, "Moderate")
+
+            # Cap confidence at 0.95
+            confidence = min(0.95, confidence)
+
+            return {
+                "level": severity,
+                "score": round(severity_score, 1),
+                "confidence": round(confidence, 2),
+            }
+
+        except Exception as e:
+            self.logger.error(f"Error assessing severity: {e}")
+            return {"level": "Unknown", "score": 0.0, "confidence": 0.0}
+
+    def extract_findings(self, text):
+        """
+        Extract key clinical findings from the report.
+
+        Args:
+            text (str): Medical report text
+
+        Returns:
+            list: List of key findings
+        """
+        try:
+            # Split text into sentences
+            sentences = re.split(r"[.!?]\s+", text)
+            findings = []
+
+            # Key phrases that often introduce findings
+            finding_markers = [
+                "finding",
+                "observed",
+                "noted",
+                "shows",
+                "reveals",
+                "demonstrates",
+                "indicates",
+                "evident",
+                "apparent",
+                "consistent with",
+                "suggestive of",
+            ]
+
+            # Negative markers
+            negation_markers = ["no", "not", "none", "negative", "without", "denies"]
+
+            for sentence in sentences:
+                # Skip very short sentences
+                if len(sentence.split()) < 3:
+                    continue
+
+                sentence = sentence.strip()
+
+                # Check if this sentence likely contains a finding
+                contains_finding_marker = any(
+                    marker in sentence.lower() for marker in finding_markers
+                )
+
+                # Check for negation
+                contains_negation = any(
+                    marker in sentence.lower().split() for marker in negation_markers
+                )
+
+                # Only include positive findings or explicitly negated findings that are important
+                if contains_finding_marker or (
+                    contains_negation
+                    and any(
+                        term in sentence.lower()
+                        for term in self.finding_severity.keys()
+                    )
+                ):
+                    findings.append(sentence)
+
+            return findings
+
+        except Exception as e:
+            self.logger.error(f"Error extracting findings: {e}")
+            return []
+
+    def suggest_followup(self, text, entities, severity):
+        """
+        Suggest follow-up actions based on report analysis.
+
+        Args:
+            text (str): Medical report text
+            entities (dict): Extracted entities
+            severity (dict): Severity assessment
+
+        Returns:
+            list: Suggested follow-up actions
+        """
+        try:
+            followups = []
+
+            # Base recommendations on severity
+            severity_level = severity.get("level", "Unknown")
+            severity_score = severity.get("score", 0)
+
+            # Extract problems from entities
+            problems = entities.get("problem", [])
+
+            # Check if follow-up is already mentioned in the text
+            followup_mentioned = any(
+                phrase in text.lower()
+                for phrase in [
+                    "follow up",
+                    "follow-up",
+                    "followup",
+                    "return",
+                    "refer",
+                    "consult",
+                ]
+            )
+
+            # Default recommendations based on severity
+            if severity_level == "Critical":
+                followups.append("Immediate specialist consultation recommended.")
+
+            elif severity_level == "Severe":
+                followups.append("Prompt follow-up with specialist is recommended.")
+
+                # Add specific recommendations for common severe conditions
+                for problem in problems:
+                    if "pneumonia" in problem.lower():
+                        followups.append(
+                            "Consider antibiotic therapy and close monitoring."
+                        )
+                    elif "fracture" in problem.lower():
+                        followups.append(
+                            "Orthopedic consultation for treatment planning."
+                        )
+                    elif "mass" in problem.lower() or "tumor" in problem.lower():
+                        followups.append(
+                            "Further imaging and possible biopsy recommended."
+                        )
+
+            elif severity_level == "Moderate":
+                followups.append("Follow-up with primary care physician recommended.")
+                if not followup_mentioned and problems:
+                    followups.append(
+                        "Consider additional imaging or tests for further evaluation."
+                    )
+
+            elif severity_level == "Mild":
+                if problems:
+                    followups.append(
+                        "Routine follow-up with primary care physician as needed."
+                    )
+                else:
+                    followups.append("No immediate follow-up required.")
+
+            else:  # Normal
+                followups.append(
+                    "No specific follow-up indicated based on this report."
+                )
+
+            # Check for specific findings that always need follow-up
+            for critical_term in ["mass", "tumor", "nodule", "opacity"]:
+                if (
+                    critical_term in text.lower()
+                    and "follow-up" not in " ".join(followups).lower()
+                ):
+                    followups.append(
+                        f"Follow-up imaging recommended to monitor {critical_term}."
+                    )
+                    break
+
+            return followups
+
+        except Exception as e:
+            self.logger.error(f"Error suggesting follow-up: {e}")
+            return ["Unable to generate follow-up recommendations."]
+
+    def analyze(self, text):
+        """
+        Perform comprehensive analysis of medical report text.
+
+        Args:
+            text (str): Medical report text
+
+        Returns:
+            dict: Complete analysis results
+        """
+        try:
+            # Extract entities
+            entities = self.extract_entities(text)
+
+            # Assess severity
+            severity = self.assess_severity(text)
+
+            # Extract key findings
+            findings = self.extract_findings(text)
+
+            # Generate follow-up suggestions
+            followups = self.suggest_followup(text, entities, severity)
+
+            # Create detailed report
+            report = {
+                "entities": entities,
+                "severity": severity,
+                "findings": findings,
+                "followup_recommendations": followups,
+            }
+
+            return report
+
+        except Exception as e:
+            self.logger.error(f"Error analyzing report: {e}")
+            return {"error": str(e)}
+
+
+# Example usage
+if __name__ == "__main__":
+    # Set up logging
+    logging.basicConfig(level=logging.INFO)
+
+    # Test on a sample report
+    analyzer = MedicalReportAnalyzer()
+
+    sample_report = """
+    CHEST X-RAY EXAMINATION
+    
+    CLINICAL HISTORY: 55-year-old male with cough and fever.
+    
+    FINDINGS: The heart size is at the upper limits of normal. The lungs are clear without focal consolidation, 
+    effusion, or pneumothorax. There is mild prominence of the pulmonary vasculature. No pleural effusion is seen. 
+    There is a small nodular opacity noted in the right lower lobe measuring approximately 8mm, which is suspicious 
+    and warrants further investigation. The mediastinum is unremarkable. The visualized bony structures show no acute abnormalities.
+    
+    IMPRESSION:
+    1. Mild cardiomegaly.
+    2. 8mm nodular opacity in the right lower lobe, recommend follow-up CT for further evaluation.
+    3. No acute pulmonary parenchymal abnormality.
+    
+    RECOMMENDATIONS: Follow-up chest CT to further characterize the nodular opacity in the right lower lobe.
+    """
+
+    results = analyzer.analyze(sample_report)
+
+    print("\nMedical Report Analysis:")
+    print(
+        f"\nSeverity: {results['severity']['level']} (Score: {results['severity']['score']})"
+    )
+
+    print("\nKey Findings:")
+    for finding in results["findings"]:
+        print(f"- {finding}")
+
+    print("\nEntities:")
+    for category, items in results["entities"].items():
+        if items:
+            print(f"- {category.capitalize()}: {', '.join(items)}")
+
+    print("\nFollow-up Recommendations:")
+    for rec in results["followup_recommendations"]:
+        print(f"- {rec}")

mediSync/utils/__init__.py

@@ -1,27 +1,38 @@
-from .preprocessing import (
-    enhance_xray_image,
-    extract_measurements,
-    extract_sections,
-    normalize_report_text,
-    preprocess_image,
-)
-from .visualization import (
-    create_heatmap_overlay,
-    figure_to_base64,
-    plot_image_prediction,
-    plot_multimodal_results,
-    plot_report_entities,
-)
-
-__all__ = [
-    "preprocess_image",
-    "normalize_report_text",
-    "enhance_xray_image",
-    "extract_sections",
-    "extract_measurements",
-    "plot_image_prediction",
-    "plot_report_entities",
-    "plot_multimodal_results",
-    "create_heatmap_overlay",
-    "figure_to_base64",
-]
+"""
+U-HRES: Utils Module
+===================
+
+This module contains utility functions for the U-HRES system:
+
+1. preprocessing: Functions for preprocessing images and text
+2. visualization: Functions for visualizing analysis results
+3. download_samples: Functions for downloading sample data
+"""
+
+from .preprocessing import (
+    enhance_xray_image,
+    extract_measurements,
+    extract_sections,
+    normalize_report_text,
+    preprocess_image,
+)
+from .visualization import (
+    create_heatmap_overlay,
+    figure_to_base64,
+    plot_image_prediction,
+    plot_multimodal_results,
+    plot_report_entities,
+)
+
+__all__ = [
+    "preprocess_image",
+    "normalize_report_text",
+    "enhance_xray_image",
+    "extract_sections",
+    "extract_measurements",
+    "plot_image_prediction",
+    "plot_report_entities",
+    "plot_multimodal_results",
+    "create_heatmap_overlay",
+    "figure_to_base64",
+]

mediSync/utils/download_samples.py

@@ -1,101 +1,135 @@
-import logging
-import urllib.request
-from pathlib import Path
-
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-)
-logger = logging.getLogger(__name__)
-
-SAMPLE_IMAGES = [
-    {
-        "url": "https://prod-images-static.radiopaedia.org/images/53448173/322830a37f0fa0852773ca2db3e8d8_big_gallery.jpeg",
-        "filename": "normal_chest_xray.jpg",
-        "description": "Normal chest X-ray",
-    },
-    {
-        "url": "https://prod-images-static.radiopaedia.org/images/52465460/e4d8791bd7502ab72af8d9e5c322db_big_gallery.jpg",
-        "filename": "pneumonia_xray.jpg",
-        "description": "X-ray with pneumonia",
-    },
-    {
-        "url": "https://prod-images-static.radiopaedia.org/images/556520/cf17c05750adb04b2a6e23afb47c7d_big_gallery.jpg",
-        "filename": "cardiomegaly_xray.jpg",
-        "description": "X-ray with cardiomegaly",
-    },
-    {
-        "url": "https://prod-images-static.radiopaedia.org/images/19972291/41eed1a2cdad06d26c3f415a6ed65a_big_gallery.jpeg",
-        "filename": "nodule_xray.jpg",
-        "description": "X-ray with lung nodule",
-    },
-]
-
-def download_sample_images(output_dir="data/sample"):
-    script_dir = Path(__file__).resolve().parent.parent
-
-    output_path = script_dir / output_dir
-    output_path.mkdir(parents=True, exist_ok=True)
-
-    downloaded_paths = []
-
-    for image in SAMPLE_IMAGES:
-        try:
-            filename = image["filename"]
-            url = image["url"]
-            output_file = output_path / filename
-
-            if output_file.exists():
-                logger.info(f"File already exists: {output_file}")
-                downloaded_paths.append(str(output_file))
-                continue
-
-            logger.info(f"Downloading {url} to {output_file}")
-
-            opener = urllib.request.build_opener()
-            opener.addheaders = [("User-Agent", "Mozilla/5.0")]
-            urllib.request.install_opener(opener)
-
-            urllib.request.urlretrieve(url, output_file)
-
-            logger.info(f"Successfully downloaded {filename}")
-            downloaded_paths.append(str(output_file))
-
-        except Exception as e:
-            logger.error(f"Error downloading {image['url']}: {e}")
-
-    logger.info(
-        f"Downloaded {len(downloaded_paths)} out of {len(SAMPLE_IMAGES)} images"
-    )
-    return downloaded_paths
-
-def create_sample_info_file(output_dir="data/sample"):
-    script_dir = Path(__file__).resolve().parent.parent
-
-    output_path = script_dir / output_dir
-    info_file = output_path / "sample_info.txt"
-
-    with open(info_file, "w") as f:
-        f.write("# Sample X-ray Images\n\n")
-
-        for image in SAMPLE_IMAGES:
-            f.write(f"## {image['filename']}\n")
-            f.write(f"Description: {image['description']}\n")
-            f.write(f"Source: {image['url']}\n\n")
-
-        f.write(
-            "\nThese images are used for testing and demonstration purposes only.\n"
-        )
-        f.write(
-            "Please note that these images are from public medical education sources.\n"
-        )
-        f.write("Do not use for clinical decision making.\n")
-
-    logger.info(f"Created sample info file: {info_file}")
-
-if __name__ == "__main__":
-    downloaded_paths = download_sample_images()
-
-    create_sample_info_file()
-
-    print(f"Downloaded {len(downloaded_paths)} sample images.")
-    print("Run the application with: python app.py")
+import logging
+import urllib.request
+from pathlib import Path
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+# Sample X-ray image URLs (from public sources)
+SAMPLE_IMAGES = [
+    # Normal chest X-ray
+    {
+        "url": "https://prod-images-static.radiopaedia.org/images/53448173/322830a37f0fa0852773ca2db3e8d8_big_gallery.jpeg",
+        "filename": "normal_chest_xray.jpg",
+        "description": "Normal chest X-ray",
+    },
+    # X-ray with pneumonia
+    {
+        "url": "https://prod-images-static.radiopaedia.org/images/52465460/e4d8791bd7502ab72af8d9e5c322db_big_gallery.jpg",
+        "filename": "pneumonia_xray.jpg",
+        "description": "X-ray with pneumonia",
+    },
+    # X-ray with cardiomegaly
+    {
+        "url": "https://prod-images-static.radiopaedia.org/images/556520/cf17c05750adb04b2a6e23afb47c7d_big_gallery.jpg",
+        "filename": "cardiomegaly_xray.jpg",
+        "description": "X-ray with cardiomegaly",
+    },
+    # X-ray with lung nodule
+    {
+        "url": "https://prod-images-static.radiopaedia.org/images/19972291/41eed1a2cdad06d26c3f415a6ed65a_big_gallery.jpeg",
+        "filename": "nodule_xray.jpg",
+        "description": "X-ray with lung nodule",
+    },
+]
+
+
+def download_sample_images(output_dir="data/sample"):
+    """
+    Download sample X-ray images for testing.
+
+    Args:
+        output_dir (str): Directory to save images
+
+    Returns:
+        list: Paths to downloaded images
+    """
+    # Get the directory of the script
+    script_dir = Path(__file__).resolve().parent.parent
+
+    # Create output directory if it doesn't exist
+    output_path = script_dir / output_dir
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    downloaded_paths = []
+
+    for image in SAMPLE_IMAGES:
+        try:
+            filename = image["filename"]
+            url = image["url"]
+            output_file = output_path / filename
+
+            # Skip if file already exists
+            if output_file.exists():
+                logger.info(f"File already exists: {output_file}")
+                downloaded_paths.append(str(output_file))
+                continue
+
+            # Download the image
+            logger.info(f"Downloading {url} to {output_file}")
+
+            # Set a user agent to avoid blocking
+            opener = urllib.request.build_opener()
+            opener.addheaders = [("User-Agent", "Mozilla/5.0")]
+            urllib.request.install_opener(opener)
+
+            # Download the file
+            urllib.request.urlretrieve(url, output_file)
+
+            logger.info(f"Successfully downloaded {filename}")
+            downloaded_paths.append(str(output_file))
+
+        except Exception as e:
+            logger.error(f"Error downloading {image['url']}: {e}")
+
+    logger.info(
+        f"Downloaded {len(downloaded_paths)} out of {len(SAMPLE_IMAGES)} images"
+    )
+    return downloaded_paths
+
+
+def create_sample_info_file(output_dir="data/sample"):
+    """
+    Create a text file with information about the sample images.
+
+    Args:
+        output_dir (str): Directory with sample images
+    """
+    # Get the directory of the script
+    script_dir = Path(__file__).resolve().parent.parent
+
+    # Output path
+    output_path = script_dir / output_dir
+    info_file = output_path / "sample_info.txt"
+
+    with open(info_file, "w") as f:
+        f.write("# Sample X-ray Images\n\n")
+
+        for image in SAMPLE_IMAGES:
+            f.write(f"## {image['filename']}\n")
+            f.write(f"Description: {image['description']}\n")
+            f.write(f"Source: {image['url']}\n\n")
+
+        f.write(
+            "\nThese images are used for testing and demonstration purposes only.\n"
+        )
+        f.write(
+            "Please note that these images are from public medical education sources.\n"
+        )
+        f.write("Do not use for clinical decision making.\n")
+
+    logger.info(f"Created sample info file: {info_file}")
+
+
+if __name__ == "__main__":
+    # Download sample images
+    downloaded_paths = download_sample_images()
+
+    # Create info file
+    create_sample_info_file()
+
+    print(f"Downloaded {len(downloaded_paths)} sample images.")
+    print("Run the application with: python app.py")

mediSync/utils/preprocessing.py

@@ -1,172 +1,262 @@
-import logging
-import os
-import re
-
-import cv2
-from PIL import Image
-
-logger = logging.getLogger(__name__)
-
-def preprocess_image(image_path, target_size=(224, 224)):
-    try:
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image file not found: {image_path}")
-
-        image = Image.open(image_path)
-
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        image = image.resize(target_size, Image.LANCZOS)
-
-        return image
-
-    except Exception as e:
-        logger.error(f"Error preprocessing image: {e}")
-        raise
-
-def enhance_xray_image(image_path, output_path=None, clahe_clip=2.0, clahe_grid=(8, 8)):
-    try:
-        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
-
-        if img is None:
-            raise ValueError(f"Failed to read image: {image_path}")
-
-        clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=clahe_grid)
-
-        enhanced = clahe.apply(img)
-
-        if output_path:
-            cv2.imwrite(output_path, enhanced)
-            return output_path
-        else:
-            return enhanced
-
-    except Exception as e:
-        logger.error(f"Error enhancing X-ray image: {e}")
-        raise
-
-def normalize_report_text(text):
-    try:
-        text = re.sub(r"\s+", " ", text)
-
-        section_patterns = {
-            r"(?i)clinical\s*(?:history|indication)": "CLINICAL HISTORY:",
-            r"(?i)technique": "TECHNIQUE:",
-            r"(?i)comparison": "COMPARISON:",
-            r"(?i)findings": "FINDINGS:",
-            r"(?i)impression": "IMPRESSION:",
-            r"(?i)recommendation": "RECOMMENDATION:",
-            r"(?i)comment": "COMMENT:",
-        }
-
-        for pattern, replacement in section_patterns.items():
-            text = re.sub(pattern + r"\s*:", replacement, text)
-
-        abbrev_patterns = {
-            r"(?i)\bw\/\b": "with",
-            r"(?i)\bw\/o\b": "without",
-            r"(?i)\bs\/p\b": "status post",
-            r"(?i)\bc\/w\b": "consistent with",
-            r"(?i)\br\/o\b": "rule out",
-            r"(?i)\bhx\b": "history",
-            r"(?i)\bdx\b": "diagnosis",
-            r"(?i)\btx\b": "treatment",
-        }
-
-        for pattern, replacement in abbrev_patterns.items():
-            text = re.sub(pattern, replacement, text)
-
-        return text.strip()
-
-    except Exception as e:
-        logger.error(f"Error normalizing report text: {e}")
-        return text
-
-def extract_sections(text):
-    try:
-        normalized_text = normalize_report_text(text)
-
-        section_headers = [
-            "CLINICAL HISTORY:",
-            "TECHNIQUE:",
-            "COMPARISON:",
-            "FINDINGS:",
-            "IMPRESSION:",
-            "RECOMMENDATION:",
-        ]
-
-        sections = {}
-        current_section = "PREAMBLE"
-        sections[current_section] = []
-
-        for line in normalized_text.split("\n"):
-            section_found = False
-
-            for header in section_headers:
-                if header in line:
-                    current_section = header.rstrip(":")
-                    sections[current_section] = []
-                    section_found = True
-                    content = line.split(header, 1)[1].strip()
-                    if content:
-                        sections[current_section].append(content)
-                    break
-
-            if not section_found and current_section:
-                sections[current_section].append(line)
-
-        for section, lines in sections.items():
-            sections[section] = " ".join(lines).strip()
-
-        sections = {k: v for k, v in sections.items() if v}
-
-        return sections
-
-    except Exception as e:
-        logger.error(f"Error extracting sections: {e}")
-        return {"FULL_TEXT": text}
-
-def extract_measurements(text):
-    try:
-        size_pattern = r"(\d+(?:\.\d+)?(?:\s*[x×]\s*\d+(?:\.\d+)?)?(?:\s*[x×]\s*\d+(?:\.\d+)?)?)\s*(mm|cm|mm2|cm2|mm3|cm3|ml|cc)"
-
-        context_pattern = (
-            r"([A-Za-z\s]+(?:mass|nodule|effusion|opacity|lesion|tumor|cyst|structure|area|region)[A-Za-z\s]*)"
-            + size_pattern
-        )
-
-        context_measurements = []
-        for match in re.finditer(context_pattern, text, re.IGNORECASE):
-            context, size, unit = match.groups()
-            context_measurements.append((context.strip(), size, unit))
-
-        all_measurements = []
-        for match in re.finditer(size_pattern, text):
-            size, unit = match.groups()
-            all_measurements.append((size, unit))
-
-        return context_measurements
-
-    except Exception as e:
-        logger.error(f"Error extracting measurements: {e}")
-        return []
-
-def prepare_sample_batch(image_paths, reports=None, target_size=(224, 224)):
-    try:
-        processed_images = []
-        processed_reports = []
-
-        for i, image_path in enumerate(image_paths):
-            image = preprocess_image(image_path, target_size)
-            processed_images.append(image)
-
-            if reports and i < len(reports):
-                normalized_report = normalize_report_text(reports[i])
-                processed_reports.append(normalized_report)
-
-        return processed_images, processed_reports if reports else None
-
-    except Exception as e:
-        logger.error(f"Error preparing sample batch: {e}")
-        raise
+import logging
+import os
+import re
+
+import cv2
+from PIL import Image
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+
+def preprocess_image(image_path, target_size=(224, 224)):
+    """
+    Preprocess X-ray image for model input.
+
+    Args:
+        image_path (str): Path to the X-ray image
+        target_size (tuple): Target size for resizing
+
+    Returns:
+        PIL.Image: Preprocessed image
+    """
+    try:
+        # Check if file exists
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+
+        # Load image
+        image = Image.open(image_path)
+
+        # Convert grayscale to RGB if needed
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        # Resize image
+        image = image.resize(target_size, Image.LANCZOS)
+
+        return image
+
+    except Exception as e:
+        logger.error(f"Error preprocessing image: {e}")
+        raise
+
+
+def enhance_xray_image(image_path, output_path=None, clahe_clip=2.0, clahe_grid=(8, 8)):
+    """
+    Enhance X-ray image contrast using CLAHE (Contrast Limited Adaptive Histogram Equalization).
+
+    Args:
+        image_path (str): Path to the X-ray image
+        output_path (str, optional): Path to save enhanced image
+        clahe_clip (float): Clip limit for CLAHE
+        clahe_grid (tuple): Grid size for CLAHE
+
+    Returns:
+        str or np.ndarray: Path to enhanced image or image array
+    """
+    try:
+        # Read image
+        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+
+        if img is None:
+            raise ValueError(f"Failed to read image: {image_path}")
+
+        # Create CLAHE object
+        clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=clahe_grid)
+
+        # Apply CLAHE
+        enhanced = clahe.apply(img)
+
+        # Save enhanced image if output path is provided
+        if output_path:
+            cv2.imwrite(output_path, enhanced)
+            return output_path
+        else:
+            return enhanced
+
+    except Exception as e:
+        logger.error(f"Error enhancing X-ray image: {e}")
+        raise
+
+
+def normalize_report_text(text):
+    """
+    Normalize medical report text for consistent processing.
+
+    Args:
+        text (str): Medical report text
+
+    Returns:
+        str: Normalized text
+    """
+    try:
+        # Remove multiple whitespaces
+        text = re.sub(r"\s+", " ", text)
+
+        # Standardize section headers
+        section_patterns = {
+            r"(?i)clinical\s*(?:history|indication)": "CLINICAL HISTORY:",
+            r"(?i)technique": "TECHNIQUE:",
+            r"(?i)comparison": "COMPARISON:",
+            r"(?i)findings": "FINDINGS:",
+            r"(?i)impression": "IMPRESSION:",
+            r"(?i)recommendation": "RECOMMENDATION:",
+            r"(?i)comment": "COMMENT:",
+        }
+
+        for pattern, replacement in section_patterns.items():
+            text = re.sub(pattern + r"\s*:", replacement, text)
+
+        # Standardize common abbreviations
+        abbrev_patterns = {
+            r"(?i)\bw\/\b": "with",
+            r"(?i)\bw\/o\b": "without",
+            r"(?i)\bs\/p\b": "status post",
+            r"(?i)\bc\/w\b": "consistent with",
+            r"(?i)\br\/o\b": "rule out",
+            r"(?i)\bhx\b": "history",
+            r"(?i)\bdx\b": "diagnosis",
+            r"(?i)\btx\b": "treatment",
+        }
+
+        for pattern, replacement in abbrev_patterns.items():
+            text = re.sub(pattern, replacement, text)
+
+        return text.strip()
+
+    except Exception as e:
+        logger.error(f"Error normalizing report text: {e}")
+        return text  # Return original text if normalization fails
+
+
+def extract_sections(text):
+    """
+    Extract sections from a medical report.
+
+    Args:
+        text (str): Medical report text
+
+    Returns:
+        dict: Dictionary of extracted sections
+    """
+    try:
+        # Normalize text first
+        normalized_text = normalize_report_text(text)
+
+        # Define section patterns
+        section_headers = [
+            "CLINICAL HISTORY:",
+            "TECHNIQUE:",
+            "COMPARISON:",
+            "FINDINGS:",
+            "IMPRESSION:",
+            "RECOMMENDATION:",
+        ]
+
+        # Find all section headers in the text
+        sections = {}
+        current_section = "PREAMBLE"  # For text before first section header
+        sections[current_section] = []
+
+        for line in normalized_text.split("\n"):
+            section_found = False
+
+            for header in section_headers:
+                if header in line:
+                    current_section = header.rstrip(":")
+                    sections[current_section] = []
+                    section_found = True
+                    # Add the rest of the line after the header
+                    content = line.split(header, 1)[1].strip()
+                    if content:
+                        sections[current_section].append(content)
+                    break
+
+            if not section_found and current_section:
+                sections[current_section].append(line)
+
+        # Join each section's lines
+        for section, lines in sections.items():
+            sections[section] = " ".join(lines).strip()
+
+        # Remove empty sections
+        sections = {k: v for k, v in sections.items() if v}
+
+        return sections
+
+    except Exception as e:
+        logger.error(f"Error extracting sections: {e}")
+        return {"FULL_TEXT": text}  # Return full text if extraction fails
+
+
+def extract_measurements(text):
+    """
+    Extract measurements from medical text (sizes, volumes, etc.).
+
+    Args:
+        text (str): Medical text
+
+    Returns:
+        list: List of tuples containing (measurement, value, unit)
+    """
+    try:
+        # Pattern for measurements like "5mm nodule" or "nodule measuring 5mm"
+        # or "8x10mm mass" or "mass of size 8x10mm"
+        size_pattern = r"(\d+(?:\.\d+)?(?:\s*[x×]\s*\d+(?:\.\d+)?)?(?:\s*[x×]\s*\d+(?:\.\d+)?)?)\s*(mm|cm|mm2|cm2|mm3|cm3|ml|cc)"
+
+        # Find measurements with context
+        context_pattern = (
+            r"([A-Za-z\s]+(?:mass|nodule|effusion|opacity|lesion|tumor|cyst|structure|area|region)[A-Za-z\s]*)"
+            + size_pattern
+        )
+
+        context_measurements = []
+        for match in re.finditer(context_pattern, text, re.IGNORECASE):
+            context, size, unit = match.groups()
+            context_measurements.append((context.strip(), size, unit))
+
+        # For measurements without clear context, just extract size and unit
+        all_measurements = []
+        for match in re.finditer(size_pattern, text):
+            size, unit = match.groups()
+            all_measurements.append((size, unit))
+
+        return context_measurements
+
+    except Exception as e:
+        logger.error(f"Error extracting measurements: {e}")
+        return []
+
+
+def prepare_sample_batch(image_paths, reports=None, target_size=(224, 224)):
+    """
+    Prepare a batch of samples for model processing.
+
+    Args:
+        image_paths (list): List of paths to images
+        reports (list, optional): List of corresponding reports
+        target_size (tuple): Target image size
+
+    Returns:
+        tuple: Batch of preprocessed images and reports
+    """
+    try:
+        processed_images = []
+        processed_reports = []
+
+        for i, image_path in enumerate(image_paths):
+            # Process image
+            image = preprocess_image(image_path, target_size)
+            processed_images.append(image)
+
+            # Process report if available
+            if reports and i < len(reports):
+                normalized_report = normalize_report_text(reports[i])
+                processed_reports.append(normalized_report)
+
+        return processed_images, processed_reports if reports else None
+
+    except Exception as e:
+        logger.error(f"Error preparing sample batch: {e}")
+        raise

mediSync/utils/visualization.py

@@ -1,402 +1,516 @@
-import base64
-import io
-import logging
-
-import cv2
-import matplotlib.pyplot as plt
-import numpy as np
-from PIL import Image
-
-logger = logging.getLogger(__name__)
-
-def plot_image_prediction(image, predictions, title=None, figsize=(10, 8)):
-    try:
-        if isinstance(image, str):
-            img = Image.open(image)
-        else:
-            img = image
-
-        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
-
-        ax1.imshow(img)
-        ax1.set_title("X-ray Image")
-        ax1.axis("off")
-
-        if predictions:
-            sorted_pred = sorted(predictions, key=lambda x: x[1], reverse=True)
-
-            top_n = min(5, len(sorted_pred))
-            labels = [pred[0] for pred in sorted_pred[:top_n]]
-            probs = [pred[1] for pred in sorted_pred[:top_n]]
-
-            y_pos = np.arange(top_n)
-            ax2.barh(y_pos, probs, align="center")
-            ax2.set_yticks(y_pos)
-            ax2.set_yticklabels(labels)
-            ax2.set_xlabel("Probability")
-            ax2.set_title("Top Predictions")
-            ax2.set_xlim(0, 1)
-
-            for i, prob in enumerate(probs):
-                ax2.text(prob + 0.02, i, f"{prob:.1%}", va="center")
-
-        if title:
-            fig.suptitle(title, fontsize=16)
-
-        fig.tight_layout()
-        return fig
-
-    except Exception as e:
-        logger.error(f"Error plotting image prediction: {e}")
-        fig, ax = plt.subplots(figsize=(8, 6))
-        ax.text(0.5, 0.5, f"Error: {str(e)}", ha="center", va="center")
-        return fig
-
-def create_heatmap_overlay(image, heatmap, alpha=0.4):
-    try:
-        if isinstance(image, str):
-            img = cv2.imread(image)
-            if img is None:
-                raise ValueError(f"Could not load image: {image}")
-        elif isinstance(image, Image.Image):
-            img = np.array(image)
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        else:
-            img = image
-
-        if len(img.shape) == 2:
-            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-
-        heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
-
-        heatmap = np.maximum(heatmap, 0)
-        heatmap = np.minimum(heatmap / np.max(heatmap), 1)
-
-        heatmap = np.uint8(255 * heatmap)
-        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-
-        overlay = cv2.addWeighted(img, 1 - alpha, heatmap, alpha, 0)
-
-        overlay = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
-        overlay_img = Image.fromarray(overlay)
-
-        return overlay_img
-
-    except Exception as e:
-        logger.error(f"Error creating heatmap overlay: {e}")
-        if isinstance(image, str):
-            return Image.open(image)
-        elif isinstance(image, Image.Image):
-            return image
-        else:
-            return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-
-def plot_report_entities(text, entities, figsize=(12, 8)):
-    try:
-        fig, ax = plt.subplots(figsize=figsize)
-        ax.axis("off")
-
-        fig.patch.set_facecolor("#f8f9fa")
-        ax.set_facecolor("#f8f9fa")
-
-        ax.text(
-            0.5,
-            0.98,
-            "Medical Report Analysis",
-            ha="center",
-            va="top",
-            fontsize=18,
-            fontweight="bold",
-            color="#2c3e50",
-        )
-
-        y_pos = 0.9
-        ax.text(
-            0.05,
-            y_pos,
-            "Extracted Entities:",
-            fontsize=14,
-            fontweight="bold",
-            color="#2c3e50",
-        )
-        y_pos -= 0.05
-
-        category_colors = {
-            "problem": "#e74c3c",
-            "test": "#3498db",
-            "treatment": "#2ecc71",
-            "anatomy": "#9b59b6",
-        }
-
-        for category, items in entities.items():
-            if items:
-                y_pos -= 0.05
-                ax.text(
-                    0.1,
-                    y_pos,
-                    f"{category.capitalize()}:",
-                    fontsize=12,
-                    fontweight="bold",
-                )
-                y_pos -= 0.05
-                ax.text(
-                    0.15,
-                    y_pos,
-                    ", ".join(items),
-                    wrap=True,
-                    fontsize=11,
-                    color=category_colors.get(category, "black"),
-                )
-
-        y_pos -= 0.1
-        ax.text(
-            0.05,
-            y_pos,
-            "Report Text (with highlighted entities):",
-            fontsize=14,
-            fontweight="bold",
-            color="#2c3e50",
-        )
-        y_pos -= 0.05
-
-        all_entities = []
-        for category, items in entities.items():
-            for item in items:
-                all_entities.append((item, category))
-
-        all_entities.sort(key=lambda x: len(x[0]), reverse=True)
-
-        highlighted_text = text
-        for entity, category in all_entities:
-            entity_escaped = (
-                entity.replace("(", r"\(")
-                .replace(")", r"\)")
-                .replace("[", r"\[")
-                .replace("]", r"\]")
-            )
-
-            pattern = r"\b" + entity_escaped + r"\b"
-            color_code = category_colors.get(category, "black")
-            replacement = f"\\textcolor{{{color_code}}}{{{entity}}}"
-            highlighted_text = highlighted_text.replace(entity, replacement)
-
-        ax.text(0.05, y_pos, highlighted_text, va="top", fontsize=10, wrap=True)
-
-        fig.tight_layout(rect=[0, 0.03, 1, 0.97])
-        return fig
-
-    except Exception as e:
-        logger.error(f"Error plotting report entities: {e}")
-        fig, ax = plt.subplots(figsize=(8, 6))
-        ax.text(0.5, 0.5, f"Error: {str(e)}", ha="center", va="center")
-        return fig
-
-def plot_multimodal_results(
-    fused_results, image=None, report_text=None, figsize=(12, 10)
-):
-    try:
-        fig = plt.figure(figsize=figsize)
-        gs = fig.add_gridspec(2, 2)
-
-        fig.suptitle(
-            "Multimodal Medical Analysis Results",
-            fontsize=18,
-            fontweight="bold",
-            y=0.98,
-        )
-
-        ax_overview = fig.add_subplot(gs[0, 0])
-        ax_overview.axis("off")
-
-        severity = fused_results.get("severity", {})
-        severity_level = severity.get("level", "Unknown")
-        severity_score = severity.get("score", 0)
-
-        primary_finding = fused_results.get("primary_finding", "Unknown")
-
-        agreement = fused_results.get("agreement_score", 0)
-
-        overview_text = [
-            "ANALYSIS OVERVIEW",
-            f"Primary Finding: {primary_finding}",
-            f"Severity Level: {severity_level} ({severity_score}/4)",
-            f"Agreement Score: {agreement:.0%}",
-        ]
-
-        severity_colors = {
-            "Normal": "#2ecc71",
-            "Mild": "#3498db",
-            "Moderate": "#f39c12",
-            "Severe": "#e74c3c",
-            "Critical": "#c0392b",
-        }
-
-        y_pos = 0.9
-        ax_overview.text(
-            0.5,
-            y_pos,
-            overview_text[0],
-            fontsize=14,
-            fontweight="bold",
-            ha="center",
-            va="center",
-        )
-        y_pos -= 0.15
-
-        ax_overview.text(
-            0.1, y_pos, overview_text[1], fontsize=12, ha="left", va="center"
-        )
-        y_pos -= 0.1
-
-        severity_color = severity_colors.get(severity_level, "black")
-        ax_overview.text(
-            0.1, y_pos, "Severity Level:", fontsize=12, ha="left", va="center"
-        )
-        ax_overview.text(
-            0.4,
-            y_pos,
-            severity_level,
-            fontsize=12,
-            color=severity_color,
-            fontweight="bold",
-            ha="left",
-            va="center",
-        )
-        ax_overview.text(
-            0.6, y_pos, f"({severity_score}/4)", fontsize=10, ha="left", va="center"
-        )
-        y_pos -= 0.1
-
-        agreement_color = (
-            "#2ecc71"
-            if agreement > 0.7
-            else "#f39c12"
-            if agreement > 0.4
-            else "#e74c3c"
-        )
-        ax_overview.text(
-            0.1, y_pos, "Agreement Score:", fontsize=12, ha="left", va="center"
-        )
-        ax_overview.text(
-            0.4,
-            y_pos,
-            f"{agreement:.0%}",
-            fontsize=12,
-            color=agreement_color,
-            fontweight="bold",
-            ha="left",
-            va="center",
-        )
-
-        ax_findings = fig.add_subplot(gs[0, 1])
-        ax_findings.axis("off")
-
-        findings = fused_results.get("findings", [])
-
-        y_pos = 0.9
-        ax_findings.text(
-            0.5,
-            y_pos,
-            "KEY FINDINGS",
-            fontsize=14,
-            fontweight="bold",
-            ha="center",
-            va="center",
-        )
-        y_pos -= 0.1
-
-        if findings:
-            for i, finding in enumerate(findings[:5]):
-                ax_findings.text(0.05, y_pos, "•", fontsize=14, ha="left", va="center")
-                ax_findings.text(
-                    0.1, y_pos, finding, fontsize=11, ha="left", va="center", wrap=True
-                )
-                y_pos -= 0.15
-        else:
-            ax_findings.text(
-                0.1,
-                y_pos,
-                "No specific findings detailed.",
-                fontsize=11,
-                ha="left",
-                va="center",
-            )
-
-        ax_image = fig.add_subplot(gs[1, 0])
-
-        if image is not None:
-            if isinstance(image, str):
-                img = Image.open(image)
-            else:
-                img = image
-
-            ax_image.imshow(img)
-            ax_image.set_title("X-ray Image", fontsize=12)
-        else:
-            ax_image.text(0.5, 0.5, "No image available", ha="center", va="center")
-
-        ax_image.axis("off")
-
-        ax_rec = fig.add_subplot(gs[1, 1])
-        ax_rec.axis("off")
-
-        recommendations = fused_results.get("followup_recommendations", [])
-
-        y_pos = 0.9
-        ax_rec.text(
-            0.5,
-            y_pos,
-            "RECOMMENDATIONS",
-            fontsize=14,
-            fontweight="bold",
-            ha="center",
-            va="center",
-        )
-        y_pos -= 0.1
-
-        if recommendations:
-            for i, rec in enumerate(recommendations):
-                ax_rec.text(0.05, y_pos, "•", fontsize=14, ha="left", va="center")
-                ax_rec.text(
-                    0.1, y_pos, rec, fontsize=11, ha="left", va="center", wrap=True
-                )
-                y_pos -= 0.15
-        else:
-            ax_rec.text(
-                0.1,
-                y_pos,
-                "No specific recommendations provided.",
-                fontsize=11,
-                ha="left",
-                va="center",
-            )
-
-        fig.text(
-            0.5,
-            0.03,
-            "DISCLAIMER: This analysis is for informational purposes only and should not replace professional medical advice.",
-            fontsize=9,
-            style="italic",
-            ha="center",
-        )
-
-        fig.tight_layout(rect=[0, 0.05, 1, 0.95])
-        return fig
-
-    except Exception as e:
-        logger.error(f"Error plotting multimodal results: {e}")
-        fig, ax = plt.subplots(figsize=(8, 6))
-        ax.text(0.5, 0.5, f"Error: {str(e)}", ha="center", va="center")
-        return fig
-
-def figure_to_base64(fig):
-    try:
-        buf = io.BytesIO()
-        fig.savefig(buf, format="png", bbox_inches="tight")
-        buf.seek(0)
-        img_str = base64.b64encode(buf.read()).decode("utf-8")
-        return img_str
-
-    except Exception as e:
-        logger.error(f"Error converting figure to base64: {e}")
-        return ""
+import base64
+import io
+import logging
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+
+def plot_image_prediction(image, predictions, title=None, figsize=(10, 8)):
+    """
+    Plot an image with its predictions.
+
+    Args:
+        image (PIL.Image or str): Image or path to image
+        predictions (list): List of (label, probability) tuples
+        title (str, optional): Plot title
+        figsize (tuple): Figure size
+
+    Returns:
+        matplotlib.figure.Figure: The figure object
+    """
+    try:
+        # Load image if path is provided
+        if isinstance(image, str):
+            img = Image.open(image)
+        else:
+            img = image
+
+        # Create figure
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
+
+        # Plot image
+        ax1.imshow(img)
+        ax1.set_title("X-ray Image")
+        ax1.axis("off")
+
+        # Plot predictions
+        if predictions:
+            # Sort predictions by probability
+            sorted_pred = sorted(predictions, key=lambda x: x[1], reverse=True)
+
+            # Get top 5 predictions
+            top_n = min(5, len(sorted_pred))
+            labels = [pred[0] for pred in sorted_pred[:top_n]]
+            probs = [pred[1] for pred in sorted_pred[:top_n]]
+
+            # Plot horizontal bar chart
+            y_pos = np.arange(top_n)
+            ax2.barh(y_pos, probs, align="center")
+            ax2.set_yticks(y_pos)
+            ax2.set_yticklabels(labels)
+            ax2.set_xlabel("Probability")
+            ax2.set_title("Top Predictions")
+            ax2.set_xlim(0, 1)
+
+            # Annotate probabilities
+            for i, prob in enumerate(probs):
+                ax2.text(prob + 0.02, i, f"{prob:.1%}", va="center")
+
+        # Set overall title
+        if title:
+            fig.suptitle(title, fontsize=16)
+
+        fig.tight_layout()
+        return fig
+
+    except Exception as e:
+        logger.error(f"Error plotting image prediction: {e}")
+        # Create empty figure if error occurs
+        fig, ax = plt.subplots(figsize=(8, 6))
+        ax.text(0.5, 0.5, f"Error: {str(e)}", ha="center", va="center")
+        return fig
+
+
+def create_heatmap_overlay(image, heatmap, alpha=0.4):
+    """
+    Create a heatmap overlay on an X-ray image to highlight areas of interest.
+
+    Args:
+        image (PIL.Image or str): Image or path to image
+        heatmap (numpy.ndarray): Heatmap array
+        alpha (float): Transparency of the overlay
+
+    Returns:
+        PIL.Image: Image with heatmap overlay
+    """
+    try:
+        # Load image if path is provided
+        if isinstance(image, str):
+            img = cv2.imread(image)
+            if img is None:
+                raise ValueError(f"Could not load image: {image}")
+        elif isinstance(image, Image.Image):
+            img = np.array(image)
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        else:
+            img = image
+
+        # Ensure image is in BGR format for OpenCV
+        if len(img.shape) == 2:  # Grayscale
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+        # Resize heatmap to match image dimensions
+        heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
+
+        # Normalize heatmap (0-1)
+        heatmap = np.maximum(heatmap, 0)
+        heatmap = np.minimum(heatmap / np.max(heatmap), 1)
+
+        # Apply colormap (jet) to heatmap
+        heatmap = np.uint8(255 * heatmap)
+        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+
+        # Create overlay
+        overlay = cv2.addWeighted(img, 1 - alpha, heatmap, alpha, 0)
+
+        # Convert back to PIL image
+        overlay = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
+        overlay_img = Image.fromarray(overlay)
+
+        return overlay_img
+
+    except Exception as e:
+        logger.error(f"Error creating heatmap overlay: {e}")
+        # Return original image if error occurs
+        if isinstance(image, str):
+            return Image.open(image)
+        elif isinstance(image, Image.Image):
+            return image
+        else:
+            return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+
+def plot_report_entities(text, entities, figsize=(12, 8)):
+    """
+    Visualize entities extracted from a medical report.
+
+    Args:
+        text (str): Report text
+        entities (dict): Dictionary of entities by category
+        figsize (tuple): Figure size
+
+    Returns:
+        matplotlib.figure.Figure: The figure object
+    """
+    try:
+        fig, ax = plt.subplots(figsize=figsize)
+        ax.axis("off")
+
+        # Set background color
+        fig.patch.set_facecolor("#f8f9fa")
+        ax.set_facecolor("#f8f9fa")
+
+        # Title
+        ax.text(
+            0.5,
+            0.98,
+            "Medical Report Analysis",
+            ha="center",
+            va="top",
+            fontsize=18,
+            fontweight="bold",
+            color="#2c3e50",
+        )
+
+        # Display entity counts
+        y_pos = 0.9
+        ax.text(
+            0.05,
+            y_pos,
+            "Extracted Entities:",
+            fontsize=14,
+            fontweight="bold",
+            color="#2c3e50",
+        )
+        y_pos -= 0.05
+
+        # Define colors for different entity categories
+        category_colors = {
+            "problem": "#e74c3c",  # Red
+            "test": "#3498db",  # Blue
+            "treatment": "#2ecc71",  # Green
+            "anatomy": "#9b59b6",  # Purple
+        }
+
+        # Display entities by category
+        for category, items in entities.items():
+            if items:
+                y_pos -= 0.05
+                ax.text(
+                    0.1,
+                    y_pos,
+                    f"{category.capitalize()}:",
+                    fontsize=12,
+                    fontweight="bold",
+                )
+                y_pos -= 0.05
+                ax.text(
+                    0.15,
+                    y_pos,
+                    ", ".join(items),
+                    wrap=True,
+                    fontsize=11,
+                    color=category_colors.get(category, "black"),
+                )
+
+        # Add the report text with highlighted entities
+        y_pos -= 0.1
+        ax.text(
+            0.05,
+            y_pos,
+            "Report Text (with highlighted entities):",
+            fontsize=14,
+            fontweight="bold",
+            color="#2c3e50",
+        )
+        y_pos -= 0.05
+
+        # Get all entities to highlight
+        all_entities = []
+        for category, items in entities.items():
+            for item in items:
+                all_entities.append((item, category))
+
+        # Sort entities by length (longest first to avoid overlap issues)
+        all_entities.sort(key=lambda x: len(x[0]), reverse=True)
+
+        # Highlight entities in text
+        highlighted_text = text
+        for entity, category in all_entities:
+            # Escape regex special characters
+            entity_escaped = (
+                entity.replace("(", r"\(")
+                .replace(")", r"\)")
+                .replace("[", r"\[")
+                .replace("]", r"\]")
+            )
+
+            # Find entity in text (word boundary)
+            pattern = r"\b" + entity_escaped + r"\b"
+            color_code = category_colors.get(category, "black")
+            replacement = f"\\textcolor{{{color_code}}}{{{entity}}}"
+            highlighted_text = highlighted_text.replace(entity, replacement)
+
+        # Display highlighted text
+        ax.text(0.05, y_pos, highlighted_text, va="top", fontsize=10, wrap=True)
+
+        fig.tight_layout(rect=[0, 0.03, 1, 0.97])
+        return fig
+
+    except Exception as e:
+        logger.error(f"Error plotting report entities: {e}")
+        # Create empty figure if error occurs
+        fig, ax = plt.subplots(figsize=(8, 6))
+        ax.text(0.5, 0.5, f"Error: {str(e)}", ha="center", va="center")
+        return fig
+
+
+def plot_multimodal_results(
+    fused_results, image=None, report_text=None, figsize=(12, 10)
+):
+    """
+    Visualize the results of multimodal analysis.
+
+    Args:
+        fused_results (dict): Results from multimodal fusion
+        image (PIL.Image or str, optional): Image or path to image
+        report_text (str, optional): Report text
+        figsize (tuple): Figure size
+
+    Returns:
+        matplotlib.figure.Figure: The figure object
+    """
+    try:
+        # Create figure with a grid layout
+        fig = plt.figure(figsize=figsize)
+        gs = fig.add_gridspec(2, 2)
+
+        # Add title
+        fig.suptitle(
+            "Multimodal Medical Analysis Results",
+            fontsize=18,
+            fontweight="bold",
+            y=0.98,
+        )
+
+        # 1. Overview panel (top left)
+        ax_overview = fig.add_subplot(gs[0, 0])
+        ax_overview.axis("off")
+
+        # Get severity info
+        severity = fused_results.get("severity", {})
+        severity_level = severity.get("level", "Unknown")
+        severity_score = severity.get("score", 0)
+
+        # Get primary finding
+        primary_finding = fused_results.get("primary_finding", "Unknown")
+
+        # Get agreement score
+        agreement = fused_results.get("agreement_score", 0)
+
+        # Create overview text
+        overview_text = [
+            "ANALYSIS OVERVIEW",
+            f"Primary Finding: {primary_finding}",
+            f"Severity Level: {severity_level} ({severity_score}/4)",
+            f"Agreement Score: {agreement:.0%}",
+        ]
+
+        # Define severity colors
+        severity_colors = {
+            "Normal": "#2ecc71",  # Green
+            "Mild": "#3498db",  # Blue
+            "Moderate": "#f39c12",  # Orange
+            "Severe": "#e74c3c",  # Red
+            "Critical": "#c0392b",  # Dark Red
+        }
+
+        # Add overview text to the panel
+        y_pos = 0.9
+        ax_overview.text(
+            0.5,
+            y_pos,
+            overview_text[0],
+            fontsize=14,
+            fontweight="bold",
+            ha="center",
+            va="center",
+        )
+        y_pos -= 0.15
+
+        ax_overview.text(
+            0.1, y_pos, overview_text[1], fontsize=12, ha="left", va="center"
+        )
+        y_pos -= 0.1
+
+        # Severity with color
+        severity_color = severity_colors.get(severity_level, "black")
+        ax_overview.text(
+            0.1, y_pos, "Severity Level:", fontsize=12, ha="left", va="center"
+        )
+        ax_overview.text(
+            0.4,
+            y_pos,
+            severity_level,
+            fontsize=12,
+            color=severity_color,
+            fontweight="bold",
+            ha="left",
+            va="center",
+        )
+        ax_overview.text(
+            0.6, y_pos, f"({severity_score}/4)", fontsize=10, ha="left", va="center"
+        )
+        y_pos -= 0.1
+
+        # Agreement score with color
+        agreement_color = (
+            "#2ecc71"
+            if agreement > 0.7
+            else "#f39c12"
+            if agreement > 0.4
+            else "#e74c3c"
+        )
+        ax_overview.text(
+            0.1, y_pos, "Agreement Score:", fontsize=12, ha="left", va="center"
+        )
+        ax_overview.text(
+            0.4,
+            y_pos,
+            f"{agreement:.0%}",
+            fontsize=12,
+            color=agreement_color,
+            fontweight="bold",
+            ha="left",
+            va="center",
+        )
+
+        # 2. Findings panel (top right)
+        ax_findings = fig.add_subplot(gs[0, 1])
+        ax_findings.axis("off")
+
+        # Get findings
+        findings = fused_results.get("findings", [])
+
+        # Add findings to the panel
+        y_pos = 0.9
+        ax_findings.text(
+            0.5,
+            y_pos,
+            "KEY FINDINGS",
+            fontsize=14,
+            fontweight="bold",
+            ha="center",
+            va="center",
+        )
+        y_pos -= 0.1
+
+        if findings:
+            for i, finding in enumerate(findings[:5]):  # Limit to 5 findings
+                ax_findings.text(0.05, y_pos, "•", fontsize=14, ha="left", va="center")
+                ax_findings.text(
+                    0.1, y_pos, finding, fontsize=11, ha="left", va="center", wrap=True
+                )
+                y_pos -= 0.15
+        else:
+            ax_findings.text(
+                0.1,
+                y_pos,
+                "No specific findings detailed.",
+                fontsize=11,
+                ha="left",
+                va="center",
+            )
+
+        # 3. Image panel (bottom left)
+        ax_image = fig.add_subplot(gs[1, 0])
+
+        if image is not None:
+            # Load image if path is provided
+            if isinstance(image, str):
+                img = Image.open(image)
+            else:
+                img = image
+
+            # Display image
+            ax_image.imshow(img)
+            ax_image.set_title("X-ray Image", fontsize=12)
+        else:
+            ax_image.text(0.5, 0.5, "No image available", ha="center", va="center")
+
+        ax_image.axis("off")
+
+        # 4. Recommendation panel (bottom right)
+        ax_rec = fig.add_subplot(gs[1, 1])
+        ax_rec.axis("off")
+
+        # Get recommendations
+        recommendations = fused_results.get("followup_recommendations", [])
+
+        # Add recommendations to the panel
+        y_pos = 0.9
+        ax_rec.text(
+            0.5,
+            y_pos,
+            "RECOMMENDATIONS",
+            fontsize=14,
+            fontweight="bold",
+            ha="center",
+            va="center",
+        )
+        y_pos -= 0.1
+
+        if recommendations:
+            for i, rec in enumerate(recommendations):
+                ax_rec.text(0.05, y_pos, "•", fontsize=14, ha="left", va="center")
+                ax_rec.text(
+                    0.1, y_pos, rec, fontsize=11, ha="left", va="center", wrap=True
+                )
+                y_pos -= 0.15
+        else:
+            ax_rec.text(
+                0.1,
+                y_pos,
+                "No specific recommendations provided.",
+                fontsize=11,
+                ha="left",
+                va="center",
+            )
+
+        # Add disclaimer
+        fig.text(
+            0.5,
+            0.03,
+            "DISCLAIMER: This analysis is for informational purposes only and should not replace professional medical advice.",
+            fontsize=9,
+            style="italic",
+            ha="center",
+        )
+
+        fig.tight_layout(rect=[0, 0.05, 1, 0.95])
+        return fig
+
+    except Exception as e:
+        logger.error(f"Error plotting multimodal results: {e}")
+        # Create empty figure if error occurs
+        fig, ax = plt.subplots(figsize=(8, 6))
+        ax.text(0.5, 0.5, f"Error: {str(e)}", ha="center", va="center")
+        return fig
+
+
+def figure_to_base64(fig):
+    """
+    Convert matplotlib figure to base64 string.
+
+    Args:
+        fig (matplotlib.figure.Figure): Figure object
+
+    Returns:
+        str: Base64 encoded string
+    """
+    try:
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight")
+        buf.seek(0)
+        img_str = base64.b64encode(buf.read()).decode("utf-8")
+        return img_str
+
+    except Exception as e:
+        logger.error(f"Error converting figure to base64: {e}")
+        return ""

run.py

@@ -1,55 +1,71 @@
-import logging
-import sys
-from pathlib import Path
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    handlers=[logging.StreamHandler(), logging.FileHandler("uhres.log")],
-)
-logger = logging.getLogger(__name__)
-
-def main():
-    logger.info("Starting U-HRES setup")
-
-    current_dir = Path(__file__).resolve().parent
-    sys.path.append(str(current_dir))
-
-    try:
-        logger.info("Checking for sample data")
-        from mediSync.utils.download_samples import (
-            create_sample_info_file,
-            download_sample_images,
-        )
-
-        sample_dir = current_dir / "mediSync" / "data" / "sample"
-        sample_dir.mkdir(parents=True, exist_ok=True)
-
-        image_files = list(sample_dir.glob("*.jpg")) + list(sample_dir.glob("*.png"))
-
-        if not image_files:
-            logger.info("No sample images found. Downloading...")
-            download_sample_images()
-            create_sample_info_file()
-        else:
-            logger.info(f"Found {len(image_files)} existing sample images")
-
-    except Exception as e:
-        logger.error(f"Error setting up sample data: {e}")
-        print(f"Warning: Could not set up sample data: {e}")
-
-    try:
-        logger.info("Launching U-HRES application")
-        from mediSync.app import create_interface
-
-        create_interface()
-
-    except Exception as e:
-        logger.error(f"Error launching application: {e}")
-        print(f"Error: Failed to launch U-HRES application: {e}")
-        return 1
-
-    return 0
-
-if __name__ == "__main__":
-    sys.exit(main())
+"""
+U-HRES: Unified Health Record Exchange System AI Layer - Runner Script
+=====================================================================
+This script downloads sample data and launches the U-HRES application.
+"""
+
+import logging
+import sys
+from pathlib import Path
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler(), logging.FileHandler("uhres.log")],
+)
+logger = logging.getLogger(__name__)
+
+
+def main():
+    """
+    Main function to set up and run the U-HRES application.
+    """
+    logger.info("Starting U-HRES setup")
+
+    # Add the current directory to Python path
+    current_dir = Path(__file__).resolve().parent
+    sys.path.append(str(current_dir))
+
+    # Download sample data if needed
+    try:
+        logger.info("Checking for sample data")
+        from mediSync.utils.download_samples import (
+            create_sample_info_file,
+            download_sample_images,
+        )
+
+        # Check if sample directory exists and contains images
+        sample_dir = current_dir / "mediSync" / "data" / "sample"
+        sample_dir.mkdir(parents=True, exist_ok=True)
+
+        image_files = list(sample_dir.glob("*.jpg")) + list(sample_dir.glob("*.png"))
+
+        if not image_files:
+            logger.info("No sample images found. Downloading...")
+            download_sample_images()
+            create_sample_info_file()
+        else:
+            logger.info(f"Found {len(image_files)} existing sample images")
+
+    except Exception as e:
+        logger.error(f"Error setting up sample data: {e}")
+        print(f"Warning: Could not set up sample data: {e}")
+
+    # Launch the application
+    try:
+        logger.info("Launching U-HRES application")
+        from mediSync.app import create_interface
+
+        create_interface()
+
+    except Exception as e:
+        logger.error(f"Error launching application: {e}")
+        print(f"Error: Failed to launch U-HRES application: {e}")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())