Spaces:
Runtime error
Runtime error
| """ | |
| Utility functions for Multi-Label Thoracic Disease Classification with DenseNet121+GRU and Score-CAM. | |
| """ | |
| import os | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import tensorflow as tf | |
| from tensorflow.keras import layers, models | |
| from config import CLASS_NAMES, AUROC_SCORES, THRESHOLD | |
| def load_and_preprocess_image(img_path, target_size=(224, 224)): | |
| """ | |
| Load an image from file, convert to RGB, resize, and normalize to [0, 1]. | |
| Returns: | |
| img_tensor (np.ndarray): Tensor of shape (1, 224, 224, 3), normalized [0, 1]. | |
| img_display (np.ndarray): Image array of shape (224, 224, 3), float32 [0, 1]. | |
| """ | |
| if not os.path.exists(img_path): | |
| raise FileNotFoundError(f"Image not found at path: {img_path}") | |
| # Read image using OpenCV | |
| img_bgr = cv2.imread(img_path) | |
| if img_bgr is None: | |
| raise ValueError(f"Could not read image file: {img_path}") | |
| # Convert BGR to RGB | |
| img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) | |
| # Resize to target size | |
| img_resized = cv2.resize(img_rgb, target_size, interpolation=cv2.INTER_AREA) | |
| # Normalize to [0, 1] | |
| img_display = img_resized.astype(np.float32) / 255.0 | |
| # Expand batch dimension | |
| img_tensor = np.expand_dims(img_display, axis=0) | |
| return img_tensor, img_display | |
| def build_densenet_gru_model(num_classes=15, input_shape=(224, 224, 3)): | |
| """ | |
| Constructs the DenseNet121 + GRU hybrid model architecture: | |
| - Backbone: DenseNet121 (pre-trained on ImageNet) | |
| - Sequential/Recurrent Head: Reshape 7x7 spatial grid -> (49, 1024) -> GRU -> Dense Sigmoid | |
| """ | |
| inputs = layers.Input(shape=input_shape, name="input_image") | |
| # DenseNet121 Feature Extractor | |
| base_model = tf.keras.applications.DenseNet121( | |
| include_top=False, | |
| weights='imagenet', | |
| input_tensor=inputs | |
| ) | |
| # Target feature maps from last conv/relu: shape (batch, 7, 7, 1024) | |
| features = base_model.output | |
| # Reshape spatial features into sequence format: (batch, 49, 1024) | |
| feature_dim = features.shape[-1] # 1024 | |
| grid_dim = features.shape[1] * features.shape[2] # 7 * 7 = 49 | |
| x = layers.Reshape((grid_dim, feature_dim), name="spatial_to_sequence")(features) | |
| # Recurrent Temporal / Spatial Sequence Modeling with GRU | |
| x = layers.GRU(256, return_sequences=False, name="gru_reasoning")(x) | |
| x = layers.Dropout(0.3, name="dropout_1")(x) | |
| x = layers.Dense(128, activation="relu", name="dense_intermediate")(x) | |
| x = layers.Dropout(0.2, name="dropout_2")(x) | |
| # Multi-label classification head (Sigmoid activations for independent label probabilities) | |
| outputs = layers.Dense(num_classes, activation="sigmoid", name="multi_label_predictions")(x) | |
| model = models.Model(inputs=inputs, outputs=outputs, name="DenseNet121_GRU_Classifier") | |
| model.compile( | |
| optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4), | |
| loss="binary_crossentropy", | |
| metrics=["binary_accuracy"] | |
| ) | |
| return model | |
| def find_target_conv_layer(model): | |
| """ | |
| Finds the last Conv2D or ReLU feature layer in DenseNet121 for CAM extraction. | |
| """ | |
| # Priority search for DenseNet121 terminal conv layers | |
| preferred_layers = ['conv5_block16_2_conv', 'conv5_block16_2_relu', 'relu'] | |
| for name in preferred_layers: | |
| for layer in model.layers: | |
| if layer.name == name: | |
| return layer | |
| # Fallback search for any layer that outputs 4D tensor before reshaping | |
| for layer in reversed(model.layers): | |
| if isinstance(layer, (layers.Conv2D, layers.Activation)) or 'conv' in layer.name.lower() or 'relu' in layer.name.lower(): | |
| if len(layer.output_shape) == 4: | |
| return layer | |
| # Secondary fallback to base model if nested | |
| for layer in model.layers: | |
| if hasattr(layer, 'layers'): | |
| for sub_layer in reversed(layer.layers): | |
| if isinstance(sub_layer, layers.Conv2D) or 'conv' in sub_layer.name.lower(): | |
| return sub_layer | |
| raise ValueError("Could not automatically locate a target 4D Conv layer in the model.") | |
| def generate_scorecam(model, img_tensor, class_idx, top_k=20): | |
| """ | |
| Gradient-free Score-CAM implementation (Wang et al., 2020). | |
| Args: | |
| model: Trained Keras multi-label classification model. | |
| img_tensor (np.ndarray): Preprocessed image tensor of shape (1, 224, 224, 3). | |
| class_idx (int): Index of the disease class to explain. | |
| top_k (int): Number of top activation maps to evaluate for efficiency. | |
| Returns: | |
| heatmap (np.ndarray): Normalized 2D Score-CAM heatmap of shape (224, 224) in [0, 1]. | |
| """ | |
| target_layer = find_target_conv_layer(model) | |
| # Sub-model extracting activation maps from target conv layer | |
| activation_model = tf.keras.models.Model( | |
| inputs=model.input, | |
| outputs=target_layer.output | |
| ) | |
| # 1. Forward pass to extract feature activation maps | |
| activations = activation_model.predict(img_tensor, verbose=0)[0] # (H, W, C) e.g., (7, 7, 1024) | |
| h, w, c = activations.shape | |
| input_h, input_w = img_tensor.shape[1], img_tensor.shape[2] | |
| # 2. Select top_k activation channels by activation variance / variance energy | |
| channel_energies = np.array([np.var(activations[:, :, i]) for i in range(c)]) | |
| top_indices = np.argsort(channel_energies)[::-1][:min(top_k, c)] | |
| weights = [] | |
| upsampled_masks = [] | |
| # 3. Upsample and normalize each selected activation channel | |
| for idx in top_indices: | |
| act_map = activations[:, :, idx] | |
| # Upsample to input resolution (224, 224) | |
| act_upsampled = cv2.resize(act_map, (input_w, input_h), interpolation=cv2.INTER_LINEAR) | |
| # Min-Max Normalization to [0, 1] | |
| act_min, act_max = np.min(act_upsampled), np.max(act_upsampled) | |
| if act_max - act_min > 1e-8: | |
| norm_mask = (act_upsampled - act_min) / (act_max - act_min) | |
| else: | |
| norm_mask = np.zeros_like(act_upsampled) | |
| upsampled_masks.append(norm_mask) | |
| # 4. Mask input image with the normalized activation map | |
| masked_img = img_tensor * np.expand_dims(np.expand_dims(norm_mask, axis=0), axis=-1) | |
| # 5. Forward pass on masked image to compute target class score | |
| score_preds = model.predict(masked_img, verbose=0) | |
| target_score = float(score_preds[0, class_idx]) | |
| weights.append(target_score) | |
| # 6. Linear combination of activation maps weighted by forward prediction score | |
| weights = np.array(weights) | |
| upsampled_masks = np.array(upsampled_masks) | |
| # Apply softmax on weights to stabilize contribution | |
| if np.sum(weights) > 0: | |
| exp_weights = np.exp(weights - np.max(weights)) | |
| norm_weights = exp_weights / np.sum(exp_weights) | |
| else: | |
| norm_weights = np.ones_like(weights) / len(weights) | |
| score_cam = np.zeros((input_h, input_w), dtype=np.float32) | |
| for i in range(len(top_indices)): | |
| score_cam += norm_weights[i] * upsampled_masks[i] | |
| # 7. Apply ReLU (keep positive influences only) | |
| score_cam = np.maximum(score_cam, 0) | |
| # 8. Final Min-Max normalization to [0, 1] | |
| cam_min, cam_max = np.min(score_cam), np.max(score_cam) | |
| if cam_max - cam_min > 1e-8: | |
| score_cam = (score_cam - cam_min) / (cam_max - cam_min) | |
| else: | |
| score_cam = np.zeros_like(score_cam) | |
| return score_cam | |
| def save_visual_report(img_np, heatmap, disease_name, prob, output_path, auroc=None): | |
| """ | |
| Generates and saves a publication-quality 3-panel visualization: | |
| 1. Original Chest X-Ray | |
| 2. Score-CAM Heatmap | |
| 3. Blended Overlay | |
| """ | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Convert heatmap to uint8 colormap | |
| heatmap_uint8 = np.uint8(255 * heatmap) | |
| jet_colormap = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET) | |
| jet_rgb = cv2.cvtColor(jet_colormap, cv2.COLOR_BGR2RGB) / 255.0 | |
| # Ensure img_np is float [0, 1] | |
| if img_np.max() > 1.0: | |
| img_np = img_np / 255.0 | |
| # Alpha blending: Overlay = 0.55 * original + 0.45 * heatmap | |
| overlay = 0.55 * img_np + 0.45 * jet_rgb | |
| overlay = np.clip(overlay, 0.0, 1.0) | |
| # Create figure | |
| fig, axes = plt.subplots(1, 3, figsize=(16, 5), dpi=150) | |
| fig.patch.set_facecolor('#0f172a') # Modern dark background | |
| # 1. Original X-Ray | |
| axes[0].imshow(img_np) | |
| axes[0].set_title("Input Chest X-Ray", fontsize=12, fontweight='bold', color='white', pad=10) | |
| axes[0].axis('off') | |
| # 2. Score-CAM Heatmap | |
| im1 = axes[1].imshow(heatmap, cmap='jet') | |
| axes[1].set_title(f"Score-CAM Activation\n({disease_name})", fontsize=12, fontweight='bold', color='white', pad=10) | |
| axes[1].axis('off') | |
| cbar = fig.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04) | |
| cbar.ax.yaxis.set_tick_params(color='white') | |
| plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='white') | |
| # 3. Superimposed Overlay | |
| axes[2].imshow(overlay) | |
| axes[2].set_title(f"Diagnostic Overlay\nProb: {prob*100:.1f}% (Threshold: {THRESHOLD*100:.0f}%)", | |
| fontsize=12, fontweight='bold', color='#38bdf8', pad=10) | |
| axes[2].axis('off') | |
| # Main Super Title | |
| auroc_str = f" | Baseline AUROC: {auroc:.4f}" if auroc is not None else "" | |
| fig.suptitle( | |
| f"Thoracic AI Diagnostic Report: {disease_name.upper()}{auroc_str}", | |
| fontsize=14, fontweight='bold', color='#f8fafc', y=0.98 | |
| ) | |
| plt.tight_layout() | |
| plt.savefig(output_path, facecolor=fig.get_facecolor(), bbox_inches='tight') | |
| plt.close(fig) | |
| def generate_synthetic_chest_xray(output_path): | |
| """ | |
| Creates an anatomically textured synthetic chest X-Ray image for testing. | |
| """ | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| img_size = 512 | |
| xray = np.zeros((img_size, img_size), dtype=np.float32) | |
| # 1. Background soft gradient | |
| y, x = np.mgrid[0:img_size, 0:img_size] | |
| center_y, center_x = img_size // 2, img_size // 2 | |
| r = np.sqrt((x - center_x)**2 + (y - center_y)**2) | |
| xray += 30 + 20 * np.sin(y / 50.0) | |
| # 2. Thoracic cage / Outer body contour | |
| body_mask = ((x - center_x)**2 / (200**2) + (y - center_y)**2 / (230**2)) <= 1.0 | |
| xray[body_mask] += 60.0 | |
| # 3. Bilateral Lung Fields (darker, air-filled regions) | |
| # Left Lung | |
| left_lung_mask = (((x - 170)**2 / (65**2) + (y - 250)**2 / (130**2)) <= 1.0) | |
| xray[left_lung_mask] = np.clip(xray[left_lung_mask] - 50.0, 15, 255) | |
| # Right Lung | |
| right_lung_mask = (((x - 340)**2 / (65**2) + (y - 250)**2 / (130**2)) <= 1.0) | |
| xray[right_lung_mask] = np.clip(xray[right_lung_mask] - 50.0, 15, 255) | |
| # 4. Mediastinum & Cardiac Silhouette (white/dense heart in left lower quadrant) | |
| heart_mask = (((x - 220)**2 / (55**2) + (y - 300)**2 / (55**2)) <= 1.0) | |
| xray[heart_mask] += 90.0 | |
| # 5. Clavicles and Ribs (horizontal elliptical density lines) | |
| for rib_y in range(160, 420, 35): | |
| rib_line = np.exp(-((y - rib_y - 0.05 * (x - center_x)**2 / 10)**2) / 12.0) | |
| xray += rib_line * 25.0 | |
| # Clavicles (top arches) | |
| left_clavicle = np.exp(-((y - 130 - 0.15 * (x - 180))**2) / 8.0) * (x < 250) * (x > 100) | |
| right_clavicle = np.exp(-((y - 130 + 0.15 * (x - 330))**2) / 8.0) * (x > 260) * (x < 410) | |
| xray += (left_clavicle + right_clavicle) * 45.0 | |
| # 6. Simulate a focal opacity / consolidation infiltrate in the right lower lung | |
| infiltrate = np.exp(-(((x - 340)**2 + (y - 310)**2) / (30.0**2))) * 70.0 | |
| xray += infiltrate | |
| # 7. Add realistic anatomical Poisson/Gaussian noise & smooth | |
| noise = np.random.normal(0, 4.0, (img_size, img_size)) | |
| xray = np.clip(xray + noise, 0, 255).astype(np.uint8) | |
| xray_blurred = cv2.GaussianBlur(xray, (3, 3), 0) | |
| # Convert to 3-channel RGB image and save | |
| xray_rgb = cv2.cvtColor(xray_blurred, cv2.COLOR_GRAY2BGR) | |
| cv2.imwrite(output_path, xray_rgb) | |
| return output_path | |