Spaces:
Sleeping
Sleeping
File size: 4,089 Bytes
d43892c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """
变化检测数据集
"""
import os
from PIL import Image
import numpy as np
from torch.utils import data
from datasets.data_utils import CDDataAugmentation
"""
CD data set with pixel-level labels；
├─image
├─image_post
├─label
└─list
"""
IMG_FOLDER_NAME = "A"
IMG_POST_FOLDER_NAME = 'B'
LIST_FOLDER_NAME = 'list'
ANNOT_FOLDER_NAME = "label"
IGNORE = 255
label_suffix='.png' # jpg for gan dataset, others : png
def load_img_name_list(dataset_path):
img_name_list = np.loadtxt(dataset_path, dtype=np.str)
if img_name_list.ndim == 2:
return img_name_list[:, 0]
return img_name_list
def load_image_label_list_from_npy(npy_path, img_name_list):
cls_labels_dict = np.load(npy_path, allow_pickle=True).item()
return [cls_labels_dict[img_name] for img_name in img_name_list]
def get_img_post_path(root_dir,img_name):
return os.path.join(root_dir, IMG_POST_FOLDER_NAME, img_name)
def get_img_path(root_dir, img_name):
return os.path.join(root_dir, IMG_FOLDER_NAME, img_name)
def get_label_path(root_dir, img_name):
return os.path.join(root_dir, ANNOT_FOLDER_NAME, img_name.replace('.jpg', label_suffix))
class ImageDataset(data.Dataset):
"""VOCdataloder"""
def __init__(self, root_dir, split='train', img_size=256, is_train=True,to_tensor=True):
super(ImageDataset, self).__init__()
self.root_dir = root_dir
self.img_size = img_size
self.split = split # train | train_aug | val
# self.list_path = self.root_dir + '/' + LIST_FOLDER_NAME + '/' + self.list + '.txt'
self.list_path = os.path.join(self.root_dir, LIST_FOLDER_NAME, self.split+'.txt')
self.img_name_list = load_img_name_list(self.list_path)
self.A_size = len(self.img_name_list) # get the size of dataset A
self.to_tensor = to_tensor
if is_train:
self.augm = CDDataAugmentation(
img_size=self.img_size,
with_random_hflip=True,
with_random_vflip=True,
with_scale_random_crop=True,
with_random_blur=True,
)
else:
self.augm = CDDataAugmentation(
img_size=self.img_size
)
def __getitem__(self, index):
name = self.img_name_list[index]
A_path = get_img_path(self.root_dir, self.img_name_list[index % self.A_size])
B_path = get_img_post_path(self.root_dir, self.img_name_list[index % self.A_size])
img = np.asarray(Image.open(A_path).convert('RGB'))
img_B = np.asarray(Image.open(B_path).convert('RGB'))
[img, img_B], _ = self.augm.transform([img, img_B],[], to_tensor=self.to_tensor)
return {'A': img, 'B': img_B, 'name': name}
def __len__(self):
"""Return the total number of images in the dataset."""
return self.A_size
class CDDataset(ImageDataset):
def __init__(self, root_dir, img_size, split='train', is_train=True, label_transform=None,
to_tensor=True):
super(CDDataset, self).__init__(root_dir, img_size=img_size, split=split, is_train=is_train,
to_tensor=to_tensor)
self.label_transform = label_transform
def __getitem__(self, index):
name = self.img_name_list[index]
A_path = get_img_path(self.root_dir, self.img_name_list[index % self.A_size])
B_path = get_img_post_path(self.root_dir, self.img_name_list[index % self.A_size])
img = np.asarray(Image.open(A_path).convert('RGB'))
img_B = np.asarray(Image.open(B_path).convert('RGB'))
L_path = get_label_path(self.root_dir, self.img_name_list[index % self.A_size])
label = np.array(Image.open(L_path), dtype=np.uint8)
# 二分类中，前景标注为255
if self.label_transform == 'norm':
label = label // 255
[img, img_B], [label] = self.augm.transform([img, img_B], [label], to_tensor=self.to_tensor)
# print(label.max())
return {'name': name, 'A': img, 'B': img_B, 'L': label}
|