Spaces:
Running on Zero
Running on Zero
File size: 3,382 Bytes
5442314 09a3fbd d8fe2bd 09a3fbd 92d49cd 09a3fbd 92d49cd 09a3fbd 92d49cd 09a3fbd 92d49cd 09a3fbd | 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 106 107 | try:
import spaces
except ImportError:
spaces = None
import os
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import gradio as gr
from models.networks import define_G
# --- 1. Model Configuration ---
class Args:
n_class = 2
net_G = 'base_transformer_pos_s4_dd8_dedim8'
gpu_ids = [] # [] for CPU inference
img_size = 256
args = Args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# --- 2. Load Model & Weights ---
print(f"Loading model on {device}...")
model = define_G(args=args, gpu_ids=args.gpu_ids)
checkpoint_path = os.path.join("checkpoints", "BIT_LEVIR", "best_ckpt.pt")
if not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Checkpoint not found at: {checkpoint_path}")
try:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
except TypeError:
checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint["model_G_state_dict"])
model.to(device)
model.eval()
# --- 3. Image Preprocessing ---
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# --- 4. Prediction Function ---
def _predict_change(img_a_pil, img_b_pil):
if img_a_pil is None or img_b_pil is None:
return None, None
# Convert to RGB
img_a_rgb = img_a_pil.convert("RGB")
img_b_rgb = img_b_pil.convert("RGB")
# Determine runtime device (cuda if ZeroGPU allocated, else cpu)
runtime_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(runtime_device)
# Transform to tensor
tensor_a = transform(img_a_rgb).unsqueeze(0).to(runtime_device)
tensor_b = transform(img_b_rgb).unsqueeze(0).to(runtime_device)
with torch.no_grad():
output = model(tensor_a, tensor_b)
pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
# Binary mask: 0 (No change, Black), 255 (Change, White)
mask = (pred * 255).astype(np.uint8)
mask_pil = Image.fromarray(mask, mode="L")
# Create a red change overlay on Image B for clear visualization
img_b_resized = img_b_rgb.resize((256, 256))
img_b_np = np.array(img_b_resized)
overlay_np = img_b_np.copy()
overlay_np[mask == 255] = [255, 50, 50] # Highlight changes in Red
blended = Image.blend(img_b_resized, Image.fromarray(overlay_np), alpha=0.5)
return mask_pil, blended
if spaces is not None:
predict_change = spaces.GPU(_predict_change)
else:
predict_change = _predict_change
# --- 5. Gradio Web UI ---
demo = gr.Interface(
fn=predict_change,
inputs=[
gr.Image(type="pil", label="Image A (Before)"),
gr.Image(type="pil", label="Image B (After)")
],
outputs=[
gr.Image(type="pil", label="Change Detection Mask"),
gr.Image(type="pil", label="Overlay on Image B (Red = Change)")
],
title="🛰️ Remote Sensing Change Detection (BIT_CD)",
description="Upload two satellite/aerial images (Before & After) to detect structural and landscape changes.",
examples=[
["samples/A/test_2_0000_0000.png", "samples/B/test_2_0000_0000.png"]
] if os.path.exists("samples/A/test_2_0000_0000.png") else None
)
if __name__ == "__main__":
demo.launch()
|