neuralcomputation/batik

1import streamlit as st2import pandas as pd3from utils.utils import create_embeddings_csv_io, create_annot_fname_dict_io, generate_embeddings_stream_io4 5if "video_embeddings" not in st.session_state:6    st.session_state.video_embeddings = None7    st.session_state.video_frames = None8    st.session_state.fnames = []9 10st.title('batik: embedding generator')11uploaded_files = st.file_uploader("Choose a video file", type=['seq', 'mp4'], accept_multiple_files=True)12with st.form('initial_settings'):13    st.header('Embedding Generation Options')14    model_select = st.selectbox('Select Model', ['SLIP', 'CLIP'])15    downsample_rate = st.number_input('Downsample Rate',value=4)16    #save_csv = st.toggle('Save Individual Results', value=False)17    submit_initial_settings = st.form_submit_button('Create Embeddings', type='secondary')18 19if submit_initial_settings and uploaded_files is not None and len(uploaded_files) > 0:20    video_embeddings, video_frames = generate_embeddings_stream_io(uploaded_files,21                                                                model_select,22                                                                downsample_rate,23                                                                False)24    fnames = [vid_file.name for vid_file in uploaded_files]25    st.session_state.video_embeddings = video_embeddings26    st.session_state.video_frames = video_frames27    st.session_state.fnames = fnames28 29if st.session_state.video_embeddings is not None:30    st.header('CSV Configuration Options')31    st.markdown('If using `.annot` files and multiple files should be grouped together, '\32                'please ensure that they share a common name and end with a number describing '\33                'the order of the files. For example:\n\n'\34                '`mouse_224_file_1.annot`, `mouse_224_file_2.annot`.')35    annot_files = st.file_uploader("Upload all annotation files", type=['.annot','.csv'], accept_multiple_files=True)36 37    annot_options = []38    if annot_files is not None and len(annot_files) > 0:39        annot_fnames = [annot_file.name for annot_file in annot_files]40        annot_fname_dict = create_annot_fname_dict_io(annot_fnames=annot_fnames,41                                                      annot_files=annot_files)42        annot_options = [str(key) for key in annot_fname_dict.keys()]43 44    if len(annot_options) > 0:45        with st.form('csv_settings'):46            csv_setting_def = pd.DataFrame(47                {48                    "File Name" : st.session_state.fnames,49                    "Annotations" : [50                        "Upload File" for _ in st.session_state.fnames51                    ],52                    "Test" : [53                        False for _ in st.session_state.fnames54                    ],55                    "View" : [56                        "Top" for _ in st.session_state.fnames57                    ],58                    "Condition" : [59                        "None" for _ in st.session_state.fnames60                    ]61 62                }63            )64 65            csv_settings = st.data_editor(66                csv_setting_def,67                column_config={68                    "Annotations" : st.column_config.SelectboxColumn(69                        "Annotations",70                        help="The annotation file(s) to use for the given video file.",71                        width="medium",72                        options=annot_options,73                        required=True74                    ),75                    "Test" : st.column_config.CheckboxColumn(76                        "Test",77                        help="Designate file(s) to use as the test set.",78                        default=False,79                        required=True80                    ),81                    "View" : st.column_config.SelectboxColumn(82                        "View",83                        help="The view used within the video (either Top or Front).",84                        options=["Top", "Front"],85                        required=True86                    ),87                    "Condition" : st.column_config.TextColumn(88                        "Condition",89                        help="A condition the video has (i.e. Control).",90                        default="None",91                        max_chars=30,92                        validate=r"[a-z]+$",93                    )94                },95                hide_index=True96            )97            save_csv_bttn = st.form_submit_button("Create CSV")98 99        if save_csv_bttn and csv_settings is not None:100            annot_chosen_options = csv_settings['Annotations'].tolist()101            annot_option = [annot_fname_dict[key] for key in annot_chosen_options]102            test_chosen_option = csv_settings['Test'].tolist()103            test_option = [st.session_state.fnames[i] for i, is_test in enumerate(test_chosen_option) if is_test]104            view_option = csv_settings['View'].tolist()105            condition_option =  csv_settings['Condition'].tolist()106 107            out_name = st.text_input("Embeddings Outpit File Name", "out.csv")108            try:109                df = create_embeddings_csv_io(out=out_name,110                                            fnames=st.session_state.fnames,111                                            embeddings=st.session_state.video_embeddings,112                                            frames=st.session_state.video_frames,113                                            annotations=annot_option,114                                            test_fnames=test_option,115                                            views=view_option,116                                            conditions=condition_option,117                                            downsample_rate=downsample_rate)118                st.success('Created Embeddings File!', icon="✅")119                st.download_button(120                    label="Download CSV",121                    data=df.to_csv().encode("utf-8"),122                    file_name=out_name,123                    mime="text/csv"124                )125            except:126                st.error('Something went wrong.')127 128    else:129        st.text('Please Upload Files')130else:131    st.text('Please Upload Files')132