AndySAnker/DeepStruc

1import sys, argparse2import streamlit as st3from tools.module import Net4import torch, random, time5import numpy as np6import pytorch_lightning as pl7from tools.utils import get_data, format_predictions, plot_ls, get_model, save_predictions8 9def main(args):10    time_start = time.time()11    data, data_name, project_name = get_data(args)12    model_path, model_arch = get_model(args.model)13    14    Net(model_arch=model_arch)15    DeepStruc = Net.load_from_checkpoint(model_path,model_arch=model_arch)16    #start_time = time.time()17    xyz_pred, latent_space, kl, mu, sigma = DeepStruc(data, mode='prior', sigma_scale=args.sigma)18    #st.write("one prediction: " , time.time() - start_time)19    #start_time = time.time()20    #for i in range(1000):21    #    xyz_pred, latent_space, kl, mu, sigma = DeepStruc(data, mode='prior', sigma_scale=args.sigma)22    #st.write("thousand predictions: " , time.time() - start_time)23 24    samling_pairs = format_predictions(latent_space, data_name, mu, sigma, args.sigma)25    26    df, mk_dir, index_highlight = samling_pairs, project_name, args.index_plot27 28    these_cords = save_predictions(xyz_pred, samling_pairs, project_name, model_arch, args)29    30    return df, index_highlight, these_cords31