leo/khqp_trainer
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

实现mode 3

Browse Source
This commit is contained in:
leo
2024-03-07 23:06:03 +08:00
parent da51178211
commit 09b90a09d7
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
main.py
View File

@@ -70,14 +70,14 @@ def prepare_train_data(mode_id) -> dict[str, list]:
conditions=(BatchData.is_validation == 1,)
)
elif mode_id == 3:
res['train_data'] = make_train_dicts(conditions=(
BatchData.is_train == 1,
BatchData.ann_file.is_not(None)
))
val_db_objs = session.query(BatchData).where(BatchData.is_validation == 1).all()
for db_obj in val_db_objs:
res['val_data'].append(DataRow(path=db_obj.ann_file,
img_prefix=db_obj.img_prefix).dict())
res['train_data'] = make_train_dicts(
with_entities=(BatchData.img_prefix, BatchData.ann_file.label('ann_file')),
conditions=(BatchData.is_train == 1, BatchData.ann_file.is_not(None))
)
res['val_data'] = make_train_dicts(
with_entities=(BatchData.img_prefix, BatchData.ann_file.label('ann_file')),
conditions=(BatchData.is_validation == 1, BatchData.ann_file.is_not(None))
)
return res