Upload 6 files

app.py

@@ -1,657 +1,735 @@
-"""Gradio Inference App for Transfer Learning Models - HuggingFace Spaces Version
-
-Features:
-- Auto-scan models directory or HuggingFace Hub for available models and approaches.
-- Dropdown selection of Model and Approach.
-- Dynamic architecture detection from filename (e.g., resnet50, densenet121, inception_v3, efficientnet_b0, resnet34).
-- Image upload and preprocessing (ImageNet normalization).
-- Top-K prediction display (configurable class labels).
-- Optional Grad-CAM visualization for interpretability.
-- Environment variable configuration for HuggingFace deployment.
-- Graceful error handling and clear user feedback.
-
-Environment Variables:
-- HF_TOKEN: HuggingFace API token for private repositories
-- MODEL_REPO_ID: HuggingFace repository containing models
-- NUM_CLASSES: Number of output classes (default: 2)
-- DEBUG: Enable debug logging
 """
-import numpy as np
-import pandas as pd
-import os
-import re
-import logging
-from typing import Dict, Tuple, List, Optional, Union
-from pathlib import Path
+Hugging Face Gradio App for Pneumonia Detection Ensemble
+Supports JPEG, PNG, and DICOM image formats
+"""
 
-import pydicom
+import gradio as gr
 import torch
-import torch.nn.functional as F
-from torchvision import models, transforms
+import torch.nn as nn
+import torchvision.transforms as transforms
+import torchvision.models as models
 from PIL import Image
-import gradio as gr
-from huggingface_hub import hf_hub_download, list_repo_files
-from dotenv import load_dotenv
+import numpy as np
+import json
+from pathlib import Path
+import io
+import os
 
-try:
-    import timm  # EfficientNet etc.
-except ImportError:
-    timm = None
+# ----------------------------------------------------------------------------
+# Debug / Diagnostics Configuration
+# Set environment variable CLINICAL_DEBUG=1 to enable verbose logging
+# ----------------------------------------------------------------------------
+DEBUG = os.getenv("CLINICAL_DEBUG", "0") in ("1", "true", "True")
+
+def _dbg(msg):
+    if DEBUG:
+        print(f"[DEBUG] {msg}")
 
+# DICOM support
 try:
-    import cv2
+    import pydicom
+    DICOM_AVAILABLE = True
 except ImportError:
-    cv2 = None
+    DICOM_AVAILABLE = False
 
-import gradio_client.utils as gc_utils
+# ============================================================================
+# Model Architectures (simplified versions for deployment)
+# ============================================================================
 
-# Hot-patch for TypeError: argument of type 'bool' is not iterable
-def _safe_json_schema_to_python_type(schema, defs=None):
-    try:
-        # Handle if schema is a boolean (Gradio 5.x bug)
-        if isinstance(schema, bool):
-            # Return readable type name to avoid crash
-            return "bool" if schema else "NoneType"
-        # If schema is a dict, proceed as usual
-        if isinstance(schema, dict) and "const" in schema:
-            pass
-        # Fall back to original internal logic if it exists
-        if hasattr(gc_utils, "_json_schema_to_python_type_original"):
-            return gc_utils._json_schema_to_python_type_original(schema, defs)
-        else:
-            # If no original stored yet, call the default one
-            return gc_utils._json_schema_to_python_type(schema, defs)
-    except Exception:
-        return "Any"
-
-# Save original before replacing (only once)
-if not hasattr(gc_utils, "_json_schema_to_python_type_original"):
-    gc_utils._json_schema_to_python_type_original = gc_utils._json_schema_to_python_type
-
-# Apply the patch
-gc_utils._json_schema_to_python_type = _safe_json_schema_to_python_type
-
-# Load environment variables
-load_dotenv()
-
-# Configuration from environment variables
-HF_TOKEN = os.getenv("HF_TOKEN")
-MODEL_REPO_ID = os.getenv("MODEL_REPO_ID")
-NUM_CLASSES = int(os.getenv("NUM_CLASSES", "2"))
-DEBUG = os.getenv("DEBUG", "False").lower() == "true"
-
-# Setup logging
-logging.basicConfig(level=logging.DEBUG if DEBUG else logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Directory paths
-MODELS_DIR = Path("models")
-MODELS_DIR.mkdir(exist_ok=True)
-
-# Placeholder class labels; customize based on your dataset
-# CLASS_LABELS = {i: f"Class_{i}" for i in range(NUM_CLASSES)}
-CLASS_LABELS = {0:'No Pneumonia',1:'Pneumonia'}
-
-# Architecture-specific input sizes
-_ARCH_INPUT_SIZE = {
-    "inception_v3": 299,
-    # Most others default to 224
-}
+class MobileNetV2Model(nn.Module):
+    def __init__(self, num_classes=2):
+        super(MobileNetV2Model, self).__init__()
+        self.model = models.mobilenet_v2(weights=None)
+        self.model.classifier[1] = nn.Linear(1280, num_classes)
 
-# Regex to parse weight filenames: Appr_<Approach>_<arch>.pt
-WEIGHT_PATTERN = re.compile(r"Appr_([A-Za-z0-9]+)_([A-Za-z0-9_]+)\.pt")
+    def forward(self, x):
+        return self.model(x)
 
-ModelMap = Dict[str, Dict[str, Dict[str, str]]]
-_model_map_cache: Optional[ModelMap] = None
-# Cache for instantiated models to avoid recreation per inference
-_model_instance_cache: Dict[Tuple[str, str], Tuple[torch.nn.Module, str]] = {}
 
+class ResNet50Model(nn.Module):
+    def __init__(self, num_classes=2):
+        super(ResNet50Model, self).__init__()
+        self.model = models.resnet50(weights=None)
+        self.model.fc = nn.Linear(2048, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class EfficientNetB0Model(nn.Module):
+    def __init__(self, num_classes=2):
+        super(EfficientNetB0Model, self).__init__()
+        try:
+            from torchvision.models import efficientnet_b0
+            self.model = efficientnet_b0(weights=None)
+        except:
+            self.model = models.efficientnet_b0(weights=None)
+        num_features = self.model.classifier[1].in_features
+        self.model.classifier[1] = nn.Linear(num_features, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class VGG19Model(nn.Module):
+    def __init__(self, num_classes=2):
+        super(VGG19Model, self).__init__()
+        self.model = models.vgg19(weights=None)
+        self.model.classifier[6] = nn.Linear(4096, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DenseNet101Model(nn.Module):
+    def __init__(self, num_classes=2):
+        super(DenseNet101Model, self).__init__()
+        self.model = models.densenet101(weights=None)
+        self.model.classifier = nn.Linear(1024, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+# ============================================================================
+# DICOM Processing Functions
+# ============================================================================
+
+def process_dicom_file(file_obj):
+    """Process DICOM file and convert to PIL Image with improved medical handling.
+
+    Adds support for:
+      - RescaleSlope / RescaleIntercept
+      - PhotometricInterpretation inversion (MONOCHROME1)
+      - Float window center/width handling
+      - Detailed pixel statistics for debugging
+    """
+    if not DICOM_AVAILABLE:
+        raise ValueError("DICOM support not available. Please install pydicom.")
 
-def download_models_from_hub() -> None:
-    """Download models from HuggingFace Hub if MODEL_REPO_ID is set."""
-    if not MODEL_REPO_ID:
-        logger.info("No MODEL_REPO_ID set, skipping Hub download")
-        return
-    
     try:
-        logger.info(f"Downloading models from {MODEL_REPO_ID}")
-        repo_files = list_repo_files(MODEL_REPO_ID, token=HF_TOKEN)
-        model_files = [f for f in repo_files if f.endswith('.pt')]
-        
-        for model_file in model_files:
-            local_path = MODELS_DIR / model_file
-            if not local_path.exists():
-                logger.info(f"Downloading {model_file}")
-                downloaded_path = hf_hub_download(
-                    MODEL_REPO_ID,
-                    model_file,
-                    cache_dir=str(MODELS_DIR),
-                    token=HF_TOKEN
-                )
-                # Move to our models directory structure
-                local_path.parent.mkdir(parents=True, exist_ok=True)
-                Path(downloaded_path).rename(local_path)
-                
-    except Exception as e:
-        logger.warning(f"Failed to download from Hub: {e}")
+        # Read DICOM
+        ds = pydicom.dcmread(file_obj.name if hasattr(file_obj, 'name') else file_obj)
+        pixel_array = ds.pixel_array.astype(np.float32)
+
+        _dbg(f"DICOM original shape={pixel_array.shape} dtype={pixel_array.dtype} min={pixel_array.min():.2f} max={pixel_array.max():.2f}")
+
+        # Apply rescale if present
+        slope = float(getattr(ds, 'RescaleSlope', 1.0))
+        intercept = float(getattr(ds, 'RescaleIntercept', 0.0))
+        if slope != 1.0 or intercept != 0.0:
+            pixel_array = pixel_array * slope + intercept
+            _dbg(f"Applied rescale slope={slope} intercept={intercept} new_min={pixel_array.min():.2f} new_max={pixel_array.max():.2f}")
+
+        # Photometric interpretation inversion (MONOCHROME1 means high values = dark)
+        photometric = getattr(ds, 'PhotometricInterpretation', '').upper()
+        if photometric == 'MONOCHROME1':
+            max_val = pixel_array.max()
+            pixel_array = max_val - pixel_array
+            _dbg("Applied MONOCHROME1 inversion")
+
+        # Windowing
+        window_center = getattr(ds, 'WindowCenter', None)
+        window_width = getattr(ds, 'WindowWidth', None)
+        if window_center is not None and window_width is not None:
+            if isinstance(window_center, (list, tuple)): window_center = float(window_center[0])
+            if isinstance(window_width, (list, tuple)): window_width = float(window_width[0])
+            window_min = window_center - window_width / 2.0
+            window_max = window_center + window_width / 2.0
+            pixel_array = np.clip(pixel_array, window_min, window_max)
+            pixel_array = (pixel_array - window_min) / max(window_max - window_min, 1e-6)
+            _dbg(f"Applied windowing center={window_center} width={window_width} -> min={window_min:.2f} max={window_max:.2f}")
+        else:
+            # Min-max normalize
+            pmin, pmax = pixel_array.min(), pixel_array.max()
+            pixel_array = (pixel_array - pmin) / max(pmax - pmin, 1e-6)
+            _dbg("Applied min-max normalization (no window tags)")
 
+        # Scale to 0-255
+        pixel_array = (pixel_array * 255.0).clip(0, 255).astype(np.uint8)
+        image = Image.fromarray(pixel_array, mode='L')
 
-def scan_local_models() -> ModelMap:
-    """Scan local models directory for available models."""
-    mapping: ModelMap = {}
-    
-    # First, try to download from Hub
-    download_models_from_hub()
-    
-    # Scan the models directory
-    if MODELS_DIR.exists():
-        for item in MODELS_DIR.iterdir():
-            if item.is_dir():
-                # Model directory structure: models/Model_X/Appr_Y_arch.pt
-                model_name = item.name
-                for model_file in item.glob("*.pt"):
-                    match = WEIGHT_PATTERN.match(model_file.name)
-                    if match:
-                        appr_code, arch = match.groups()
-                        mapping.setdefault(model_name, {})[appr_code] = {
-                            "path": str(model_file),
-                            "arch": arch.lower(),
-                        }
-            elif item.suffix == ".pt":
-                # Flat structure: models/Appr_Y_arch.pt
-                match = WEIGHT_PATTERN.match(item.name)
-                if match:
-                    appr_code, arch = match.groups()
-                    model_name = f"Model_{arch}"
-                    mapping.setdefault(model_name, {})[appr_code] = {
-                        "path": str(item),
-                        "arch": arch.lower(),
-                    }
-    
-    return mapping
+        # Optional mild contrast enhancement
+        try:
+            from PIL import ImageEnhance
+            enhancer = ImageEnhance.Contrast(image)
+            image = enhancer.enhance(1.15)
+        except Exception:
+            pass
 
+        image = image.convert('RGB')
 
-def scan_weights(refresh: bool = False) -> ModelMap:
-    """Scan for available models, with caching."""
-    global _model_map_cache
-    if _model_map_cache is not None and not refresh:
-        return _model_map_cache
+        # Log summary stats
+        arr = np.array(image)  # RGB
+        _dbg(f"Post-process RGB stats: mean={arr.mean():.2f} std={arr.std():.2f} min={arr.min()} max={arr.max()}")
+
+        return image
+    except Exception as e:
+        raise ValueError(f"Error processing DICOM file: {str(e)}")
+
+
+def process_uploaded_image(file_obj):
+    """
+    Process uploaded image file (JPEG, PNG, or DICOM)
     
