UNET-5-Classes

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

=======================

1. Configuration and Constants

=======================

CONFIG = { "HEIGHT": 512, "WIDTH": 512, "N_CLASSES": 5, # Class labels are 1-5 "N_CHANNELS": 10 # 10 Sentinel-2 bands }

--- MODEL SOURCE ---

We are now loading the model directly from the Hugging Face Hub

HF_REPO_ID = "FYP-PROJECT/UNET-5-Classes"

--- DEBUG CHANGE: Trying alternative filename for debugging 404 error ---

HF_MODEL_FILE = "best_multiclass_model.keras"

Please check your Hugging Face repo and set this to the EXACT file name.

=== USER INPUT REQUIRED ===

CHANGE THIS TO THE PATH OF YOUR INPUT TIF FILE (512x512, 10 bands)

TIF_FILE_PATH = "/content/drive/MyDrive/new-data-2025/test/images/9.tif"

===========================

Class mapping (1-indexed)

CLASS_MAPPING = { 1: "Water", 2: "Built-up", 3: "Forest", 4: "Soil", 5: "Crop" }

Color mapping for visualization (1-indexed keys, RGB 0-255)

COLOR_MAP = { 1: [0, 0, 255], 2: [255, 0, 0], 3: [0, 100, 0], 4: [139, 69, 19], 5: [255, 255, 0] }

=======================

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)