forked from huawei/mindspore2022
89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
# Copyright 2022 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.
|
|
# ==============================================================================
|
|
"""
|
|
This module contains common utility functions for minddataset tests.
|
|
"""
|
|
import os
|
|
import pytest
|
|
from mindspore.mindrecord import FileWriter
|
|
|
|
FILES_NUM = 4
|
|
CV_DIR_NAME = "../data/mindrecord/testImageNetData"
|
|
|
|
|
|
def get_data(dir_name):
|
|
"""
|
|
usage: get data from imagenet dataset
|
|
params:
|
|
dir_name: directory containing folder images and annotation information
|
|
|
|
"""
|
|
if not os.path.isdir(dir_name):
|
|
raise IOError("Directory {} does not exist".format(dir_name))
|
|
img_dir = os.path.join(dir_name, "images")
|
|
ann_file = os.path.join(dir_name, "annotation.txt")
|
|
with open(ann_file, "r") as file_reader:
|
|
lines = file_reader.readlines()
|
|
|
|
data_list = []
|
|
for i, line in enumerate(lines):
|
|
try:
|
|
filename, label = line.split(",")
|
|
label = label.strip("\n")
|
|
with open(os.path.join(img_dir, filename), "rb") as file_reader:
|
|
img = file_reader.read()
|
|
data_json = {"id": i,
|
|
"file_name": filename,
|
|
"data": img,
|
|
"label": int(label)}
|
|
data_list.append(data_json)
|
|
except FileNotFoundError:
|
|
continue
|
|
return data_list
|
|
|
|
|
|
@pytest.fixture
|
|
def add_and_remove_cv_file():
|
|
"""add/remove cv file"""
|
|
file_name = os.environ.get('PYTEST_CURRENT_TEST').split(':')[-1].split(' ')[0]
|
|
paths = ["{}{}".format(file_name, str(x).rjust(1, '0'))
|
|
for x in range(FILES_NUM)]
|
|
try:
|
|
for x in paths:
|
|
if os.path.exists("{}".format(x)):
|
|
os.remove("{}".format(x))
|
|
if os.path.exists("{}.db".format(x)):
|
|
os.remove("{}.db".format(x))
|
|
writer = FileWriter(file_name, FILES_NUM)
|
|
data = get_data(CV_DIR_NAME)
|
|
cv_schema_json = {"id": {"type": "int32"},
|
|
"file_name": {"type": "string"},
|
|
"label": {"type": "int32"},
|
|
"data": {"type": "bytes"}}
|
|
writer.add_schema(cv_schema_json, "img_schema")
|
|
writer.add_index(["file_name", "label"])
|
|
writer.write_raw_data(data)
|
|
writer.commit()
|
|
yield "yield_cv_data"
|
|
except Exception as error:
|
|
for x in paths:
|
|
os.remove("{}".format(x))
|
|
os.remove("{}.db".format(x))
|
|
raise error
|
|
else:
|
|
for x in paths:
|
|
os.remove("{}".format(x))
|
|
os.remove("{}.db".format(x))
|