2_train.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. """train"""
  2. from datetime import datetime
  3. import os
  4. import streamlit as st
  5. from db import subject
  6. from db import train
  7. from components.remove_style import hide_footer
  8. import page_utils
  9. def _create_train(conn, subjects):
  10. with st.form("train_form"):
  11. st.write("创建训练")
  12. position = st.selectbox("训练部位", ['左手', '右手'])
  13. finger_model = st.selectbox("气动手手势", list(page_utils.fingermodel_trans.keys()))
  14. trial_num = st.number_input("训练次数", value=10, step=5)
  15. owner_name = st.selectbox("用户", subjects.name.to_list())
  16. virtual_feedback_rate = st.number_input("假反馈比例", min_value=0., max_value=1., value=0., step=0.2)
  17. model_path = page_utils.file_selector(os.path.join(f'./static/models/{owner_name}'))
  18. submitted = st.form_submit_button("开始训练")
  19. if submitted:
  20. start_time = datetime.strptime(datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "%Y-%m-%d %H:%M:%S")
  21. train_new = {"position": position,
  22. "finger_model": page_utils.fingermodel_trans[finger_model],
  23. "trial_num": int(trial_num),
  24. "start_time": start_time,
  25. "owner_name": owner_name,
  26. "virtual_feedback_rate": float(virtual_feedback_rate),
  27. "model_path": model_path.name if model_path is not None else None}
  28. train_id = train.create_train(conn, train_new)
  29. # run psychopy process
  30. os.system(f'python general_grasp_training.py --subj {owner_name} --train-id {train_id} --finger-model {finger_model}, --virtual-feedback-rate {virtual_feedback_rate} --model-path {model_path}')
  31. return owner_name
  32. def render():
  33. st.set_page_config(
  34. page_title="手势训练", page_icon=":chart_with_upwards_trend:"
  35. )
  36. hide_footer()
  37. st.markdown("# 手势训练")
  38. st.sidebar.success("训练")
  39. conn = st.connection("sql_app", type="sql")
  40. train.create_table(conn)
  41. subjects = subject.get_subjects(conn)
  42. sub_name = _create_train(conn, subjects)
  43. if sub_name:
  44. trains = train.get_trains(conn, sub_name)
  45. st.write("# 训练列表")
  46. st.dataframe(trains)
  47. render()