azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-23 03:26:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

small fix

Browse Source 
add code for revert dataset
add get_recent_model
This commit is contained in:
Oleksandr Bezdieniezhnykh
2024-06-06 01:46:48 +03:00
parent 07ea67746a
commit 1f7b4a95d2
1 changed files with 28 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+28 -19
View File
@@ -36,8 +36,11 @@ def move_annotations(images, folder):
image_path = os.path.join(current_images_dir, image_name)
label_name = f'{Path(image_name).stem}.txt'
label_path = os.path.join(current_labels_dir, label_name)
os.replace(image_path, os.path.join(destination_images, image_name))
os.replace(label_path, os.path.join(destination_labels, label_name))
if not os.path.exists(label_path):
os.remove(image_path)
else:
os.replace(image_path, os.path.join(destination_images, image_name))
os.replace(label_path, os.path.join(destination_labels, label_name))
def create_yaml():
@@ -55,18 +58,17 @@ def create_yaml():
def get_recent_model():
date_sets = []
datasets = [next((file for file in os.listdir(os.path.join('datasets', d)) if file.endswith('pt')), None)
for d in os.listdir('datasets')]
# date_str = d.replace(prefix, '')
# if date_str == 'current' or date_str == f'{datetime.now():%Y-%m-%d}':
# continue
# if len(date_sets) == 0:
# return None
recent = max(date_sets)
return os.path.join('datasets', f'{prefix}{recent}', f'{prefix}{recent}.pt')
cur_model = None
cur_date = None
for d in os.listdir('datasets'):
date_str = d.replace(prefix, '')
if date_str == 'current' or date_str == f'{datetime.now():%Y-%m-%d}':
continue
date = datetime.strptime(date_str, '%Y-%m-%d')
for file in os.listdir(os.path.join('datasets', d)):
if file.endswith('pt') and cur_date is None or cur_date < date:
cur_model = os.path.join('datasets', d, file)
return cur_model
def retrain():
@@ -75,12 +77,19 @@ def retrain():
def revert_to_current(date):
def revert_dir(dir):
os.listdir(os.path.join(current_images_dir, 'images'))
def revert_dir(src_dir, dest_dir):
for file in os.listdir(src_dir):
s = os.path.join(src_dir, file)
d = os.path.join(dest_dir, file)
os.replace(s, d)
date_dataset = os.path.join('datasets', f'{prefix}{date}')
current_dataset = os.path.join('datasets', f'{prefix}current')
for subset in ['test', 'train', 'valid']:
revert_dir(os.path.join(date_dataset, subset, 'images'), os.path.join(current_dataset, 'images'))
revert_dir(os.path.join(date_dataset, subset, 'labels'), os.path.join(current_dataset, 'labels'))
shutil.rmtree(date_dataset)
date_dataset = f'{prefix}{date}'
revert_dir(os.path.join(date_dataset, 'test'))
form_dataset()
create_yaml()
# revert_to_current('2024-06-06')
retrain()