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
__version__ = '0.0.6'
__version__ = '0.0.7'

View File

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