fix bug when test_size not set

This commit is contained in:
greatx 2022-12-06 14:35:15 +08:00
parent 2794fbe31c
commit b59e6d3d49
2 changed files with 3 additions and 1 deletions

View File

@ -2,4 +2,4 @@
# #
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
__version__ = '0.0.6' __version__ = '0.0.7'

View File

@ -179,6 +179,8 @@ class Labelme2YOLO(object):
test_size=val_size) test_size=val_size)
tmp_train_len = len(train_idxs) tmp_train_len = len(train_idxs)
test_idxs = [] test_idxs = []
if test_size is None:
test_size = 0.0
if test_size > 1e-8: if test_size > 1e-8:
train_idxs, test_idxs = train_test_split( train_idxs, test_idxs = train_test_split(
range(tmp_train_len), test_size=test_size / (1 - val_size)) range(tmp_train_len), test_size=test_size / (1 - val_size))