-    mapping = scan_local_models()
+    Args:
+        file_obj: File object from Gradio upload
+        
+    Returns:
+        PIL Image object
+    """
+    if file_obj is None:
+        return None
     
-    # Fallback: create demo models if none found
-    if not mapping:
-        logger.warning("No models found, creating demo entries")
-        mapping = {
-            "ResNet34": {
-                "G1": {"path": "", "arch": "resnet34"},
+    try:
+        # Check file extension
+        file_name = getattr(file_obj, 'name', '').lower()
+        
+        if file_name.endswith(('.dcm', '.dicom')):
+            # Process as DICOM
+            return process_dicom_file(file_obj)
+        else:
+            # Process as regular image
+            if hasattr(file_obj, 'name'):
+                # File path provided
+                image = Image.open(file_obj.name)
+            else:
+                # File object provided
+                image = Image.open(file_obj)
+            
+            # Ensure RGB
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            return image
+            
+    except Exception as e:
+        raise ValueError(f"Error processing image file: {str(e)}")
+
+
+# ============================================================================
+# Model Classes
+# ============================================================================
+
+class PneumoniaModelSystem:
+    """Flexible model system supporting both individual models and ensemble"""
+    def __init__(self, device='cpu'):
+        self.device = device
+        self.models = {}
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                               std=[0.229, 0.224, 0.225])
+        ])
+        
+        # Model definitions with their architectures and weights
+        self.model_definitions = {
+            'Model_A1_7CB_Appr_D': {
+                'architecture': 'VGG19',
+                'file': 'Model_A1_7CB_Appr_D.pt',
+                'description': 'VGG19 - Model A1 with 7CB approach D'
+            },
+            'Model_C_Appr_B': {
+                'architecture': 'MobileNetV2', 
+                'file': 'Model_C_Appr_B.pt',
+                'description': 'MobileNetV2 - Model C with approach B'
+            },
+            'Model_F_Appr_B': {
+                'architecture': 'ResNet50',
+                'file': 'Model_F_Appr_B.pt', 
+                'description': 'ResNet50 - Model F with approach B'
             },
-            "InceptionV3": {
-                "H1": {"path": "", "arch": "inception_v3"},
+            'Model_G_Appr_B': {
+                'architecture': 'EfficientNet-B0',
+                'file': 'Model_G_Appr_B.pt',
+                'description': 'EfficientNet-B0 - Model G with approach B'
+            },
+            'Model_H_Appr_B': {
+                'architecture': 'DenseNet101',
+                'file': 'Model_H_Appr_B.pt',
+                'description': 'DenseNet101 - Model H with approach B'
             }
         }
+        
+        # Ensemble configuration
+        self.ensemble_weights = {
+            'Model_A1_7CB_Appr_D': 0.30,  # VGG19 - Higher weight
+            'Model_C_Appr_B': 0.175,      # MobileNetV2
+            'Model_F_Appr_B': 0.175,      # ResNet50
+            'Model_G_Appr_B': 0.175,      # EfficientNet-B0
+            'Model_H_Appr_B': 0.175       # DenseNet101
+        }
     
-    _model_map_cache = mapping
-    logger.info(f"Found models: {list(mapping.keys())}")
-    return mapping
-
-
-def _create_model(arch: str, num_classes: int) -> torch.nn.Module:
-    """Instantiate a model given architecture string."""
-    arch = arch.lower()
-    logger.debug(f"Creating model: {arch} with {num_classes} classes")
     
-    if arch.startswith("resnet"):
-        base = getattr(models, arch)(weights=None)
-        base.fc = torch.nn.Linear(base.fc.in_features, num_classes)
-        return base
-    elif arch.startswith("densenet"):
-        base = getattr(models, arch)(weights=None)
-        base.classifier = torch.nn.Linear(base.classifier.in_features, num_classes)
-        return base
-    elif arch.startswith("inception_v3"):
-        # Disable aux logits for lighter inference path
-        base = models.inception_v3(weights=None, aux_logits=False)
-        base.fc = torch.nn.Linear(base.fc.in_features, num_classes)
-        return base
-    elif arch.startswith("efficientnet"):
-        if timm is None:
-            raise RuntimeError("timm not installed; cannot create EfficientNet model.")
-        base = timm.create_model(arch, pretrained=False, num_classes=num_classes)
-        return base
+    def _create_model(self, architecture):
+        """Create a model instance based on architecture type"""
+        if architecture == 'MobileNetV2':
+            return MobileNetV2Model(num_classes=2).to(self.device)
+        elif architecture == 'ResNet50':
+            return ResNet50Model(num_classes=2).to(self.device)
+        elif architecture == 'EfficientNet-B0':
+            return EfficientNetB0Model(num_classes=2).to(self.device)
+        elif architecture == 'VGG19':
+            return VGG19Model(num_classes=2).to(self.device)
+        elif architecture == 'DenseNet101':
+            return DenseNet101Model(num_classes=2).to(self.device)
+        else:
+            raise ValueError(f"Unknown architecture: {architecture}")
     
-    # Fallback using timm if available
-    if timm is not None:
-        try:
-            return timm.create_model(arch, pretrained=False, num_classes=num_classes)
-        except Exception as e:
-            logger.warning(f"timm failed to create {arch}: {e}")
+    def load_models(self, model_dir='models'):
+        """Load all available models from directory"""
+        model_dir = Path(model_dir)
+        loaded_models = {}
+        
+        load_kwargs = {"map_location": self.device, "weights_only": False}
+        
+        for model_name, model_info in self.model_definitions.items():
+            model_path = model_dir / model_info['file']
+            if model_path.exists():
+                try:
+                    model = self._create_model(model_info['architecture'])
+                    model.load_state_dict(torch.load(model_path, **load_kwargs))
+                    model.eval()
+                    loaded_models[model_name] = {
+                        'model': model,
+                        'info': model_info
+                    }
+                    _dbg(f"Loaded {model_name} ({model_info['architecture']})")
+                except Exception as e:
+                    print(f"Warning: Could not load {model_name}: {str(e)}")
+            else:
+                print(f"Warning: Model file not found: {model_path}")
+        
+        self.models = loaded_models
+        return self
     
-    raise ValueError(f"Unsupported architecture: {arch}")
-
-
-def _resolve_device() -> torch.device:
-    """Resolve device based on DEVICE env var (cpu|cuda|auto)."""
-    device_pref = os.getenv("DEVICE", "auto").lower()
-    if device_pref == "cpu":
-        return torch.device("cpu")
-    if device_pref == "cuda":
-        if torch.cuda.is_available():
-            return torch.device("cuda")
-        logger.warning("DEVICE=cuda requested but CUDA not available; falling back to CPU")
-        return torch.device("cpu")
-    # auto
-    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-def load_model(model_name: str, approach: str, use_cache: bool = True) -> Tuple[torch.nn.Module, str]:
-    """Load a model given model name and approach, optionally using cache."""
-    if use_cache:
-        cache_key = (model_name, approach)
-        if cache_key in _model_instance_cache:
-            return _model_instance_cache[cache_key]
-
-    mapping = scan_weights()
-    if model_name not in mapping:
-        raise ValueError(f"Model '{model_name}' not found in {list(mapping.keys())}")
-    if approach not in mapping[model_name]:
-        raise ValueError(
-            f"Approach '{approach}' not found for model '{model_name}'. Available: {list(mapping[model_name].keys())}"
-        )
+    def get_available_models(self):
+        """Get list of available model names"""
+        return list(self.models.keys())
+    
+    def predict_single_model(self, image, model_name):
+        """
+        Predict using a single specified model
+        
+        Args:
+            image: PIL Image
+            model_name: Name of the model to use
+            
+        Returns:
+            dict with predictions and probabilities
+        """
+        if model_name not in self.models:
+            raise ValueError(f"Model {model_name} not available")
+            
+        # Convert to PIL if needed
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        # Ensure RGB
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        # Transform
+        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
+        
+        model = self.models[model_name]['model']
+        model_info = self.models[model_name]['info']
+        
+        with torch.no_grad():
+            outputs = model(img_tensor)
+            probs = torch.softmax(outputs, dim=1)
+            
+        probs_numpy = probs[0].cpu().numpy()
+        prediction_index = int(probs_numpy.argmax())
+        predicted_label = 'PNEUMONIA' if prediction_index == 1 else 'NORMAL'
+        
+        _dbg(f"{model_name} logits={outputs.cpu().numpy()} probs={probs_numpy} label={predicted_label}")
+        
+        result = {
+            'prediction': predicted_label,
+            'confidence': float(probs_numpy[prediction_index]),
+            'pneumonia_probability': float(probs_numpy[1]),
+            'normal_probability': float(probs_numpy[0]),
+            'model_used': model_name,
+            'model_architecture': model_info['architecture'],
+            'model_description': model_info['description']
+        }
+        return result
 
-    info = mapping[model_name][approach]
-    arch = info["arch"]
-    weight_path = info["path"]
+    def predict_ensemble(self, image, selected_models=None):
+        """
+        Predict using ensemble of models
+        
+        Args:
+            image: PIL Image
+            selected_models: List of model names to include in ensemble, or None for all
+            
+        Returns:
+            dict with predictions and probabilities
+        """
+        if selected_models is None:
+            selected_models = list(self.ensemble_weights.keys())
+        
+        # Filter to only available models
+        available_models = [m for m in selected_models if m in self.models]
+        if not available_models:
+            raise ValueError("No valid models available for ensemble")
+        
+        # Convert to PIL if needed
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
 
-    device = _resolve_device()
-    logger.info(f"Loading {model_name}/{approach} ({arch}) on {device}")
+        # Ensure RGB
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
 
-    model = _create_model(arch, NUM_CLASSES)
+        # Transform
+        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
 
-    # Load weights if path exists
-    if weight_path and Path(weight_path).exists():
-        try:
-            state = torch.load(weight_path, map_location=device)
-            if isinstance(state, dict) and "state_dict" in state:
-                state = state["state_dict"]
+        # Get predictions from each model
+        per_model = []
+        with torch.no_grad():
+            ensemble_probs = torch.zeros(1, 2).to(self.device)
+            total_weight = 0
+            
+            for model_name in available_models:
+                model = self.models[model_name]['model']
+                weight = self.ensemble_weights.get(model_name, 1.0)
+                
+                outputs = model(img_tensor)
+                probs = torch.softmax(outputs, dim=1)
+                ensemble_probs += weight * probs
+                total_weight += weight
+                
+                per_model.append({
+                    'model_name': model_name,
+                    'architecture': self.models[model_name]['info']['architecture'],
+                    'weight': weight,
+                    'logits': outputs.detach().cpu().numpy().tolist(),
+                    'probs': probs.detach().cpu().numpy().tolist()
+                })
+                _dbg(f"{model_name} weight={weight} logits={outputs.cpu().numpy()} probs={probs.cpu().numpy()}")
+
+            # Normalize by total weight
+            if total_weight > 0:
+                ensemble_probs /= total_weight
+
+        probs_numpy = ensemble_probs[0].cpu().numpy()
+        prediction_index = int(probs_numpy.argmax())
+        predicted_label = 'PNEUMONIA' if prediction_index == 1 else 'NORMAL'
+        _dbg(f"Ensemble probs={probs_numpy} predicted_index={prediction_index} label={predicted_label}")
+
+        result = {
+            'prediction': predicted_label,
+            'confidence': float(probs_numpy[prediction_index]),
+            'pneumonia_probability': float(probs_numpy[1]),
+            'normal_probability': float(probs_numpy[0]),
+            'models_used': available_models,
+            'per_model': per_model,
+            'prediction_type': 'ensemble'
+        }
+        return result
 
-            # Remove any DistributedDataParallel prefixes
-            new_state = {k.replace("module.", ""): v for k, v in state.items()}
 
-            missing, unexpected = model.load_state_dict(new_state, strict=False)
-            if missing:
-                logger.warning(f"Missing keys: {missing[:5]}...")
-            if unexpected:
-                logger.warning(f"Unexpected keys: {unexpected[:5]}...")
+# ============================================================================
+# Global Model System Instance
+# ============================================================================
 
-        except Exception as e:
-            logger.warning(f"Failed to load weights from {weight_path}: {e}")
-    else:
-        logger.warning(f"No weights file found at {weight_path}, using random weights")
+# Initialize model system
+model_system = PneumoniaModelSystem(device='cpu')
 
-    model.to(device)
-    model.eval()
+# Try to load models
+try:
+    model_system.load_models('models')
+    available_models = model_system.get_available_models()
+    print(f"Loaded {len(available_models)} models: {available_models}")
+except Exception as e:
+    print(f"Warning: Could not load models - {str(e)}")
+    print("   Running in demo mode")
+    available_models = []
 
-    if use_cache:
-        _model_instance_cache[cache_key] = (model, arch)
 
-    return model, arch
+# ============================================================================
+# Gradio Interface Functions
+# ============================================================================
 
-def load_image_any(file_obj):
+def predict_pneumonia(file_path, selected_model_option):
     """
