Spaces:
Runtime error
Runtime error
| # 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 math | |
| import cv2 | |
| import numpy as np | |
| import random | |
| import paddle | |
| from paddleseg.cvlibs import manager | |
| import matting.transforms as T | |
| class MattingDataset(paddle.io.Dataset): | |
| """ | |
| Pass in a dataset that conforms to the format. | |
| matting_dataset/ | |
| |--bg/ | |
| | | |
| |--train/ | |
| | |--fg/ | |
| | |--alpha/ | |
| | | |
| |--val/ | |
| | |--fg/ | |
| | |--alpha/ | |
| | |--trimap/ (if existing) | |
| | | |
| |--train.txt | |
| | | |
| |--val.txt | |
| See README.md for more information of dataset. | |
| Args: | |
| dataset_root(str): The root path of dataset. | |
| transforms(list): Transforms for image. | |
| mode (str, optional): which part of dataset to use. it is one of ('train', 'val', 'trainval'). Default: 'train'. | |
| train_file (str|list, optional): File list is used to train. It should be `foreground_image.png background_image.png` | |
| or `foreground_image.png`. It shold be provided if mode equal to 'train'. Default: None. | |
| val_file (str|list, optional): File list is used to evaluation. It should be `foreground_image.png background_image.png` | |
| or `foreground_image.png` or ``foreground_image.png background_image.png trimap_image.png`. | |
| It shold be provided if mode equal to 'val'. Default: None. | |
| get_trimap (bool, optional): Whether to get triamp. Default: True. | |
| separator (str, optional): The separator of train_file or val_file. If file name contains ' ', '|' may be perfect. Default: ' '. | |
| """ | |
| def __init__(self, | |
| dataset_root, | |
| transforms, | |
| mode='train', | |
| train_file=None, | |
| val_file=None, | |
| get_trimap=True, | |
| separator=' '): | |
| super().__init__() | |
| self.dataset_root = dataset_root | |
| self.transforms = T.Compose(transforms) | |
| self.mode = mode | |
| self.get_trimap = get_trimap | |
| self.separator = separator | |
| # check file | |
| if mode == 'train' or mode == 'trainval': | |
| if train_file is None: | |
| raise ValueError( | |
| "When `mode` is 'train' or 'trainval', `train_file must be provided!" | |
| ) | |
| if isinstance(train_file, str): | |
| train_file = [train_file] | |
| file_list = train_file | |
| if mode == 'val' or mode == 'trainval': | |
| if val_file is None: | |
| raise ValueError( | |
| "When `mode` is 'val' or 'trainval', `val_file must be provided!" | |
| ) | |
| if isinstance(val_file, str): | |
| val_file = [val_file] | |
| file_list = val_file | |
| if mode == 'trainval': | |
| file_list = train_file + val_file | |
| # read file | |
| self.fg_bg_list = [] | |
| for file in file_list: | |
| file = os.path.join(dataset_root, file) | |
| with open(file, 'r') as f: | |
| lines = f.readlines() | |
| for line in lines: | |
| line = line.strip() | |
| self.fg_bg_list.append(line) | |
| def __getitem__(self, idx): | |
| data = {} | |
| fg_bg_file = self.fg_bg_list[idx] | |
| fg_bg_file = fg_bg_file.split(self.separator) | |
| data['img_name'] = fg_bg_file[0] # using in save prediction results | |
| fg_file = os.path.join(self.dataset_root, fg_bg_file[0]) | |
| alpha_file = fg_file.replace('/fg', '/alpha') | |
| fg = cv2.imread(fg_file) | |
| alpha = cv2.imread(alpha_file, 0) | |
| data['alpha'] = alpha | |
| data['gt_fields'] = [] | |
| # line is: fg [bg] [trimap] | |
| if len(fg_bg_file) >= 2: | |
| bg_file = os.path.join(self.dataset_root, fg_bg_file[1]) | |
| bg = cv2.imread(bg_file) | |
| data['img'], data['bg'] = self.composite(fg, alpha, bg) | |
| data['fg'] = fg | |
| if self.mode in ['train', 'trainval']: | |
| data['gt_fields'].append('fg') | |
| data['gt_fields'].append('bg') | |
| data['gt_fields'].append('alpha') | |
| if len(fg_bg_file) == 3 and self.get_trimap: | |
| if self.mode == 'val': | |
| trimap_path = os.path.join(self.dataset_root, fg_bg_file[2]) | |
| if os.path.exists(trimap_path): | |
| data['trimap'] = trimap_path | |
| data['gt_fields'].append('trimap') | |
| data['ori_trimap'] = cv2.imread(trimap_path, 0) | |
| else: | |
| raise FileNotFoundError( | |
| 'trimap is not Found: {}'.format(fg_bg_file[2])) | |
| else: | |
| data['img'] = fg | |
| if self.mode in ['train', 'trainval']: | |
| data['fg'] = fg.copy() | |
| data['bg'] = fg.copy() | |
| data['gt_fields'].append('fg') | |
| data['gt_fields'].append('bg') | |
| data['gt_fields'].append('alpha') | |
| data['trans_info'] = [] # Record shape change information | |
| # Generate trimap from alpha if no trimap file provided | |
| if self.get_trimap: | |
| if 'trimap' not in data: | |
| data['trimap'] = self.gen_trimap( | |
| data['alpha'], mode=self.mode).astype('float32') | |
| data['gt_fields'].append('trimap') | |
| if self.mode == 'val': | |
| data['ori_trimap'] = data['trimap'].copy() | |
| data = self.transforms(data) | |
| # When evaluation, gt should not be transforms. | |
| if self.mode == 'val': | |
| data['gt_fields'].append('alpha') | |
| data['img'] = data['img'].astype('float32') | |
| for key in data.get('gt_fields', []): | |
| data[key] = data[key].astype('float32') | |
| if 'trimap' in data: | |
| data['trimap'] = data['trimap'][np.newaxis, :, :] | |
| if 'ori_trimap' in data: | |
| data['ori_trimap'] = data['ori_trimap'][np.newaxis, :, :] | |
| data['alpha'] = data['alpha'][np.newaxis, :, :] / 255. | |
| return data | |
| def __len__(self): | |
| return len(self.fg_bg_list) | |
| def composite(self, fg, alpha, ori_bg): | |
| fg_h, fg_w = fg.shape[:2] | |
| ori_bg_h, ori_bg_w = ori_bg.shape[:2] | |
| wratio = fg_w / ori_bg_w | |
| hratio = fg_h / ori_bg_h | |
| ratio = wratio if wratio > hratio else hratio | |
| # Resize ori_bg if it is smaller than fg. | |
| if ratio > 1: | |
| resize_h = math.ceil(ori_bg_h * ratio) | |
| resize_w = math.ceil(ori_bg_w * ratio) | |
| bg = cv2.resize( | |
| ori_bg, (resize_w, resize_h), interpolation=cv2.INTER_LINEAR) | |
| else: | |
| bg = ori_bg | |
| bg = bg[0:fg_h, 0:fg_w, :] | |
| alpha = alpha / 255 | |
| alpha = np.expand_dims(alpha, axis=2) | |
| image = alpha * fg + (1 - alpha) * bg | |
| image = image.astype(np.uint8) | |
| return image, bg | |
| def gen_trimap(alpha, mode='train', eval_kernel=7): | |
| if mode == 'train': | |
| k_size = random.choice(range(2, 5)) | |
| iterations = np.random.randint(5, 15) | |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, | |
| (k_size, k_size)) | |
| dilated = cv2.dilate(alpha, kernel, iterations=iterations) | |
| eroded = cv2.erode(alpha, kernel, iterations=iterations) | |
| trimap = np.zeros(alpha.shape) | |
| trimap.fill(128) | |
| trimap[eroded > 254.5] = 255 | |
| trimap[dilated < 0.5] = 0 | |
| else: | |
| k_size = eval_kernel | |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, | |
| (k_size, k_size)) | |
| dilated = cv2.dilate(alpha, kernel) | |
| trimap = np.zeros(alpha.shape) | |
| trimap.fill(128) | |
| trimap[alpha >= 250] = 255 | |
| trimap[dilated <= 5] = 0 | |
| return trimap | |