File size: 2,236 Bytes
2cd560a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: [email protected]
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
import argparse
import os
import os.path as osp
import pickle
from tqdm import tqdm
import numpy as np
def clean_fn(fn, output_folder='output'):
with open(fn, 'rb') as body_file:
body_data = pickle.load(body_file)
output_dict = {}
for key, data in body_data.iteritems():
if 'chumpy' in str(type(data)):
output_dict[key] = np.array(data)
else:
output_dict[key] = data
out_fn = osp.split(fn)[1]
out_path = osp.join(output_folder, out_fn)
with open(out_path, 'wb') as out_file:
pickle.dump(output_dict, out_file)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input-models', dest='input_models', nargs='+',
required=True, type=str,
help='The path to the model that will be processed')
parser.add_argument('--output-folder', dest='output_folder',
required=True, type=str,
help='The path to the output folder')
args = parser.parse_args()
input_models = args.input_models
output_folder = args.output_folder
if not osp.exists(output_folder):
print('Creating directory: {}'.format(output_folder))
os.makedirs(output_folder)
for input_model in input_models:
clean_fn(input_model, output_folder=output_folder)
|