You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

136 lines
5.2 KiB

# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import copy
import cv2
import numpy as np
from paddlers.models.ppseg.cvlibs import manager
from paddlers.models.ppseg.transforms import Compose
from paddlers.models.ppseg.datasets import Dataset
from paddlers.models.ppseg.utils.download import download_file_and_uncompress
from paddlers.models.ppseg.utils import seg_env
import paddlers.models.ppseg.transforms.functional as F
URL = "https://paddleseg.bj.bcebos.com/dataset/Supervisely_face.zip"
@manager.DATASETS.add_component
class SUPERVISELY(Dataset):
"""
Supervise.ly dataset `https://supervise.ly/`.
Args:
common_transforms (list): A list of common image transformations for two inputs of portrait net.
transforms1 (list): A list of image transformations for the first input of portrait net.
transforms2 (list): A list of image transformations for the second input of portrait net.
dataset_root (str, optional): The Supervise.ly dataset directory. Default: None.
mode (str, optional): A subset of the entire dataset. It should be one of ('train', 'val'). Default: 'train'.
edge (bool, optional): Whether to compute edge while training. Default: False
"""
NUM_CLASSES = 2
def __init__(self,
common_transforms,
transforms1,
transforms2,
dataset_root=None,
mode='train',
edge=False):
self.dataset_root = dataset_root
self.common_transforms = Compose(common_transforms)
self.transforms = self.common_transforms
if transforms1 is not None:
self.transforms1 = Compose(transforms1, to_rgb=False)
if transforms2 is not None:
self.transforms2 = Compose(transforms2, to_rgb=False)
mode = mode.lower()
self.ignore_index = 255
self.mode = mode
self.num_classes = self.NUM_CLASSES
self.input_width = 224
self.input_height = 224
if self.dataset_root is None:
self.dataset_root = download_file_and_uncompress(
url=URL,
savepath=seg_env.DATA_HOME,
extrapath=seg_env.DATA_HOME)
elif not os.path.exists(self.dataset_root):
self.dataset_root = os.path.normpath(self.dataset_root)
savepath, extraname = self.dataset_root.rsplit(
sep=os.path.sep, maxsplit=1)
self.dataset_root = download_file_and_uncompress(
url=URL,
savepath=savepath,
extrapath=savepath,
extraname=extraname)
if mode == 'train':
path = os.path.join(dataset_root, 'supervisely_face_train_easy.txt')
else:
path = os.path.join(dataset_root, 'supervisely_face_test_easy.txt')
with open(path, 'r') as f:
files = f.readlines()
files = ["/".join(file.split('/')[1:]) for file in files]
img_files = [os.path.join(dataset_root, file).strip() for file in files]
label_files = [
os.path.join(dataset_root, file.replace('/img/', '/ann/')).strip()
for file in files
]
self.file_list = [
[img_path, label_path]
for img_path, label_path in zip(img_files, label_files)
]
def __getitem__(self, item):
image_path, label_path = self.file_list[item]
im = cv2.imread(image_path)
label = cv2.imread(label_path, 0)
label[label > 0] = 1
if self.mode == "val":
common_im, label = self.common_transforms(im=im, label=label)
im = np.float32(common_im[::-1, :, :]) # RGB => BGR
im_aug = copy.deepcopy(im)
else:
common_im, label = self.common_transforms(im=im, label=label)
common_im = np.transpose(common_im, [1, 2, 0])
# add augmentation
im, _ = self.transforms1(common_im)
im_aug, _ = self.transforms2(common_im)
im = np.float32(im[::-1, :, :]) # RGB => BGR
im_aug = np.float32(im_aug[::-1, :, :]) # RGB => BGR
label = cv2.resize(
np.uint8(label), (self.input_width, self.input_height),
interpolation=cv2.INTER_NEAREST)
# add mask blur
label = np.uint8(cv2.blur(label, (5, 5)))
label[label >= 0.5] = 1
label[label < 0.5] = 0
edge_mask = F.mask_to_binary_edge(
label, radius=4, num_classes=self.num_classes)
edge_mask = np.transpose(edge_mask, [1, 2, 0]).squeeze(axis=-1)
im = np.concatenate([im_aug, im])
if self.mode == "train":
return im, label, edge_mask
else:
return im, label