From 2ef6e629da40049626c1493880bcfb70aeace27f Mon Sep 17 00:00:00 2001 From: yaochaorui Date: Thu, 18 Feb 2021 18:13:59 +0800 Subject: [PATCH 1/3] Generate training_list.json. --- setup.cfg | 2 +- tools/data/matting/comp1k/extend_fg.py | 16 +++++++++------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/setup.cfg b/setup.cfg index 6aec562119..da38d8823f 100644 --- a/setup.cfg +++ b/setup.cfg @@ -17,6 +17,6 @@ line_length = 79 multi_line_output = 0 known_standard_library = pkg_resources,setuptools known_first_party = mmedit -known_third_party =PIL,cv2,lmdb,mmcv,numpy,onnx,onnxruntime,pymatting,pytest,scipy,titlecase,torch,torchvision,tqdm +known_third_party =PIL,cv2,lmdb,mmcv,numpy,onnx,onnxruntime,pymatting,pytest,scipy,titlecase,torch,torchvision no_lines_before = STDLIB,LOCALFOLDER default_section = THIRDPARTY diff --git a/tools/data/matting/comp1k/extend_fg.py b/tools/data/matting/comp1k/extend_fg.py index d3627904cb..8133bdd52e 100644 --- a/tools/data/matting/comp1k/extend_fg.py +++ b/tools/data/matting/comp1k/extend_fg.py @@ -5,10 +5,10 @@ import subprocess from multiprocessing import Pool +import mmcv import numpy as np from PIL import Image from pymatting import estimate_foreground_ml, load_image -from tqdm import tqdm def fix_png_file(filename, folder): @@ -70,6 +70,10 @@ def extend(self, fg_name): fg = Image.fromarray(np.uint8(F * 255)) fg.save(extended_path) fix_png_file(osp.basename(extended_path), osp.dirname(extended_path)) + data_info = dict() + data_info['alpha_path'] = alpha_path + data_info['fg_extended_path'] = extended_path + return data_info def parse_args(): @@ -112,16 +116,14 @@ def main(): os.makedirs(p, exist_ok=True) fg_names = osp.join(dir_prefix, f'{fname_prefix}_fg_names.txt') + save_json_path = f'{fname_prefix}_list.json' fg_names = open(osp.join(data_root, fg_names)).readlines() fg_iter = iter(fg_names) - num = len(fg_names) extend_fg = ExtendFg(data_root, fg_dirs, alpha_dirs) - with Pool(processes=args.nproc) as p: - with tqdm(total=num) as pbar: - for i, _ in tqdm( - enumerate(p.imap_unordered(extend_fg.extend, fg_iter))): - pbar.update() + data_infos = mmcv.track_parallel_progress(extend_fg.extend, list(fg_iter), + args.nproc) + mmcv.dump(data_infos, osp.join(data_root, save_json_path)) print('train done') From 96bd06db380da1751fe90eecaf8d933cbc19e3a6 Mon Sep 17 00:00:00 2001 From: yaochaorui Date: Thu, 18 Feb 2021 18:14:36 +0800 Subject: [PATCH 2/3] Format. --- tools/data/matting/comp1k/extend_fg.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tools/data/matting/comp1k/extend_fg.py b/tools/data/matting/comp1k/extend_fg.py index 8133bdd52e..a2b3a36445 100644 --- a/tools/data/matting/comp1k/extend_fg.py +++ b/tools/data/matting/comp1k/extend_fg.py @@ -3,7 +3,6 @@ import os.path as osp import re import subprocess -from multiprocessing import Pool import mmcv import numpy as np From 2e1824ef28056a937ad47e98fdc593026398ef11 Mon Sep 17 00:00:00 2001 From: yaochaorui Date: Thu, 18 Feb 2021 20:50:40 +0800 Subject: [PATCH 3/3] Tiny. --- tools/data/matting/comp1k/extend_fg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/data/matting/comp1k/extend_fg.py b/tools/data/matting/comp1k/extend_fg.py index a2b3a36445..938554fea9 100644 --- a/tools/data/matting/comp1k/extend_fg.py +++ b/tools/data/matting/comp1k/extend_fg.py @@ -115,7 +115,7 @@ def main(): os.makedirs(p, exist_ok=True) fg_names = osp.join(dir_prefix, f'{fname_prefix}_fg_names.txt') - save_json_path = f'{fname_prefix}_list.json' + save_json_path = f'{fname_prefix}_list_fba.json' fg_names = open(osp.join(data_root, fg_names)).readlines() fg_iter = iter(fg_names)