222 lines
10 KiB
Python
222 lines
10 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 cropland_module.data_utils import get_transform, calMetric_iou, TestDatasetFromFolder
|
||
from cropland_module.loss.losses import cross_entropy
|
||
from cropland_module.model.network import CDNet
|
||
# from model.network import CDNet
|
||
# from data_utils import DA_DatasetFromFolder, LoadDatasetFromFolder, TestDatasetFromFolder, calMetric_iou, \
|
||
# TestDatasetFromFolderWithoutLabel, get_transform
|
||
# from loss.losses import cross_entropy
|
||
from PIL import Image
|
||
|
||
class ChangeDetectionModule:
|
||
|
||
|
||
|
||
|
||
def __init__(self, gpu_id,img_size,lr):
|
||
# 参数完整性检查
|
||
# 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"] = 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=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=lr, betas=(0.9, 0.999))
|
||
|
||
# def train(self):
|
||
# print("Starting training...")
|
||
# train_set = DA_DatasetFromFolder(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,save_dir):
|
||
|
||
print("Starting prediction...")
|
||
if not os.path.exists(save_dir):
|
||
os.makedirs(save_dir)
|
||
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()
|
||
inter, union = 0, 0
|
||
with torch.no_grad():
|
||
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, os.path.basename(image_names[i])))
|
||
|
||
iou = inter / union if union != 0 else 0
|
||
print(f"Test IoU: {iou:.4f}")
|
||
|
||
# def predict_without_label(self):
|
||
# print("Starting prediction...")
|
||
# if not os.path.exists(self.args.save_dir):
|
||
# os.makedirs(self.args.save_dir)
|
||
#
|
||
# # 使用无标签的数据集
|
||
# test_set = TestDatasetFromFolderWithoutLabel(self.args, self.args.path_img1, self.args.path_img2)
|
||
# test_loader = DataLoader(test_set, num_workers=24, batch_size=self.args.batch_size, shuffle=False)
|
||
#
|
||
# self.model.eval()
|
||
# with torch.no_grad():
|
||
# for image1, image2, image_names in tqdm(test_loader, desc='Testing'):
|
||
# image1 = image1.to(self.device, dtype=torch.float)
|
||
# image2 = image2.to(self.device, dtype=torch.float)
|
||
#
|
||
# output, _, _ = self.model(image1, image2)
|
||
# pred = torch.argmax(output, 1).unsqueeze(1)
|
||
# prob_np = pred.squeeze().cpu().detach().numpy()
|
||
#
|
||
# # 二值化并保存结果
|
||
# binary_result = np.where(prob_np > 0.5, 255, 0).astype('uint8')
|
||
# result = Image.fromarray(binary_result)
|
||
# result.save(os.path.join(self.args.save_dir, image_names))
|
||
|
||
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 |