|
- # Copyright 2021 Huawei Technologies Co., Ltd
- #
- # 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.
- # ============================================================================
- """train resnet."""
- import os
- import argparse
- import json
- import numpy as np
- from moco.loader import create_dataset_lincls
-
- parser = argparse.ArgumentParser('preprocess')
- parser.add_argument('--dataset_name', type=str, default="imagenet2012")
- parser.add_argument('--data_path', type=str, default='', help='eval data dir')\
- #parser.add_argument('--result_path', type=str, default='./preprocess_Result/', help='result path')
-
- args = parser.parse_args()
-
- def search_dir(save_path):
- '''
- search_dir
- '''
- if not os.path.exists(save_path):
- os.makedirs(save_path)
- bin_path = save_path+"images"
- if not os.path.exists(bin_path):
- os.makedirs(bin_path)
-
- def data_to_bin(val_dataset, save_path):
- '''
- data_to_bin
- '''
- i = 0
- ids = []
- for item in val_dataset.create_dict_iterator():
- inputs = item['image'].asnumpy()
- d = item['label'].asnumpy()
- if not os.path.exists(save_path+"images/"+str(d[0])):
- os.makedirs(save_path+"images/"+str(d[0]))
- inputs.tofile(save_path+"images/"+str(d[0]) + "/" + str(i)+".bin")
- #ids.append(d.astype(np.float32))
- i = i+1
- np.save(os.path.join(save_path, "label.npy"), np.array(ids, dtype=np.float32))
-
- def create_label(result_path, dir_path):
- """create label json for imagenet"""
- print("[WARNING] Create imagenet label. Currently only use for Imagenet2012!")
- dirs = os.listdir(dir_path)
- file_list = []
- for file in dirs:
- file_list.append(file)
- file_list = sorted(file_list)
-
- total = 0
- img_label = {}
- for i, file_dir in enumerate(file_list):
- files = os.listdir(os.path.join(dir_path, file_dir))
- for f in files:
- img_label[f] = i
- total += len(files)
-
- json_file = os.path.join(result_path, "imagenet_label.json")
- with open(json_file, "w+") as label:
- json.dump(img_label, label)
-
- print("[INFO] Completed! Total {} data.".format(total))
-
- def main():
- #create_label('./preprocess_Result/', args.data_path)
- valid_dataset = create_dataset_lincls(dataset_path=args.data_path, do_train=False,
- batch_size=1)
- save_path = "data_bin//"
- search_dir(save_path)
- data_to_bin(valid_dataset, save_path)
-
- if __name__ == '__main__':
- main()
|