kraken/backend/pages/2_train.py

"""train"""
from datetime import datetime
import os

import streamlit as st

from db.models import subject
from db.models import train
from components.remove_style import hide_footer


def _create_train(conn, subjects):
    with st.form("train_form"):
        st.write("创建训练")
        position = st.text_input("训练部位")
        trial_num = st.number_input("训练次数", value=1, step=1)
        owner_name = st.selectbox("用户", subjects.name.to_list())
        submitted = st.form_submit_button("开始训练")
        if submitted:
            start_time = datetime.strptime(datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "%Y-%m-%d %H:%M:%S")
            train_new = {"position": position, "trial_num": int(trial_num), "start_time": start_time, "owner_name": owner_name}
            train.create_train(conn, train_new)
            # TODO: set absolute path
            os.system("python train_1.py")
            return owner_name


def render():
    st.set_page_config(
        page_title="train", page_icon=":chart_with_upwards_trend:"
    )
    hide_footer()

    st.markdown("# Train")
    st.sidebar.success("训练")
    conn = st.connection("sql_app", type="sql")
    train.create_table(conn)
    subjects = subject.get_subjects(conn)
    sub_name = _create_train(conn, subjects)
    if sub_name:
        trains = train.get_trains(conn, sub_name)
        st.write("# 训练列表")
        st.dataframe(trains)


render()