Spaces:
Build error
Build error
File size: 3,186 Bytes
01b9c69 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image
# -------------------------------
# Load Model and Data
# -------------------------------
MODEL_PATH = "satellite_segmentation_unet.h5"
X_TEST_PATH = "X_test.npy"
Y_TEST_PATH = "y_test.npy"
TOTAL_CLASSES_PATH = "total_classes.npy"
print("Loading model and test data...")
model = load_model(MODEL_PATH, compile=False)
X_test = np.load(X_TEST_PATH)
y_test = np.load(Y_TEST_PATH)
total_classes = int(np.load(TOTAL_CLASSES_PATH)[0])
# -------------------------------
# Class color map
# -------------------------------
CLASS_COLORS = {
0: [226, 169, 41], # water
1: [132, 41, 246], # land
2: [110, 193, 228], # road
3: [60, 16, 152], # building
4: [254, 221, 58], # vegetation
5: [155, 155, 155], # unlabeled
}
# -------------------------------
# Helper Functions
# -------------------------------
def decode_segmentation(mask):
"""Convert model output or label mask to RGB."""
if mask.ndim == 4: # Model output shape (1, H, W, num_classes)
mask = np.argmax(mask[0], axis=-1)
elif mask.ndim == 3 and mask.shape[-1] == total_classes:
mask = np.argmax(mask, axis=-1)
elif mask.ndim == 3:
mask = mask.squeeze()
color_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
for class_idx, color in CLASS_COLORS.items():
color_mask[mask == class_idx] = color
return color_mask
def predict_segmentation(index):
"""Run model inference and return all visualizations."""
try:
index = int(index)
image = (X_test[index] * 255).astype(np.uint8)
# Model prediction
pred = model.predict(np.expand_dims(X_test[index], axis=0))
pred_mask = decode_segmentation(pred)
# Ground truth
gt_mask = decode_segmentation(y_test[index])
return image, gt_mask, pred_mask
except Exception as e:
print("Error during prediction:", e)
return None, None, None
# -------------------------------
# Gradio Interface
# -------------------------------
title = "🛰️ Satellite Segmentation (U-Net)"
description = """
Select one of the **test images** from the dataset to visualize:
- The original input image
- The ground truth mask
- The predicted segmentation mask from the trained **U-Net**
"""
indices = list(range(min(50, len(X_test))))
interface = gr.Interface(
fn=predict_segmentation,
inputs=gr.Dropdown(
choices=[str(i) for i in indices],
label="Select Test Image Index",
value="0",
interactive=True,
),
outputs=[
gr.Image(label="Original Image"),
gr.Image(label="Ground Truth Mask"),
gr.Image(label="Predicted Mask"),
],
title=title,
description=description,
allow_flagging="never",
theme="gradio/soft",
)
# -------------------------------
# Launch
# -------------------------------
if __name__ == "__main__":
interface.launch(server_name="0.0.0.0", server_port=7860)
|