Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!16196 If Imagefolder input directory is not right, will occur out of memory error 

From: @xiefangqi
Reviewed-by: @ms_yan,@heleiwang,@jonyguo
Signed-off-by: @jonyguo
This commit is contained in:
mindspore-ci-bot 2021-05-11 17:36:28 +08:00 committed by Gitee
parent 909b0191c3 e0148f7181
commit ebed057e72
2 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mindspore/ccsrc/minddata/dataset/core/tensor.cc
View File

@ -260,6 +260,10 @@ Status Tensor::CreateFromByteList(const dataengine::BytesList &bytes_list, const
#endif #endif
Status Tensor::CreateFromFile(const std::string &path, std::shared_ptr<Tensor> *out) { Status Tensor::CreateFromFile(const std::string &path, std::shared_ptr<Tensor> *out) {
Path file(path);
if (file.IsDirectory()) {
RETURN_STATUS_UNEXPECTED("Invalid file found: " + path + ", should be file, but got directory.");
}
std::ifstream fs; std::ifstream fs;
fs.open(path, std::ios::binary | std::ios::in); fs.open(path, std::ios::binary | std::ios::in);
CHECK_FAIL_RETURN_UNEXPECTED(!fs.fail(), "Fail to open file: " + path); CHECK_FAIL_RETURN_UNEXPECTED(!fs.fail(), "Fail to open file: " + path);

9
tests/ut/python/dataset/test_datasets_imagefolder.py
View File

@ -759,6 +759,15 @@ def test_imagefolder_exception():
except RuntimeError as e: except RuntimeError as e:
assert "map operation: [PyFunc] failed. The corresponding data files" in str(e) assert "map operation: [PyFunc] failed. The corresponding data files" in str(e)
data_dir_invalid = "../data/dataset/testPK"
try:
data = ds.ImageFolderDataset(data_dir_invalid)
for _ in data.__iter__():
pass
assert False
except RuntimeError as e:
assert "should be file, but got directory" in str(e)
if __name__ == '__main__': if __name__ == '__main__':
test_imagefolder_basic() test_imagefolder_basic()