364 lines
16 KiB
Python
364 lines
16 KiB
Python
|
# ChangeDetection.py
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import torch
|
|||
|
|
import numpy as np
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
import itertools
|
|||
|
|
import torch.optim as optim
|
|||
|
|
from torch.utils.data import DataLoader
|
|||
|
|
from model.network import CDNet
|
|||
|
|
from data_utils import DynamicPatchDataset, LoadDatasetFromCoords, TestDatasetFromCoords, calMetric_iou, get_transform, \
|
|||
|
|
TestDatasetFromPic
|
|||
|
|
from loss.losses import cross_entropy
|
|||
|
|
from PIL import Image
|
|||
|
|
from cood_csv import PatchIndexer, PredictionAggregator
|
|||
|
|
import torchvision.transforms as transforms
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChangeDetection:
|
|||
|
|
|
|||
|
|
def __init__(self, args):
|
|||
|
|
# 参数完整性检查
|
|||
|
|
required_args = ['gpu_id', 'img_size', 'lr']
|
|||
|
|
for arg in required_args:
|
|||
|
|
if not hasattr(args, arg):
|
|||
|
|
raise ValueError(f"Missing required argument: {arg}")
|
|||
|
|
self.args = args
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
|
|||
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|||
|
|
self.best_iou = 0.0
|
|||
|
|
|
|||
|
|
# Initialize model
|
|||
|
|
self.model = CDNet(img_size=self.args.img_size).to(self.device, dtype=torch.float)
|
|||
|
|
if torch.cuda.device_count() > 1:
|
|||
|
|
print(f"Using {torch.cuda.device_count()} GPUs!")
|
|||
|
|
self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
|
|||
|
|
|
|||
|
|
# Loss and optimizer
|
|||
|
|
self.criterion = cross_entropy().to(self.device, dtype=torch.float)
|
|||
|
|
self.optimizer = optim.Adam(itertools.chain(self.model.parameters()),
|
|||
|
|
lr=self.args.lr, betas=(0.9, 0.999))
|
|||
|
|
|
|||
|
|
def train(self):
|
|||
|
|
print("Starting training...")
|
|||
|
|
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
|
|||
|
|
|
|||
|
|
# 自动推断 save_dir = 'val/', 自动保存为 'coords.csv'
|
|||
|
|
# 1训练scv文件生成
|
|||
|
|
indexer_train = PatchIndexer(
|
|||
|
|
img_dir=self.args.hr1_train,
|
|||
|
|
patch_size=256,
|
|||
|
|
stride=256
|
|||
|
|
)
|
|||
|
|
indexer_train.index_all()
|
|||
|
|
# 1验证scv文件生成
|
|||
|
|
indexer_val = PatchIndexer(
|
|||
|
|
img_dir=self.args.hr1_val,
|
|||
|
|
patch_size=256,
|
|||
|
|
stride=256
|
|||
|
|
)
|
|||
|
|
indexer_val.index_all()
|
|||
|
|
|
|||
|
|
# 调用相对路径
|
|||
|
|
# 1获取上一级文件路径
|
|||
|
|
parent_train_cood = os.path.dirname(self.args.hr1_train)
|
|||
|
|
train_cood = os.path.join(parent_train_cood, "coords.csv")
|
|||
|
|
parent_val_cood = os.path.dirname(self.args.hr1_val)
|
|||
|
|
val_cood = os.path.join(parent_val_cood, "coords.csv")
|
|||
|
|
|
|||
|
|
train_set = DynamicPatchDataset(coords_csv=train_cood, img1_dir=self.args.hr1_train,
|
|||
|
|
img2_dir=self.args.hr2_train, label_dir=self.args.lab_train, crop=True,
|
|||
|
|
augment=True, angle=30, num_classes=2)
|
|||
|
|
val_set = LoadDatasetFromCoords(coords_csv=val_cood, hr1_dir=self.args.hr1_val, hr2_dir=self.args.hr2_val,
|
|||
|
|
label_dir=self.args.lab_val)
|
|||
|
|
|
|||
|
|
# train_set = DynamicPatchDataset(self.args.hr1_train, self.args.hr2_train, self.args.lab_train, crop=False)
|
|||
|
|
# val_set = LoadDatasetFromFolder(self.args, self.args.hr1_val, self.args.hr2_val, self.args.lab_val)
|
|||
|
|
train_loader = DataLoader(train_set, num_workers=self.args.num_workers,
|
|||
|
|
batch_size=self.args.batchsize, shuffle=True)
|
|||
|
|
val_loader = DataLoader(val_set, num_workers=self.args.num_workers,
|
|||
|
|
batch_size=self.args.val_batchsize, shuffle=False)
|
|||
|
|
|
|||
|
|
for epoch in range(1, self.args.num_epochs + 1):
|
|||
|
|
self.model.train()
|
|||
|
|
train_bar = tqdm(train_loader, desc=f"Epoch [{epoch}/{self.args.num_epochs}]")
|
|||
|
|
for hr_img1, hr_img2, label in train_bar:
|
|||
|
|
hr_img1 = hr_img1.to(self.device, dtype=torch.float)
|
|||
|
|
hr_img2 = hr_img2.to(self.device, dtype=torch.float)
|
|||
|
|
label = label.to(self.device, dtype=torch.float)
|
|||
|
|
label = torch.argmax(label, 1).unsqueeze(1).float()
|
|||
|
|
|
|||
|
|
self.optimizer.zero_grad()
|
|||
|
|
out1, out2, out3 = self.model(hr_img1, hr_img2)
|
|||
|
|
cd_loss = (self.criterion(out1, label) +
|
|||
|
|
self.criterion(out2, label) +
|
|||
|
|
self.criterion(out3, label))
|
|||
|
|
cd_loss.backward()
|
|||
|
|
self.optimizer.step()
|
|||
|
|
train_bar.set_postfix(loss=cd_loss.item())
|
|||
|
|
|
|||
|
|
# Evaluate only on best epoch (IoU improved)
|
|||
|
|
val_iou = self.validate(val_loader)
|
|||
|
|
print(f"Validation IoU: {val_iou:.4f}")
|
|||
|
|
|
|||
|
|
if val_iou > self.best_iou:
|
|||
|
|
self.best_iou = val_iou
|
|||
|
|
self._save_best_model()
|
|||
|
|
print(f"New best model saved with IoU: {val_iou:.4f}")
|
|||
|
|
|
|||
|
|
def validate(self, val_loader):
|
|||
|
|
self.model.eval()
|
|||
|
|
inter, union = 0, 0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for hr_img1, hr_img2, label in tqdm(val_loader, desc='Validating'):
|
|||
|
|
hr_img1 = hr_img1.to(self.device, dtype=torch.float)
|
|||
|
|
hr_img2 = hr_img2.to(self.device, dtype=torch.float)
|
|||
|
|
label = label.to(self.device, dtype=torch.float)
|
|||
|
|
label = torch.argmax(label, 1).unsqueeze(1).float()
|
|||
|
|
|
|||
|
|
output, _, _ = self.model(hr_img1, hr_img2)
|
|||
|
|
pred = torch.argmax(output, 1).unsqueeze(1).float()
|
|||
|
|
gt = (label > 0).float()
|
|||
|
|
prob = (pred > 0).float()
|
|||
|
|
|
|||
|
|
gt_np = np.squeeze(gt.cpu().detach().numpy())
|
|||
|
|
pred_np = np.squeeze(prob.cpu().detach().numpy())
|
|||
|
|
intr, unn = calMetric_iou(gt_np, pred_np)
|
|||
|
|
inter += intr
|
|||
|
|
union += unn
|
|||
|
|
|
|||
|
|
iou = inter / union if union != 0 else 0
|
|||
|
|
return iou
|
|||
|
|
|
|||
|
|
def _save_best_model(self):
|
|||
|
|
if not os.path.exists(self.args.model_dir):
|
|||
|
|
os.makedirs(self.args.model_dir)
|
|||
|
|
# Remove any existing .pth files
|
|||
|
|
for f in os.listdir(self.args.model_dir):
|
|||
|
|
if f.endswith('.pth'):
|
|||
|
|
os.remove(os.path.join(self.args.model_dir, f))
|
|||
|
|
save_path = os.path.join(self.args.model_dir, 'best_model.pth')
|
|||
|
|
torch.save(self.model.state_dict(), save_path)
|
|||
|
|
print(f"Best model saved to {save_path}")
|
|||
|
|
|
|||
|
|
def load_model(self, model_path):
|
|||
|
|
self.model.load_state_dict(torch.load(model_path, map_location=self.device))
|
|||
|
|
self.model.eval()
|
|||
|
|
print(f"Loaded model from {model_path}")
|
|||
|
|
|
|||
|
|
def predict(self):
|
|||
|
|
print(f"结果将保存到目录: {self.args.save_dir}")
|
|||
|
|
print("Starting prediction...")
|
|||
|
|
indexer_test = PatchIndexer(
|
|||
|
|
img_dir=self.args.path_img1,
|
|||
|
|
patch_size=256,
|
|||
|
|
stride=256
|
|||
|
|
)
|
|||
|
|
indexer_test.index_all()
|
|||
|
|
# 调用scv相对路径
|
|||
|
|
parent_test_cood = os.path.dirname(self.args.path_img1)
|
|||
|
|
test_cood = os.path.join(parent_test_cood, "coords.csv")
|
|||
|
|
|
|||
|
|
if not os.path.exists(self.args.save_dir):
|
|||
|
|
os.makedirs(self.args.save_dir)
|
|||
|
|
|
|||
|
|
transform = transforms.Compose([
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
img1Url = r"D:\project\2025-7-10_data+model\val\image1\1.png"
|
|||
|
|
img2Url = r"D:\project\2025-7-10_data+model\val\image2\1.png"
|
|||
|
|
|
|||
|
|
test_set = TestDatasetFromPic(img1Url, img2Url)
|
|||
|
|
|
|||
|
|
# test_set = TestDatasetFromCoords(coords_csv=test_cood, hr1_dir=self.args.path_img1, hr2_dir=self.args.path_img2,
|
|||
|
|
# label_dir=None, transform=transform)
|
|||
|
|
# test_set = TestDatasetFromFolder(self.args, self.args.path_img1,
|
|||
|
|
# self.args.path_img2, self.args.path_lab)
|
|||
|
|
test_loader = DataLoader(test_set, num_workers=24,
|
|||
|
|
batch_size=self.args.batch_size, shuffle=False)
|
|||
|
|
|
|||
|
|
self.model.eval()
|
|||
|
|
aggregator = PredictionAggregator(img_dir=self.args.path_img1)
|
|||
|
|
# aggregator=PredictionAggregator()
|
|||
|
|
test_bar = tqdm(test_loader)
|
|||
|
|
inter, union = 0, 0
|
|||
|
|
|
|||
|
|
for img1, img2, label, patch_name, fname_list, x_list, y_list, h_list, w_list in test_bar:
|
|||
|
|
img1 = img1.to(self.device)
|
|||
|
|
img2 = img2.to(self.device)
|
|||
|
|
label = label.to(self.device)
|
|||
|
|
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output = self.model(img1, img2)
|
|||
|
|
if isinstance(output, tuple): # ← 检查是否是 tuple
|
|||
|
|
output = output[0]
|
|||
|
|
|
|||
|
|
pred = torch.argmax(output, 1).unsqueeze(1).float()
|
|||
|
|
# label = torch.argmax(label, 1).unsqueeze(1).float()
|
|||
|
|
label = label.float()
|
|||
|
|
|
|||
|
|
for i in range(img1.size(0)):
|
|||
|
|
# prob = pred[i].cpu().numpy()
|
|||
|
|
# gt_value = label[i].cpu().numpy()
|
|||
|
|
patch = pred[i][0].cpu().numpy().astype(np.uint8)
|
|||
|
|
gt_value = label[i][0].cpu().numpy().astype(np.uint8)
|
|||
|
|
|
|||
|
|
# intr, unn = calMetric_iou(gt_value.squeeze(), prob.squeeze())
|
|||
|
|
intr, unn = calMetric_iou(gt_value, patch)
|
|||
|
|
inter += intr
|
|||
|
|
union += unn
|
|||
|
|
fname = fname_list[i]
|
|||
|
|
x, y = x_list[i].item(), y_list[i].item()
|
|||
|
|
# patch = pred[i][0].cpu().numpy().astype(np.uint8)
|
|||
|
|
h, w = h_list[i].item(), w_list[i].item()
|
|||
|
|
aggregator.add_patch(fname, x, y, patch, h, w)
|
|||
|
|
# aggregator.add_patch(fname, x, y, patch)
|
|||
|
|
|
|||
|
|
test_bar.set_description('IoU: %.4f' % (inter / union if union > 0 else 0))
|
|||
|
|
aggregator.save_all(self.args.save_dir)
|
|||
|
|
iou = inter / union if union != 0 else 0
|
|||
|
|
print(f"Test IoU: {iou:.4f}")
|
|||
|
|
|
|||
|
|
def predict_from_imgurl(self,img1Url,img2Url,save_dir):
|
|||
|
|
print(f"结果将保存到目录: {self.args.save_dir}")
|
|||
|
|
print("Starting prediction...")
|
|||
|
|
indexer_test = PatchIndexer(
|
|||
|
|
img_dir=self.args.path_img1,
|
|||
|
|
patch_size=256,
|
|||
|
|
stride=256
|
|||
|
|
)
|
|||
|
|
indexer_test.index_all()
|
|||
|
|
# 调用scv相对路径
|
|||
|
|
parent_test_cood = os.path.dirname(self.args.path_img1)
|
|||
|
|
# test_cood = os.path.join(parent_test_cood, "coords.csv")
|
|||
|
|
|
|||
|
|
if not os.path.exists(self.args.save_dir):
|
|||
|
|
os.makedirs(self.args.save_dir)
|
|||
|
|
|
|||
|
|
transform = transforms.Compose([
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
img1Url = r"D:\project\2025-7-10_data+model\val\image1\1.png"
|
|||
|
|
img2Url = r"D:\project\2025-7-10_data+model\val\image2\1.png"
|
|||
|
|
|
|||
|
|
test_set = TestDatasetFromPic(img1Url, img2Url, None, transform)
|
|||
|
|
|
|||
|
|
# test_set = TestDatasetFromCoords(coords_csv=test_cood, hr1_dir=self.args.path_img1, hr2_dir=self.args.path_img2,
|
|||
|
|
# label_dir=None, transform=transform)
|
|||
|
|
# test_set = TestDatasetFromFolder(self.args, self.args.path_img1,
|
|||
|
|
# self.args.path_img2, self.args.path_lab)
|
|||
|
|
test_loader = DataLoader(test_set, num_workers=24,
|
|||
|
|
batch_size=self.args.batch_size, shuffle=False)
|
|||
|
|
|
|||
|
|
self.model.eval()
|
|||
|
|
aggregator = PredictionAggregator(img_dir=self.args.path_img1)
|
|||
|
|
# aggregator=PredictionAggregator()
|
|||
|
|
test_bar = tqdm(test_loader)
|
|||
|
|
inter, union = 0, 0
|
|||
|
|
|
|||
|
|
for img1, img2, label, patch_name, fname_list, x_list, y_list, h_list, w_list in test_bar:
|
|||
|
|
img1 = img1.to(self.device)
|
|||
|
|
img2 = img2.to(self.device)
|
|||
|
|
label = label.to(self.device)
|
|||
|
|
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output = self.model(img1, img2)
|
|||
|
|
if isinstance(output, tuple): # ← 检查是否是 tuple
|
|||
|
|
output = output[0]
|
|||
|
|
|
|||
|
|
pred = torch.argmax(output, 1).unsqueeze(1).float()
|
|||
|
|
# label = torch.argmax(label, 1).unsqueeze(1).float()
|
|||
|
|
label = label.float()
|
|||
|
|
|
|||
|
|
for i in range(img1.size(0)):
|
|||
|
|
# prob = pred[i].cpu().numpy()
|
|||
|
|
# gt_value = label[i].cpu().numpy()
|
|||
|
|
patch = pred[i][0].cpu().numpy().astype(np.uint8)
|
|||
|
|
gt_value = label[i][0].cpu().numpy().astype(np.uint8)
|
|||
|
|
|
|||
|
|
# intr, unn = calMetric_iou(gt_value.squeeze(), prob.squeeze())
|
|||
|
|
intr, unn = calMetric_iou(gt_value, patch)
|
|||
|
|
inter += intr
|
|||
|
|
union += unn
|
|||
|
|
fname = fname_list[i]
|
|||
|
|
x, y = x_list[i].item(), y_list[i].item()
|
|||
|
|
# patch = pred[i][0].cpu().numpy().astype(np.uint8)
|
|||
|
|
h, w = h_list[i].item(), w_list[i].item()
|
|||
|
|
aggregator.add_patch(fname, x, y, patch, h, w)
|
|||
|
|
# aggregator.add_patch(fname, x, y, patch)
|
|||
|
|
|
|||
|
|
test_bar.set_description('IoU: %.4f' % (inter / union if union > 0 else 0))
|
|||
|
|
out_save_path=aggregator.save_all(save_dir)
|
|||
|
|
print(f"out_save_pathout_save_path {out_save_path}")
|
|||
|
|
iou = inter / union if union != 0 else 0
|
|||
|
|
print(f"Test IoU: {iou:.4f}")
|
|||
|
|
return out_save_path
|
|||
|
|
|
|||
|
|
# for image1, image2, label, image_names in tqdm(test_loader, desc='Testing'):
|
|||
|
|
# image1 = image1.to(self.device, dtype=torch.float)
|
|||
|
|
# image2 = image2.to(self.device, dtype=torch.float)
|
|||
|
|
# label = label.to(self.device, dtype=torch.float)
|
|||
|
|
|
|||
|
|
# output, _, _ = self.model(image1, image2)
|
|||
|
|
# pred = torch.argmax(output, 1).unsqueeze(1)
|
|||
|
|
# label = torch.argmax(label, 1).unsqueeze(1)
|
|||
|
|
|
|||
|
|
# for i in range(pred.size(0)):
|
|||
|
|
# gt_value = (label[i] > 0).float()
|
|||
|
|
# prob = (pred[i] > 0).float()
|
|||
|
|
# gt_np = np.squeeze(gt_value.cpu().detach().numpy())
|
|||
|
|
# prob_np = np.squeeze(prob.cpu().detach().numpy())
|
|||
|
|
# intr, unn = calMetric_iou(gt_np, prob_np)
|
|||
|
|
# inter += intr
|
|||
|
|
# union += unn
|
|||
|
|
# binary_result = np.where(prob_np > 0.5, 255, 0).astype('uint8') # 将预测值转换为0和255
|
|||
|
|
|
|||
|
|
# result = Image.fromarray(binary_result) # 使用转换后的二值图像
|
|||
|
|
# result.save(os.path.join(self.args.save_dir, image_names[i]))
|
|||
|
|
|
|||
|
|
# iou = inter / union if union != 0 else 0
|
|||
|
|
# print(f"Test IoU: {iou:.4f}")
|
|||
|
|
|
|||
|
|
def predict_without_label(self, img1, img2):
|
|||
|
|
"""
|
|||
|
|
直接对输入的 image1 和 image2 进行预测。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
image1 (torch.Tensor 或 np.ndarray): 输入图像1（需与模型输入兼容）。
|
|||
|
|
image2 (torch.Tensor 或 np.ndarray): 输入图像2（需与模型输入兼容）。
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
torch.Tensor: 模型的输出（未经过 argmax 处理）。
|
|||
|
|
torch.Tensor: 二值化后的预测结果（0 或 255）。
|
|||
|
|
"""
|
|||
|
|
print("Starting prediction...")
|
|||
|
|
# if not os.path.exists(self.args.save_dir):
|
|||
|
|
# os.makedirs(self.args.save_dir)
|
|||
|
|
transform = get_transform(convert=True, normalize=True)
|
|||
|
|
# image1 = transform(img1.convert('RGB'))
|
|||
|
|
# image2 = transform(img2.convert('RGB'))
|
|||
|
|
image1 = transform(img1.convert('RGB')).to(self.device) # 移动到设备
|
|||
|
|
image2 = transform(img2.convert('RGB')).to(self.device) # 移动到设备
|
|||
|
|
# # 确保输入是 Tensor，并移动到指定设备
|
|||
|
|
# if isinstance(image1, np.ndarray):
|
|||
|
|
# image1 = torch.from_numpy(image1).to(self.device, dtype=torch.float)
|
|||
|
|
# if isinstance(image2, np.ndarray):
|
|||
|
|
# image2 = torch.from_numpy(image2).to(self.device, dtype=torch.float)
|
|||
|
|
|
|||
|
|
# 模型推理
|
|||
|
|
self.model.eval()
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output, _, _ = self.model(image1.unsqueeze(0), image2.unsqueeze(0)) # 添加 batch 维度
|
|||
|
|
pred = torch.argmax(output, 1).squeeze(0) # 移除 batch 维度
|
|||
|
|
prob_np = pred.cpu().detach().numpy()
|
|||
|
|
|
|||
|
|
# 二值化
|
|||
|
|
binary_result = np.where(prob_np > 0.5, 255, 0).astype('uint8')
|
|||
|
|
|
|||
|
|
return output, binary_result
|