jonasmaltebecker/vae_drilling

1import numpy as np2import streamlit as st3import torch4import matplotlib.pyplot as plt5 6import disvae7import transforms as trans8 9P_MODEL = "model/drilling_ds_btcvae"10SAMPLING_TIME = 0.1511 12st.set_page_config(page_title="Drilling VAE")13 14@st.cache_resource15def load_decode_function():16    17    sorter = trans.LatentSorter(disvae.get_kl_dict(P_MODEL))18    vae = disvae.load_model(P_MODEL)19    scaler = trans.MinMaxScaler(_min=torch.tensor([1.3]),_max=torch.tensor([4.0]),min_norm=0.3,max_norm=0.6)20    imaging = trans.SumField()21 22    _dec = trans.sequential_function(23        sorter.inv,24        vae.decoder,25        scaler.inv,26        imaging.inv27    )28 29    def decode(latent):30        with torch.no_grad():31            return trans.np_sample(_dec)(latent)32    33    return decode34    35decode = load_decode_function()36 37col1,col2 = st.columns(2)38 39with col1:40    st.markdown("**Latent Space Parameters**")41    latent_vector = np.array([st.slider(f"Latent Dimension {l}",min_value=-3.0,max_value=3.0,value=0.0) for l in range(3)])42    latent_vector = np.concatenate([latent_vector,np.zeros(7)],axis=0)43 44    ts = decode(latent_vector)45 46with col2:47    st.markdown("**Generated Time Series**")48 49    fig, ax = plt.subplots(figsize=(4,3))50 51    time = np.arange(0,len(ts)*SAMPLING_TIME,SAMPLING_TIME)52    ax.plot(time,ts.ravel())53    ax.set_xlabel("Time t [s]")54    ax.set_ylabel("Spindle torque t [Nm]")55    ax.set_ylim([0,4])56    ax.grid()57 58    st.pyplot(fig)