From 543cead191ff0267716064550f6de302d29a871a Mon Sep 17 00:00:00 2001 From: zxsanny Date: Fri, 23 Aug 2024 13:47:54 +0300 Subject: [PATCH] refactoring, fixes --- train.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/train.py b/train.py index 534bd3a..9a30991 100644 --- a/train.py +++ b/train.py @@ -19,7 +19,7 @@ valid_set = 20 test_set = 10 -def form_dataset(set_date: datetime): +def form_dataset(from_date: datetime): makedirs(today_dataset, exist_ok=True) images = [] with scandir(processed_images_dir) as imd: @@ -27,9 +27,9 @@ def form_dataset(set_date: datetime): if not image_file.is_file(): continue mod_time = datetime.fromtimestamp(image_file.stat().st_mtime) - if set_date is None: + if from_date is None: images.append(image_file) - elif mod_time > set_date: + elif mod_time > from_date: images.append(image_file) print('shuffling images') @@ -142,9 +142,7 @@ if __name__ == '__main__': cur_folder = today_dataset yaml = abspath(path.join(cur_folder, 'data.yaml')) - results = model.train(data=yaml, epochs=100, batch=60, imgsz=640, save_period=1) + results = model.train(data=yaml, epochs=100, batch=60, imgsz=640) - shutil.copy(f'{results.save_dir}/weights/best.pt', latest_model) - shutil.copytree(results.save_dir, path.join(models_dir, cur_folder)) + shutil.copytree(results.save_dir, path.join(models_dir, today_folder)) shutil.rmtree('runs') - shutil.rmtree(path.join(models_dir, f'{prefix}latest'))