Spaces:
Build error
Build error
| # coding:utf-8 | |
| import os | |
| import argparse | |
| import time | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import DataLoader | |
| from sklearn.decomposition import IncrementalPCA | |
| from MaskLoader import MaskLoader | |
| from utils import inverse_sigmoid | |
| VALUE_MAX = 0.05 | |
| VALUE_MIN = 0.01 | |
| def mask_encoding(masks, n_components=60, class_agnostic=True, whiten=True, sigmoid=True, batch_size=1024): | |
| components_c = [] | |
| mean_c = [] | |
| ratio_c = [] | |
| explained_variance_c = [] | |
| if class_agnostic: | |
| if sigmoid: | |
| value_random = VALUE_MAX * np.random.rand(masks.shape[0], masks.shape[1]) | |
| value_random = np.maximum(value_random, VALUE_MIN) | |
| masks = np.where(masks > value_random, 1-value_random, value_random) | |
| masks = inverse_sigmoid(masks) | |
| pca = IncrementalPCA(n_components=n_components, copy=False, whiten=whiten, batch_size=batch_size) | |
| pca.fit(masks) | |
| components_c.append(pca.components_[np.newaxis, :, :]) | |
| mean_c.append(pca.mean_[np.newaxis, :]) | |
| ratio_c.append(pca.explained_variance_ratio_[np.newaxis, :]) | |
| explained_variance_c.append(pca.explained_variance_[np.newaxis, :]) | |
| ratio = pca.explained_variance_ratio_.sum() | |
| else: | |
| # TODO: We have not achieve the function in class-specific. | |
| raise NotImplemented | |
| return components_c, mean_c, ratio_c, explained_variance_c, ratio | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description='PCA Mask Encoding for local mask.') | |
| parser.add_argument('--root', default='datasets', type=str) | |
| parser.add_argument('--dataset', default='coco_2017_train', type=str) | |
| parser.add_argument('--output', default='./projects/LME', type=str) | |
| # mask encoding params. | |
| parser.add_argument('--mask_size', default=28, type=int) | |
| parser.add_argument('--n_components', default=128, type=int) | |
| parser.add_argument('--class_agnostic', default=True, type=bool) | |
| parser.add_argument('--whiten', default=True, type=bool) | |
| parser.add_argument('--sigmoid', default=True, type=bool) | |
| parser.add_argument('--batch-size', default=1024, type=int) | |
| args = parser.parse_args() | |
| return args | |
| if __name__ == "__main__": | |
| args = parse_args() | |
| # parse args. | |
| mask_size = args.mask_size | |
| n_components = args.n_components | |
| class_agnostic = args.class_agnostic | |
| whiten = args.whiten | |
| sigmoid = args.sigmoid | |
| dataset_root = './datasets/' | |
| output_dir = args.output | |
| os.makedirs(output_dir, exist_ok=True) | |
| # build data loader. | |
| mask_data = MaskLoader(root=dataset_root, dataset=args.dataset, size=mask_size) | |
| mask_loader = DataLoader(mask_data, batch_size=args.batch_size, shuffle=False, num_workers=4) | |
| # loading masks. | |
| masks = list() | |
| print("Start Loading Masks.") | |
| tic = time.time() | |
| for mask in mask_loader: | |
| masks.append(mask.squeeze(1)) | |
| toc = time.time() - tic | |
| print("Finish Loading Masks in {}s.".format(toc)) | |
| masks = torch.cat(masks, 0) | |
| masks = masks.view(masks.shape[0], -1).numpy() | |
| masks = masks.astype(np.float32) | |
| # mask encoding. | |
| print("Start to mask encoding ...") | |
| print("It may take several times, please wait patiently ...") | |
| tic = time.time() | |
| components_c, mean_c, ratio_c, explained_variance_c, ratio = \ | |
| mask_encoding(masks, n_components, class_agnostic, whiten, sigmoid, args.batch_size) | |
| toc = time.time() - tic | |
| print("Finish the mask encoding in {}s.".format(toc)) | |
| components_c = np.concatenate(components_c).mean(0)[np.newaxis, :, :].astype(np.float32) | |
| mean_c = np.concatenate(mean_c).mean(0)[np.newaxis, :].astype(np.float32) | |
| ratio_c = np.concatenate(ratio_c).mean(0)[np.newaxis, :].astype(np.float32) | |
| explained_variance_c = np.concatenate(explained_variance_c).mean(0)[np.newaxis, :].astype(np.float32) | |
| print("The mean variance_ratio for all categories is {}".format(np.mean(ratio))) | |
| # save the parameters. | |
| output_path = os.path.join(output_dir, args.dataset + '_class_agnostic' + str(class_agnostic) | |
| + '_whiten' + str(whiten) + '_sigmoid' + str(sigmoid) + '_' + str(n_components) + '_siz' + str(mask_size) | |
| + '.npz') | |
| print("Save the local mask encoding matrix: " + output_path) | |
| np.savez(output_path, | |
| components_c=components_c, | |
| mean_c=mean_c, | |
| ratio_c=ratio_c, | |
| explained_variance_c=explained_variance_c) | |