Medical AI Development: Computer Vision for Healthcare Image Analysis

This preprocessing pipeline handles medical images with proper normalization, augmentation, and quality control for AI model training.

import torch
import torchvision.transforms as transforms
import numpy as np
import pydicom
from typing import Dict, List, Tuple
import cv2
from skimage import exposure, filters

class MedicalImagePreprocessor:
    def __init__(self, config: Dict):
        self.config = config
        self.target_size = config.get('target_size', (512, 512))
        self.normalization_method = config.get('normalization', 'z_score')
        
    def load_dicom_image(self, file_path: str) -> np.ndarray:
        """Load DICOM medical image with proper handling"""
        try:
            dicom = pydicom.dcmread(file_path)
            image = dicom.pixel_array.astype(np.float32)
            
            # Handle different photometric interpretations
            if dicom.PhotometricInterpretation == 'MONOCHROME1':
                image = np.max(image) - image
            
            # Apply windowing if available
            if hasattr(dicom, 'WindowCenter') and hasattr(dicom, 'WindowWidth'):
                image = self.apply_windowing(image, dicom.WindowCenter, dicom.WindowWidth)
            
            return image
            
        except Exception as e:
            raise ValueError(f"Failed to load DICOM image: {str(e)}")
    
    def apply_windowing(self, image: np.ndarray, center: float, width: float) -> np.ndarray:
        """Apply medical image windowing for optimal contrast"""
        min_val = center - width / 2
        max_val = center + width / 2
        
        image = np.clip(image, min_val, max_val)
        image = (image - min_val) / (max_val - min_val)
        
        return image
    
    def normalize_image(self, image: np.ndarray) -> np.ndarray:
        """Normalize medical image using appropriate method"""
        if self.normalization_method == 'z_score':
            # Z-score normalization
            mean = np.mean(image)
            std = np.std(image)
            if std > 0:
                image = (image - mean) / std
        elif self.normalization_method == 'min_max':
            # Min-max normalization
            min_val = np.min(image)
            max_val = np.max(image)
            if max_val > min_val:
                image = (image - min_val) / (max_val - min_val)
        elif self.normalization_method == 'histogram_equalization':
            # Histogram equalization for better contrast
            image = exposure.equalize_hist(image)
        
        return image
    
    def augment_medical_image(self, image: np.ndarray, mask: np.ndarray = None) -> Tuple[np.ndarray, np.ndarray]:
        """Apply medical-specific data augmentation"""
        augmented_image = image.copy()
        augmented_mask = mask.copy() if mask is not None else None
        
        # Random rotation (small angles for medical images)
        if np.random.random() > 0.5:
            angle = np.random.uniform(-15, 15)
            augmented_image = self.rotate_image(augmented_image, angle)
            if augmented_mask is not None:
                augmented_mask = self.rotate_image(augmented_mask, angle)
        
        # Random brightness adjustment
        if np.random.random() > 0.5:
            factor = np.random.uniform(0.8, 1.2)
            augmented_image = np.clip(augmented_image * factor, 0, 1)
        
        # Random noise addition (simulate acquisition noise)
        if np.random.random() > 0.5:
            noise_level = np.random.uniform(0.01, 0.05)
            noise = np.random.normal(0, noise_level, augmented_image.shape)
            augmented_image = np.clip(augmented_image + noise, 0, 1)
        
        # Random elastic deformation (simulate tissue deformation)
        if np.random.random() > 0.5:
            augmented_image = self.elastic_deformation(augmented_image)
            if augmented_mask is not None:
                augmented_mask = self.elastic_deformation(augmented_mask)
        
        return augmented_image, augmented_mask
    
    def elastic_deformation(self, image: np.ndarray, alpha: float = 1000, sigma: float = 30) -> np.ndarray:
        """Apply elastic deformation for data augmentation"""
        random_state = np.random.RandomState(None)
        
        shape = image.shape
        dx = filters.gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha
        dy = filters.gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha
        
        x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1))
        
        return cv2.remap(image, indices[1].astype(np.float32), indices[0].astype(np.float32), cv2.INTER_LINEAR)
    
    def create_training_batch(self, image_paths: List[str], labels: List[int]) -> Dict:
        """Create training batch with proper preprocessing"""
        batch_images = []
        batch_labels = []
        
        for image_path, label in zip(image_paths, labels):
            # Load and preprocess image
            image = self.load_dicom_image(image_path)
            image = self.normalize_image(image)
            
            # Resize to target size
            image = cv2.resize(image, self.target_size)
            
            # Apply augmentation during training
            if self.config.get('augment', True):
                image, _ = self.augment_medical_image(image)
            
            batch_images.append(image)
            batch_labels.append(label)
        
        return {
            'images': torch.tensor(np.array(batch_images)).unsqueeze(1),  # Add channel dimension
            'labels': torch.tensor(batch_labels)
        }

