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"
|
||
)
|