580 lines
22 KiB
Python
580 lines
22 KiB
Python
|
#coding=utf-8
|
|||
|
|
from os.path import join
|
|||
|
|
import torch
|
|||
|
|
from PIL import Image, ImageEnhance
|
|||
|
|
from torch.utils.data.dataset import Dataset
|
|||
|
|
|
|||
|
|
from torchvision.transforms import Compose, RandomCrop, ToTensor, ToPILImage, CenterCrop, Resize
|
|||
|
|
import numpy as np
|
|||
|
|
import torchvision.transforms as transforms
|
|||
|
|
import os
|
|||
|
|
from argparse import Namespace
|
|||
|
|
import imageio
|
|||
|
|
import pandas as pd
|
|||
|
|
import rasterio
|
|||
|
|
from rasterio.windows import Window
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
Image.MAX_IMAGE_PIXELS = None # 禁用像素限制警告
|
|||
|
|
def is_image_file(filename):
|
|||
|
|
return any(filename.endswith(extension) for extension in ['.tif','.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])
|
|||
|
|
|
|||
|
|
def calMetric_iou(predict, label):
|
|||
|
|
tp = np.sum(np.logical_and(predict == 1, label == 1))
|
|||
|
|
fp = np.sum(predict==1)
|
|||
|
|
fn = np.sum(label == 1)
|
|||
|
|
return tp,fp+fn-tp
|
|||
|
|
|
|||
|
|
|
|||
|
|
def getDataList(img_path):
|
|||
|
|
dataline = open(img_path, 'r').readlines()
|
|||
|
|
datalist =[]
|
|||
|
|
for line in dataline:
|
|||
|
|
temp = line.strip('\n')
|
|||
|
|
datalist.append(temp)
|
|||
|
|
return datalist
|
|||
|
|
|
|||
|
|
|
|||
|
|
def make_one_hot(input, num_classes):
|
|||
|
|
"""Convert class index tensor to one hot encoding tensor.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
input: A tensor of shape [N, 1, *]
|
|||
|
|
num_classes: An int of number of class
|
|||
|
|
Returns:
|
|||
|
|
A tensor of shape [N, num_classes, *]
|
|||
|
|
"""
|
|||
|
|
shape = np.array(input.shape)
|
|||
|
|
shape[1] = num_classes
|
|||
|
|
shape = tuple(shape)
|
|||
|
|
result = torch.zeros(shape)
|
|||
|
|
result = result.scatter_(1, input.cpu(), 1)
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_transform(convert=True, normalize=False):
|
|||
|
|
transform_list = []
|
|||
|
|
if convert:
|
|||
|
|
transform_list += [transforms.ToTensor()]
|
|||
|
|
if normalize:
|
|||
|
|
transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
|
|||
|
|
(0.5, 0.5, 0.5))]
|
|||
|
|
return transforms.Compose(transform_list)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# class LoadDatasetFromFolder(Dataset):
|
|||
|
|
# def __init__(self, args, hr1_path, hr2_path, lab_path):
|
|||
|
|
# super(LoadDatasetFromFolder, self).__init__()
|
|||
|
|
# # 获取图片列表
|
|||
|
|
# datalist = [name for name in os.listdir(hr1_path) for item in args.suffix if
|
|||
|
|
# os.path.splitext(name)[1] == item]
|
|||
|
|
|
|||
|
|
# self.hr1_filenames = [join(hr1_path, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.hr2_filenames = [join(hr2_path, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.lab_filenames = [join(lab_path, x) for x in datalist if is_image_file(x)]
|
|||
|
|
|
|||
|
|
# self.transform = get_transform(convert=True, normalize=True) # convert to tensor and normalize to [-1,1]
|
|||
|
|
# self.label_transform = get_transform() # only convert to tensor
|
|||
|
|
|
|||
|
|
# def __getitem__(self, index):
|
|||
|
|
# hr1_img = self.transform(Image.open(self.hr1_filenames[index]).convert('RGB'))
|
|||
|
|
# # lr2_img = self.transform(Image.open(self.lr2_filenames[index]).convert('RGB'))
|
|||
|
|
# hr2_img = self.transform(Image.open(self.hr2_filenames[index]).convert('RGB'))
|
|||
|
|
|
|||
|
|
# label = self.label_transform(Image.open(self.lab_filenames[index]))
|
|||
|
|
# label = make_one_hot(label.unsqueeze(0).long(), 2).squeeze(0)
|
|||
|
|
|
|||
|
|
# return hr1_img, hr2_img, label
|
|||
|
|
|
|||
|
|
# def __len__(self):
|
|||
|
|
# return len(self.hr1_filenames)
|
|||
|
|
#验证
|
|||
|
|
class LoadDatasetFromCoords(Dataset):
|
|||
|
|
def __init__(self, coords_csv, hr1_dir, hr2_dir, label_dir, num_classes=2):
|
|||
|
|
super(LoadDatasetFromCoords, self).__init__()
|
|||
|
|
self.coords = pd.read_csv(coords_csv)
|
|||
|
|
self.hr1_dir = hr1_dir
|
|||
|
|
self.hr2_dir = hr2_dir
|
|||
|
|
self.label_dir = label_dir
|
|||
|
|
self.num_classes = num_classes
|
|||
|
|
|
|||
|
|
self.transform = get_transform(convert=True, normalize=True)
|
|||
|
|
self.label_transform = get_transform()
|
|||
|
|
self.open_images = {}
|
|||
|
|
|
|||
|
|
def _load_image(self, path):
|
|||
|
|
if path not in self.open_images:
|
|||
|
|
self.open_images[path] = rasterio.open(path)
|
|||
|
|
return self.open_images[path]
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.coords)
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
row = self.coords.iloc[idx]
|
|||
|
|
fname = row['orig_name']
|
|||
|
|
x, y, h, w = int(row['x']), int(row['y']), int(row['h']), int(row['w'])
|
|||
|
|
|
|||
|
|
# 加载影像
|
|||
|
|
img1 = self._load_image(os.path.join(self.hr1_dir, fname))
|
|||
|
|
img2 = self._load_image(os.path.join(self.hr2_dir, fname))
|
|||
|
|
label = self._load_image(os.path.join(self.label_dir, fname))
|
|||
|
|
|
|||
|
|
window = Window(x, y, w, h)
|
|||
|
|
|
|||
|
|
patch1 = img1.read(window=window) # [C, H, W]
|
|||
|
|
patch2 = img2.read(window=window)
|
|||
|
|
patch_label = label.read(1, window=window) # [H, W]
|
|||
|
|
|
|||
|
|
# 转为 PIL.Image 以使用 transform
|
|||
|
|
patch1 = Image.fromarray(np.transpose(patch1, (1, 2, 0)).astype(np.uint8))
|
|||
|
|
patch2 = Image.fromarray(np.transpose(patch2, (1, 2, 0)).astype(np.uint8))
|
|||
|
|
patch_label = Image.fromarray(patch_label.astype(np.uint8))
|
|||
|
|
|
|||
|
|
patch1 = self.transform(patch1)
|
|||
|
|
patch2 = self.transform(patch2)
|
|||
|
|
patch_label = self.label_transform(patch_label)
|
|||
|
|
|
|||
|
|
patch_label = make_one_hot(patch_label.unsqueeze(0).long(), self.num_classes).squeeze(0)
|
|||
|
|
|
|||
|
|
return patch1, patch2, patch_label
|
|||
|
|
|
|||
|
|
|
|||
|
|
#测试
|
|||
|
|
# class TestDatasetFromCoords(Dataset):
|
|||
|
|
# def __init__(self, coords_csv, hr1_dir, hr2_dir, label_dir=None, transform=None):
|
|||
|
|
# self.coords = pd.read_csv(coords_csv)
|
|||
|
|
# self.hr1_dir = hr1_dir
|
|||
|
|
# self.hr2_dir = hr2_dir
|
|||
|
|
# self.label_dir = label_dir
|
|||
|
|
# self.transform = transform
|
|||
|
|
# self.label_transform = transforms.ToTensor()
|
|||
|
|
|
|||
|
|
# def __len__(self):
|
|||
|
|
# return len(self.coords)
|
|||
|
|
|
|||
|
|
# def __getitem__(self, idx):
|
|||
|
|
# row = self.coords.iloc[idx]
|
|||
|
|
# x, y, h, w = row['x'], row['y'], row['h'], row['w']
|
|||
|
|
# fname = row['orig_name']
|
|||
|
|
|
|||
|
|
# img1 = Image.open(os.path.join(self.hr1_dir, fname)).convert('RGB')
|
|||
|
|
# img2 = Image.open(os.path.join(self.hr2_dir, fname)).convert('RGB')
|
|||
|
|
|
|||
|
|
# img1 = img1.crop((x, y, x + w, y + h))
|
|||
|
|
# img2 = img2.crop((x, y, x + w, y + h))
|
|||
|
|
|
|||
|
|
# if self.label_dir:
|
|||
|
|
# label = Image.open(os.path.join(self.label_dir, fname)).convert('L')
|
|||
|
|
# label = label.crop((x, y, x + w, y + h))
|
|||
|
|
# label = self.label_transform(label)
|
|||
|
|
# else:
|
|||
|
|
# label = torch.zeros((1, h, w)) # dummy label
|
|||
|
|
|
|||
|
|
# if self.transform:
|
|||
|
|
# img1 = self.transform(img1)
|
|||
|
|
# img2 = self.transform(img2)
|
|||
|
|
|
|||
|
|
# patch_name = f"{os.path.splitext(fname)[0]}_{x}_{y}"
|
|||
|
|
# return img1, img2, label, patch_name, fname, x, y
|
|||
|
|
class TestDatasetFromCoords(Dataset):
|
|||
|
|
def __init__(self, coords_csv, hr1_dir, hr2_dir, label_dir=None, transform=None):
|
|||
|
|
self.coords = pd.read_csv(coords_csv)
|
|||
|
|
self.hr1_dir = hr1_dir
|
|||
|
|
self.hr2_dir = hr2_dir
|
|||
|
|
self.label_dir = label_dir
|
|||
|
|
self.transform = transform
|
|||
|
|
self.label_transform = transforms.ToTensor()
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.coords)
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
row = self.coords.iloc[idx]
|
|||
|
|
x, y, h, w = int(row['x']), int(row['y']), int(row['h']), int(row['w'])
|
|||
|
|
fname = row['orig_name']
|
|||
|
|
|
|||
|
|
# 读取图像并裁剪 patch
|
|||
|
|
img1 = Image.open(os.path.join(self.hr1_dir, fname)).convert('RGB')
|
|||
|
|
img2 = Image.open(os.path.join(self.hr2_dir, fname)).convert('RGB')
|
|||
|
|
img1 = img1.crop((x, y, x + w, y + h))
|
|||
|
|
img2 = img2.crop((x, y, x + w, y + h))
|
|||
|
|
|
|||
|
|
# 标签
|
|||
|
|
if self.label_dir:
|
|||
|
|
label = Image.open(os.path.join(self.label_dir, fname)).convert('L')
|
|||
|
|
label = label.crop((x, y, x + w, y + h))
|
|||
|
|
label = self.label_transform(label)
|
|||
|
|
else:
|
|||
|
|
label = torch.zeros((1, h, w)) # dummy label
|
|||
|
|
|
|||
|
|
# 图像 transform
|
|||
|
|
if self.transform:
|
|||
|
|
img1 = self.transform(img1)
|
|||
|
|
img2 = self.transform(img2)
|
|||
|
|
|
|||
|
|
# padding 到 256x256，并记录原始 h, w
|
|||
|
|
pad_h = 256 - h
|
|||
|
|
pad_w = 256 - w
|
|||
|
|
if pad_h > 0 or pad_w > 0:
|
|||
|
|
img1 = F.pad(img1, (0, pad_w, 0, pad_h)) # pad right & bottom
|
|||
|
|
img2 = F.pad(img2, (0, pad_w, 0, pad_h))
|
|||
|
|
label = F.pad(label, (0, pad_w, 0, pad_h))
|
|||
|
|
|
|||
|
|
patch_name = f"{os.path.splitext(fname)[0]}_{x}_{y}"
|
|||
|
|
return img1, img2, label, patch_name, fname, x, y, h, w # 返回原始大小
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TestDatasetFromImagePath(Dataset):
|
|||
|
|
def __init__(self, img1_path,img2_path, transform=None):
|
|||
|
|
# self.coords = pd.read_csv(coords_csv)
|
|||
|
|
# self.hr1_dir = hr1_dir
|
|||
|
|
# self.hr2_dir = hr2_dir
|
|||
|
|
# self.label_dir = label_dir
|
|||
|
|
self.img1_path = img1_path
|
|||
|
|
self.img2_path = img2_path
|
|||
|
|
self.transform = transform
|
|||
|
|
self.label_transform = transforms.ToTensor()
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
# row = self.coords.iloc[idx]
|
|||
|
|
# x, y, h, w = int(row['x']), int(row['y']), int(row['h']), int(row['w'])
|
|||
|
|
# fname = row['orig_name']
|
|||
|
|
|
|||
|
|
# 读取图像并裁剪 patch
|
|||
|
|
img1 = Image.open(self.img1_path).convert('RGB')
|
|||
|
|
img2 = Image.open(self.img2_path).convert('RGB')
|
|||
|
|
img1 = img1.crop((x, y, x + w, y + h))
|
|||
|
|
img2 = img2.crop((x, y, x + w, y + h))
|
|||
|
|
|
|||
|
|
# 标签
|
|||
|
|
if self.label_dir:
|
|||
|
|
label = Image.open(os.path.join(self.label_dir, fname)).convert('L')
|
|||
|
|
label = label.crop((x, y, x + w, y + h))
|
|||
|
|
label = self.label_transform(label)
|
|||
|
|
else:
|
|||
|
|
label = torch.zeros((1, h, w)) # dummy label
|
|||
|
|
|
|||
|
|
# 图像 transform
|
|||
|
|
if self.transform:
|
|||
|
|
img1 = self.transform(img1)
|
|||
|
|
img2 = self.transform(img2)
|
|||
|
|
|
|||
|
|
# padding 到 256x256，并记录原始 h, w
|
|||
|
|
pad_h = 256 - h
|
|||
|
|
pad_w = 256 - w
|
|||
|
|
if pad_h > 0 or pad_w > 0:
|
|||
|
|
img1 = F.pad(img1, (0, pad_w, 0, pad_h)) # pad right & bottom
|
|||
|
|
img2 = F.pad(img2, (0, pad_w, 0, pad_h))
|
|||
|
|
label = F.pad(label, (0, pad_w, 0, pad_h))
|
|||
|
|
|
|||
|
|
patch_name = f"{os.path.splitext(fname)[0]}_{x}_{y}"
|
|||
|
|
return img1, img2, label, patch_name, fname, x, y, h, w # 返回原始大小
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# class TestDatasetFromFolder(Dataset):
|
|||
|
|
# def __init__(self, args, Time1_dir, Time2_dir, Label_dir):
|
|||
|
|
# super(TestDatasetFromFolder, self).__init__()
|
|||
|
|
|
|||
|
|
# # 确保 args.suffix 存在，否则使用默认值
|
|||
|
|
# suffixes = getattr(args, 'suffix', ['.png', '.jpg'])
|
|||
|
|
|
|||
|
|
# # 筛选符合后缀的文件
|
|||
|
|
# datalist = [
|
|||
|
|
# name for name in os.listdir(Time1_dir)
|
|||
|
|
# if os.path.splitext(name)[1].lower() in suffixes
|
|||
|
|
# ]
|
|||
|
|
|
|||
|
|
# print(f"找到 {len(datalist)} 个文件: {datalist[:5]}...") # 调试输出
|
|||
|
|
|
|||
|
|
# self.image1_filenames = [join(Time1_dir, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.image2_filenames = [join(Time2_dir, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.image3_filenames = [join(Label_dir, x) for x in datalist if is_image_file(x)]
|
|||
|
|
|
|||
|
|
# self.transform = get_transform(convert=True, normalize=True) # convert to tensor and normalize to [-1,1]
|
|||
|
|
# self.label_transform = get_transform()
|
|||
|
|
|
|||
|
|
# def __getitem__(self, index):
|
|||
|
|
# image1 = self.transform(Image.open(self.image1_filenames[index]).convert('RGB'))
|
|||
|
|
# image2 = self.transform(Image.open(self.image2_filenames[index]).convert('RGB'))
|
|||
|
|
|
|||
|
|
# label = self.label_transform(Image.open(self.image3_filenames[index]))
|
|||
|
|
# label = make_one_hot(label.unsqueeze(0).long(), 2).squeeze(0)
|
|||
|
|
|
|||
|
|
# image_name = self.image1_filenames[index].split('/', -1)
|
|||
|
|
# image_name = image_name[len(image_name)-1]
|
|||
|
|
|
|||
|
|
# return image1, image2, label, image_name
|
|||
|
|
|
|||
|
|
# def __len__(self):
|
|||
|
|
# return len(self.image1_filenames)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class trainImageAug(object):
|
|||
|
|
def __init__(self, crop = True, augment = True, angle = 30):
|
|||
|
|
self.crop =crop
|
|||
|
|
self.augment = augment
|
|||
|
|
self.angle = angle
|
|||
|
|
|
|||
|
|
def __call__(self, image1, image2, mask):
|
|||
|
|
if self.crop:
|
|||
|
|
w = np.random.randint(0,256)
|
|||
|
|
h = np.random.randint(0,256)
|
|||
|
|
box = (w, h, w+256, h+256)
|
|||
|
|
image1 = image1.crop(box)
|
|||
|
|
image2 = image2.crop(box)
|
|||
|
|
mask = mask.crop(box)
|
|||
|
|
if self.augment:
|
|||
|
|
prop = np.random.uniform(0, 1)
|
|||
|
|
if prop < 0.15:
|
|||
|
|
image1 = image1.transpose(Image.FLIP_LEFT_RIGHT)
|
|||
|
|
image2 = image2.transpose(Image.FLIP_LEFT_RIGHT)
|
|||
|
|
mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
|
|||
|
|
elif prop < 0.3:
|
|||
|
|
image1 = image1.transpose(Image.FLIP_TOP_BOTTOM)
|
|||
|
|
image2 = image2.transpose(Image.FLIP_TOP_BOTTOM)
|
|||
|
|
mask = mask.transpose(Image.FLIP_TOP_BOTTOM)
|
|||
|
|
elif prop < 0.5:
|
|||
|
|
image1 = image1.rotate(transforms.RandomRotation.get_params([-self.angle, self.angle]))
|
|||
|
|
image2 = image2.rotate(transforms.RandomRotation.get_params([-self.angle, self.angle]))
|
|||
|
|
mask = mask.rotate(transforms.RandomRotation.get_params([-self.angle, self.angle]))
|
|||
|
|
|
|||
|
|
return image1, image2, mask
|
|||
|
|
|
|||
|
|
def get_transform(convert=True, normalize=False):
|
|||
|
|
transform_list = []
|
|||
|
|
if convert:
|
|||
|
|
transform_list += [
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
]
|
|||
|
|
if normalize:
|
|||
|
|
transform_list += [
|
|||
|
|
# transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
|
|||
|
|
transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))]
|
|||
|
|
return transforms.Compose(transform_list)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DynamicPatchDataset(Dataset):
|
|||
|
|
def __init__(self, coords_csv, img1_dir, img2_dir, label_dir,
|
|||
|
|
crop=True, augment=True, angle=30, num_classes=2):
|
|||
|
|
self.coords = pd.read_csv(coords_csv)
|
|||
|
|
self.img1_dir = img1_dir
|
|||
|
|
self.img2_dir = img2_dir
|
|||
|
|
self.label_dir = label_dir
|
|||
|
|
self.num_classes = num_classes
|
|||
|
|
|
|||
|
|
# 与原 DA_DatasetFromFolder 一致的设置
|
|||
|
|
self.data_augment = trainImageAug(crop=crop, augment=augment, angle=angle)
|
|||
|
|
self.img_transform = get_transform(convert=True, normalize=True)
|
|||
|
|
self.lab_transform = get_transform()
|
|||
|
|
|
|||
|
|
self.open_images = {} # 缓存图像
|
|||
|
|
|
|||
|
|
def _load_image(self, path):
|
|||
|
|
if path not in self.open_images:
|
|||
|
|
self.open_images[path] = rasterio.open(path)
|
|||
|
|
return self.open_images[path]
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.coords)
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
row = self.coords.iloc[idx]
|
|||
|
|
fname = row['orig_name']
|
|||
|
|
x, y, h, w = int(row['x']), int(row['y']), int(row['h']), int(row['w'])
|
|||
|
|
|
|||
|
|
# 打开大图并裁剪 patch
|
|||
|
|
img1 = self._load_image(os.path.join(self.img1_dir, fname))
|
|||
|
|
img2 = self._load_image(os.path.join(self.img2_dir, fname))
|
|||
|
|
label = self._load_image(os.path.join(self.label_dir, fname))
|
|||
|
|
|
|||
|
|
window = Window(x, y, w, h)
|
|||
|
|
|
|||
|
|
patch1 = img1.read(window=window) # [C, H, W]
|
|||
|
|
patch2 = img2.read(window=window)
|
|||
|
|
patch_label = label.read(1, window=window) # [H, W]
|
|||
|
|
|
|||
|
|
# 转换为 PIL Image 以便和原有增强方式兼容
|
|||
|
|
patch1 = Image.fromarray(np.transpose(patch1, (1, 2, 0)).astype(np.uint8)) # [H, W, C]
|
|||
|
|
patch2 = Image.fromarray(np.transpose(patch2, (1, 2, 0)).astype(np.uint8))
|
|||
|
|
patch_label = Image.fromarray(patch_label.astype(np.uint8))
|
|||
|
|
|
|||
|
|
# 数据增强（裁剪、旋转等）
|
|||
|
|
patch1, patch2, patch_label = self.data_augment(patch1, patch2, patch_label)
|
|||
|
|
|
|||
|
|
# 图像 transform（归一化、转 tensor）
|
|||
|
|
patch1 = self.img_transform(patch1)
|
|||
|
|
patch2 = self.img_transform(patch2)
|
|||
|
|
patch_label = self.lab_transform(patch_label)
|
|||
|
|
|
|||
|
|
# one-hot label
|
|||
|
|
patch_label = make_one_hot(patch_label.unsqueeze(0).long(), self.num_classes).squeeze(0)
|
|||
|
|
|
|||
|
|
return patch1, patch2, patch_label
|
|||
|
|
# class DA_DatasetFromFolder(Dataset):
|
|||
|
|
# def __init__(self, Image_dir1, Image_dir2, Label_dir, crop=True, augment = True, angle = 30):
|
|||
|
|
# super(DA_DatasetFromFolder, self).__init__()
|
|||
|
|
# # 获取图片列表
|
|||
|
|
# datalist = os.listdir(Image_dir1)
|
|||
|
|
# self.image_filenames1 = [join(Image_dir1, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.image_filenames2 = [join(Image_dir2, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.label_filenames = [join(Label_dir, x) for x in datalist if is_image_file(x)]
|
|||
|
|
# self.data_augment = trainImageAug(crop=crop, augment = augment, angle=angle)
|
|||
|
|
# self.img_transform = get_transform(convert=True, normalize=True)
|
|||
|
|
# self.lab_transform = get_transform()
|
|||
|
|
|
|||
|
|
# def __getitem__(self, index):
|
|||
|
|
# image1 = Image.open(self.image_filenames1[index]).convert('RGB')
|
|||
|
|
# image2 = Image.open(self.image_filenames2[index]).convert('RGB')
|
|||
|
|
# label = Image.open(self.label_filenames[index])
|
|||
|
|
# image1, image2, label = self.data_augment(image1, image2, label)
|
|||
|
|
# image1, image2 = self.img_transform(image1), self.img_transform(image2)
|
|||
|
|
# label = self.lab_transform(label)
|
|||
|
|
# label = make_one_hot(label.unsqueeze(0).long(), 2).squeeze(0)
|
|||
|
|
# return image1, image2, label
|
|||
|
|
|
|||
|
|
# def __len__(self):
|
|||
|
|
# return len(self.image_filenames1)
|
|||
|
|
|
|||
|
|
class TestDatasetFromPic(Dataset):
|
|||
|
|
def __init__(self, img1_path, img2_path, label_path=None, transform=None,
|
|||
|
|
patch_size=256, stride=256, pad_to_patch=False):
|
|||
|
|
"""
|
|||
|
|
动态滑动窗口测试数据集（单张大图版本）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
img1_path (str): 第一期影像的完整路径
|
|||
|
|
img2_path (str): 第二期影像的完整路径
|
|||
|
|
label_path (str, optional): 标签路径（测试时可能不需要）
|
|||
|
|
transform (callable, optional): 图像变换（如归一化）
|
|||
|
|
patch_size (int): 裁剪的patch尺寸（默认256）
|
|||
|
|
stride (int): 滑动步长（默认256，无重叠）
|
|||
|
|
pad_to_patch (bool): 是否将图像边缘不足patch_size的部分填充到完整patch
|
|||
|
|
"""
|
|||
|
|
self.img1_path = img1_path
|
|||
|
|
self.img2_path = img2_path
|
|||
|
|
self.label_path = label_path
|
|||
|
|
self.transform = transform or transforms.Compose([
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
|
|||
|
|
])
|
|||
|
|
self.patch_size = patch_size
|
|||
|
|
self.stride = stride
|
|||
|
|
self.pad_to_patch = pad_to_patch
|
|||
|
|
self.label_transform = transforms.ToTensor()
|
|||
|
|
|
|||
|
|
# 提取文件名（用于保存结果时命名）
|
|||
|
|
self.img1_name = os.path.basename(img1_path)
|
|||
|
|
self.img2_name = os.path.basename(img2_path)
|
|||
|
|
|
|||
|
|
# 动态生成坐标列表
|
|||
|
|
self.coords = self._generate_coords()
|
|||
|
|
|
|||
|
|
def _generate_coords(self):
|
|||
|
|
"""为单张图像生成滑动窗口坐标"""
|
|||
|
|
coords = []
|
|||
|
|
|
|||
|
|
# 获取图像尺寸
|
|||
|
|
with Image.open(self.img1_path) as img:
|
|||
|
|
width, height = img.size
|
|||
|
|
|
|||
|
|
# 计算需要裁剪的patch数量（考虑边缘填充）
|
|||
|
|
if self.pad_to_patch:
|
|||
|
|
pad_w = (self.patch_size - width % self.patch_size) % self.patch_size
|
|||
|
|
pad_h = (self.patch_size - height % self.patch_size) % self.patch_size
|
|||
|
|
n_patches_w = (width + pad_w) // self.stride
|
|||
|
|
n_patches_h = (height + pad_h) // self.stride
|
|||
|
|
else:
|
|||
|
|
n_patches_w = (width - self.patch_size) // self.stride + 1
|
|||
|
|
n_patches_h = (height - self.patch_size) // self.stride + 1
|
|||
|
|
|
|||
|
|
# 生成所有patch的坐标
|
|||
|
|
for i in range(n_patches_w):
|
|||
|
|
for j in range(n_patches_h):
|
|||
|
|
x = i * self.stride
|
|||
|
|
y = j * self.stride
|
|||
|
|
# 确保不超出原始图像边界（不填充时）
|
|||
|
|
if not self.pad_to_patch:
|
|||
|
|
if x + self.patch_size > width or y + self.patch_size > height:
|
|||
|
|
continue
|
|||
|
|
coords.append({
|
|||
|
|
'orig_name': self.img1_name, # 保持与原代码兼容
|
|||
|
|
'x': x,
|
|||
|
|
'y': y,
|
|||
|
|
'h': self.patch_size,
|
|||
|
|
'w': self.patch_size
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return coords
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.coords)
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
"""获取单个patch的数据"""
|
|||
|
|
row = self.coords[idx]
|
|||
|
|
x, y, h, w = row['x'], row['y'], row['h'], row['w']
|
|||
|
|
|
|||
|
|
# 加载图像并裁剪
|
|||
|
|
img1 = Image.open(self.img1_path).convert('RGB')
|
|||
|
|
img2 = Image.open(self.img2_path).convert('RGB')
|
|||
|
|
|
|||
|
|
# 处理标签（如果有）
|
|||
|
|
if self.label_path:
|
|||
|
|
label = Image.open(self.label_path).convert('L')
|
|||
|
|
label = label.crop((x, y, x + w, y + h))
|
|||
|
|
label = self.label_transform(label)
|
|||
|
|
else:
|
|||
|
|
label = torch.zeros((1, h, w)) # 虚拟标签
|
|||
|
|
|
|||
|
|
# 裁剪图像
|
|||
|
|
img1 = img1.crop((x, y, x + w, y + h))
|
|||
|
|
img2 = img2.crop((x, y, x + w, y + h))
|
|||
|
|
|
|||
|
|
# 应用变换
|
|||
|
|
if self.transform:
|
|||
|
|
img1 = self.transform(img1)
|
|||
|
|
img2 = self.transform(img2)
|
|||
|
|
|
|||
|
|
# 填充到统一尺寸（如果启用）
|
|||
|
|
if self.pad_to_patch:
|
|||
|
|
pad_h = self.patch_size - h
|
|||
|
|
pad_w = self.patch_size - w
|
|||
|
|
if pad_h > 0 or pad_w > 0:
|
|||
|
|
img1 = F.pad(img1, (0, pad_w, 0, pad_h))
|
|||
|
|
img2 = F.pad(img2, (0, pad_w, 0, pad_h))
|
|||
|
|
label = F.pad(label, (0, pad_w, 0, pad_h))
|
|||
|
|
|
|||
|
|
# 返回结果（保持与原代码兼容）
|
|||
|
|
patch_name = f"{os.path.splitext(self.img1_name)[0]}_{x}_{y}.png"
|
|||
|
|
return img1, img2, label, patch_name, self.img1_name, x, y, h, w
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 模拟 args（可以是 argparse.Namespace 或硬编码的 Namespace）
|
|||
|
|
args = Namespace()
|
|||
|
|
args.suffix = ['.png', '.jpg'] # 确保 suffix 存在
|
|||
|
|
|
|||
|
|
dataset = TestDatasetFromFolder(
|
|||
|
|
args,
|
|||
|
|
Time1_dir="path/to/Time1",
|
|||
|
|
Time2_dir="path/to/Time2",
|
|||
|
|
Label_dir="path/to/Label"
|
|||
|
|
)
|