class MedicalImageQualityControl:
    def __init__(self):
        self.quality_metrics = {}
    
    def assess_image_quality(self, image: np.ndarray) -> Dict:
        """Assess medical image quality for AI training"""
        quality_report = {
            'contrast_ratio': self.calculate_contrast_ratio(image),
            'sharpness': self.calculate_sharpness(image),
            'noise_level': self.calculate_noise_level(image),
            'brightness': np.mean(image),
            'overall_quality': 0
        }
        
        # Calculate overall quality score
        quality_score = 0
        if quality_report['contrast_ratio'] > 0.3:
            quality_score += 25
        if quality_report['sharpness'] > 0.1:
            quality_score += 25
        if quality_report['noise_level'] < 0.1:
            quality_score += 25
        if 0.2 < quality_report['brightness'] < 0.8:
            quality_score += 25
        
        quality_report['overall_quality'] = quality_score
        
        return quality_report
    
    def calculate_contrast_ratio(self, image: np.ndarray) -> float:
        """Calculate contrast ratio of medical image"""
        return np.std(image) / (np.mean(image) + 1e-8)
    
    def calculate_sharpness(self, image: np.ndarray) -> float:
        """Calculate image sharpness using Laplacian variance"""
        laplacian = cv2.Laplacian(image, cv2.CV_64F)
        return np.var(laplacian)
    
    def calculate_noise_level(self, image: np.ndarray) -> float:
        """Estimate noise level in medical image"""
        # Use high-frequency content as noise estimate
        kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]])
        filtered = cv2.filter2D(image, -1, kernel)
        return np.std(filtered)

This medical AI model includes uncertainty quantification, attention mechanisms, and clinical decision support features required for healthcare applications.

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, List
import numpy as np

