From a5a8f920594b4f175cac264e1f3f5417f8816d0f Mon Sep 17 00:00:00 2001
From: Arshadoid <50555286+Arshadoid@users.noreply.github.com>
Date: Wed, 29 Mar 2023 10:17:02 -0400
Subject: [PATCH] Update l2y.py (#16)

Fixed _train_test_split() for correct split of test and train.
---
 src/labelme2yolo/l2y.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/labelme2yolo/l2y.py b/src/labelme2yolo/l2y.py
index 7cc3f3e..7dfb17a 100644
--- a/src/labelme2yolo/l2y.py
+++ b/src/labelme2yolo/l2y.py
@@ -181,7 +181,7 @@ class Labelme2YOLO(object):
             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))
+                train_idxs, test_size=test_size / (1 - val_size))
         train_json_names = [json_names[train_idx] for train_idx in train_idxs]
         val_json_names = [json_names[val_idx] for val_idx in val_idxs]
         test_json_names = [json_names[test_idx] for test_idx in test_idxs]