Temporarily introduce image-mask pattern data sets , by manual segmentation data set EGTEA Gaze + Example .
Prepare the data folder
Needs to Image and Mask stored separately , corresponding to the file name must be consistent . Reminder : Mask image generally png single channel
EGTEA Gaze + dataset download extract obtained after following directory , without treatment
hand14k
┣━ Images
┃ ┣━ OP01-R01-PastaSalad_000014.jpg
┃ ┣━ OP01-R01-PastaSalad_000015.jpg
┃ ┣━ OP01-R01-PastaSalad_000016.jpg
┃ ┗━ ···
┗━ Masks
┣━ OP01-R01-pastasalad_000014.png
┣━ OP01-R01-pastasalad_000015.png
┣━ OP01-R01-pastasalad_000016.png
┗━ ···
Generating a file path , Data Partitioning
Script is as follows : import cv2 as cv
import numpy as np
import PIL.Image as Image
import os
np.random.seed(42)
def split_dataset():
# Read image files
images_path = "./Images/"
images_list = os.listdir (images_path) # returned each time the list of files order is inconsistent
images_list.sort () # required sorting process
# Reads the tag / Mask image
labels_path = "./Masks/"
labels_list = os.listdir(labels_path)
labels_list.sort()
# Create a path to a file ( using binary encoding , to avoid the operating system does not match )
train_file = "./train.data"
test_file = "./test.data"
if os.path.isfile(train_file) and os.path.isfile(test_file):
return
train_file = open(train_file, "wb")
test_file = open(test_file, "wb")
split_ratio = 0.8
for image, label in zip(images_list, labels_list):
image = os.path.join(images_path, image)
label = os.path.join(labels_path, label)
if os.path.basename(image).split('.')[0] != os.path.basename(label).split('.')[0]:
continue
file = train_file if np.random.rand() < split_ratio else test_file
file.write((image + "\t" + label + "\n").encode("utf-8"))
train_file.close()
test_file.close()
print ( " successfully divided the data set !")
def read_image(path):
img = np.array(Image.open(path))
if img.ndim == 2:
img = cv.merge([img, img, img])
return img
def test_read():
train_file = "./test.data"
with open(train_file, 'rb') as f:
datalist = f.readlines()
datalist = [(k, v) for k, v in map(lambda x: x.decode('utf-8').strip('\n').split('\t'), datalist)]
item = datalist[np.random.randint(42)]
image = read_image(item[0])
mask = read_image(item[1])
cv.imshow("image", image)
cv.imshow("mask", mask)
cv.waitKey(0)
cv.destroyAllWindows()
if __name__ == '__main__':
split_dataset()
test_read()
Derived Dataset class
class MyDataset(Dataset):
def __init__(
self, data_file, data_dir, transform_trn=None, transform_val=None
):
"""
Args:
data_file (string): Path to the data file with annotations.
data_dir (string): Directory with all the images.
transform_{trn, val} (callable, optional): Optional transform to be applied
on a sample.
"""
with open(data_file, 'rb') as f:
datalist = f.readlines()
self.datalist = [(k, v) for k, v in map(lambda x: x.decode('utf-8').strip('\n').split('\t'), datalist)]
self.root_dir = data_dir
self.transform_trn = transform_trn
self.transform_val = transform_val
self.stage = 'train'
def set_stage(self, stage):
self.stage = stage
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
img_name = os.path.join(self.root_dir, self.datalist[idx][0])
msk_name = os.path.join(self.root_dir, self.datalist[idx][1])
def read_image(x):
img_arr = np.array(Image.open(x))
if len(img_arr.shape) == 2: # grayscale
img_arr = np.tile(img_arr, [3, 1, 1]).transpose(1, 2, 0)
return img_arr
image = read_image(img_name)
mask = np.array(Image.open(msk_name))
if img_name != msk_name:
assert len(mask.shape) == 2, 'Masks must be encoded without colourmap'
sample = {'image': image, 'mask': mask}
if self.stage == 'train':
if self.transform_trn:
sample = self.transform_trn(sample)
elif self.stage == 'val':
if self.transform_val:
sample = self.transform_val(sample)
return sample
Construction DataLoader
# Define Transform
composed_trn = transforms.Compose([ResizeShorterScale(shorter_side, low_scale, high_scale),
Pad(crop_size, [123.675, 116.28, 103.53], ignore_label),
RandomMirror(),
RandomCrop(crop_size),
Normalise(*normalise_params),
ToTensor ()])
composed_val = transforms.Compose([Normalise(*normalise_params),
ToTensor ()])
# Import data set
trainset = MyDataset(data_file=train_list,
data_dir = train_dir,
transform_trn=composed_trn,
transform_val=composed_val)
rolled = MyDataset (data_file = val_list,
data_dir = val_dir,
transform_trn=None,
transform_val=composed_val)
# Build generator
train_loader = DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
drop_last=True)
val_loader = Data Loader (rolled,
batch_size=1,
shuffle=False,
num_workers=num_workers,
pin_memory=True)
training
for i, sample in enumerate(train_loader):
image = sample['image'].cuda()
target = sample['mask'].cuda()
image_var = torch.autograd.Variable(image).float()
target_var = torch.autograd.Variable(target).long()
# Compute output
output = net(image_var)
...
Original link: https://blog.csdn.net/Augurlee/article/details/103652444