class MedicalAIDiagnosticModel(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_classes = config['num_classes']
        self.uncertainty_threshold = config.get('uncertainty_threshold', 0.8)
        
        # Backbone network (e.g., ResNet, DenseNet, or Vision Transformer)
        self.backbone = self._build_backbone(config['backbone'])
        
        # Attention mechanism for interpretability
        self.attention = SpatialAttentionModule()
        
        # Uncertainty quantification
        self.uncertainty_head = UncertaintyHead(self.backbone.feature_dim)
        
        # Clinical decision support
        self.clinical_head = ClinicalDecisionHead(self.backbone.feature_dim)
        
        # Main classification head
        self.classifier = nn.Linear(self.backbone.feature_dim, self.num_classes)
        
    def _build_backbone(self, backbone_type: str) -> nn.Module:
        """Build backbone network for medical image analysis"""
        if backbone_type == 'resnet50':
            return MedicalResNet50()
        elif backbone_type == 'densenet121':
            return MedicalDenseNet121()
        elif backbone_type == 'vit':
            return MedicalVisionTransformer()
        else:
            raise ValueError(f"Unsupported backbone: {backbone_type}")
    
    def forward(self, x: torch.Tensor) -> Dict:
        """Forward pass with uncertainty quantification"""
        batch_size = x.size(0)
        
        # Extract features
        features = self.backbone(x)
        
        # Apply attention mechanism
        attended_features, attention_map = self.attention(features, x)
        
        # Calculate uncertainty
        uncertainty = self.uncertainty_head(attended_features)
        
        # Generate clinical insights
        clinical_insights = self.clinical_head(attended_features)
        
        # Main classification
        logits = self.classifier(attended_features)
        probabilities = F.softmax(logits, dim=1)
        
        # Calculate confidence
        confidence = torch.max(probabilities, dim=1)[0]
        
        return {
            'logits': logits,
            'probabilities': probabilities,
            'confidence': confidence,
            'uncertainty': uncertainty,
            'attention_map': attention_map,
            'clinical_insights': clinical_insights,
            'prediction': torch.argmax(logits, dim=1)
        }
    
    def predict_with_uncertainty(self, x: torch.Tensor) -> Dict:
        """Make prediction with uncertainty quantification"""
        with torch.no_grad():
            outputs = self.forward(x)
            
            # Filter low-confidence predictions
            high_confidence_mask = outputs['confidence'] > self.uncertainty_threshold
            
            results = {
                'predictions': outputs['prediction'],
                'probabilities': outputs['probabilities'],
                'confidence': outputs['confidence'],
                'uncertainty': outputs['uncertainty'],
                'attention_map': outputs['attention_map'],
                'clinical_insights': outputs['clinical_insights'],
                'high_confidence': high_confidence_mask
            }
            
            return results

class SpatialAttentionModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention_conv = nn.Sequential(
            nn.Conv2d(2048, 1024, 1),
            nn.ReLU(),
            nn.Conv2d(1024, 1, 1),
            nn.Sigmoid()
        )
    
    def forward(self, features: torch.Tensor, original_image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Apply spatial attention to features"""
        # Generate attention map
        attention_map = self.attention_conv(features)
        
        # Apply attention to features
        attended_features = features * attention_map
        
        # Global average pooling
        attended_features = F.adaptive_avg_pool2d(attended_features, (1, 1))
        attended_features = attended_features.flatten(1)
        
        return attended_features, attention_map

class UncertaintyHead(nn.Module):
    def __init__(self, feature_dim: int):
        super().__init__()
        self.uncertainty_net = nn.Sequential(
            nn.Linear(feature_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    
    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Estimate prediction uncertainty"""
        return self.uncertainty_net(features)

class ClinicalDecisionHead(nn.Module):
    def __init__(self, feature_dim: int):
        super().__init__()
        self.clinical_net = nn.Sequential(
            nn.Linear(feature_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64)  # Clinical features
        )
    
    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Generate clinical decision support features"""
        return self.clinical_net(features)

class MedicalResNet50(nn.Module):
    def __init__(self):
        super().__init__()
        import torchvision.models as models
        
        # Load pre-trained ResNet50
        resnet = models.resnet50(pretrained=True)
        
        # Remove final classification layer
        self.features = nn.Sequential(*list(resnet.children())[:-1])
        self.feature_dim = 2048
        
        # Medical-specific modifications
        self.medical_adaptation = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through ResNet backbone"""
        return self.features(x)

class MedicalTrainingPipeline:
    def __init__(self, model: MedicalAIDiagnosticModel, config: Dict):
        self.model = model
        self.config = config
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=config.get('learning_rate', 1e-4),
            weight_decay=config.get('weight_decay', 1e-5)
        )
        
        # Medical-specific loss function
        self.criterion = MedicalLossFunction(
            class_weights=config.get('class_weights'),
            uncertainty_weight=config.get('uncertainty_weight', 0.1)
        )
    
    def train_epoch(self, dataloader) -> Dict:
        """Train model for one epoch with medical-specific metrics"""
        self.model.train()
        epoch_loss = 0
        epoch_accuracy = 0
        epoch_uncertainty = 0
        
        for batch_idx, batch in enumerate(dataloader):
            images = batch['images'].cuda()
            labels = batch['labels'].cuda()
            
            self.optimizer.zero_grad()
            
            # Forward pass
            outputs = self.model(images)
            
            # Calculate loss
            loss = self.criterion(outputs, labels)
            
            # Backward pass
            loss.backward()
            self.optimizer.step()
            
            # Calculate metrics
            epoch_loss += loss.item()
            epoch_accuracy += self.calculate_accuracy(outputs['prediction'], labels)
            epoch_uncertainty += outputs['uncertainty'].mean().item()
        
        return {
            'loss': epoch_loss / len(dataloader),
            'accuracy': epoch_accuracy / len(dataloader),
            'uncertainty': epoch_uncertainty / len(dataloader)
        }
    
    def calculate_accuracy(self, predictions: torch.Tensor, labels: torch.Tensor) -> float:
        """Calculate classification accuracy"""
        correct = (predictions == labels).float()
        return correct.mean().item()

class MedicalLossFunction(nn.Module):
    def __init__(self, class_weights: List[float] = None, uncertainty_weight: float = 0.1):
        super().__init__()
        self.class_weights = class_weights
        self.uncertainty_weight = uncertainty_weight
        
        # Cross-entropy loss for classification
        self.classification_loss = nn.CrossEntropyLoss(weight=self.class_weights)
        
        # Uncertainty loss (encourage model to be uncertain on difficult cases)
        self.uncertainty_loss = nn.MSELoss()
    
    def forward(self, outputs: Dict, labels: torch.Tensor) -> torch.Tensor:
        """Calculate combined medical loss"""
        # Classification loss
        cls_loss = self.classification_loss(outputs['logits'], labels)
        
        # Uncertainty loss (penalize overconfident wrong predictions)
        uncertainty_loss = self.uncertainty_loss(
            outputs['uncertainty'],
            (outputs['prediction'] != labels).float()
        )
        
        # Combined loss
        total_loss = cls_loss + self.uncertainty_weight * uncertainty_loss
        
        return total_loss