建筑物变化检测算法baseline工程
使用PyTorch框架，并选择U-Net来进行二分类任务（变化/不变）Baseline工程将基于深度学习方法来检测建筑物的变化
备注：博客 所有文章代码仅供参考！
如何使用建筑物变化检测算法的Baseline工程，一个详细的步骤和代码示例。这个Baseline工程将基于深度学习方法来检测建筑物的变化。我们将使用PyTorch框架，并选择一个常用的模型架构（如U-Net）来进行二分类任务（变化/不变）。
项目结构
building_change_detection/
├── main.py
├── train.py
├── evaluate.py
├── infer.py
├── visualize.py
├── datasets/
│ ├── building_images/
│ │ ├── images/
│ │ │ ├── before/
│ │ │ └── after/
│ │ └── masks/
│ │ ├── change/
│ │ └── nochange/
├── best_building_model.pth
├── requirements.txt
└── config.yaml
文件内容
requirements.txt
torch==1.9.0+cu111
torchvision==0.10.0+cu111
matplotlib
numpy
pandas
albumentations
pyyaml
scikit-image
config.yaml
配置文件用于存储训练参数：
train:
dataset_dir: ./datasets/building_images/images
mask_dir: ./datasets/building_images/masks
batch_size: 4
epochs: 50
learning_rate: 0.001
img_size: 256
num_classes: 1
test:
dataset_dir: ./datasets/building_images/images
mask_dir: ./datasets/building_images/masks
batch_size: 4
img_size: 256
数据准备
-
确认数据集目录结构:
确保你的数据集已经按照上述结构组织好：datasets/ └── building_images/ ├── images/ │ ├── before/ │ │ ├── image_0001_before.png │ │ ├── image_0002_before.png │ │ └── ... │ └── after/ │ ├── image_0001_after.png │ ├── image_0002_after.png │ └── ... └── masks/ ├── change/ │ ├── image_0001_change_mask.png │ ├── image_0002_change_mask.png │ └── ... └── nochange/ ├── image_0001_nochange_mask.png ├── image_0002_nochange_mask.png └── ... -
检查
config.yaml文件:
确认config.yaml文件中的路径和参数正确无误。
训练脚本
train.py
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import os
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import matplotlib.pyplot as plt
import yaml
import numpy as np
# 加载配置文件
with open('config.yaml', 'r') as f:
config = yaml.safe_load(f)
# 定义数据预处理
data_transforms = {
'train': A.Compose([
A.Resize(height=config['train']['img_size'], width=config['train']['img_size']),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.Rotate(limit=180, p=0.7),
A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.3),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
]),
'test': A.Compose([
A.Resize(height=config['test']['img_size'], width=config['test']['img_size']),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
])
}
class BuildingChangeDataset(Dataset):
def __init__(self, root_dir, mask_dir, transform=None):
self.root_dir = root_dir
self.mask_dir = mask_dir
self.transform = transform
self.before_files = sorted(os.listdir(os.path.join(root_dir, 'before')))
self.after_files = sorted(os.listdir(os.path.join(root_dir, 'after')))
def __len__(self):
return len(self.before_files)
def __getitem__(self, index):
before_img_path = os.path.join(self.root_dir, 'before', self.before_files[index])
after_img_path = os.path.join(self.root_dir, 'after', self.after_files[index])
mask_path = os.path.join(self.mask_dir, 'change', self.before_files[index])
before_img = Image.open(before_img_path).convert("RGB")
after_img = Image.open(after_img_path).convert("RGB")
mask = Image.open(mask_path).convert("L")
sample = {'before': before_img, 'after': after_img, 'mask': mask}
if self.transform:
augmented = self.transform(image=np.array(sample['before']), image1=np.array(sample['after']), mask=np.array(sample['mask']))
sample['before'] = augmented['image']
sample['after'] = augmented['image1']
sample['mask'] = augmented['mask']
return sample['before'], sample['after'], sample['mask'].unsqueeze(0).float() / 255.0
# 选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载数据集
train_dataset = BuildingChangeDataset(root_dir=config['train']['dataset_dir'],
mask_dir=config['train']['mask_dir'],
transform=data_transforms['train'])
val_dataset = BuildingChangeDataset(root_dir=config['test']['dataset_dir'],
mask_dir=config['test']['mask_dir'],
transform=data_transforms['test'])
train_loader = DataLoader(train_dataset, batch_size=config['train']['batch_size'], shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=config['test']['batch_size'], shuffle=False, num_workers=4)
# 定义UNet模型
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
def CBR2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
layers = []
layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)]
layers += [nn.BatchNorm2d(num_features=out_channels)]
layers += [nn.ReLU()]
cbr = nn.Sequential(*layers)
return cbr
# Contracting path
self.enc1_1 = CBR2d(in_channels=in_channels, out_channels=64)
self.enc1_2 = CBR2d(in_channels=64, out_channels=64)
self.pool1 = nn.MaxPool2d(kernel_size=2)
self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
self.enc2_2 = CBR2d(in_channels=128, out_channels=128)
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
self.enc3_2 = CBR2d(in_channels=256, out_channels=256)
self.pool3 = nn.MaxPool2d(kernel_size=2)
self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
self.enc4_2 = CBR2d(in_channels=512, out_channels=512)
self.pool4 = nn.MaxPool2d(kernel_size=2)
self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)
# Expanding path
self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)
self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec4_2 = CBR2d(in_channels=1024, out_channels=512)
self.dec4_1 = CBR2d(in_channels=512, out_channels=512)
self.unpool3 = nn.ConvTranspose2d(in_channels=512, out_channels=256,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec3_2 = CBR2d(in_channels=512, out_channels=256)
self.dec3_1 = CBR2d(in_channels=256, out_channels=256)
self.unpool2 = nn.ConvTranspose2d(in_channels=256, out_channels=128,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec2_2 = CBR2d(in_channels=256, out_channels=128)
self.dec2_1 = CBR2d(in_channels=128, out_channels=128)
self.unpool1 = nn.ConvTranspose2d(in_channels=128, out_channels=64,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec1_2 = CBR2d(in_channels=128, out_channels=64)
self.dec1_1 = CBR2d(in_channels=64, out_channels=64)
self.fc = nn.Conv2d(in_channels=64, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=True)
def forward(self, x):
enc1_1 = self.enc1_1(x)
enc1_2 = self.enc1_2(enc1_1)
pool1 = self.pool1(enc1_2)
enc2_1 = self.enc2_1(pool1)
enc2_2 = self.enc2_2(enc2_1)
pool2 = self.pool2(enc2_2)
enc3_1 = self.enc3_1(pool2)
enc3_2 = self.enc3_2(enc3_1)
pool3 = self.pool3(enc3_2)
enc4_1 = self.enc4_1(pool3)
enc4_2 = self.enc4_2(enc4_1)
pool4 = self.pool4(enc4_2)
enc5_1 = self.enc5_1(pool4)
dec5_1 = self.dec5_1(enc5_1)
unpool4 = self.unpool4(dec5_1)
cat4 = torch.cat((unpool4, enc4_2), dim=1)
dec4_2 = self.dec4_2(cat4)
dec4_1 = self.dec4_1(dec4_2)
unpool3 = self.unpool3(dec4_1)
cat3 = torch.cat((unpool3, enc3_2), dim=1)
dec3_2 = self.dec3_2(cat3)
dec3_1 = self.dec3_1(dec3_2)
unpool2 = self.unpool2(dec3_1)
cat2 = torch.cat((unpool2, enc2_2), dim=1)
dec2_2 = self.dec2_2(cat2)
dec2_1 = self.dec2_1(dec2_2)
unpool1 = self.unpool1(dec2_1)
cat1 = torch.cat((unpool1, enc1_2), dim=1)
dec1_2 = self.dec1_2(cat1)
dec1_1 = self.dec1_1(dec1_2)
x = self.fc(dec1_1)
return x
# 初始化模型
model = UNet(in_channels=6, out_channels=config['train']['num_classes']).to(device)
# 定义损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=config['train']['learning_rate'])
# 训练函数
def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10):
model.train()
metric_logger = utils.MetricLogger(delimiter=" ")
header = f'Epoch: [{epoch}]'
for before_imgs, after_imgs, masks in metric_logger.log_every(data_loader, print_freq, header):
inputs = torch.cat((before_imgs, after_imgs), dim=1).to(device)
targets = masks.to(device)
outputs = model(inputs)
loss = criterion(outputs, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
metric_logger.update(loss=loss.item())
# 验证函数
@torch.no_grad()
def validate_one_epoch(model, data_loader, device, print_freq=10):
model.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
header = "Test:"
for before_imgs, after_imgs, masks in metric_logger.log_every(data_loader, print_freq, header):
inputs = torch.cat((before_imgs, after_imgs), dim=1).to(device)
targets = masks.to(device)
outputs = model(inputs)
loss = criterion(outputs, targets)
metric_logger.update(loss=loss.item())
return metric_logger.meters['loss'].global_avg
# 训练模型
for epoch in range(config['train']['epochs']):
train_one_epoch(model, optimizer, train_loader, device, epoch, print_freq=10)
val_loss = validate_one_epoch(model, val_loader, device, print_freq=10)
print(f'Validation Loss: {val_loss:.4f}')
# 保存最佳模型
torch.save(model.state_dict(), 'best_building_model.pth')
评估脚本
evaluate.py
import torch
from torchvision import models
from PIL import Image
import os
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import matplotlib.pyplot as plt
import yaml
import numpy as np
import random
# 加载配置文件
with open('config.yaml', 'r') as f:
config = yaml.safe_load(f)
# 定义数据预处理
data_transforms = {
'test': A.Compose([
A.Resize(height=config['test']['img_size'], width=config['test']['img_size']),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
])
}
class BuildingChangeDataset(torch.utils.data.Dataset):
def __init__(self, root_dir, mask_dir, transform=None):
self.root_dir = root_dir
self.mask_dir = mask_dir
self.transform = transform
self.before_files = sorted(os.listdir(os.path.join(root_dir, 'before')))
self.after_files = sorted(os.listdir(os.path.join(root_dir, 'after')))
def __len__(self):
return len(self.before_files)
def __getitem__(self, index):
before_img_path = os.path.join(self.root_dir, 'before', self.before_files[index])
after_img_path = os.path.join(self.root_dir, 'after', self.after_files[index])
mask_path = os.path.join(self.mask_dir, 'change', self.before_files[index])
before_img = Image.open(before_img_path).convert("RGB")
after_img = Image.open(after_img_path).convert("RGB")
mask = Image.open(mask_path).convert("L")
sample = {'before': before_img, 'after': after_img, 'mask': mask}
if self.transform:
augmented = self.transform(image=np.array(sample['before']), image1=np.array(sample['after']), mask=np.array(sample['mask']))
sample['before'] = augmented['image']
sample['after'] = augmented['image1']
sample['mask'] = augmented['mask']
return sample['before'], sample['after'], sample['mask'].unsqueeze(0).float() / 255.0
# 选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载数据集
val_dataset = BuildingChangeDataset(root_dir=config['test']['dataset_dir'],
mask_dir=config['test']['mask_dir'],
transform=data_transforms['test'])
val_loader = DataLoader(val_dataset, batch_size=config['test']['batch_size'], shuffle=False, num_workers=4)
# 定义UNet模型
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
def CBR2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
layers = []
layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)]
layers += [nn.BatchNorm2d(num_features=out_channels)]
layers += [nn.ReLU()]
cbr = nn.Sequential(*layers)
return cbr
# Contracting path
self.enc1_1 = CBR2d(in_channels=in_channels, out_channels=64)
self.enc1_2 = CBR2d(in_channels=64, out_channels=64)
self.pool1 = nn.MaxPool2d(kernel_size=2)
self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
self.enc2_2 = CBR2d(in_channels=128, out_channels=128)
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
self.enc3_2 = CBR2d(in_channels=256, out_channels=256)
self.pool3 = nn.MaxPool2d(kernel_size=2)
self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
self.enc4_2 = CBR2d(in_channels=512, out_channels=512)
self.pool4 = nn.MaxPool2d(kernel_size=2)
self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)
# Expanding path
self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)
self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec4_2 = CBR2d(in_channels=1024, out_channels=512)
self.dec4_1 = CBR2d(in_channels=512, out_channels=512)
self.unpool3 = nn.ConvTranspose2d(in_channels=512, out_channels=256,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec3_2 = CBR2d(in_channels=512, out_channels=256)
self.dec3_1 = CBR2d(in_channels=256, out_channels=256)
self.unpool2 = nn.ConvTranspose2d(in_channels=256, out_channels=128,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec2_2 = CBR2d(in_channels=256, out_channels=128)
self.dec2_1 = CBR2d(in_channels=128, out_channels=128)
self.unpool1 = nn.ConvTranspose2d(in_channels=128, out_channels=64,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec1_2 = CBR2d(in_channels=128, out_channels=64)
self.dec1_1 = CBR2d(in_channels=64, out_channels=64)
self.fc = nn.Conv2d(in_channels=64, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=True)
def forward(self, x):
enc1_1 = self.enc1_1(x)
enc1_2 = self.enc1_2(enc1_1)
pool1 = self.pool1(enc1_2)
enc2_1 = self.enc2_1(pool1)
enc2_2 = self.enc2_2(enc2_1)
pool2 = self.pool2(enc2_2)
enc3_1 = self.enc3_1(pool2)
enc3_2 = self.enc3_2(enc3_1)
pool3 = self.pool3(enc3_2)
enc4_1 = self.enc4_1(pool3)
enc4_2 = self.enc4_2(enc4_1)
pool4 = self.pool4(enc4_2)
enc5_1 = self.enc5_1(pool4)
dec5_1 = self.dec5_1(enc5_1)
unpool4 = self.unpool4(dec5_1)
cat4 = torch.cat((unpool4, enc4_2), dim=1)
dec4_2 = self.dec4_2(cat4)
dec4_1 = self.dec4_1(dec4_2)
unpool3 = self.unpool3(dec4_1)
cat3 = torch.cat((unpool3, enc3_2), dim=1)
dec3_2 = self.dec3_2(cat3)
dec3_1 = self.dec3_1(dec3_2)
unpool2 = self.unpool2(dec3_1)
cat2 = torch.cat((unpool2, enc2_2), dim=1)
dec2_2 = self.dec2_2(cat2)
dec2_1 = self.dec2_1(dec2_2)
unpool1 = self.unpool1(dec2_1)
cat1 = torch.cat((unpool1, enc1_2), dim=1)
dec1_2 = self.dec1_2(cat1)
dec1_1 = self.dec1_1(dec1_2)
x = self.fc(dec1_1)
return x
# 加载预训练的UNet模型并加载权重
model = UNet(in_channels=6, out_channels=config['train']['num_classes'])
model.load_state_dict(torch.load('best_building_model.pth'))
model.to(device)
model.eval() # 设置模型为评估模式
# 评估模型
def evaluate(model, data_loader, device):
model.eval()
total_iou = 0.0
total_samples = 0
with torch.no_grad():
for before_imgs, after_imgs, masks in data_loader:
inputs = torch.cat((before_imgs, after_imgs), dim=1).to(device)
targets = masks.to(device)
outputs = model(inputs)
outputs = torch.sigmoid(outputs)
outputs = (outputs > 0.5).float()
intersection = (outputs * targets).sum(dim=(2, 3))
union = (outputs + targets - outputs * targets).sum(dim=(2, 3))
iou = intersection / union
total_iou += iou.sum().item()
total_samples += iou.size(0) * iou.size(1)
mean_iou = total_iou / total_samples
print(f'Mean IoU: {mean_iou:.4f}')
evaluate(model, val_loader, device)
推理脚本
infer.py
import torch
from torchvision import models
from PIL import Image
import os
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import matplotlib.pyplot as plt
import yaml
import numpy as np
# 加载配置文件
with open('config.yaml', 'r') as f:
config = yaml.safe_load(f)
# 定义数据预处理
data_transforms = {
'test': A.Compose([
A.Resize(height=config['test']['img_size'], width=config['test']['img_size']),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
])
}
# 选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 定义UNet模型
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
def CBR2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
layers = []
layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)]
layers += [nn.BatchNorm2d(num_features=out_channels)]
layers += [nn.ReLU()]
cbr = nn.Sequential(*layers)
return cbr
# Contracting path
self.enc1_1 = CBR2d(in_channels=in_channels, out_channels=64)
self.enc1_2 = CBR2d(in_channels=64, out_channels=64)
self.pool1 = nn.MaxPool2d(kernel_size=2)
self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
self.enc2_2 = CBR2d(in_channels=128, out_channels=128)
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
self.enc3_2 = CBR2d(in_channels=256, out_channels=256)
self.pool3 = nn.MaxPool2d(kernel_size=2)
self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
self.enc4_2 = CBR2d(in_channels=512, out_channels=512)
self.pool4 = nn.MaxPool2d(kernel_size=2)
self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)
# Expanding path
self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)
self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec4_2 = CBR2d(in_channels=1024, out_channels=512)
self.dec4_1 = CBR2d(in_channels=512, out_channels=512)
self.unpool3 = nn.ConvTranspose2d(in_channels=512, out_channels=256,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec3_2 = CBR2d(in_channels=512, out_channels=256)
self.dec3_1 = CBR2d(in_channels=256, out_channels=256)
self.unpool2 = nn.ConvTranspose2d(in_channels=256, out_channels=128,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec2_2 = CBR2d(in_channels=256, out_channels=128)
self.dec2_1 = CBR2d(in_channels=128, out_channels=128)
self.unpool1 = nn.ConvTranspose2d(in_channels=128, out_channels=64,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec1_2 = CBR2d(in_channels=128, out_channels=64)
self.dec1_1 = CBR2d(in_channels=64, out_channels=64)
self.fc = nn.Conv2d(in_channels=64, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=True)
def forward(self, x):
enc1_1 = self.enc1_1(x)
enc1_2 = self.enc1_2(enc1_1)
pool1 = self.pool1(enc1_2)
enc2_1 = self.enc2_1(pool1)
enc2_2 = self.enc2_2(enc2_1)
pool2 = self.pool2(enc2_2)
enc3_1 = self.enc3_1(pool2)
enc3_2 = self.enc3_2(enc3_1)
pool3 = self.pool3(enc3_2)
enc4_1 = self.enc4_1(pool3)
enc4_2 = self.enc4_2(enc4_1)
pool4 = self.pool4(enc4_2)
enc5_1 = self.enc5_1(pool4)
dec5_1 = self.dec5_1(enc5_1)
unpool4 = self.unpool4(dec5_1)
cat4 = torch.cat((unpool4, enc4_2), dim=1)
dec4_2 = self.dec4_2(cat4)
dec4_1 = self.dec4_1(dec4_2)
unpool3 = self.unpool3(dec4_1)
cat3 = torch.cat((unpool3, enc3_2), dim=1)
dec3_2 = self.dec3_2(cat3)
dec3_1 = self.dec3_1(dec3_2)
unpool2 = self.unpool2(dec3_1)
cat2 = torch.cat((unpool2, enc2_2), dim=1)
dec2_2 = self.dec2_2(cat2)
dec2_1 = self.dec2_1(dec2_2)
unpool1 = self.unpool1(dec2_1)
cat1 = torch.cat((unpool1, enc1_2), dim=1)
dec1_2 = self.dec1_2(cat1)
dec1_1 = self.dec1_1(dec1_2)
x = self.fc(dec1_1)
return x
# 加载预训练的UNet模型并加载权重
model = UNet(in_channels=6, out_channels=config['train']['num_classes'])
model.load_state_dict(torch.load('best_building_model.pth'))
model.to(device)
model.eval() # 设置模型为评估模式
def predict_image(before_img_path, after_img_path):
before_img = Image.open(before_img_path).convert("RGB")
after_img = Image.open(after_img_path).convert("RGB")
original_before_img = before_img.copy()
original_after_img = after_img.copy()
transformed = data_transforms['test'](image=np.array(before_img), image1=np.array(after_img))
before_img = transformed['image'].unsqueeze(0).to(device)
after_img = transformed['image1'].unsqueeze(0).to(device)
inputs = torch.cat((before_img, after_img), dim=1)
with torch.no_grad():
output = model(inputs)[0]
output = torch.sigmoid(output)
output = (output > 0.5).float().squeeze().cpu().numpy()
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
axes[0].imshow(original_before_img)
axes[0].set_title('Before Image')
axes[0].axis('off')
axes[1].imshow(original_after_img)
axes[1].set_title('After Image')
axes[1].axis('off')
axes[2].imshow(output, cmap='gray')
axes[2].set_title('Change Mask')
axes[2].axis('off')
plt.tight_layout()
plt.show()
if __name__ == "__main__":
before_img_path = 'path/to/your/before/image.jpg' # 替换为你的“之前”图像路径
after_img_path = 'path/to/your/after/image.jpg' # 替换为你的“之后”图像路径
predict_image(before_img_path, after_img_path)
可视化脚本
visualize.py
import torch
from torchvision import models
from PIL import Image
import os
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import matplotlib.pyplot as plt
import yaml
import numpy as np
import random
# 加载配置文件
with open('config.yaml', 'r') as f:
config = yaml.safe_load(f)
# 定义数据预处理
data_transforms = {
'test': A.Compose([
A.Resize(height=config['test']['img_size'], width=config['test']['img_size']),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
])
}
class BuildingChangeDataset(torch.utils.data.Dataset):
def __init__(self, root_dir, mask_dir, transform=None):
self.root_dir = root_dir
self.mask_dir = mask_dir
self.transform = transform
self.before_files = sorted(os.listdir(os.path.join(root_dir, 'before')))
self.after_files = sorted(os.listdir(os.path.join(root_dir, 'after')))
def __len__(self):
return len(self.before_files)
def __getitem__(self, index):
before_img_path = os.path.join(self.root_dir, 'before', self.before_files[index])
after_img_path = os.path.join(self.root_dir, 'after', self.after_files[index])
mask_path = os.path.join(self.mask_dir, 'change', self.before_files[index])
before_img = Image.open(before_img_path).convert("RGB")
after_img = Image.open(after_img_path).convert("RGB")
mask = Image.open(mask_path).convert("L")
sample = {'before': before_img, 'after': after_img, 'mask': mask}
if self.transform:
augmented = self.transform(image=np.array(sample['before']), image1=np.array(sample['after']), mask=np.array(sample['mask']))
sample['before'] = augmented['image']
sample['after'] = augmented['image1']
sample['mask'] = augmented['mask']
return sample['before'], sample['after'], sample['mask'].unsqueeze(0).float() / 255.0
# 选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载数据集
val_dataset = BuildingChangeDataset(root_dir=config['test']['dataset_dir'],
mask_dir=config['test']['mask_dir'],
transform=data_transforms['test'])
val_loader = DataLoader(val_dataset, batch_size=config['test']['batch_size'], shuffle=False, num_workers=4)
# 定义UNet模型
class UNet(nn.Module):
def __init__(