-    Accepts .dcm, .png, .jpg, .jpeg, etc.
-    Returns a Pillow RGB image.
+    Main prediction function for Gradio
+
+    Args:
+        file_path: File path from Gradio file upload
+        selected_model_option: Selected model from radio button
+
+    Returns:
+        tuple: (processed_image, result_text, probability_dict, confidence_html)
     """
-    # Handle case when file is removed
-    if file_obj is None:
-        return None
+    if file_path is None:
+        return None, "Please upload an X-ray image", {}, ""
 
-    if isinstance(file_obj, str):
-        filepath = file_obj
-    else:
-        filepath = file_obj.name  # if it's a tempfile
-
-    ext = os.path.splitext(filepath)[1].lower()
-
-    if ext == ".dcm":
-        # Read DICOM file
-        ds = pydicom.dcmread(filepath)
-        img_array = ds.pixel_array.astype(float)
-
-        # Normalize to 0-255
-        img_array = (np.maximum(img_array, 0) / img_array.max()) * 255.0
-        img_array = np.uint8(img_array)
-        # Convert grayscale to RGB
-        img_rgb = Image.fromarray(img_array).convert("RGB")
-        return img_rgb
-    else:
-        # Standard formats
-        return Image.open(filepath).convert("RGB")
-
-
-_image_cache_transform: Dict[int, transforms.Compose] = {}
-
-
-def get_transform(arch: str) -> transforms.Compose:
-    """Get image transforms for the given architecture."""
-    size = _ARCH_INPUT_SIZE.get(arch, 224)
-    if size in _image_cache_transform:
-        return _image_cache_transform[size]
-
-    t = transforms.Compose([
-        transforms.Resize((size, size)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    _image_cache_transform[size] = t
-    return t
-
-
-# Grad-CAM utilities
-class GradCAM:
-    """Grad-CAM implementation for generating attention heatmaps."""
-
-    def __init__(self, model: torch.nn.Module, target_layer: Optional[str] = None):
-        self.model = model
-        self.model.eval()
-        self.target_layer = target_layer
-        self.activations = None
-        self.gradients = None
-
-        # Try to automatically pick a layer if not provided
-        if target_layer is None:
-            layer = None
-            # Common patterns for different architectures
-            layer_candidates = ["layer4", "features.denseblock4", "blocks.6", "conv_head", "features"]
-            for cand in layer_candidates:
-                parts = cand.split('.')
-                current = model
-                try:
-                    for part in parts:
-                        current = getattr(current, part)
-                    layer = current
-                    break
-                except AttributeError:
-                    continue
-            self.target_module = layer if layer is not None else model
+    try:
+        # Process uploaded image (handles JPEG, PNG, DICOM)
+        class FileObj:
+            def __init__(self, path):
+                self.name = path
+        
+        file_obj = FileObj(file_path)
+        processed_image = process_uploaded_image(file_obj)
+        
+        if processed_image is None:
+            return None, "Error: Could not process the uploaded image", {}, ""
+
+        # Get prediction based on selected model option
+        if selected_model_option == "Ensemble (All Models)":
+            result = model_system.predict_ensemble(processed_image)
+            model_info = f"Ensemble of {len(result['models_used'])} models"
         else:
-            self.target_module = dict(model.named_modules()).get(target_layer, model)
+            # Individual model prediction
+            result = model_system.predict_single_model(processed_image, selected_model_option)
+            model_info = f"{result['model_architecture']} - {result['model_description']}"
+
+        # Format result text
+        result_text = f"## Prediction: {result['prediction']}\n\n"
+        result_text += f"**Confidence:** {result['confidence']*100:.2f}%\n\n"
+        result_text += f"**Model Used:** {model_info}\n\n"
+
+        if result['prediction'] == 'PNEUMONIA':
+            result_text += "⚠️ **Pneumonia detected**\n\n"
+            result_text += "This X-ray shows signs consistent with pneumonia. "
+            result_text += "Please consult a qualified radiologist for confirmation."
+        else:
+            result_text += "✓ **No pneumonia detected**\n\n"
+            result_text += "This X-ray appears normal. "
+            result_text += "However, always consult a healthcare professional for accurate diagnosis."
+
+        # Create probability dictionary for bar chart
+        prob_dict = {
+            "Normal": result['normal_probability'],
+            "Pneumonia": result['pneumonia_probability']
+        }
 
-        def fwd_hook(_, __, output):
-            self.activations = output.detach()
+        # Create confidence HTML with color coding
+        confidence_pct = result['confidence'] * 100
+        if confidence_pct >= 90:
+            color = "green"
+            level = "Very High"
+        elif confidence_pct >= 75:
+            color = "blue"
+            level = "High"
+        elif confidence_pct >= 60:
+            color = "orange"
+            level = "Moderate"
+        else:
+            color = "red"
+            level = "Low"
+
+        confidence_html = f"""
+        <div style="padding: 20px; border-radius: 10px; background-color: #f0f0f0;">
+            <h3 style="color: {color};">Confidence Level: {level}</h3>
+            <p style="font-size: 24px; color: {color}; font-weight: bold;">{confidence_pct:.1f}%</p>
+            <p style="font-size: 12px; color: #666;">
+                Model: {model_info}
+            </p>
+        </div>
+        """
 
-        def bwd_hook(_, grad_input, grad_output):
-            if grad_output[0] is not None:
-                self.gradients = grad_output[0].detach()
+        return processed_image, result_text, prob_dict, confidence_html
 
-        self.target_module.register_forward_hook(fwd_hook)
-        # Use full backward hook (non-deprecated) for gradient capture
-        try:
-            self.target_module.register_full_backward_hook(bwd_hook)
-        except AttributeError:
-            # Fallback if running older torch without full hook
-            self.target_module.register_backward_hook(bwd_hook)
-
-    def generate(self, tensor: torch.Tensor, class_idx: Optional[int] = None) -> torch.Tensor:
-        """Generate Grad-CAM heatmap."""
-        tensor = tensor.requires_grad_(True)
-        logits = self.model(tensor)
-        if isinstance(logits, tuple):  # Inception may return (logits, aux)
-            logits = logits[0]
-
-        if class_idx is None:
-            class_idx = logits.argmax(dim=1).item()
-
-        score = logits[:, class_idx]
-        score.backward(retain_graph=True)
-
-        # Compute weights
-        grads = self.gradients  # [B, C, H, W]
-        acts = self.activations
-
-        if grads is None or acts is None:
-            raise RuntimeError("GradCAM hooks did not capture activations/gradients")
-
-        weights = grads.mean(dim=(2, 3), keepdim=True)  # [B, C, 1, 1]
-        cam = (weights * acts).sum(dim=1, keepdim=True)
-        cam = F.relu(cam)
-        cam = F.interpolate(cam, size=tensor.shape[2:], mode="bilinear", align_corners=False)
-
-        # Normalize
-        cam_min, cam_max = cam.min(), cam.max()
-        cam = (cam - cam_min) / (cam_max - cam_min + 1e-8)
-        return cam.squeeze(0).squeeze(0)  # [H, W]
-
-
-def predict(image: Image.Image, model_name: str, approach: str, grad_cam: bool = False, top_k: int = 5):
-    """Run inference on the uploaded image."""
-    if image is None:
-        return [], None, None
-    else:
-        image = load_image_any(image)
-    try:
-        model, arch = load_model(model_name, approach, use_cache=True)
     except Exception as e:
-        logger.error(f"Model loading failed: {e}")
-        error_df = [["Error", 0.0, str(e)]]
-        return error_df, image, None
+        error_msg = f"Error processing image: {str(e)}"
+        return None, error_msg, {}, f"<p style='color: red;'>{error_msg}</p>"
 
-    try:
-        transform = get_transform(arch)
-        tensor = transform(image).unsqueeze(0)
-        device = next(model.parameters()).device
-        tensor = tensor.to(device)
 
-        with torch.no_grad():
-            out = model(tensor)
-            if isinstance(out, tuple):  # Inception
-                out = out[0]
-            probs = F.softmax(out, dim=1).cpu().squeeze(0)
-
-        top_k = min(top_k, probs.shape[0])
-        top_probs, top_indices = torch.topk(probs, top_k)
-
-        results = []
-        for p, idx in zip(top_probs.tolist(), top_indices.tolist()):
-            label = CLASS_LABELS.get(idx, f"Class_{idx}")
-            results.append([label, f"{p * 100:.2f}%"])
-
-        cam_img = None
-        if grad_cam:
-            try:
-                gcam = GradCAM(model)
-                cam = gcam.generate(tensor)
-
-                # Convert cam to PIL heatmap overlay
-                if cv2 is not None:
-                    import numpy as np
-                    base_img = image.resize((cam.shape[1], cam.shape[0]))
-                    base_arr = np.array(base_img)
-                    heatmap = (cam.cpu().numpy() * 255).astype('uint8')
-                    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-                    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
-                    overlay = (0.4 * heatmap + 0.6 * base_arr).astype('uint8')
-                    cam_img = Image.fromarray(overlay)
-                else:
-                    # Fallback without OpenCV
-                    import numpy as np
-                    cam_np = cam.cpu().numpy()
-                    cam_img = Image.fromarray((cam_np * 255).astype('uint8'))
-
-            except Exception as e:
-                logger.warning(f"Grad-CAM failed: {e}")
-                cam_img = Image.new('RGB', (224, 224), color=(255, 100, 100))
-
-        return results, image, cam_img
+def get_model_system_info():
+    """Return information about the model system"""
+    info = f"""
+    # Model System Information
 
