| import os |
| import numpy as np |
| import matplotlib.pyplot as plt |
| import tensorflow as tf |
| from data_generator import RockGlacierDataGenerator |
| from models.swin_unet import build_model |
|
|
| |
| np.Inf = np.inf |
|
|
| TEST_FEAT = "dataset/test/features" |
| TEST_MASK = "dataset/test/masks" |
| CKPT_PATH = "checkpoints/best_model.weights.h5" |
| OUTPUT_DIR = "reuslt/comparisons" |
|
|
| os.makedirs(OUTPUT_DIR, exist_ok=True) |
|
|
| |
| def dice_loss(y_true, y_pred): |
| smooth = 1e-6 |
| y_true_f = tf.cast(y_true, tf.float32) |
| y_pred_f = tf.cast(y_pred, tf.float32) |
| intersection = tf.reduce_sum(y_true_f * y_pred_f, axis=[1, 2, 3]) |
| union = tf.reduce_sum(y_true_f, axis=[1, 2, 3]) + tf.reduce_sum(y_pred_f, axis=[1, 2, 3]) |
| dice = (2. * intersection + smooth) / (union + smooth) |
| return 1 - tf.reduce_mean(dice) |
|
|
| def combined_loss(y_true, y_pred): |
| bce = tf.keras.losses.binary_crossentropy(y_true, y_pred) |
| bce = tf.reduce_mean(bce) |
| weighted_bce = bce * 2.0 |
| dl = dice_loss(y_true, y_pred) |
| return weighted_bce + dl |
|
|
| print("Building Model and Loading Weights...") |
| model = build_model(input_shape=(128, 128, 12), num_classes=1, num_heads=8) |
| model.compile(optimizer='adam', loss=combined_loss) |
| model.load_weights(CKPT_PATH) |
|
|
| print("Initializing Test Data Generator...") |
| test_gen = RockGlacierDataGenerator(TEST_FEAT, TEST_MASK, batch_size=1, augment=False, shuffle=False) |
|
|
| print("Generating Comparison Plots...") |
| |
| num_samples = len(test_gen) |
|
|
| for i in range(num_samples): |
| x, y_true = test_gen[i] |
| y_pred = model.predict(x, verbose=0) |
| |
| y_pred_bin = (y_pred > 0.5).astype(np.uint8)[0, :, :, 0] |
| y_true_bin = y_true[0, :, :, 0] |
| |
| |
| |
| |
| |
| |
| r = x[0, :, :, 2] |
| g = x[0, :, :, 1] |
| b = x[0, :, :, 0] |
| |
| rgb = np.dstack((r, g, b)) |
| |
| |
| rgb_min, rgb_max = np.percentile(rgb, (2, 98)) |
| if rgb_max > rgb_min: |
| rgb_disp = np.clip((rgb - rgb_min) / (rgb_max - rgb_min), 0, 1) |
| else: |
| rgb_disp = np.zeros_like(rgb) |
| |
| file_path = test_gen.feature_files[i] |
| base_name = os.path.basename(file_path).replace('feature_', '').replace('.tif', '') |
| |
| fig, axes = plt.subplots(1, 3, figsize=(15, 5)) |
| |
| axes[0].imshow(rgb_disp) |
| axes[0].set_title(f"Optical (RGB) - {base_name}") |
| axes[0].axis('off') |
| |
| axes[1].imshow(y_true_bin, cmap='gray', vmin=0, vmax=1) |
| axes[1].set_title("Ground Truth Mask") |
| axes[1].axis('off') |
| |
| axes[2].imshow(y_pred_bin, cmap='gray', vmin=0, vmax=1) |
| axes[2].set_title("Predicted Mask") |
| axes[2].axis('off') |
| |
| plt.tight_layout() |
| out_path = os.path.join(OUTPUT_DIR, f"compare_{base_name}.png") |
| plt.savefig(out_path, dpi=150, bbox_inches='tight') |
| plt.close() |
|
|
| print(f"Successfully saved {num_samples} comparison plots in '{OUTPUT_DIR}'") |
|
|
