|
- # 单独运行生成 mindrecord 文件
- """create mindrecord for training EfficientDet."""
-
-
- from mindspore.mindrecord import FileWriter
- import argparse
- from src.dataset import create_mindrecord
- import moxing as mox
- import os
- from src.dataset import create_coco_label
- import numpy as np
-
-
- if __name__ == "__main__":
- parser = argparse.ArgumentParser(description="EfficientDet dataset create")
- parser.add_argument("--dataset", type=str, default="coco", help="Dataset, default is coco.")
- parser.add_argument("--data_url", type=str, default=None, help="mindrecord dir")
- parser.add_argument("--train_url", type=str, default=None, help="ckpt output dir in obs")
-
- args_opt = parser.parse_args()
-
- device_id = int(os.getenv('DEVICE_ID'), 0)
- local_data_url = "/cache/data/" + str(device_id)
- mox.file.make_dirs(local_data_url)
-
- mindrecord_dir = "/cache/mr/" + str(device_id)
- mox.file.make_dirs(mindrecord_dir)
-
- mox.file.copy_parallel(args_opt.data_url, local_data_url)
-
- prefix = "EfficientDet.mindrecord"
-
- # ----------------------------------------------
-
-
- mindrecord_path = os.path.join(mindrecord_dir, prefix)
- writer = FileWriter(mindrecord_path, 8)
-
- images, image_path_dict, image_anno_dict = create_coco_label(True, local_data_url) # 调用创建标签代码(下一个)
-
- EfficientDet_json = {
- "image": {"type": "bytes"},
- "annotation": {"type": "float32", "shape": [-1, 5]},
- }
- writer.add_schema(EfficientDet_json, "EfficientDet_json")
-
- for img_id in images: # 对于每一张图片的操作
-
- image_path = image_path_dict[img_id] # 取出这一张图片的完整的路径
- with open(image_path, 'rb') as f:
- img = f.read() # 根据图片的id,读出图片,保存在img中
- annos = np.array(image_anno_dict[img_id], dtype=np.int32) # 这一张图片所有的anno
- img_id = np.array([img_id], dtype=np.int32)
- row = {"img_id": img_id, "image": img, "annotation": annos}
- writer.write_raw_data([row])
-
- writer.commit()
-
- # ----------------------------------------------
- print("Create Mindrecord Done, at {}".format(mindrecord_dir))
-
- mox.file.copy_parallel(mindrecord_dir, args_opt.train_url)
-
- print("transfer OK. ")
-
-
-
|