-    except Exception as e:
-        logger.error(f"Prediction failed: {e}")
-        error_df = [["Error", 0.0, str(e)]]
-        return error_df, image, None
-
-def build_interface():
-    """Build the Gradio interface."""
-    mapping = scan_weights()
-    model_choices = sorted(mapping.keys())
-
-    with gr.Blocks(
-        title="Transfer Learning Inference",
-        theme=gr.themes.Soft(),
-        css="""
-        .gradio-container {
-            max-width: 1200px;
-            margin: auto;
-        }
+    **Available Models:** {len(available_models)}
+
+    **Individual Models:**
+    """
+    for model_name in available_models:
+        model_info = model_system.model_definitions[model_name]
+        info += f"\n- **{model_name}**: {model_info['architecture']} - {model_info['description']}"
+
+    info += f"""
+
+    **Ensemble Configuration:**
+    - Uses weighted voting from multiple models
+    - Ensemble weights: {model_system.ensemble_weights}
+
+    **Performance Expectations:**
+    - Individual models may have varying strengths
+    - Ensemble typically provides more robust predictions
+    - Model selection allows comparing different approaches
+
+    **Best Practices:**
+    - Use ensemble for most reliable results
+    - Compare individual models to understand prediction confidence
+    - Consider model architecture strengths for specific cases
+    """
+
+    return info
+
+
+# ============================================================================
+# Gradio Interface
+# ============================================================================
+
+# Custom CSS
+custom_css = """
+.gradio-container {
+    font-family: 'Arial', sans-serif;
+}
+.output-markdown h2 {
+    color: #2c3e50;
+}
+"""
+
+# Create Gradio interface
+with gr.Blocks(css=custom_css, title="Pneumonia Detection AI") as demo:
+    gr.Markdown(
         """
-    ) as demo:
-        gr.Markdown(
-            """
-            # Transfer Learning Inference
+        # Pneumonia Detection from Chest X-rays
+        ### AI-Powered Individual Models and Ensemble for Medical Screening
+
+        Upload a chest X-ray image and select a model to detect signs of pneumonia using our
+        state-of-the-art deep learning models.
+
+        **DISCLAIMER:** This tool is for research and educational purposes only.
+        It should not be used as a substitute for professional medical diagnosis.
+        Always consult qualified healthcare professionals for medical advice.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Input section - Use File upload to handle DICOM properly
+            input_file = gr.File(
+                label="Upload Chest X-Ray Image (JPEG, PNG, DICOM)",
+                file_types=[".jpg", ".jpeg", ".png", ".dcm", ".dicom"],
+                type="filepath"
+            )
             
-            Select a pre-trained model and approach, upload an image, and view predictions with optional attention visualization.
+            # Model selection radio button
+            model_options = ["Ensemble (All Models)"] + available_models
+            model_selector = gr.Radio(
+                choices=model_options,
+                value="Ensemble (All Models)",
+                label="Select Model",
+                info="Choose an individual model or use the ensemble of all models"
+            )
+
+            predict_btn = gr.Button("Analyze X-Ray", variant="primary", size="lg")
+
+            gr.Markdown("### Supported Formats")
+            gr.Markdown("**JPEG, PNG, DICOM** (.dcm, .dicom files)")
+            gr.Markdown("Try with your own chest X-ray images")
+
+        with gr.Column(scale=1):
+            # Preview section for processed image
+            preview_image = gr.Image(
+                label="Processed Image Preview",
+                height=400,
+                interactive=False
+            )
             
-            **Available Models:** ResNet, Inception V3, and Ensemble
+            gr.Markdown("*This shows how the image appears to the AI model*")
+
+        with gr.Column(scale=1):
+            # Output section
+            output_text = gr.Markdown(label="Diagnosis Result")
+            confidence_html = gr.HTML(label="Confidence Level")
+            prob_chart = gr.Label(label="Probability Distribution", num_top_classes=2)
+
+    # Model system info accordion
+    with gr.Accordion("🤖 About These Models", open=False):
+        model_info = gr.Markdown(get_model_system_info())
+
+    # Technical details accordion
+    with gr.Accordion("🔬 Technical Details", open=False):
+        gr.Markdown(
+            """
+            ### Model Architecture
+
+            This system uses an ensemble of multiple deep learning models:
+            - **VGG19** (Transfer Learning): Deep convolutional network with 19 layers
+            - **VGG16** (Transfer Learning): 16-layer convolutional network
+            - **EfficientNet-B0**: Efficient architecture with compound scaling
+
+            ### How It Works
+
+            1. **Image Preprocessing**: X-ray image is resized to 224×224 and normalized
+            2. **Ensemble Prediction**: Each model independently analyzes the image
+            3. **Weighted Voting**: Predictions are combined using learned weights
+            4. **Confidence Score**: Final probability based on ensemble agreement
+
+            ### Performance Metrics
+
+            The ensemble has been validated on 33 different model configurations
+            and ranked based on clinical utility metrics.
+
+            ### Dataset
+
+            Models trained on chest X-ray dataset with thousands of images
+            from actual clinical cases.
             """
-        )
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                model_dd = gr.Dropdown(
-                    choices=model_choices,
-                    label="Model",
-                    value=model_choices[0] if model_choices else None,
-                    info="Select the model architecture"
-                )
-                approach_dd = gr.Dropdown(
-                    choices=[],
-                    label="Approach",
-                    info="Select the training approach/variant"
-                )
-                grad_cam_cb = gr.Checkbox(
-                    label="Generate Grad-CAM",
-                    value=False,
-                    info="Show attention heatmap overlay"
-                )
-                
-            with gr.Column(scale=2):
-                image_in = gr.File(
-                    file_types=["image", ".dcm"],
-                    label="Upload a chest x-Ray image for classification (.png, .jpg, .jpeg, .dcm)",
-                )
-                image_preview = gr.Image(interactive=False, size=(224, 224))
-                original_out = gr.Image(interactive=False, size=(224, 224))
-                cam_out = gr.Image(interactive=False, size=(224, 224))
-                #image_in.change(fn=load_image_any, inputs=image_in, outputs=image_preview)
-                image_in.change(
-                    fn=load_image_any,
-                    inputs=image_in,
-                    outputs=[image_preview, original_out, cam_out],
-                    show_progress=False
-                )
-        
-        submit_btn = gr.Button("Run Inference", variant="primary", size="lg")
-        
-        with gr.Row():
-            with gr.Column():
-                results_out = gr.Dataframe(
-                    headers=["Label", "Probability"],
-                    datatype=["str", "number"],
-                    label="Top Predictions",
-                    interactive=False
-                )
-            
-            with gr.Column():
-                with gr.Row():
-                    original_out = gr.Image(label="Original Image", interactive=False)
-                    cam_out = gr.Image(label="Grad-CAM Overlay", interactive=False)
-        
-        # Add model info
-        with gr.Accordion("Model Information", open=False):
-            gr.Markdown(f"""
-            - **Number of Classes:** {NUM_CLASSES}
-            - **Available Models:** {len(model_choices)}
-            - **Environment:** {'HuggingFace Spaces' if MODEL_REPO_ID else 'Local'}
-            """)
-
-        def update_approaches(selected_model):
-            if not selected_model:
-                return gr.update(choices=[], value=None)
-            mapping_local = scan_weights()
-            apprs = sorted(mapping_local.get(selected_model, {}).keys())
-            value = apprs[0] if apprs else None
-            return gr.update(choices=apprs, value=value)
-
-        model_dd.change(fn=update_approaches, inputs=model_dd, outputs=approach_dd)
-
-        submit_btn.click(
-            fn=predict,
-            inputs=[image_in, model_dd, approach_dd, grad_cam_cb],
-            outputs=[results_out, original_out, cam_out],
         )
 
-        def init_approaches():
-            if not model_choices:
-                return gr.update(choices=[], value=None)
-            mapping_local = scan_weights()
-            apprs = sorted(mapping_local.get(model_choices[0], {}).keys())
-            value = apprs[0] if apprs else None
-            return gr.update(choices=apprs, value=value)
-
-        # Initialize approaches for first model
-        if model_choices:
-            demo.load(fn=init_approaches, inputs=[], outputs=approach_dd)
-    
-    return demo
+    # Citation accordion
+    with gr.Accordion("📚 Citation & Credits", open=False):
+        gr.Markdown(
+            """
+            ### Citation
 
-def safe_gradio_launch(demo):
-    try:
-        demo.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)
-    except TypeError as e:
-        if "argument of type 'bool' is not iterable" in str(e):
-            logger.warning("Gradio schema bug detected, restarting with share=True fallback.")
-            demo.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=False)
-        else:
-            raise
+            If you use this model in your research, please cite:
 
-def main():
-    """Main function to launch the Gradio app."""
-    logger.info("Starting Transfer Learning Inference App")
-    demo = build_interface()
-    
-    # Configuration for different environments
-    server_name = os.getenv("GRADIO_SERVER_NAME", "0.0.0.0")
-    server_port = int(os.getenv("GRADIO_SERVER_PORT", "7860"))
-    share = os.getenv("GRADIO_SHARE", "False").lower() == "true"
-    
-    # demo.launch(
-    #     server_name=server_name,
-    #     server_port=server_port,
-    #     share=share,
-    #     show_error=True
-    # )
+            ```
+            @software{pneumonia_ensemble_2025,
+              title={Deep Learning Ensemble for Pneumonia Detection},
+              author={Prabakaran Thangamani},
+              year={2025},
+              url={https://huggingface.co/spaces/papsofts/pneumonia-detection}
+            }
+            ```
+
+            ### Credits
 
-    safe_gradio_launch(demo)
+            - **Model Development**: Based on 33 model architectures including VGG, EfficientNet, ResNet families
+            - **Framework**: PyTorch, torchvision
+            - **Interface**: Gradio
+            - **Deployment**: Hugging Face Spaces
 
+            ### License
+
+            This model is released for research and educational purposes.
+            """
+        )
+
+    # Connect the prediction function
+    predict_btn.click(
+        fn=predict_pneumonia,
+        inputs=[input_file, model_selector],
+        outputs=[preview_image, output_text, prob_chart, confidence_html]
+    )
+
+    # Examples (you can add actual example images here)
+    gr.Examples(
+        examples=[],
+        inputs=input_file,
+        label="Example X-Ray Images"
+    )
+
+
+# ============================================================================
+# Launch
+# ============================================================================
 
 if __name__ == "__main__":
-    # Quick internal test if requested
-    if os.environ.get("HEADLESS_TEST") == "1":
-        logger.info("Running headless test")
-        mapping = scan_weights()
-        if mapping:
-            first_model = next(iter(mapping))
-            first_appr = next(iter(mapping[first_model]))
-            try:
-                model, arch = load_model(first_model, first_appr)
-                size = _ARCH_INPUT_SIZE.get(arch, 224)
-                x = torch.randn(1, 3, size, size)
-                out = model(x)
-                if isinstance(out, tuple):
-                    out = out[0]
-                logger.info(f"Test forward output shape: {out.shape}")
-            except Exception as e:
-                logger.error(f"Test failed: {e}")
-        else:
-            logger.warning("No models found for testing")
-    else:
-        main()
+    demo.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860
+    )

models/Model_A1_7CB_Appr_D.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ee1e63db1ee140d6dd425400e3745918ff0aed42d1fe65c85a520d1967b5e2e
+size 558329434

models/Model_C_Appr_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4018a7db58cf03fa3965b5195fdaef3ffee034eb445d98b6ee70ecbf4e104153
+size 9163946

models/Model_F_Appr_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69bf003475478056556f558cc9f3f511f99869c0ba57617a2c918776207f1dab
+size 94379802

models/Model_G_Appr_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05bb0b26d6bc63fe82f47e3c5ffc4c6dc0d34bd483d17b74480930f9c6583fa8
+size 16359314

models/Model_H_Appr_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a6c44c009fde3eb81e55dd87d1a8606a325ed5fb82bcc5635f6a7c0a7fbdde2
+size 28459058