import os
import numpy as np
import tensorflow as tf
import keras
from tensorflow.keras import layers # Used only to satisfy model loading structure
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import rasterio # REQUIRED for reading TIF files
from huggingface_hub import hf_hub_download # NEW IMPORT for reliable download
2. Custom Metric Definition (CRITICAL for loading the model)
=======================
def sparse_mean_iou(y_true, y_pred):
"""Mean IoU across all classes (Re-defining the custom metric for model loading)."""
# y_true and y_pred are (B, H, W, 5)
y_pred_labels = tf.argmax(y_pred, axis=-1)
y_true_labels = tf.argmax(y_true, axis=-1)
ious = []
for i in range(CONFIG["N_CLASSES"]):
y_true_class = tf.cast(y_true_labels == i, tf.float32)
y_pred_class = tf.cast(y_pred_labels == i, tf.float32)
intersection = tf.reduce_sum(y_true_class * y_pred_class)
union = tf.reduce_sum(y_true_class) + tf.reduce_sum(y_pred_class) - intersection
iou = intersection / (union + 1e-7)
ious.append(iou)
return tf.reduce_mean(ious)
=======================
3. Utility Functions
=======================
def mask_to_rgb(mask_1_indexed):
"""Convert 1-indexed class mask (1-5) to an RGB image."""
h, w = mask_1_indexed.shape
rgb = np.zeros((h, w, 3), dtype=np.uint8)
for class_id, color in COLOR_MAP.items():
rgb[mask_1_indexed == class_id] = color
return rgb
def load_and_preprocess_patch(image_data, normalization_factor=10000.0):
"""
1. Normalizes the input patch (H, W, C) to 0-1 range.
2. Adds the batch dimension (1, H, W, C).
"""
if image_data.ndim != 3 or image_data.shape[-1] != CONFIG["N_CHANNELS"]:
print(f"Error: Input shape must be (512, 512, 10), but got {image_data.shape}")
return None, None
# 1. Normalize the data (Crucial step to match training: reflectance / 10000)
normalized_data = np.clip(image_data.astype(np.float32) / normalization_factor, 0.0, 1.0)
# 2. Add batch dimension: (H, W, C) -> (1, H, W, C)
input_tensor = np.expand_dims(normalized_data, axis=0)
return input_tensor, normalized_data
def visualize_prediction(input_image, pred_labels_1_indexed):
"""Displays the input RGB composite and the final segmented mask."""
# Show RGB composite (bands 3,2,1 for natural color - indices 2, 1, 0)
# The input image here is the normalized data (0-1 range)
rgb_display = input_image[..., [2, 1, 0]]
# Convert prediction to RGB
pred_rgb = mask_to_rgb(pred_labels_1_indexed)
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Input Image (Normalized RGB Composite)")
plt.imshow(rgb_display)
plt.axis('off')
plt.subplot(1, 2, 2)
plt.title("Predicted Segmentation")
plt.imshow(pred_rgb)
plt.axis('off')
plt.tight_layout()
plt.show()
=======================
4. Main Prediction Logic
=======================
if name == "main":
# --- MODEL LOADING ---
model = None
custom_objects = {"sparse_mean_iou": sparse_mean_iou}
try:
print(f"Downloading model file '{HF_MODEL_FILE}' from Hub: {HF_REPO_ID}...")
# 1. Download the file locally first
# CRITICAL CHECK: The 404 error means the file or repo is not found.
# Ensure the filename and REPO_ID are EXACTLY correct in HF_MODEL_FILE and HF_REPO_ID.
# Ensure your token (if required for private repo) has 'read' access.
local_model_path = hf_hub_download(
repo_id=HF_REPO_ID,
filename=HF_MODEL_FILE,
repo_type="model"
)
print(f"Model downloaded to temporary path: {local_model_path}")
# 2. Load the model from the local path
model = keras.saving.load_model(
local_model_path,
custom_objects=custom_objects
)
print("Model loaded successfully.")
except Exception as e:
print(f"\nFATAL ERROR: Model Loading Failed: {e}")
print("-" * 50)
print(f"DEBUG STEP 1: VERIFY REPOSITORY AND FILENAME.")
print(f"The download attempt failed for file '{HF_MODEL_FILE}'.")
print(f"ACTION: Go to https://huggingface.co/{HF_REPO_ID}/tree/main in your browser and check the exact case-sensitive filename.")
print("-" * 50)
print("DEBUG STEP 2: CHECK READ ACCESS.")
print("If the repository is private or belongs to an organization, ensure your token has 'read' permissions (you may need to run `from huggingface_hub import notebook_login; notebook_login()` and paste your token).")
print("-" * 50)
# Exit the script immediately if loading failed, preventing the NameError later
exit()
# --- LOAD INPUT DATA from TIF file ---
try:
with rasterio.open(TIF_FILE_PATH) as src:
if src.height != CONFIG["HEIGHT"] or src.width != CONFIG["WIDTH"] or src.count != CONFIG["N_CHANNELS"]:
print(f"Error: Input TIF shape ({src.height}, {src.width}, {src.count}) does not match expected model input.")
exit()
# Read all 10 bands. Shape: (10, 512, 512)
real_input_patch = src.read()
# Transpose to (512, 512, 10) to match model input shape
real_input_patch = np.transpose(real_input_patch, (1, 2, 0))
print(f"\nSuccessfully loaded TIF patch with shape: {real_input_patch.shape}")
except rasterio.RasterioIOError:
print(f"Error: Could not open TIF file at {TIF_FILE_PATH}. Check file path and Google Drive mounting.")
exit()
except Exception as e:
print(f"An unexpected error occurred during TIF loading: {e}")
exit()
# --- PREPROCESS ---
input_tensor, normalized_input_for_display = load_and_preprocess_patch(real_input_patch)
if input_tensor is None:
exit()
# --- PREDICT ---
print("Generating prediction...")
# 'model' is guaranteed to be defined here because we exit() on failure above.
prediction_output = model.predict(input_tensor, verbose=0)
# --- POST-PROCESS & VISUALIZE ---
prediction_mask_ohe = prediction_output[0]
pred_labels_0_indexed = np.argmax(prediction_mask_ohe, axis=-1)
pred_labels_1_indexed = pred_labels_0_indexed + 1
print("Prediction complete. Displaying results...")
visualize_prediction(normalized_input_for_display, pred_labels_1_indexed)
