add main program files

app.py

@@ -0,0 +1,21 @@
+import os
+import streamlit as st
+
+pages = {
+    "home": [
+        st.Page("home.py",title="about",icon=":material/home:")
+    ],
+    "generate embeddings": [
+        st.Page("generate_embeddings.py", title="generate",icon=":material/dataset:")
+    ],
+    "annotation": [
+        st.Page("train_model.py", title="train model",icon=":material/model_training:"),
+        st.Page("apply_model.py", title="apply model",icon=":material/grade:"),
+    ],
+    "behavior discovery": [
+        st.Page("explore.py", title="explore",icon=":material/search:")
+    ],
+}
+
+pg = st.navigation(pages)
+pg.run()

apply_model.py

@@ -0,0 +1,195 @@
+import os
+import pickle
+from random import random
+import streamlit as st
+import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap
+import numpy as np
+import pandas as pd
+import torch
+from utils.mp4Io import mp4Io_reader
+from utils.seqIo import seqIo_reader
+import pandas as pd
+from PIL import Image
+from pathlib import Path
+from transformers import AutoProcessor, AutoModel
+from tempfile import NamedTemporaryFile
+from tqdm import tqdm
+from sklearn.metrics import accuracy_score, classification_report
+from utils.utils import create_embeddings_csv_io, process_dataset_in_mem, multiclass_merge_and_filter_bouts, generate_embeddings_stream_io
+
+# --server.maxUploadSize 3000
+
+def get_io_reader(uploaded_file):
+    if uploaded_file.name[-3:]=='seq':
+        with NamedTemporaryFile(suffix="seq", delete=False) as temp:
+            temp.write(uploaded_file.getvalue())
+            sr = seqIo_reader(temp.name)
+    else:
+        with NamedTemporaryFile(suffix="mp4", delete=False) as temp:
+            temp.write(uploaded_file.getvalue())
+            sr = mp4Io_reader(temp.name)
+    return sr
+
+def get_unique_labels(label_list: list[str]):
+    label_set = set()
+    for label in label_list:
+        individual_labels = label.split('||')
+        for individual_label in individual_labels:
+            label_set.add(individual_label)
+    return list(label_set)
+
+def get_smoothed_predictions(svm_model, test_embeds):
+    test_pred = svm_model.predict(test_embeds)
+    test_prob = svm_model.predict_proba(test_embeds)
+
+    bout_threshold = 5
+    proximity_threshold = 2
+
+    predictions = multiclass_merge_and_filter_bouts(test_pred, bout_threshold, proximity_threshold)
+    return predictions
+
+if "embeddings_df" not in st.session_state:
+    st.session_state.embeddings_df = None
+
+if "smoothed_predictions" not in st.session_state:
+    st.session_state.smoothed_predictions = None
+    st.session_state.test_labels = []
+
+st.title('batik: frame classifier')
+
+st.text("Upload files to apply trained classifier on.")
+with st.form('embedding_generation_settings'):
+    seq_file = st.file_uploader("Choose a video file", type=['seq', 'mp4'], accept_multiple_files=False)
+    annot_files = st.file_uploader("Choose an annotation File", type=['annot','csv'], accept_multiple_files=True)
+    downsample_rate = st.number_input('Downsample Rate',value=4)
+    submit_embed_settings = st.form_submit_button('Create Embeddings', type='secondary')
+
+st.markdown("**(Optional)** Upload embeddings if not generating above.")
+embeddings_csv = st.file_uploader("Choose a .csv File", type=['csv'])
+
+if submit_embed_settings and seq_file is not None and annot_files is not None:
+    video_embeddings, video_frames = generate_embeddings_stream_io([seq_file],
+                                                                "SLIP",
+                                                                downsample_rate,
+                                                                False)
+    
+    fnames = [seq_file.name]
+    embeddings_df = create_embeddings_csv_io(out="file",
+                                fnames=fnames,
+                                embeddings=video_embeddings,
+                                frames=video_frames,
+                                annotations=[annot_files],
+                                test_fnames=None,
+                                views=None,
+                                conditions=None,
+                                downsample_rate=downsample_rate)
+    st.session_state.embeddings_df = embeddings_df
+
+elif embeddings_csv is not None:
+    embeddings_df = pd.read_csv(embeddings_csv)
+    st.session_state.embeddings_df = embeddings_df
+else:
+    st.text('Please upload file(s).')
+
+st.divider()
+st.markdown("Upload classifier model.")
+pickled_file = st.file_uploader("Choose a .pkl File", type=['pkl'])
+
+if pickled_file is not None:
+    with NamedTemporaryFile(suffix='pkl', delete=False) as temp:
+        temp.write(pickled_file.getvalue())
+        with open(temp.name, 'rb') as pickled_model:
+            svm_clf = pickle.load(pickled_model)
+else:
+    svm_clf = None
+
+st.divider()
+if st.session_state.embeddings_df is not None and svm_clf is not None:
+    st.subheader("specify dataset labels")
+    label_list = st.session_state.embeddings_df['Label'].to_list()
+    unique_label_list = get_unique_labels(label_list)
+
+    with st.form('apply_model_settings'):
+        st.text("Select label(s):")
+        specified_classes = st.multiselect("Label(s) included:", options=unique_label_list)
+
+
+        apply_model = st.form_submit_button("Apply Model")
+
+    if apply_model:
+        kwargs = {'embeddings_df' : st.session_state.embeddings_df, 
+                'specified_classes' : specified_classes,
+                'classes_to_remove' : None,
+                'max_class_size' : None,
+                'animal_state' : None,
+                'view' : None,
+                'shuffle_data' : False,
+                'test_videos' : list(set(st.session_state.embeddings_df['Source'].to_list()))}
+        train_embeds, train_labels, train_images, test_embeds, test_labels, test_images =\
+            process_dataset_in_mem(**kwargs)
+        
+        # get predictions from embeddings
+        with st.spinner("Model application in progress..."):
+            smoothed_predictions = get_smoothed_predictions(svm_clf, test_embeds)
+
+        # save variables to state
+        st.session_state.smoothed_predictions = smoothed_predictions
+        st.session_state.test_labels = test_labels
+
+    if st.session_state.smoothed_predictions is not None:
+        # Convert labels to numerical values
+        label_to_appear_first = 'other'
+        unique_labels = set(st.session_state.test_labels)
+        unique_labels.discard(label_to_appear_first)
+
+        label_to_index = {label_to_appear_first: 0}
+
+        label_to_index.update({label: idx + 1 for idx, label in enumerate(unique_labels)})
+        index_to_label = {idx: label for label, idx in label_to_index.items()}
+
+        numerical_labels_test = np.array([label_to_index[label] for label in st.session_state.test_labels])
+        print("Label Valence: ", label_to_index)
+
+        #smoothed_predictions test labels
+        if len(st.session_state.smoothed_predictions) > 0:
+            test_accuracy = accuracy_score(numerical_labels_test, st.session_state.smoothed_predictions)
+        else:
+            test_accuracy = 0  # If no predictions meet the threshold, set accuracy to 0
+
+        # test_accuracy = accuracy_score(numerical_labels_test, test_pred)
+        report = classification_report(numerical_labels_test,
+                                       st.session_state.smoothed_predictions,
+                                       target_names=[index_to_label[idx] for idx in range(len(index_to_label))],
+                                       output_dict=True)
+        report_df = pd.DataFrame(report).transpose()
+
+        st.text(f"Eval Accuracy: {test_accuracy}")
+        st.subheader("Classification Report:")
+        st.dataframe(report_df)
+
+        # create figure (behavior raster)
+        fig, ax = plt.subplots()
+        raster = ax.imshow(st.session_state.smoothed_predictions.reshape((1,st.session_state.smoothed_predictions.size)),
+                  aspect='auto',
+                  interpolation='nearest',
+                  cmap=ListedColormap(['white'] + [(random(),random(),random()) for i in range(len(index_to_label) - 1)]))
+        ax.set_yticklabels([])
+        ax.set_xlabel('frames')
+        cbar = fig.colorbar(raster)
+        labels = [label_to_appear_first] + list(unique_labels)
+        spacing = (len(labels) - 1)/len(labels)
+        start = spacing/2
+        ticks = [start] + [start + spacing*i for i in range(1,len(labels))]
+        cbar.set_ticks(ticks=ticks, labels = labels)
+
+        st.pyplot(fig)
+
+         # save generated annotations
+        annotations = [labels[x] for x in st.session_state.smoothed_predictions]
+        annotations_df = pd.DataFrame(annotations, columns=['label'])
+        csv = annotations_df.to_csv(header=False).encode("utf-8")
+        output_file_name = st.text_input("Output File Name:","output")
+        st.download_button("Download annotations as .csv",
+                           data=csv,
+                           file_name=f"{output_file_name}.csv")

explore.py

@@ -0,0 +1,337 @@
+import streamlit as st
+import plotly.express as px
+import numpy as np
+import pandas as pd
+import torch
+from utils.mp4Io import mp4Io_reader
+from utils.seqIo import seqIo_reader
+import pandas as pd
+from PIL import Image
+from pathlib import Path
+from transformers import AutoProcessor, AutoModel
+from tempfile import NamedTemporaryFile
+from tqdm import tqdm
+from utils.utils import create_embeddings_csv_io, process_dataset_in_mem, generate_embeddings_stream_io
+from get_llava_response import get_llava_response, load_llava_checkpoint_hf
+from sklearn.manifold import TSNE
+from openai import OpenAI
+import cv2
+import base64
+from hdbscan import HDBSCAN, all_points_membership_vectors
+import random
+
+# --server.maxUploadSize 3000
+REPO_NAME = 'ncoria/llava-lora-vicuna-clip-5-epochs-merge'
+
+def load_llava_model(hf_token):
+    return load_llava_checkpoint_hf(REPO_NAME, hf_token)
+
+def get_unique_labels(label_list: list[str]):
+    label_set = set()
+    for label in label_list:
+        individual_labels = label.split('||')
+        for individual_label in individual_labels:
+            label_set.add(individual_label)
+    return list(label_set)
+
+SYSTEM_PROMPT = """You are a researcher studying mice interactions from videos of the inside of a resident 
+intruder box where there is either just the resident mouse (the black one) or the resident and the intruder mouse (the white one).
+Your job is to answer questions about the behavior of the mice in the image given the context that each image is a frame of a continuous video.
+Thus, you should use the visual information about the mice in the image to try to provide a detailed behavioral description of the image."""
+
+@st.cache_resource
+def get_io_reader(uploaded_file):
+    if uploaded_file.name[-3:]=='seq':
+        with NamedTemporaryFile(suffix="seq", delete=False) as temp:
+            temp.write(uploaded_file.getvalue())
+            sr = seqIo_reader(temp.name)
+    else:
+        with NamedTemporaryFile(suffix="mp4", delete=False) as temp:
+            temp.write(uploaded_file.getvalue())
+            sr = mp4Io_reader(temp.name)
+    return sr
+
+def get_image(sr, frame_no: int):
+    image, _ = sr.getFrame(frame_no)
+    return image
+
+@st.cache_data
+def get_2d_embedding(embeddings: pd.DataFrame):
+    tsne = TSNE(n_jobs=4, n_components=2, random_state=42, perplexity=50)
+    embedding_2d = tsne.fit_transform(np.array(embeddings))
+    return embedding_2d
+
+HDBSCAN_PARAMS = {
+    'min_samples': 1
+}
+
+@st.cache_data
+def hdbscan_classification(umap_embeddings, embeddings_2d, cluster_range):
+    max_num_clusters = -np.infty
+    num_clusters = []
+    min_cluster_size = np.linspace(cluster_range[0], cluster_range[1], 4)
+    for min_c in min_cluster_size:
+        learned_hierarchy = HDBSCAN(
+            prediction_data=True, min_cluster_size=int(round(min_c * 0.01 *umap_embeddings.shape[0])),
+            cluster_selection_method='leaf' ,
+            **HDBSCAN_PARAMS).fit(umap_embeddings)
+        num_clusters.append(len(np.unique(learned_hierarchy.labels_)))
+        if num_clusters[-1] > max_num_clusters:
+            max_num_clusters = num_clusters[-1]
+            retained_hierarchy = learned_hierarchy
+    assignments = retained_hierarchy.labels_
+    assign_prob = all_points_membership_vectors(retained_hierarchy)
+    soft_assignments = np.argmax(assign_prob, axis=1)
+    retained_hierarchy.fit(embeddings_2d)
+    return retained_hierarchy, assignments, assign_prob, soft_assignments
+
+def upload_image(frame: np.ndarray):
+    """returns the file ID."""
+    _, encoded_image = cv2.imencode('.png', frame)
+    return base64.b64encode(encoded_image.tobytes()).decode('utf-8')
+
+def ask_question_with_image_gpt(file_id, system_prompt, question, api_key):
+    """Asks a question about the uploaded image."""
+    client = OpenAI(api_key=api_key)
+
+    if file_id != None:
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": [
+                    {"type": "text", "text": question},
+                    {"type": "image_url",  "image_url": {"url": f"data:image/jpg:base64, {file_id}"}}]
+                }
+            ]
+        )
+    else:
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": question}
+            ]
+        )
+    return response.choices[0].message.content
+
+def ask_question_with_image_llava(image, system_prompt, question,
+                                  tokenizer, model, image_processor):
+    outputs = get_llava_response([question],
+                                 [image],
+                                 system_prompt,
+                                 tokenizer,
+                                 model,
+                                 image_processor,
+                                 REPO_NAME,
+                                 stream_output=False)
+    return outputs[0]
+
+def ask_summary_question(image_array, label_array, api_key):
+    # load llava model
+    tokenizer, model, image_processor = load_llava_model(hf_token)
+
+    # global variable
+    system_prompt = SYSTEM_PROMPT
+
+    # collect responses
+    responses = []
+
+    # create progress bar
+    j = 0
+    pbar_text = lambda j: f'Creating llava response {j}/{len(label_array)}.'
+    pbar = st.progress(0, text=pbar_text(0))
+
+    for i, image in enumerate(image_array):
+        label = label_array[i]
+        question = f"The frame is annotated by a human observer with the label: {label}. Give evidence for this label using the posture of the mice and their current behavior. "
+        question += "Also, designate a behavioral subtype of the given label that describes the current social interaction based on what you see about the posture of the mice and "\
+                    "how they are positioned with respect to each other. Usually, the body parts (i.e., tail, genitals, face, body, ears, paws)"\
+                    "of the mice that are closest to each other will give some clue. Please limit behavioral subtype to a 1-4 word phrase. limit your response to 4 sentences."
+        response = ask_question_with_image_llava(image, system_prompt, question,
+                                                 tokenizer, model, image_processor)
+        responses.append(response)
+        # update progress bar
+        j += 1
+        pbar.progress(j/len(label_array), pbar_text(j))
+
+    system_prompt_summarize = "You are a researcher studying mice interactions from videos of the inside of a resident "\
+                            "intruder box where there is either just the resident mouse (the black one) or the resident and the intruder mouse (the white one). "\
+                            "You will be given a question about a list of descriptions from frames of these videos. "\
+                            "Your job is to answer the question by focusing on the behaviors of the mice and their postures "\
+                            "as well as any other aspects of the descriptions that may be relevant to the class label associated with them"
+    user_prompt_summarize = "Here are several descriptions of individual frames from a mouse behavior video. Please summarize these descriptions and provide a suggestion for a "\
+                            "behavior label which captures what is described in the descriptions: \n\n"
+    user_prompt_summarize = user_prompt_summarize + '\n'.join(responses)
+    summary_response = ask_question_with_image_gpt(None, system_prompt_summarize, user_prompt_summarize, api_key)
+    return summary_response
+
+if "embeddings_df" not in st.session_state:
+    st.session_state.embeddings_df = None
+
+st.title('batik: frame classifier')
+
+api_key = st.text_input("OpenAI API Key:","")
+hf_token = st.text_input("HuggingFace Token:","")
+st.subheader("generate or import embeddings")
+
+st.text("Upload files to generate embeddings.")
+with st.form('embedding_generation_settings'):
+    seq_file = st.file_uploader("Choose a video file", type=['seq', 'mp4'], accept_multiple_files=True)
+    annot_files = st.file_uploader("Choose an annotation File", type=['annot','csv'], accept_multiple_files=True)
+    downsample_rate = st.number_input('Downsample Rate',value=4)
+    submit_embed_settings = st.form_submit_button('Create Embeddings', type='secondary')
+
+st.markdown("**(Optional)** Upload embeddings.")
+embeddings_csv = st.file_uploader("Choose a .csv File", type=['csv'])
+
+if submit_embed_settings and seq_file is not None and annot_files is not None:
+    video_embeddings, video_frames = generate_embeddings_stream_io([seq_file],
+                                                                "SLIP",
+                                                                downsample_rate,
+                                                                False)
+    
+    fnames = [seq_file.name]
+    embeddings_df = create_embeddings_csv_io(out="file",
+                                fnames=fnames,
+                                embeddings=video_embeddings,
+                                frames=video_frames,
+                                annotations=[annot_files],
+                                test_fnames=None,
+                                views=None,
+                                conditions=None,
+                                downsample_rate=downsample_rate)
+    st.session_state.embeddings_df = embeddings_df
+elif embeddings_csv is not None:
+    embeddings_df = pd.read_csv(embeddings_csv)
+    st.session_state.embeddings_df = embeddings_df
+else:
+    st.text('Please upload file(s).')
+
+st.divider()
+st.subheader("provide video file if not yet already provided")
+
+uploaded_file = st.file_uploader("Choose a video file", type=['seq', 'mp4'], accept_multiple_files=True)
+
+st.divider()
+if st.session_state.embeddings_df is not None and (uploaded_file is not None or seq_file is not None):
+    if seq_file is not None:
+        uploaded_file = seq_file
+    io_reader = get_io_reader(uploaded_file)
+    print("CONVERTED SEQ")
+    label_list = st.session_state.embeddings_df['Label'].to_list()
+    unique_label_list = get_unique_labels(label_list)
+    print(f"unique_labels: {unique_label_list}")
+    #unique_label_list = ['check_genital', 'wiggle', 'lordose', 'stay', 'turn', 'top_up', 'dart', 'sniff', 'approach', 'into_male_cage']
+    #unique_label_list = ['into_male_cage', 'intromission', 'male_sniff', 'mount']
+    kwargs = {'embeddings_df' : st.session_state.embeddings_df, 
+                'specified_classes' : unique_label_list,
+                'classes_to_remove' : None,
+                'max_class_size' : None,
+                'animal_state' : None,
+                'view' : None,
+                'shuffle_data' : False,
+                'test_videos' : None}
+    train_embeds, train_labels, train_images, _, _, _ = process_dataset_in_mem(**kwargs)
+    print("PROCESSED DATASET")
+    if "Images" in st.session_state.embeddings_df.keys():
+        train_images = [i for i in range(len(train_images))]
+    embedding_2d = get_2d_embedding(train_embeds)
+else:
+    st.text('Please generate embeddings and provide video file.')
+    print("GOT 2D EMBEDS")
+
+if uploaded_file is not None and st.session_state.embeddings_df is not None:
+    st.subheader("t-SNE Projection")
+    option = st.selectbox(
+        "Select Color Option",
+        ("By Label", "By Time", "By Cluster")
+    )
+    if embedding_2d is not None:
+        if option is not None:
+            if option == "By Label":
+                color = 'label'
+            elif option == "By Time":
+                color = 'frame_no'
+            else:
+                color = 'cluster_label'
+            
+            if option in ["By Label", "By Time"]:
+                edf = pd.DataFrame(embedding_2d,columns=['tsne_dim_1', 'tsne_dim_2'])
+                edf.insert(2,'frame_no',np.array([int(x) for x in train_images]))
+                edf.insert(3, 'label', train_labels)
+                fig = px.scatter(
+                    edf,
+                    x="tsne_dim_1",
+                    y="tsne_dim_2",
+                    color=color,
+                    hover_data=["frame_no"],
+                    color_discrete_sequence=px.colors.qualitative.Dark24
+                )
+            else:
+                r, _, _, _ = hdbscan_classification(train_embeds, embedding_2d, [4, 6])
+                edf = pd.DataFrame(embedding_2d,columns=['tsne_dim_1', 'tsne_dim_2'])
+                edf.insert(2,'frame_no',np.array([int(x) for x in train_images]))
+                edf.insert(3, 'label', train_labels)
+                edf.insert(4, 'cluster_label', [str(c_id) for c_id in r.labels_.tolist()])
+                fig = px.scatter(
+                    edf,
+                    x="tsne_dim_1",
+                    y="tsne_dim_2",
+                    color=color,
+                    hover_data=["frame_no"],
+                    color_discrete_sequence=px.colors.qualitative.Dark24
+                )
+
+            event = st.plotly_chart(fig, key="df", on_select="rerun")
+        else:
+            st.text("No Color Option Selected")
+    else:
+        st.text('No Embeddings Loaded')
+
+    event_dict = event.selection
+
+    if event_dict is not None:
+        custom_data = []
+        for point in event_dict['points']:
+            data = point["customdata"][0]
+            custom_data.append(int(data))
+        
+        if len(custom_data) > 10:
+            custom_data = random.sample(custom_data, 10)
+        if len(custom_data) > 1:
+            col_1, col_2 = st.columns(2)
+            with col_1:
+                for frame_no in custom_data[::2]:
+                    st.image(get_image(io_reader, frame_no))
+                    st.caption(f"Frame {frame_no}, {train_labels[frame_no]}")
+            with col_2:
+                for frame_no in custom_data[1::2]:
+                    st.image(get_image(io_reader, frame_no))
+                    st.caption(f"Frame {frame_no}, {train_labels[frame_no]}")
+        elif len(custom_data) == 1:
+            frame_no = custom_data[0]
+            st.image(get_image(io_reader, frame_no))
+            st.caption(f"Frame {frame_no}, {train_labels[frame_no]}")
+        else:
+            st.text('No Points Selected')
+        
+        if len(custom_data) == 1:
+            frame_no = custom_data[0]
+            image = get_image(io_reader, frame_no)
+            system_prompt = SYSTEM_PROMPT
+            label = train_labels[frame_no]
+            question = f"The frame is annotated by a human observer with the label: {label}. Give evidence for this label using the posture of the mice and their current behavior. "\
+                        "Also, designate a behavioral subtype of the given label that describes the current social interaction based on what you see about the posture of the mice and "\
+                        "how they are positioned with respect to each other. Usually, the body parts (i.e., tail, genitals, face, body, ears, paws)" \
+                        "of the mice that are closest to each other will give some clue. Please limit behavioral subtype to a 1-4 word phrase. limit your response to 4 sentences."
+            tokenizer, model, image_processor = load_llava_model(hf_token)
+            response = ask_question_with_image_llava(image, system_prompt, question,
+                                                     tokenizer, model, image_processor)
+            st.markdown(response)
+            
+        elif len(custom_data) > 1:
+            image_array = [get_image(io_reader, f_no) for f_no in custom_data]
+            label_array = [train_labels[f_no] for f_no in custom_data]
+            response = ask_summary_question(image_array, label_array, api_key)
+            st.markdown(response)

generate_embeddings.py

@@ -0,0 +1,131 @@
+import streamlit as st
+import pandas as pd
+from utils.utils import create_embeddings_csv_io, create_annot_fname_dict_io, generate_embeddings_stream_io
+
+if "video_embeddings" not in st.session_state:
+    st.session_state.video_embeddings = None
+    st.session_state.video_frames = None
+    st.session_state.fnames = []
+
+st.title('batik: embedding generator')
+uploaded_files = st.file_uploader("Choose a video file", type=['seq', 'mp4'], accept_multiple_files=True)
+with st.form('initial_settings'):
+    st.header('Embedding Generation Options')
+    model_select = st.selectbox('Select Model', ['SLIP', 'CLIP'])
+    downsample_rate = st.number_input('Downsample Rate',value=4)
+    save_csv = st.toggle('Save Individual Results', value=False)
+    submit_initial_settings = st.form_submit_button('Create Embeddings', type='secondary')
+
+if submit_initial_settings and uploaded_files is not None and len(uploaded_files) > 0:
+    video_embeddings, video_frames = generate_embeddings_stream_io(uploaded_files,
+                                                                model_select,
+                                                                downsample_rate,
+                                                                save_csv)
+    fnames = [vid_file.name for vid_file in uploaded_files]
+    st.session_state.video_embeddings = video_embeddings
+    st.session_state.video_frames = video_frames
+    st.session_state.fnames = fnames
+
+if st.session_state.video_embeddings is not None:
+    st.header('CSV Configuration Options')
+    st.markdown('If using `.annot` files and multiple files should be grouped together, '\
+                'please ensure that they share a common name and end with a number describing '\
+                'the order of the files. For example:\n\n'\
+                '`mouse_224_file_1.annot`, `mouse_224_file_2.annot`.')
+    annot_files = st.file_uploader("Upload all annotation files", type=['.annot','.csv'], accept_multiple_files=True)
+
+    annot_options = []
+    if annot_files is not None and len(annot_files) > 0:
+        annot_fnames = [annot_file.name for annot_file in annot_files]
+        annot_fname_dict = create_annot_fname_dict_io(annot_fnames=annot_fnames,
+                                                      annot_files=annot_files)
+        annot_options = [str(key) for key in annot_fname_dict.keys()]
+
+    if len(annot_options) > 0:
+        with st.form('csv_settings'):
+            csv_setting_def = pd.DataFrame(
+                {
+                    "File Name" : st.session_state.fnames,
+                    "Annotations" : [
+                        "Upload File" for _ in st.session_state.fnames
+                    ],
+                    "Test" : [
+                        False for _ in st.session_state.fnames
+                    ],
+                    "View" : [
+                        "Top" for _ in st.session_state.fnames
+                    ],
+                    "Condition" : [
+                        "None" for _ in st.session_state.fnames
+                    ]
+
+                }
+            )
+
+            csv_settings = st.data_editor(
+                csv_setting_def,
+                column_config={
+                    "Annotations" : st.column_config.SelectboxColumn(
+                        "Annotations",
+                        help="The annotation file(s) to use for the given video file.",
+                        width="medium",
+                        options=annot_options,
+                        required=True
+                    ),
+                    "Test" : st.column_config.CheckboxColumn(
+                        "Test",
+                        help="Designate file(s) to use as the test set.",
+                        default=False,
+                        required=True
+                    ),
+                    "View" : st.column_config.SelectboxColumn(
+                        "View",
+                        help="The view used within the video (either Top or Front).",
+                        options=["Top", "Front"],
+                        required=True
+                    ),
+                    "Condition" : st.column_config.TextColumn(
+                        "Condition",
+                        help="A condition the video has (i.e. Control).",
+                        default="None",
+                        max_chars=30,
+                        validate=r"[a-z]+$",
+                    )
+                },
+                hide_index=True
+            )
+            save_csv_bttn = st.form_submit_button("Create CSV")
+
+        if save_csv_bttn and csv_settings is not None:
+            annot_chosen_options = csv_settings['Annotations'].tolist()
+            annot_option = [annot_fname_dict[key] for key in annot_chosen_options]
+            test_chosen_option = csv_settings['Test'].tolist()
+            test_option = [st.session_state.fnames[i] for i, is_test in enumerate(test_chosen_option) if is_test]
+            view_option = csv_settings['View'].tolist()
+            condition_option =  csv_settings['Condition'].tolist()
+
+            out_name = st.text_input("Embeddings Outpit File Name", "out.csv")
+            try:
+                df = create_embeddings_csv_io(out=out_name,
+                                            fnames=st.session_state.fnames,
+                                            embeddings=st.session_state.video_embeddings,
+                                            frames=st.session_state.video_frames,
+                                            annotations=annot_option,
+                                            test_fnames=test_option,
+                                            views=view_option,
+                                            conditions=condition_option,
+                                            downsample_rate=downsample_rate)
+                st.success('Created Embeddings File!', icon="✅")
+                st.download_button(
+                    label="Download CSV",
+                    data=df.to_csv().encode("utf-8"),
+                    file_name=out_name,
+                    mime="text/csv"
+                )
+            except:
+                st.error('Something went wrong.')
+
+    else:
+        st.text('Please Upload Files')
+else:
+    st.text('Please Upload Files')

get_llava_response.py

@@ -0,0 +1,184 @@
+import argparse
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
+import numpy as np
+
+from huggingface_hub import whoami
+
+import llava
+from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN
+from llava.conversation import conv_templates, SeparatorStyle
+from llava.model.builder import load_pretrained_model
+from llava.utils import disable_torch_init
+from llava.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+from tqdm import tqdm
+
+import warnings
+warnings.filterwarnings('ignore')
+
+REPO_NAME = 'ncoria/llava-lora-vicuna-clip-5-epochs-merge'
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+def load_llava_checkpoint(model_path: str):
+    model_name = get_model_name_from_path(model_path)
+    return load_pretrained_model(model_path, None, model_name, load_4bit=True, device="cuda")
+
+def load_llava_checkpoint_hf(model_path, hf_token):
+    user = whoami(token=hf_token)
+    kwargs = {"device_map": "auto"}
+    kwargs['load_in_4bit'] = True
+    kwargs['quantization_config'] = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type='nf4'
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+    model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+    mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+    mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
+    if mm_use_im_patch_token:
+        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+    if mm_use_im_start_end:
+        tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+    model.resize_token_embeddings(len(tokenizer))
+
+    vision_tower = model.get_vision_tower()
+    if not vision_tower.is_loaded:
+        vision_tower.load_model(device_map="auto")
+    image_processor = vision_tower.image_processor
+    return tokenizer, model, image_processor
+
+def get_llava_response(user_prompts: list[str],
+                       images: list,
+                       sys_prompt: str,
+                       tokenizer,
+                       model,
+                       image_processor,
+                       model_path = REPO_NAME,
+                       stream_output = True):
+    """
+    This function returns the response from the given model. It creates a one turn conversation in which
+    the only content is a system prompt and the given user message applied to each image.
+
+    Parameters:
+    ----------
+    user_prompt : str
+        The prompt sent by the user.
+    images : str
+        List of images from file.
+    sys_prompt : str
+        The prompt that sets the tone for the conversation.
+    model_path : str
+        The path to the merged checkpoint or base model.
+
+    Returns:
+    --------
+    """
+    # set up and load model
+    model_name = get_model_name_from_path(model_path)
+    temperature = 0.2 # default
+    max_new_tokens = 512 # default
+
+    # determine conversation type
+    if "llama-2" in model_name.lower():
+        conv_mode = "llava_llama_2"
+    elif "mistral" in model_name.lower():
+        conv_mode = "mistral_instruct"
+    elif "v1.6-34b" in model_name.lower():
+        conv_mode = "chatml_direct"
+    elif "v1" in model_name.lower():
+        conv_mode = "llava_v1"
+    elif "mpt" in model_name.lower():
+        conv_mode = "mpt"
+    else:
+        conv_mode = "llava_v0"
+
+    # run clean conversation for each image
+    llm_outputs = []
+    for i, img in tqdm(enumerate(images)):
+        # set up clean conversation
+        conv = conv_templates[conv_mode].copy()
+        if "mpt" in model_name.lower():
+            roles = ('user', 'assistant')
+        else:
+            roles = conv.roles
+    
+        conv.system = sys_prompt
+    
+        # load image
+        # image = load_image("../images/mouse.png") # previous method
+        if isinstance(img, np.ndarray) and len(img.shape) == 2:
+            img = Image.fromarray(img, 'L')
+
+        image = img.convert('RGB')
+        image_size = image.size
+
+        # NOTE: image is simply PIL Image (.convert('RGB')), no need for temp files!
+        
+        # Similar operation in model_worker.py
+        image_tensor = process_images([image], image_processor, model.config)
+        if type(image_tensor) is list:
+            image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
+        else:
+            image_tensor = image_tensor.to(model.device, dtype=torch.float16)
+
+        # execute conversation
+        inp = user_prompts[i]
+        if image is not None:
+            # first message
+            if model.config.mm_use_im_start_end:
+                inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp
+            else:
+                inp = DEFAULT_IMAGE_TOKEN + '\n' + inp
+            image = None
+        conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+        input_ids = tokenizer_image_token(prompt,
+                                          tokenizer,
+                                          IMAGE_TOKEN_INDEX,
+                                          return_tensors='pt').unsqueeze(0).to(model.device)
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
+        if stream_output:
+            streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        else:
+            streamer = None
+        
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids,
+                images=image_tensor,
+                image_sizes=[image_size],
+                do_sample=True if temperature > 0 else False,
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                streamer=streamer,
+                use_cache=True)
+
+        outputs = tokenizer.decode(output_ids[0]).strip()
+        llm_outputs.append(outputs)
+    return llm_outputs
+        
+
+        
+    
+    
+
+    
+                       

home.py

@@ -0,0 +1,5 @@
+import streamlit as st
+
+st.title('batik')
+st.subheader('Abstract')
+st.markdown('Quantitative analysis of animal behavior represents a burgeoning frontier in neuroscience and ethology. Recent years have witnessed a proliferation of computational methods aimed at identifying behavioral subtypes, or "syllables," from video data. However, while significant advances have been made in behavior segmentation, comparatively few approaches address the interpretation of these behavior syllables, leaving researchers to spend considerable time curating and interpreting the characteristics of the behavioral subtype. Furthermore, most current techniques rely heavily on pose estimation—a prerequisite that, while useful, can introduce limitations concerning generalization in behavioral classification and discovery. Here, we introduce Batik, a system leveraging pre-trained and fine-tuned multimodal transformers to perform end-to-end behavior analysis directly from raw video. Batik excels at supervised behavior annotation, utilizing lightweight models trained on the transformer-extracted feature space to achieve state-of-the-art performance. By integrating a pre-trained vision transformer with a custom fine-tuned language model, Batik not only discovers behavior syllables but also provides expert-level interpretations of mouse behavior, directly from visual data. This comprehensive platform empowers researchers with automated behavior discovery and interpretation, significantly reducing the time burden on experimentalists. Coupled with an intuitive user interface, Batik offers a transformative tool for the next generation of behavioral analysis, showcasing the potential of what is possible with transformer-based language models for behavior.')

pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "llava"
+version = "1.2.2.post1"
+description = "Towards GPT-4 like large language and visual assistant."
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+]
+dependencies = [
+    "torch==2.1.2", "torchvision==0.16.2",
+    "transformers==4.37.2", "tokenizers==0.15.1", "sentencepiece==0.1.99", "shortuuid",
+    "accelerate==0.21.0", "peft", "bitsandbytes",
+    "pydantic", "markdown2[all]", "numpy", "scikit-learn==1.2.2",
+    "gradio==4.16.0", "gradio_client==0.8.1",
+    "requests", "httpx==0.24.0", "uvicorn", "fastapi",
+    "einops==0.6.1", "einops-exts==0.0.4", "timm==0.6.13",
+    "protobuf", "timecode", "sortedcontainers", "qtpy", "pyqt5-tools",
+    "scipy", "matplotlib", "colour_demosaicing", "sk-video",
+    "opencv-python", "progressbar", "openai",
+    "clip @ git+https://github.com/openai/CLIP@main",
+    "scikit-learn", "tensorflow", "sentencepiece", "streamlit",
+    "hdbscan", "plotly", "ipywidgets"
+]
+
+[project.optional-dependencies]
+train = ["deepspeed==0.12.6", "ninja", "wandb"]
+build = ["build", "twine"]
+
+[project.urls]
+"Homepage" = "https://llava-vl.github.io"
+"Bug Tracker" = "https://github.com/haotian-liu/LLaVA/issues"
+
+[tool.setuptools.packages.find]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]
+
+[tool.wheel]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]

train_model.py

@@ -0,0 +1,159 @@
+import os
+import io
+import pickle
+import regex
+import streamlit as st
+import plotly.express as px
+import numpy as np
+import pandas as pd
+import torch
+from utils.seqIo import seqIo_reader
+import pandas as pd
+from PIL import Image
+from pathlib import Path
+from transformers import AutoProcessor, AutoModel
+from tqdm import tqdm
+from sklearn.svm import SVC
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, classification_report
+from utils.utils import create_embeddings_csv_io, process_dataset_in_mem, generate_embeddings_stream_io
+
+# --server.maxUploadSize 3000
+
+def get_unique_labels(label_list: list[str]):
+    label_set = set()
+    for label in label_list:
+        individual_labels = label.split('||')
+        for individual_label in individual_labels:
+            label_set.add(individual_label)
+    return list(label_set)
+
+@st.cache_data
+def get_train_test_split(train_embeds, numerical_labels, test_size=0.05, random_state=42):
+    return train_test_split(train_embeds, numerical_labels, test_size=test_size, random_state=random_state)
+
+@st.cache_resource
+def train_model(X_train, y_train, random_state=42):
+    # Train SVM Classifier
+    svm_clf = SVC(kernel='rbf', random_state=random_state, probability=True)
+    svm_clf.fit(X_train, y_train)
+    return svm_clf
+
+def pickle_model(model):
+    pickled = io.BytesIO()
+    pickle.dump(model, pickled)
+    return pickled
+
+if "embeddings_df" not in st.session_state:
+    st.session_state.embeddings_df = None
+
+if "svm_clf" not in st.session_state:
+    st.session_state.svm_clf = None
+    st.session_state.report_df = None
+    st.session_state.accuracy = None
+
+st.title('batik: frame classifier training')
+
+st.text("Upload files to train classifier on.")
+with st.form('embedding_generation_settings'):
+    seq_file = st.file_uploader("Choose a .seq File", type=['seq'])
+    annot_files = st.file_uploader("Choose an annotation File", type=['annot','csv'], accept_multiple_files=True)
+    downsample_rate = st.number_input('Downsample Rate',value=4)
+    submit_embed_settings = st.form_submit_button('Create Embeddings', type='secondary')
+
+st.markdown("**(Optional)** Upload embeddings.")
+embeddings_csv = st.file_uploader("Choose a .csv File", type=['csv'])
+
+if submit_embed_settings and seq_file is not None and annot_files is not None:
+    video_embeddings, video_frames = generate_embeddings_stream_io([seq_file],
+                                                                "SLIP",
+                                                                downsample_rate,
+                                                                False)
+    
+    fnames = [seq_file.name]
+    embeddings_df = create_embeddings_csv_io(out="file",
+                                fnames=fnames,
+                                embeddings=video_embeddings,
+                                frames=video_frames,
+                                annotations=[annot_files],
+                                test_fnames=None,
+                                views=None,
+                                conditions=None,
+                                downsample_rate=downsample_rate)
+    st.session_state.embeddings_df = embeddings_df
+
+elif embeddings_csv is not None:
+    embeddings_df = pd.read_csv(embeddings_csv)
+    st.session_state.embeddings_df = embeddings_df
+else:
+    st.text('Please upload file(s).')
+
+st.divider()
+
+if st.session_state.embeddings_df is not None:
+    st.subheader("specify dataset preprocessing options")
+    st.text("Select frames with label(s) to include:")
+
+    with st.form('train_settings'):
+        label_list = st.session_state.embeddings_df['Label'].to_list()
+        unique_label_list = get_unique_labels(label_list)
+        specified_classes = st.multiselect("Label(s) included:", options=unique_label_list)
+
+        st.text("Select label(s) that should be removed:")
+        classes_to_remove = st.multiselect("Label(s) excluded:", options=unique_label_list)
+
+        max_class_size = st.number_input("(Optional) Specify max class size:", value=None)
+
+        shuffle_data = st.toggle("Shuffle data:")
+
+        train_model_clicked = st.form_submit_button("Train Model")
+
+    if train_model_clicked:
+        kwargs = {'embeddings_df' : st.session_state.embeddings_df, 
+                'specified_classes' : specified_classes,
+                'classes_to_remove' : classes_to_remove,
+                'max_class_size' : max_class_size,
+                'animal_state' : None,
+                'view' : None,
+                'shuffle_data' : shuffle_data,
+                'test_videos' : None}
+        train_embeds, train_labels, train_images, _, _, _ = process_dataset_in_mem(**kwargs)
+        # Convert labels to numerical values
+        label_to_appear_first = 'other'
+        unique_labels = set(train_labels)
+        unique_labels.discard(label_to_appear_first)
+
+        label_to_index = {label_to_appear_first: 0}
+
+        label_to_index.update({label: idx + 1 for idx, label in enumerate(unique_labels)})
+        index_to_label = {idx: label for label, idx in label_to_index.items()}
+        numerical_labels = np.array([label_to_index[label] for label in train_labels])
+
+        print("Label Valence: ", label_to_index)
+        # Split data into train and test sets
+        X_train, X_test, y_train, y_test = get_train_test_split(train_embeds, numerical_labels, test_size=0.05, random_state=42)
+        with st.spinner("Model training in progress..."):
+            svm_clf = train_model(X_train, y_train)
+
+        # Predict on the test set
+        with st.spinner("In progress..."):
+            y_pred = svm_clf.predict(X_test)
+        accuracy = accuracy_score(y_test, y_pred)
+        report = classification_report(y_test, y_pred, target_names=[index_to_label[idx] for idx in range(len(label_to_index))], output_dict=True)
+        report_df = pd.DataFrame(report).transpose()
+
+        # save results to session state
+        st.session_state.svm_clf = svm_clf
+        st.session_state.report_df = report_df
+        st.session_state.accuracy = accuracy
+    
+    if st.session_state.svm_clf is not None:
+        pickled_model = pickle_model(st.session_state.svm_clf)
+
+        st.text(f"Eval Accuracy: {st.session_state.accuracy}")
+        st.subheader("Classification Report:")
+        st.dataframe(st.session_state.report_df)
+
+        st.download_button("Download model as .pkl file",
+                           data=pickled_model,
+                           file_name=f"{'_'.join(specified_classes)}_classifier.pkl")

utils/annot.py

@@ -0,0 +1,641 @@
+# annot.py
+from io import StringIO
+from random import sample
+from collections import OrderedDict
+import timecode as tc
+from .behavior import Behavior
+from sortedcontainers import SortedKeyList
+from qtpy.QtCore import QObject, QRectF, Signal, Slot
+from qtpy.QtGui import QColor
+from qtpy.QtWidgets import QGraphicsItem
+
+class Bout(object):
+    """
+    """
+
+    def __init__(self, start, end, behavior):
+        self._start = start
+        self._end = end
+        self._behavior = behavior
+
+    def __lt__(self, b):
+        if type(b) is tc.Timecode:
+            return self._start.float < b.float
+        elif type(b) is Bout:
+            return self._start < b._start
+        else:
+            raise NotImplementedError(f"Comparing Bout with {type(b)} not supported")
+
+    def __le__(self, b):
+        if type(b) is tc.Timecode:
+            return self._start.float <= b.is_float
+        elif type(b) is Bout:
+            return self._start <= b._start
+        else:
+            raise NotImplementedError(f'Comparing Bout with {type(b)} not supported')
+
+    def is_at(self, t):
+        return self._start <= t and self._end >= t
+
+    def start(self):
+        return self._start
+
+    def set_start(self, start):
+        self._start = start
+
+    def end(self):
+        return self._end
+
+    def set_end(self, end):
+        self._end = end
+
+    def len(self):
+        return self._end - self._start + tc.Timecode(self._start.framerate, frames=1)
+
+    def behavior(self):
+        return self._behavior
+
+    def name(self):
+        return self._behavior.get_name()
+
+    def color(self):
+        return self._behavior.get_color()
+
+    def is_active(self):
+        return self._behavior.is_active()
+
+    def is_visible(self):
+        return self._behavior.is_visible() and self._behavior.is_active()
+
+    def __repr__(self):
+        return f"Bout: start = {self.start()}, end = {self.end()}, behavior: {self.behavior()}"
+
+class Channel(QGraphicsItem):
+    """
+    """
+
+    contentChanged = Signal()
+
+    def __init__(self, chan = None):
+        super().__init__()
+        if not chan is None:
+            self._bouts_by_start = chan._bouts_by_start
+            self._bouts_by_end = chan._bouts_by_end
+            self._top = chan._top
+        else:
+            self._bouts_by_start = SortedKeyList(key=lambda bout: bout.start().float)
+            self._bouts_by_end = SortedKeyList(key=lambda bout: bout.end().float)
+            self._top = 0.
+        self.cur_ix = 0
+        self.fakeFirstBout = Bout(
+            tc.Timecode('30.0', '0:0:0:0'),
+            tc.Timecode('30.0', '0:0:0:0'),
+            Behavior('none', '', QColor.fromRgbF(0., 0., 0.), active = True)
+        )
+        self.fakeLastBout = Bout(
+            tc.Timecode('30.0', '23:59:59:29'),
+            tc.Timecode('30.0', '23:59:59:29'),
+            Behavior('none', '', QColor.fromRgbF(0., 0., 0.), active = True)
+        )
+
+    def add(self, b):
+        if isinstance(b, Bout):
+            self._bouts_by_start.add(b)
+            self._bouts_by_end.add(b)
+        else:
+            raise TypeError("Can only add Bouts to Channel")
+
+    def remove(self, b):
+        # can raise ValueError if b is not in the channel
+        self._bouts_by_start.remove(b)
+        self._bouts_by_end.remove(b)
+
+    def update_start(self, b, new_start):
+        """
+        Update the starting time of a bout while
+        preserving the _bouts_by_start access order
+        """
+        self._bouts_by_start.remove(b)
+        b.set_start(new_start)
+        self._bouts_by_start.add(b)
+
+    def update_end(self, b, new_end):
+        """
+        Update the ending time of a bout while
+        preserving the _bouts_by_end access order
+        """
+        self._bouts_by_end.remove(b)
+        b.set_end(new_end)
+        self._bouts_by_end.add(b)
+
+    def __add__(self, b):
+        self.add(b)
+
+    def _get_next(self, t, sortedlist):
+        ix = sortedlist.bisect_key_right(t.float)
+        l = len(sortedlist)
+        if ix == l:
+            return self.fakeLastBout, t.next()
+        return sortedlist[ix], t.next()
+
+    def _get_prev(self, t, sortedlist):
+        ix = sortedlist.bisect_key_left(t.float)
+        if ix == 0:
+            return self.fakeFirstBout, t.back()
+        return sortedlist[ix-1], t.back()
+
+    def _get_inner(self, t, sortedList, op):
+        t_local = t + 0 # kludgy copy constructor!
+        visible = False
+        while not visible:
+            # no need to check for the end, because the fake first and last bouts are visible
+            bout, t_local = op(t_local, sortedList)
+            visible = bout.is_visible()
+        return bout
+
+    def get_next_start(self, t):
+        return self._get_inner(t, self._bouts_by_start, self._get_next)
+
+    def get_next_end(self, t):
+        return self._get_inner(t, self._bouts_by_end, self._get_next)
+
+    def get_prev_start(self, t):
+        return self._get_inner(t, self._bouts_by_start, self._get_prev)
+
+    def get_prev_end(self, t):
+        return self._get_inner(t, self._bouts_by_end, self._get_prev)
+
+    def get_in_range(self, start, end):
+        """
+        get all bouts that intersect the range [start, end]
+        """
+        return [bout for bout in self._bouts_by_start
+            if bout.start().float <= end.float and bout.end().float >= start.float]
+
+    def get_at(self, t):
+        """
+        get all bouts that span time t
+        """
+        return self.get_in_range(t, t)
+
+    def __iter__(self):
+        return iter(self._bouts_by_start)
+
+    def irange(self, start_time, end_time):
+        if isinstance(start_time, tc.Timecode):
+            start_time = start_time.float
+        if isinstance(end_time, tc.Timecode):
+            end_time = end_time.float
+        if not isinstance(start_time, float):
+            raise TypeError(f"Can't handle start_time of type {type(start_time)}")
+        if not isinstance(end_time, float):
+            raise TypeError(f"Can't handle end_time of type {type(end_time)}")
+        return self._bouts_by_start.irange_key(start_time, end_time)
+
+    def top(self):
+        return self._top
+
+    def set_top(self, top):
+        self._top = top
+
+    def boundingRect(self):
+        width = self.fakeLastBout.end().float
+        return QRectF(0., self.top(), width, 1.)
+
+    def paint(self, painter, option, widget=None):
+        boundingRect = option.rect
+        in_range_bouts = self._bouts_by_start.irange_key(boundingRect.left(), boundingRect.right())
+        while True:
+            try:
+                bout = next(in_range_bouts)
+            except StopIteration:
+                break
+            if bout.is_visible():
+                painter.fillRect(
+                    QRectF(bout.start().float, self.top(), bout.len().float, 1.),
+                    bout.color()
+                    )
+
+    def _delete_all_inner(self, predicate):
+        to_delete = list()
+        # can't alter the bouts within the iterator
+        for bout in iter(self):
+            if predicate(bout):
+                to_delete.append(bout)
+        deleted_names = set()
+        for bout in to_delete:
+            deleted_names.add(bout.name())
+            self.remove(bout)
+        return deleted_names
+
+    def delete_bouts_by_name(self, behavior_name):
+        return self._delete_all_inner(lambda bout: bout.name() == behavior_name)
+
+    def delete_inactive_bouts(self):
+        return self._delete_all_inner(lambda bout: not bout.is_active())
+
+    def truncate_or_remove_bouts(self, behavior, start, end, delete_all=False):
+        """
+        Delete bouts entirely within the range [start, end], and
+        truncate bouts that extend outside the range.
+        If behavior matches _deleteBehavior, the activity affects
+        all bouts.  Otherwise, it only affects bouts with matching behavior.
+        """
+        items = self.get_in_range(start, end)
+        for item in items:
+            if not delete_all and behavior.get_name() != item.name():
+                continue
+            # Delete bouts that are entirely within the range
+            if item.start() >= start and item.end() < end:
+                print(f"removing {item} from active channel")
+                self.remove(item)
+
+            # Truncate and duplicate bouts that extend out both sides of the range
+            if item.start() < start and item.end() > end:
+                self.add(Bout(end, item.end(), item.behavior()))
+                self.update_end(item, start)
+
+            # Truncate bouts at the start boundary that start before the range
+            elif item.start() < start and item.end() <= end:
+                self.update_end(item, start)
+
+            # Truncate bouts at the end boundary that end after the range
+            elif item.start() >= start and item.end() > end:
+                self.update_start(item, end)
+
+            else:
+                print(f"truncate_or_delete_bouts: Unexpected bout {item}")
+
+    def coalesce_bouts(self, start, end):
+        """
+        combine overlapping bouts of the same behavior within [start, end]
+        """
+        to_delete = []
+        items = self.get_in_range(start, end)
+        # items will be ordered by start time
+        for ix, first in enumerate(items):
+            if first in to_delete:
+                # previously coalesced
+                continue
+            if ix == len(items)-1:
+                break
+            for second in items[ix+1:]:
+                if (first.name() == second.name() and
+                    first.end() >= second.start()):
+                    if first.end() < second.end():
+                        self.update_end(first, second.end())
+                    to_delete.append(second)
+        for item in to_delete:
+            self.remove(item)
+
+class Annotations(QObject):
+    """
+    """
+
+    # Signals
+    annotations_changed = Signal()
+    active_annotations_changed = Signal()
+
+    def __init__(self, behaviors):
+        super().__init__()
+        self._channels = OrderedDict()
+        self._behaviors = behaviors
+        self._movies = []
+        self._start_frame = None
+        self._end_frame = None
+        self._sample_rate = None
+        self._stimulus = None
+        self._format = None
+        self.annotation_names = []
+        behaviors.behaviors_changed.connect(self.note_annotations_changed)
+
+    def read(self, fn):
+        with open(fn, "r") as f:
+            line = f.readline()
+            line = line.strip().lower()
+            if line.endswith("annotation file"):
+                self._format = 'Caltech'
+                self._read_caltech(f)
+            elif line.startswith("scorevideo log"):
+                self._format = 'Ethovision'
+                self._read_ethovision(f)
+            else:
+                print("Unsupported annotation file format")
+
+    def read_io(self, uploaded_file):
+        text_str = uploaded_file.getvalue().decode("utf-8")
+        with StringIO(text_str) as f:
+            f.__setattr__('name', uploaded_file.name)
+            line = f.readline()
+            line = line.strip().lower()
+            if line.endswith("annotation file"):
+                self._format = 'Caltech'
+                self._read_caltech(f)
+            elif line.startswith("scorevideo log"):
+                self._format = 'Ethovision'
+                self._read_ethovision(f)
+            else:
+                print("Unsupported annotation file format")
+
+    def _read_caltech(self, f):
+        found_movies = False
+        found_timecode = False
+        found_stimulus = False
+        found_channel_names = False
+        found_annotation_names = False
+        found_all_channels = False
+        found_all_annotations = False
+        new_behaviors_activated = False
+        reading_channel = False
+        to_activate = []
+        channel_names = []
+        current_channel = None
+        current_bout = None
+
+        self._format = 'Caltech'
+
+        line = f.readline()
+        while line:
+            if found_annotation_names and not new_behaviors_activated:
+                self.ensure_and_activate_behaviors(to_activate)
+                new_behaviors_activated = True
+
+            line.strip()
+
+            if not line:
+                if reading_channel:
+                    reading_channel = False
+                    current_channel = None
+                    current_bout = None
+            elif line.lower().startswith("movie file"):
+                items = line.split()
+                for item in items:
+                    if item.lower().startswith("movie"):
+                        continue
+                    if item.lower().startswith("file"):
+                        continue
+                    self._movies.append(item)
+                found_movies = True
+            elif line.lower().startswith("stimulus name"):
+                # TODO: do something when we know what one of these looks like
+                found_stimulus = True
+            elif line.lower().startswith("annotation start frame") or line.lower().startswith("annotation start time"):
+                items = line.split()
+                if len(items) > 3:
+                    try:
+                        self._start_frame = int(items[3])
+                    except:
+                        self._start_frame = int(float(items[3]))
+                    if self._end_frame and self._sample_rate:
+                        found_timecode = True
+            elif line.lower().startswith("annotation stop frame") or  line.lower().startswith("annotation stop time"):
+                items = line.split()
+                if len(items) > 3:
+                    try:
+                        self._end_frame = int(items[3])
+                    except:
+                        self._end_frame = int(float(items[3]))
+                    if self._start_frame and self._sample_rate:
+                        found_timecode = True
+            elif line.lower().startswith("annotation framerate"):
+                items = line.split()
+                if len(items) > 2:
+                    self._sample_rate = float(items[2])
+                    if self._start_frame and self._end_frame:
+                        found_timecode = True
+            elif line.lower().startswith("list of channels"):
+                line = f.readline()
+                while line:
+                    line = line.strip()
+                    if not line:
+                        break # blank line -- end of section
+                    channel_names.append(line)
+                    line = f.readline()
+                found_channel_names = True
+            elif line.lower().startswith("list of annotations"):
+                line = f.readline()
+                while line:
+                    line = line.strip()
+                    if not line:
+                        break # blank line -- end of section
+                    to_activate.append(line)
+                    line = f.readline().strip()
+                found_annotation_names = True
+            elif line.strip().lower().endswith("---"):
+                for ch_name in channel_names:
+                    if line.startswith(ch_name):
+                        self._channels[ch_name] = Channel()
+                        current_channel = ch_name
+                        reading_channel = True
+                        break
+                if reading_channel:
+                    reading_annot = False
+                    line = f.readline()
+                    while line:
+                        line = line.strip()
+                        if not line: # blank line
+                            if reading_annot:
+                                reading_annot = False
+                                current_bout = None
+                            else:
+                                # second blank line, so done with channel
+                                reading_channel = False
+                                current_channel = None
+                                break
+                        elif line.startswith(">"):
+                            current_bout = line[1:]
+                            reading_annot = True
+                        elif line.lower().startswith("start"):
+                            pass # skip a header line
+                        else:
+                            items = line.split()
+                            is_float = '.' in items[0] or '.' in items[1] or '.' in items[2]
+                            bout_start = items[0]
+                            if float(bout_start) < 1:
+                                bout_start = 1
+                            bout_end = items[1]
+                            if float(bout_end) < 1:
+                                bout_end = 1
+                            self.add_bout(
+                                Bout(
+                                    tc.Timecode(self._sample_rate, start_seconds=float(bout_start)) if is_float
+                                        else tc.Timecode(self._sample_rate, frames=int(bout_start)),
+                                    tc.Timecode(self._sample_rate, start_seconds=float(bout_end)) if is_float
+                                        else tc.Timecode(self._sample_rate, frames=int(bout_end)),
+                                    self._behaviors.get(current_bout)),
+                                current_channel)
+                        line = f.readline()
+            line = f.readline()
+        print(f"Done reading Caltech annotation file {f.name}")
+        self.note_annotations_changed()
+
+    def write_caltech(self, f, video_files, stimulus):
+        if not f.writable():
+            raise RuntimeError("File not writable")
+        f.write("Bento annotation file\n")
+        f.write("Movie file(s):")
+        for file in video_files:
+            f.write(' ' + file)
+        f.write('\n\n')
+
+        f.write(f"Stimulus name: {stimulus}\n")
+        f.write(f"Annotation start frame: {self._start_frame}\n")
+        f.write(f"Annotation stop frame: {self._end_frame}\n")
+        f.write(f"Annotation framerate: {self._sample_rate}\n")
+        f.write("\n")
+
+        f.write("List of Channels:\n")
+        for ch in self.channel_names():
+            f.write(ch + "\n")
+        f.write("\n")
+
+        f.write("List of annotations:\n")
+        for annot in self.annotation_names:
+            f.write(annot + "\n")
+        f.write("\n")
+
+        for ch in self.channel_names():
+            by_name = {}
+            f.write(f"{ch}----------\n")
+
+            for bout in self.channel(ch):
+                if not by_name.get(bout.name()):
+                    by_name[bout.name()] = []
+                by_name[bout.name()].append(bout)
+
+            for annot in by_name:
+                f.write(f">{annot}\n")
+                f.write("Start\tStop\tDuration\n")
+                for bout in by_name[annot]:
+                    start = bout.start().frames
+                    end = bout.end().frames
+                    f.write(f"{start}\t{end}\t{end - start}\n")
+                f.write("\n")
+
+            f.write("\n")
+
+        f.close()
+        print(f"Done writing Caltech annotation file {f.name}")
+
+    def _read_ethovision(self, f):
+        print("Ethovision annotations not yet supported")
+
+    def clear_channels(self):
+        self._channels.clear()
+
+    def channel_names(self):
+        return list(self._channels.keys())
+
+    def channel(self, ch: str) -> Channel:
+        return self._channels[ch]
+
+    def addEmptyChannel(self, ch: str):
+        if ch not in self.channel_names():
+            self._channels[ch] = Channel()
+
+    def add_bout(self, bout, channel):
+        if bout.name() not in self.annotation_names:
+            self.annotation_names.append(bout.name())
+        self._channels[channel].add(bout)
+        if bout.end() > self.end_time():
+            self.set_end_frame(bout.end())
+
+    def start_time(self):
+        """
+        At some point we will need to support a start time distinct from
+        frame number, perhaps derived from the OS file modify time
+        or the start time of the corresponding video (or other media) file
+        """
+        if not self._start_frame or not self._sample_rate:
+            return tc.Timecode('30.0', '0:0:0:0')
+        return tc.Timecode(self._sample_rate, frames=self._start_frame)
+
+    def start_frame(self):
+        return self._start_frame
+
+    def set_start_frame(self, t):
+        if isinstance(t, int):
+            self._start_frame = t
+        elif isinstance(t, tc.Timecode):
+            self._start_frame = t.frames
+        else:
+            raise TypeError("Expected a frame number or Timecode")
+
+    def end_time(self):
+        if not self._end_frame or not self._sample_rate:
+            return tc.Timecode('30.0', '23:59:59:29')
+        return tc.Timecode(self._sample_rate, frames=self._end_frame)
+
+    def end_frame(self):
+        return self._end_frame
+
+    def set_end_frame(self, t):
+        if isinstance(t, int):
+            self._end_frame = t
+        elif isinstance(t, tc.Timecode):
+            self._end_frame = t.frames
+        else:
+            raise TypeError("Expected a frame number or Timecode")
+
+    def sample_rate(self):
+        return self._sample_rate
+
+    def set_sample_rate(self, sample_rate):
+        self._sample_rate = sample_rate
+
+    def format(self):
+        return self._format
+
+    def set_format(self, format):
+        self._format = format
+
+    def delete_bouts_by_name(self, behavior_name):
+        deleted_names = set()
+        for chan_name in self.channel_names():
+            deleted_names.update(self.channel(chan_name).delete_bouts_by_name(behavior_name))
+        for name in deleted_names:
+            self.annotation_names.remove(name)
+
+    def delete_inactive_bouts(self):
+        deleted_names = set()
+        for chan_name in self.channel_names():
+            deleted_names.update(self.channel(chan_name).delete_inactive_bouts())
+        for name in deleted_names:
+            self.annotation_names.remove(name)
+
+    def ensure_and_activate_behaviors(self, toActivate):
+        behaviorSetUpdated = False
+        for behaviorName in toActivate:
+            behaviorSetUpdated |= self._behaviors.addIfMissing(behaviorName)
+            self.annotation_names.append(behaviorName)
+            self._behaviors.get(behaviorName).set_active(True)
+        if behaviorSetUpdated:
+            self.annotations_changed.emit()
+        self.active_annotations_changed.emit()
+
+    def ensure_active_behaviors(self):
+        for behavior in self._behaviors:
+            if behavior.is_active() and behavior.get_name() not in self.annotation_names:
+                self.annotation_names.append(behavior.get_name())
+
+    def truncate_or_remove_bouts(self, behavior, start, end, chan):
+        """
+        Delete bouts entirely within the range [start, end], or
+        truncate bouts that extend outside the range.
+        If behavior matches _deleteBehavior, the activity affects
+        all bouts.  Otherwise, it only affects bouts with matching behavior.
+        """
+        delete_all = (behavior.get_name() == self._behaviors.getDeleteBehavior().get_name())
+        self._channels[chan].truncate_or_remove_bouts(behavior, start, end, delete_all)
+        self.note_annotations_changed()
+
+    def coalesce_bouts(self, start, end, chan):
+        """
+        combine overlapping bouts of the same behavior within [start, end]
+        """
+        self._channels[chan].coalesce_bouts(start, end)
+        self.note_annotations_changed()
+
+    @Slot()
+    def note_annotations_changed(self):
+        self.annotations_changed.emit()

utils/behavior.py

@@ -0,0 +1,291 @@
+# behavior.py
+"""
+Overview comment here
+"""
+
+from qtpy.QtCore import QAbstractItemModel, QAbstractTableModel, QModelIndex, QObject, Qt, Signal, Slot
+from qtpy.QtGui import QColor
+import os
+
+class Behavior(QObject):
+    """
+    An annotation behavior, which is quite simple.  It comprises:
+    name        The name of the behavior, which is displayed on various UI widgets
+    hot_key     The case-sensitive key stroke used to start and stop instances of the behavior
+    color       The color with which to display this behavior
+    """
+
+    def __init__(self, name: str, hot_key: str = '', color: QColor = QColor('gray'), active = False, visible = True):
+        super().__init__()
+        self._name = name
+        self._hot_key = '' if hot_key == '_' else hot_key
+        self._color = color
+        self._visible = visible
+        self._active = active
+        self._get_functions = {
+            'hot_key': self.get_hot_key,
+            'name': self.get_name,
+            'color': self.get_color,
+            'active': self.is_active,
+            'visible': self.is_visible
+            }
+        self._set_functions = {
+            'hot_key': self.set_hot_key,
+            'name': self.set_name,
+            'color': self.set_color,
+            'active': self.set_active,
+            'visible': self.set_visible
+            }
+
+    def __repr__(self):
+        return f"Behavior: name={self._name}, hot_key={self._hot_key}, color={self._color}, active={self._active}, visible={self._visible}"
+
+    def get(self, key):
+        # may raise KeyError
+        return self._get_functions[key]()
+
+    def set(self, key, value):
+        # may raise KeyError
+        self._set_functions[key](value)
+
+    def get_hot_key(self):
+        return self._hot_key
+
+    def set_hot_key(self, hot_key: str):
+        self._hot_key = hot_key
+
+    def get_color(self):
+        return self._color
+
+    def set_color(self, color: QColor):
+        self._color = color
+
+    def is_active(self):
+        return self._active
+
+    @Slot(bool)
+    def set_active(self, active):
+        self._active = active
+
+    def is_visible(self):
+        return self._visible
+
+    @Slot(bool)
+    def set_visible(self, visible):
+        self._visible = visible
+
+    def get_name(self):
+        return self._name
+
+    def set_name(self, name):
+        self._name = name
+
+    def toDict(self):
+        return {
+            'hot_key': '_' if self._hot_key == '' else self._hot_key,
+            'color': self._color,
+            'name': self._name,
+            'active': self._active,
+            'visible': self._visible
+            }
+
+class Behaviors(QAbstractTableModel):
+    """
+    A set of behaviors, which represent "all" possible behaviors.
+    The class supports reading from and writing to profile files that specify
+    the default hot_key and color for each behavior, along with the name.
+
+    Derives from QAbstractTableModel so that it can be viewed and edited
+    directly in a QTableView widget.
+
+    Use getattr(name) to get the Behavior instance for a given name.
+    Use from_hot_key(key) to get the behavior(s) given the hot key.
+        Returns None if the hot_key isn't defined.
+    """
+
+    behaviors_changed = Signal()
+    layout_changed = Signal()
+
+    def __init__(self):
+        super().__init__()
+        self._items = []
+        self._by_name = {}
+        self._by_hot_key = {}
+        self._header = ['hot_key', 'color', 'name', 'active', 'visible']
+        self._searchList = [self._by_hot_key, None, self._by_name, None, None]
+        self._delete_behavior = Behavior('_delete', color = QColor('black'))
+        self._immutableColumns = set()
+        self._booleanColumns = set([self._header.index('active'), self._header.index('visible')])
+        self._role_to_str = {
+            Qt.DisplayRole: "DisplayRole",
+            Qt.DecorationRole: "DecorationRole",
+            Qt.EditRole: "EditRole",
+            Qt.ToolTipRole: "ToolTipRole",
+            Qt.StatusTipRole: "StatusTipRole",
+            Qt.WhatsThisRole: "WhatsThisRole",
+            Qt.SizeHintRole: "SizeHintRole",
+            Qt.FontRole: "FontRole",
+            Qt.TextAlignmentRole: "TextAlignmentRole",
+            Qt.BackgroundRole: "BackgroundRole",
+            Qt.ForegroundRole: "ForegroundRole",
+            Qt.CheckStateRole: "CheckStateRole",
+            Qt.InitialSortOrderRole: "InitialSortOrderRole",
+            Qt.AccessibleTextRole: "AccessibleTextRole",
+            Qt.UserRole: "UserRole"
+        }
+
+    def add(self, beh: Behavior, row=-1):
+        if row < 0:
+            row = self.rowCount()
+        self.beginInsertRows(QModelIndex(), row, row)
+        self._items.insert(row, beh)
+        self._by_name[beh.get_name()] = beh
+        hot_key = beh.get_hot_key()
+        if hot_key:
+            if hot_key not in self._by_hot_key.keys():
+                self._by_hot_key[hot_key] = []
+            assert(isinstance(self._by_hot_key[hot_key], list))
+            self._by_hot_key[hot_key].append(beh)
+        self.endInsertRows()
+        self.dataChanged.emit(
+            self.index(row, 0, QModelIndex()),
+            self.index(row, self.columnCount()-1, QModelIndex()),
+            [Qt.DisplayRole, Qt.EditRole])
+        self.behaviors_changed.emit()
+
+    def load(self, f):
+        line = f.readline()
+        while line:
+            hot_key, name, r, g, b = line.strip().split(' ')
+            if hot_key == '_':
+                hot_key = ''
+            self.add(Behavior(name, hot_key, QColor.fromRgbF(float(r), float(g), float(b))))
+            line = f.readline()
+
+    def save(self, f):
+        for beh in self._items:
+            h = beh.get_hot_key()
+            if h == '':
+                h = '_'
+            color = beh.get_color()
+            f.write(f"{h} {beh.get_name()} {color.redF()} {color.greenF()} {color.blueF()}" + os.linesep)
+
+    def get(self, name):
+        if name not in self._by_name.keys():
+            return None
+        return self._by_name[name]
+
+    def from_hot_key(self, key):
+        """
+        Return the list of behaviors associated with this hot key, if any
+        """
+        try:
+            return self._by_hot_key[key]
+        except KeyError:
+            return None
+
+    def len(self):
+        return len(self._items)
+
+    def header(self):
+        return self._header
+
+    def colorColumns(self):
+        return [self._header.index('color')]
+
+    def __iter__(self):
+        return iter(self._items)
+
+    def getDeleteBehavior(self):
+        return self._delete_behavior
+
+    def addIfMissing(self, nameToAdd):
+        if nameToAdd not in self._by_name:
+            self.add(Behavior(nameToAdd, '', QColor('gray')))
+            return True
+        return False
+
+    def isImmutable(self, index):
+        return index.column() in self._immutableColumns
+
+    def setImmutable(self, column):
+        self._immutableColumns.add(column)
+
+    # QAbstractTableModel API methods
+
+    def headerData(self, col, orientation, role):
+        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
+            return self._header[col]
+        return None
+
+    def rowCount(self, parent=None):
+        return len(self._items)
+
+    def columnCount(self, parent=None):
+        return len(self._header)
+
+    def data(self, index, role=Qt.DisplayRole):
+        datum = self._items[index.row()].get(self._header[index.column()])
+        if isinstance(datum, bool):
+            if role in [Qt.CheckStateRole, Qt.EditRole]:
+                return Qt.Checked if datum else Qt.Unchecked
+            return None
+        if role in [Qt.DisplayRole, Qt.EditRole]:
+            return self._items[index.row()].get(self._header[index.column()])
+        return None
+
+    def setData(self, index, value, role=Qt.EditRole):
+        if not role in [Qt.CheckStateRole, Qt.EditRole]:
+            return False
+        if role == Qt.CheckStateRole:
+            value = bool(value)
+        beh = self._items[index.row()]
+        key = self._header[index.column()]
+        name = beh.get_name()
+        hot_key = beh.get_hot_key()
+        beh.set(key, value)
+        if key == 'hot_key' and value != hot_key:
+            # disassociate this behavior from the hot_key
+            # and associate with the new hot_key if not ''
+            if hot_key != '':
+                del(self._by_hot_key[hot_key])
+            if value != '':
+                if value not in self._by_hot_key.keys():
+                    self._by_hot_key[value] = []
+                assert(isinstance(self._by_hot_key[value], list))
+                self._by_hot_key[value].append(beh)
+        elif key == 'name' and value != name:
+            if name in self._by_name.keys():
+                del(self._by_name[name])
+            self._by_name[value] = beh
+        self.behaviors_changed.emit()
+        self.dataChanged.emit(index, index, [role])
+        return True
+
+    def insertRows(self, row, count, parent):
+        if count < 1 or row < 0 or row > self.rowCount():
+            return False
+        self.beginInsertRows(QModelIndex(), row, row)
+        for r in range(count):
+            self._items.insert(row, Behavior('', active=True))
+        self.endInsertRows()
+        return True
+
+    def removeRows(self, row, count, parent=QModelIndex()):
+        if count <= 0 or row < 0 or row + count > self.rowCount(parent):
+            return False
+        self.beginRemoveRows(parent, row, row + count - 1)
+        for item in self._items[row:row+count-1]:
+            self._by_name.pop(item.name)
+            self._by_hot_key.pop(item.hot_key)
+        for i in range(count):
+            self._items.pop(row)
+        self.endRemoveRows()
+
+    def flags(self, index):
+        f = super().flags(index)
+        if index.column() not in self._immutableColumns:
+            f |= Qt.ItemIsEditable
+        if index.column() in self._booleanColumns:
+            f = (f & ~(Qt.ItemIsSelectable | Qt.ItemIsEditable)) | Qt.ItemIsUserCheckable
+        return f

utils/data_loading.py

@@ -0,0 +1,198 @@
+"""File for loading data into AnimalEditor"""
+import io
+from random import random
+from os.path import splitext
+from collections import OrderedDict
+import numpy as np
+from tempfile import NamedTemporaryFile
+
+from .annot import Annotations
+from .behavior import Behaviors
+
+def has_extension(fname:str, extension:str|list[str]) -> bool:
+    """
+    Checks to see if the passed in file name ends with an expected extension.
+    """
+    _, ext = splitext(fname)
+    if isinstance(extension, str):
+        return ext == extension
+    elif isinstance(extension, list):
+        return ext in extension
+
+def _clean_annotations(annotations):
+    """
+    While reading in behaviors from an .annot file, sometimes channels without normally
+    callable keys appear (i.e. keys that are strings which name a behavior), thus this
+    code only accepts keys which are strings.
+    """
+    if not annotations:
+        raise ValueError('No annotations found.')
+    clean_annot = OrderedDict()
+    for channel in annotations.keys():
+        channel_dict = OrderedDict()
+        for behavior_name in annotations[channel].keys():
+            if isinstance(behavior_name, str):
+                channel_dict.update({behavior_name : annotations[channel][behavior_name]})
+        clean_annot.update({channel: channel_dict})
+    return clean_annot
+
+def load_annot_sheet_txt(fname, offset = 0):
+    """
+    Generated a dictionary for retrieving the beginning and end frames of behaviors from
+    an .annot file.
+
+    Note that 0:00:00 is frame 1
+
+    Args:
+        fname       - the path to the .annot file to be read (must be Caltech format)\n
+        offset      - a value which offsets the start and end frame of each bout in
+                      the sheet, as well as the absolute start and end frame of the file.
+                      This value is optional, and is set to 0 by default
+
+    Returns:
+        annotations - dictionary of beginning and end frames for behaviors\n
+        start_time  - the frame the movie started at (0:00:00 is 1)\n
+        end_time    - the frame the movie ended at (0:00:00 is 1)\n
+        sample_rate - the sample rate reported within the file
+    """
+    # from bento for python
+    behaviors   = Behaviors()
+    annot_sheet = Annotations(behaviors)
+    annot_sheet.read(fname)
+
+    sample_rate = annot_sheet.sample_rate()
+    annotations = OrderedDict()
+    for key in annot_sheet.channel_names():
+        annot_behaviors = OrderedDict()
+        bout_names      = set()
+        for bout in annot_sheet.channel(key): #._bouts_by_start:
+            bout_names.add(bout.name())
+        for name in bout_names:
+            annot_behaviors.update({name : []})
+        for bout in annot_sheet.channel(key): #._bouts_by_start:
+            start_frame = bout.start().frames + offset
+            end_frame   = bout.end().frames + offset
+            bout_frames = [start_frame, end_frame]
+            curr_table  = annot_behaviors.get(bout.name())
+            new_table   = curr_table.append(bout_frames)
+            annot_behaviors.update({bout.name : new_table})
+        for name in bout_names:
+            curr_table  = annot_behaviors.get(name)
+            beh_array   = np.array(curr_table)
+            annot_behaviors.update({name : beh_array})
+
+        annotations.update({key : annot_behaviors})
+    annotations = _clean_annotations(annotations)
+    start_time = annot_sheet.start_frame() + offset
+    end_time   = annot_sheet.end_frame() + offset
+    return annotations, start_time, end_time, sample_rate
+
+def load_multiple_annotations(fnames):
+    """
+    Generates a single dictionary given multiple .annot files.
+    """
+    if not isinstance(fnames, list):
+        raise TypeError(f'Expected list[str], got {type(fnames)} instead.')
+    if not fnames:
+        raise ValueError('No file names passed in.')
+    if len(fnames) == 1:
+        return load_annot_sheet_txt(fnames[0])
+    head_annot, head_start_frame, head_end_frame, sample_rate = load_annot_sheet_txt(fnames[0])
+    end_frame = head_end_frame
+    for fname in fnames[1:]:
+        curr_annot, _, curr_end_frame, _ = load_annot_sheet_txt(fname, end_frame)
+        end_frame = curr_end_frame
+        for channel in curr_annot.keys():
+            if channel not in head_annot:
+                channel_dict = {}
+                head_annot.update({channel : channel_dict})
+            for behavior in curr_annot[channel].keys():
+                curr_behavior_bout_array = curr_annot[channel][behavior]
+                if channel in head_annot and behavior in head_annot[channel]:
+                    new_bout_array = np.vstack((head_annot[channel][behavior],
+                                                curr_behavior_bout_array))
+                else:
+                    new_bout_array = curr_behavior_bout_array
+                head_annot[channel].update({behavior : new_bout_array})
+    return head_annot, head_start_frame, end_frame, sample_rate
+
+def load_annot_sheet_txt_io(uploaded_file, offset = 0):
+    """
+    Generated a dictionary for retrieving the beginning and end frames of behaviors from
+    an .annot file.
+
+    Note that 0:00:00 is frame 1
+
+    Args:
+        fname       - the path to the .annot file to be read (must be Caltech format)\n
+        offset      - a value which offsets the start and end frame of each bout in
+                      the sheet, as well as the absolute start and end frame of the file.
+                      This value is optional, and is set to 0 by default
+
+    Returns:
+        annotations - dictionary of beginning and end frames for behaviors\n
+        start_time  - the frame the movie started at (0:00:00 is 1)\n
+        end_time    - the frame the movie ended at (0:00:00 is 1)\n
+        sample_rate - the sample rate reported within the file
+    """
+    # from bento for python
+    behaviors   = Behaviors()
+    annot_sheet = Annotations(behaviors)
+    
+    annot_sheet.read_io(uploaded_file)
+
+    sample_rate = annot_sheet.sample_rate()
+    annotations = OrderedDict()
+    for key in annot_sheet.channel_names():
+        annot_behaviors = OrderedDict()
+        bout_names      = set()
+        for bout in annot_sheet.channel(key): #._bouts_by_start:
+            bout_names.add(bout.name())
+        for name in bout_names:
+            annot_behaviors.update({name : []})
+        for bout in annot_sheet.channel(key): #._bouts_by_start:
+            start_frame = bout.start().frames + offset
+            end_frame   = bout.end().frames + offset
+            bout_frames = [start_frame, end_frame]
+            curr_table  = annot_behaviors.get(bout.name())
+            new_table   = curr_table.append(bout_frames)
+            annot_behaviors.update({bout.name : new_table})
+        for name in bout_names:
+            curr_table  = annot_behaviors.get(name)
+            beh_array   = np.array(curr_table)
+            annot_behaviors.update({name : beh_array})
+
+        annotations.update({key : annot_behaviors})
+    annotations = _clean_annotations(annotations)
+    start_time = annot_sheet.start_frame() + offset
+    end_time   = annot_sheet.end_frame() + offset
+    return annotations, start_time, end_time, sample_rate
+
+def load_multiple_annotations_io(uploaded_files):
+    """
+    Generates a single dictionary given multiple .annot files.
+    """
+    if not isinstance(uploaded_files, list):
+        raise TypeError(f'Expected list, got {type(uploaded_files)} instead.')
+    if not uploaded_files:
+        raise ValueError('No file names passed in.')
+    if len(uploaded_files) == 1:
+        return load_annot_sheet_txt_io(uploaded_files[0])
+    head_annot, head_start_frame, head_end_frame, sample_rate = load_annot_sheet_txt_io(uploaded_files[0])
+    end_frame = head_end_frame
+    for uploaded_file in uploaded_files[1:]:
+        curr_annot, _, curr_end_frame, _ = load_annot_sheet_txt_io(uploaded_file, end_frame)
+        end_frame = curr_end_frame
+        for channel in curr_annot.keys():
+            if channel not in head_annot:
+                channel_dict = {}
+                head_annot.update({channel : channel_dict})
+            for behavior in curr_annot[channel].keys():
+                curr_behavior_bout_array = curr_annot[channel][behavior]
+                if channel in head_annot and behavior in head_annot[channel]:
+                    new_bout_array = np.vstack((head_annot[channel][behavior],
+                                                curr_behavior_bout_array))
+                else:
+                    new_bout_array = curr_behavior_bout_array
+                head_annot[channel].update({behavior : new_bout_array})
+    return head_annot, head_start_frame, end_frame, sample_rate

utils/data_processing.py

@@ -0,0 +1,384 @@
+"""Provides the standard data processing functions performed on CNMFe and annotation data"""
+import numpy as np
+from scipy.signal import correlate
+from scipy.stats import zscore
+
+def smooth(data: np.ndarray, window_size=5):
+    """
+    Returns a smoothed version of response data using a moving average filter.
+
+    Parameters:
+    ----------
+    data : np.ndarray
+        A numpy 1-D array containing data to be smoothed.
+    window_size : int 
+        Number of data points for calculating the smoothed value. If an even number is
+        passed in, window_size is autmoatically reduced by 1.
+
+    Returns:
+    --------
+    smooth_data : np.ndarray
+        Smoothed data, returned as a 1-D array of the same size as ``data``.
+    
+    Notes:
+    ------
+    Implements MATLAB's smooth function.
+    """
+    if window_size == 0:
+        raise ValueError('window_size can not be 0.')
+    if window_size == 1:
+        return data
+    if window_size > data.size:
+        window_size = data.size
+    if window_size%2 == 0:
+        window_size = window_size - 1
+    outside_valid_window_size = int((window_size-1)/2)
+    start = np.array([np.sum(data[0:(2*k+1)]/(2*k+1)) for k in range(outside_valid_window_size)])
+    end   = np.array([np.sum(data[-(2*k+1):]/(2*k+1)) for k in range(outside_valid_window_size)])[::-1]
+    smoothed_data = np.convolve(data,np.ones(window_size,dtype=int),'valid')/window_size
+    return np.hstack((start,smoothed_data,end))
+
+def corr(x: np.ndarray, y: np.ndarray):
+    """
+    Returns a matrix of the pairwise correlation coefficient between each pair of columns
+    in the input matrices x and y.
+
+    Parameters:
+    -----------
+    x : np.ndarray
+        Input matrix, specified as an n x k_1 matrix. Its rows correspond to
+        observations, and the columns correspond to variables.
+    y : np.ndarray
+        Input matrix, specified as an n x k_2 matrix. Its rows correspond to
+        observations, and the columns correspond to variables.
+
+    Returns:
+    --------
+    rho - Pairwise linear correlation coefficient, returned as a matrix.
+    
+    Notes:
+    ------
+    Implements MATLAB's corr function.
+    """
+    return np.corrcoef(x,y)[0][1]
+
+def autocorr(x:np.ndarray,
+             max_lags=10):
+    """
+    Returns the correlations and associated lags of the univariate time series x.
+
+    Parameters:
+    -----------
+    x : np.ndarray
+        Observed univariate time series.
+    max_lags : int 
+        Number of lags, specified as a positive integer.
+
+    Returns:
+    acf : np.ndarray
+        Correlations, returned as a numeric vector of length ``max_lags`` + 1.
+    lags : np.ndarray
+        Autocorrelation lags.
+
+    Notes:
+    ------
+    Modified version of matplotlib's acorr function.
+    """
+    Nx = len(x)
+
+    correls = correlate(x, x, mode="full")
+    correls = correls / np.dot(x, x)
+
+    if max_lags is None:
+        max_lags = Nx - 1
+
+    if max_lags >= Nx or max_lags < 1:
+        raise ValueError('maxlags must be None or strictly '
+                            'positive < %d' % Nx)
+
+    lags = np.arange(-max_lags, max_lags + 1)
+    acf = correls[Nx - 1 - max_lags:Nx + max_lags]
+
+    return acf, lags
+
+def convert_to_rast(behavior_ts, time_max):
+    """
+    Converts a list of behavior time stamps to a one-hot vector where 0 indicates no
+    presence of the given behavior, and 1 indicates presence of it.
+
+    Args:
+        behavior_ts   - a list of time stamps (start and end) for a particular behavior\n
+        time_max      - the length in frames of the vector
+    
+    Returns:
+        behavior_rast - a one-hot vector
+    """
+    behavior_rast = np.zeros(time_max)
+    for time_stamps in behavior_ts:
+        start = int(round(time_stamps[0]))
+        end   = int(round(time_stamps[1] + 1))
+        if start > time_max:
+            break
+        if end > time_max:
+            end = time_max
+        np.put(behavior_rast,range(start,end),np.ones(end-start))
+    return behavior_rast
+
+def convert_to_raster(bouts: list,
+                      neural_activity_sr: float,
+                      observation_sr: float,
+                      max_frame: int):
+    """
+    Converts bouts into a behavior raster, a one hot encoding of a behavior describing
+    when it is active.
+
+    It is often the case that the start and stop timestamps found in ``bouts`` are
+    collected at a different sample rate than ``neural_activity``, which are often what
+    behavior rasters align to. In order to align the two, a ratio between the sample
+    rates of ``neural_activity`` and the bouts of behavior, which are observations,
+    is calculated and then multiplied to the timestamps.
+
+    Parameters:
+    -----------
+    bouts : np.ndarray
+        An array where each element is a pair of integers where the first integer denotes
+        the beginning of a bout of behavior, and the second integer denotes the end of
+        the bout.
+    neural_activity_sr : float
+        Sample rate of ``neural_activity``.
+    observation_sr : float
+        Sample rate for the ``bouts`` used.
+    max_frame : int
+        The length of the behavior raster, often set to the number of frames of
+        ``neural_activity``.
+    
+    Returns:
+    --------
+    behavior_raster : np.ndarray
+        A raster (a one hot encoding) of a behavior, describing when it is active.
+    """
+    sr_ratio             = neural_activity_sr/observation_sr
+    behavior_ts_adjusted = bouts*sr_ratio
+    behavior_raster      = np.zeros(max_frame)
+    for time_stamps in behavior_ts_adjusted:
+        start = int(round(time_stamps[0]))
+        end   = int(round(time_stamps[1] + 1))
+        if start > max_frame:
+            break
+        if end > max_frame:
+            end = max_frame
+        np.put(behavior_raster,range(start,end),np.ones(end-start))
+    return behavior_raster
+
+def convert_to_bouts(behavior_raster: np.ndarray):
+    """
+    Converts a behavior raster into behavior bouts, an array where each element is a
+    pair of timestamps (int) where the first timestamp denotes the beginning of a bout of
+    behavior, and the second timestamp denotes the end of the bout.
+
+    Parameters:
+    -----------
+    behavior_raster : np.ndarray
+        A raster (a one hot encoding) of a behavior, describing when it is active.
+
+    Returns:
+    --------
+    bouts : np.ndarray
+        An array where each element is a pair of timestamps (int) where the first
+        timestamp denotes the beginning of a bout of behavior, and the second timestamp
+        denotes the end of the bout.
+    """
+    dt = behavior_raster[1:] - behavior_raster[:-1]
+    start = np.where(dt==1)[0] + 1
+    stop  = np.where(dt==-1)[0]
+    if behavior_raster[0]:
+        start = np.concatenate((np.array([0]),start))
+    if behavior_raster[-1]:
+        stop = np.concatenate((stop,[behavior_raster.size]))
+    bouts = np.hstack((np.reshape(start,(len(start),1)),
+                       np.reshape(stop,(len(stop),1))))
+    return bouts
+
+def merge_rasters_down(behavior_raster_array: np.ndarray)-> np.ndarray:
+    """
+    For a behavior raster, merges down all rasters to one array in such a way that no
+    two behaviors are occuring at the same time.
+
+    It determines which behavior should remain 'on top' by determening which behavior
+    has the least amount of active frames.
+
+    This method should only be used on behavior rasters where all behaviors come from a
+    single channel.
+
+    Parameters:
+    -----------
+    behavior_raster_array : np.ndarray
+        An array where each row is a behavior raster, a one hot encoding of behaviors,
+        describing when that behavior is active. Each row of this array must use a
+        different value to indicate that a behavior is active (for example, if one
+        row uses 1s, another row must not use 1 as well).
+    
+    Returns:
+    --------
+    single_track : np.ndarray
+        An array which is the length of a behavior raster in ``behavior_raster_array``,
+        where each entry is either 0 indicating that no behavior is active, or a value
+        indicating that a specific behavior is active.
+    """
+    # single track
+    single_track = np.zeros((1,behavior_raster_array.shape[1]))
+
+    # determine order to insert row values
+    num_active_frames = [np.sum(np.where(row > 0, 1, 0)) for row in behavior_raster_array]
+
+    for i in range(behavior_raster_array.shape[0]):
+        max_i = np.argmax(num_active_frames)
+        num_active_frames[max_i] = -1
+
+        unique_values = np.unique(behavior_raster_array[max_i])
+        if len(unique_values) > 1: value = unique_values[1]
+        else: value = 0
+        active_inds = np.where(behavior_raster_array[max_i] == value)[0]
+
+        single_track[:,active_inds] = value
+    return single_track
+
+def separate_tracks(single_track: np.ndarray,
+                    behavior_values: list):
+    """
+    For a single track, separates each unique value (except for 0) into its own raster
+    within a 2-D array.
+
+    Parameters:
+    -----------
+    single_track : np.ndarray
+        An array which is the length of a behavior raster in ``behavior_raster_array``,
+        where each entry is either 0 indicating that no behavior is active, or a value
+        indicating that a specific behavior is active.
+    behavior_values : list
+        A list of values corresponding to the specific behaviors within ``single_track``.
+    
+    Returns:
+    --------
+    behavior_raster_array : np.ndarray
+        An array where each row is a behavior raster, a one hot encoding of behaviors,
+        describing when that behavior is active.
+    """
+    if len(behavior_values) < np.unique(single_track).size - 1:
+        raise KeyError("There are not sufficient values within ``behavior_values`` to "
+                       "accomodate those present in ``single_track``.")
+    tracks = []
+    for value in behavior_values:
+        tracks.append(np.where(single_track == value, value, 0))
+    return np.vstack(tracks)
+
+def config_neural_activity(config: dict, neural_activity: np.ndarray):
+    """
+    Configures `neural_activity` according to parameters set in config.
+
+    Parameters:
+    -----------
+    config : dict
+        A dictionary which specifies the following parameters: 'smooth_window',
+        'baseline_frame', and 'zscore_method'. 'zscore_method' is one of "All Data",
+        "Baseline", or "No Z-Score".
+    neural_activity : np.ndarray
+        Neural activity being used.
+    
+    Returns:
+    --------
+    mod_neural_activity : np.ndarray
+        Modified `neural_activity`, accodring to `config`.
+    """
+    smooth_window  = config['smooth_window']
+    zscore_method  = config['zscore_method']
+    baseline_frame = config['baseline_frame']
+
+    # smooth
+    if len(neural_activity.shape) > 1:
+        neural_data_smooth = np.zeros(neural_activity.shape)
+        for i in range(neural_activity.shape[0]):
+            neural_data_smooth[i] = smooth(neural_activity[i], int(smooth_window))
+        mod_neural_activity = neural_data_smooth
+    else:
+        mod_neural_activity = smooth(neural_activity, int(smooth_window))
+
+    # z-score
+    if zscore_method == 'Baseline' and (not baseline_frame is None or baseline_frame == 0):
+        if len(neural_activity.shape)> 1:
+            mean = mod_neural_activity[:,:baseline_frame].mean(axis=1,keepdims=True)
+            std  = mod_neural_activity[:,:baseline_frame].std(axis=1,keepdims=True)
+        else:
+            mean = mod_neural_activity[:baseline_frame].mean()
+            std  = mod_neural_activity[:baseline_frame].std()
+        mod_neural_activity = (mod_neural_activity - mean) / std
+    elif zscore_method == 'No Z-Score':
+        mod_neural_activity = mod_neural_activity
+    else:
+        if len(neural_activity.shape) > 1:
+            mod_neural_activity = zscore(mod_neural_activity,axis=1)
+        else:
+            mod_neural_activity = zscore(mod_neural_activity)
+    return mod_neural_activity
+
+def compress_annotations(annot: dict, downsample_rate: int, max_frame: int)-> dict:
+    """
+    Takes in an annotation dictionary and creates a single raster per channel, where the
+    raster contains the behaviors from their respective channel.
+
+    annot : dict
+        Dictionary of beginning and end frames for behaviors.
+    downsample_rate : int
+        The rate at which samples should be taken. Divides bout timing (in frames) by
+        value.
+    max_frame : int
+        The last frame for annotations from `annot`.
+    """
+    annot_single_track = {}
+    channel_behavior_map = {}
+    for channel in annot:
+        channel_rasters = []
+        behavior_map = {}
+        behavior_map.update({0: 'None'})
+        for i, behavior in enumerate(annot[channel]):
+            bouts = annot[channel][behavior]
+            raster = convert_to_raster(bouts, 1, downsample_rate, max_frame)
+            channel_rasters.append(raster*(i+1))
+            behavior_map.update({(i+1) : behavior})
+        channel_raster = merge_rasters_down(np.array(channel_rasters))[0]
+        annot_single_track.update({channel : channel_raster})
+        channel_behavior_map.update({channel : behavior_map})
+    return annot_single_track, channel_behavior_map
+
+def compress_compressed_annotations(annot_single_track: dict,
+                                    channel_behavior_map: dict,
+                                    max_frame: int):
+    """
+    Further compresses the results from `compress_annotations` to get a single array
+    where each entry is a list of the behaviors present at that frame across all channels.
+    """
+    labels = []
+    for frame in range(max_frame):
+        labels_at_frame = []
+        for channel in annot_single_track:
+            channel_raster = annot_single_track[channel]
+            behavior_map = channel_behavior_map[channel]
+            behavior_value = int(channel_raster[frame])
+            behavior_label = behavior_map.get(behavior_value)
+            labels_at_frame.append(behavior_label)
+        labels.append('||'.join(labels_at_frame))
+    return labels
+
+def generate_label_array(annot: dict,
+                         downsample_rate: int,
+                         max_frame: int)-> list[str]:
+    """
+    Generates an array of lists of labels, where each entry is a video frame, and the
+    labels come from each channel in `annot`.
+    """
+    annot_single_track,\
+    channel_behavior_map = compress_annotations(annot, downsample_rate, max_frame)
+    labels = compress_compressed_annotations(annot_single_track,
+                                             channel_behavior_map,
+                                             max_frame)
+    return labels

utils/mp4Io.py

@@ -0,0 +1,60 @@
+import cv2
+import numpy as np
+import os
+from qtpy.QtGui import QImage, QPixmap
+
+class mp4Io_reader():
+	def __init__(self, filename, info=[]):
+		self.filename = filename
+		self.file = cv2.VideoCapture(filename)
+
+		if self.file.isOpened()==False:
+			print("Error in opening video file.")
+
+		self.header={}
+		if info==[]:
+			self.readHeader()
+
+	def readHeader(self):
+
+		self.header = {
+			'width': int(self.file.get(cv2.CAP_PROP_FRAME_WIDTH)),
+			'height': int(self.file.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+			'fps': self.file.get(cv2.CAP_PROP_FPS),
+			'numFrames': int(self.file.get(cv2.CAP_PROP_FRAME_COUNT))
+		}
+
+	def seek(self, index):
+
+		self.file.set(cv2.CAP_PROP_POS_FRAMES, index)
+
+	def getTs(self,n=None):
+		if n==None:
+			n = self.header['numFrames']
+
+		ts = np.zeros(n+1)
+		for i in np.arange(1,n+1):
+			self.seek(i)
+			self.file.read()
+			ts[i] = self.file.get(cv2.CAP_PROP_POS_MSEC)/1000.
+
+		self.ts = ts[1:]
+		return self.ts
+
+	def getFrame(self, index, decode=True):
+
+		self.seek(index)
+		ret, frame = self.file.read()
+
+		ts = self.file.get(cv2.CAP_PROP_POS_MSEC)/1000.
+		return frame, ts
+
+	def getFrameAsQPixmap(self, index, decode=True):
+		image, _ = self.getFrame(index, decode)
+		h, w, ch = image.shape
+		bytes_per_line = ch * w
+		convert_to_Qt_format = QImage(image.data, w, h, bytes_per_line, QImage.Format_BGR888)
+		return QPixmap.fromImage(convert_to_Qt_format)
+
+	def close(self):
+		self.file.release()

utils/seqIo.py

@@ -0,0 +1,1189 @@
+import os, sys
+import numpy as np
+import PIL
+from PIL import Image
+import io
+from datetime import datetime, timedelta,date
+import time
+from matplotlib.dates import date2num, num2date
+# import colour_demosaicing
+import skvideo.io
+import re
+import pickle
+import cv2
+import progressbar as pb
+
+# Create interface sr for reading seq files.
+#   sr = seqIo_reader( fName )
+# Create interface sw for writing seq files.
+#   sw = seqIo_Writer( fName, header )
+# Crop sub-sequence from seq file.
+#   seqIo_crop( fName, 'crop', tName, frames )
+# Extract images from seq file to target directory or array.
+#   Is = seqIo_toImgs( fName, tDir=[],skip=1,f0=0,f1=np.inf,ext='' )
+# Create seq file from an array or directory of images or from an AVI file. DONE
+#   seqIo_frImgs( fName, fName,header,aviName=[],Is=[],sDir=[],name='I',ndig=5,f0=0,f1=1e6 )
+# Convert seq file by applying imgFun(I) to each frame I.
+#   seqIo( fName, 'convert', tName, imgFun, varargin )
+# Replace header of seq file with provided info.
+#   seqIo( fName, 'newHeader', info )
+# Create interface sr for reading dual seq files.
+#   sr = seqIo( fNames, 'readerDual', [cache] )
+
+FRAME_FORMAT_RAW_GRAY = 100 #RAW
+FRAME_FORMAT_RAW_COLOR = 200 #RAW
+FRAME_FORMAT_JPEG_GRAY = 102 #JPG
+FRAME_FORMAT_JPEG_COLOR = 201 #JPG
+FRAME_FORMAT_MONOB = 101 #BRGB8
+FRAME_FORMAT_MONOB_JPEG = 103 #JBRGB
+FRAME_FORMAT_PNG_GRAY = 0x001 #PNG
+FRAME_FORMAT_PNG_COLOR = 0x002 #PNG
+
+#matlab equivalent fread
+def fread(fid, nelements, dtype):
+
+    """Equivalent to Matlab fread function"""
+
+    if dtype is np.str_:
+        dt = np.uint8  # WARNING: assuming 8-bit ASCII for np.str!
+    else:
+        dt = dtype
+
+    data_array = np.fromfile(fid, dt, nelements)
+    if data_array.size == 1:
+        data_array = data_array[0]
+    return data_array
+
+def fwrite(fid,a,dtype=np.str_):
+    # assuming 8but ASCII for string
+    if dtype is np.str_:
+        dt = np.uint8  # WARNING: assuming 8-bit ASCII for np.str!
+    else:
+        dt = dtype
+    if isinstance(a,np.ndarray):
+        data_array = a.astype(dt)
+    else:
+        data_array = np.array(a).astype(dt)
+    data_array.tofile(fid)
+
+def tsSync(video_path, srTop, srFront):
+    #tsSync
+    name = video_path.split('/')[-1]
+    srDepth=seqIo_reader(video_path + '/' + name + '_DepGr_Raw.seq')
+
+    #read timestamp of individual frames
+    tsTop = srTop.getTs()
+    tsFront = srFront.getTs()
+    tsDepth = srDepth.getTs()
+
+    #check the version of the videos
+    videoDateStr = re.search('[0-9]+_[0-9]+-[0-9]+-[0-9]+',name).group(0)
+    videoDateNum =date2num(datetime.strptime(videoDateStr ,'%Y%m%d_%H-%M-%S'))
+    videoRefNum = date2num(datetime.strptime('20150401_00-00-00','%Y%m%d_%H-%M-%S'))
+    seqV = 1 if videoDateNum < videoRefNum else 2
+
+    ##correlate timestamps from one view to another
+    mapTs = {}
+    if seqV ==1:
+        for f in range(len(tsDepth)):
+            if tsDepth[f] - np.floor(tsDepth[f])>=.5: # Santiago's bug acquisistion software
+                tsDepth[f]-= 1
+        tsDepth-= 0.03 # substract the systematic timeshift
+
+        #load front and top view and convert from UTC to PST
+        hourShift = np.round((tsDepth[0]-tsTop[0])/3600.)*3600
+        timeShift = tsDepth[0] - tsTop[0] - .066
+        # print('tsDepth[0] - tsTop[0] = timeShift: %s sec'% str(timeShift - hourShift))
+        tsTop +=  timeShift
+        tsFront += timeShift
+
+        srTop.ts = tsTop
+        srFront.ts = tsFront
+        srDepth.ts = tsDepth
+
+        # Convert timestamps from left to right
+        mapTs['T2F'] = transformTs(tsTop, tsFront)
+        mapTs['F2T'] = transformTs(tsFront, tsTop)
+        mapTs['T2D'] = transformTs(tsTop, tsDepth)
+        mapTs['D2T'] = transformTs(tsDepth, tsTop)
+        mapTs['F2D'] = transformTs(tsFront, tsDepth)
+        mapTs['D2F'] = transformTs(tsDepth, tsFront)
+    else:
+        T = len(tsTop)
+        F = len(tsFront)
+        D = len(tsDepth)
+        mapTs['T2F'] = resizeTs(T,F)
+        mapTs['F2T'] = resizeTs(F,T)
+        mapTs['T2D'] = resizeTs(T,D)
+        mapTs['D2T'] = resizeTs(D,T)
+        mapTs['F2D'] = resizeTs(F,D)
+        mapTs['D2F'] = resizeTs(D,F)
+
+    #display the time in string format
+    # fTp=5
+    # fFr = mapTs['T2F'][fTp]
+    # fDp = mapTs['T2D'][fTp]
+    # tF = ts2str(tsFront[fFr])
+    # tT = ts2str(tsTop[fTp])
+    # tD = ts2str(tsDepth[fDp])
+    #
+    # print('mapping to tTop:')
+    # print('tTop   = ' + str(fTp) + ' - ' + tT)
+    # print('tFront = ' + str(fFr) + ' - ' + tF)
+    # print('tDepth = ' + str(fDp) + ' - ' + tD)
+
+    # fDp = np.round(len(tsDepth)/1.5).astype(int)
+    # fTp = mapTs['D2T'][fDp]
+    # tT = ts2str(tsTop[fTp])
+    # tD = ts2str(tsDepth[fDp])
+    # print('mapping to tTop:')
+    # print('tTop   = ' + str(fTp) + '. ' + tT)
+    # print('tDepth = ' + str(fDp) + '. ' + tD)
+
+    return mapTs,srTop,srFront
+
+def transformTs(ts1, ts2):
+    # map ts1 to ts2 where ts2 is the reference
+    rankTs = np.zeros((len(ts1),4))
+    for f in range(len(ts1)):
+        tsDiff = ts2-ts1[f]
+        tsRank = np.sort(abs(tsDiff))
+        ind = np.argsort(abs(tsDiff))
+        rankTs[f,:] = [ind[0]+1,ind[1]+1,tsRank[0],tsRank[1]]
+
+    mapTs = np.round(smooth(rankTs[:,0],7)).astype(int)
+    mapTs = np.round(smooth(mapTs,7)).astype(int)
+    return mapTs
+
+def smooth(a, WSZ):
+    out0 = np.convolve(a, np.ones(WSZ, dtype=int), 'valid') / WSZ
+    r = np.arange(1, WSZ - 1, 2)
+    start = np.cumsum(a[:WSZ - 1])[::2] / r
+    stop = (np.cumsum(a[:-WSZ:-1])[::2] / r)[::-1]
+    return np.concatenate((start, out0, stop))
+
+def resizeTs(t1, t2):
+    if t1>t2:
+        mapTs = np.hstack((np.array(range(t2))+1,np.ones((t1-t2),int)*t2))
+    else:
+        mapTs = np.array(range(t1))+1
+
+    return mapTs
+
+def ts2str(ts):
+    t = ts / 86400. + date.toordinal(date(1971, 1, 2))
+    # datetime.fromtimestamp(t)
+    str_time = (datetime.fromordinal(int(t)) + timedelta(days=t % 1) - timedelta(days=366)).strftime(
+        "%Y-%m-%d %H:%M:%S") + '.%03d' % np.round((ts - np.floor(ts)) * 1000)
+    return str_time
+
+def parse_ann(f_ann):
+    header = 'Caltech Behavior Annotator - Annotation File'
+    conf = 'Configuration file:'
+    fid = open(f_ann)
+    ann = fid.read().splitlines()
+    fid.close()
+    NFrames = []
+    # check the header
+    assert ann[0].rstrip() == header
+    assert ann[1].rstrip() == ''
+    assert ann[2].rstrip() == conf
+    # parse action list
+    l = 3
+    names = [None] * 1000
+    keys = [None] * 1000
+    types = []
+    bnds = []
+    k = -1
+
+    # get config keys and names
+    while True:
+        ann[l] = ann[l].rstrip()
+        if not isinstance(ann[l], str) or not ann[l]:
+            l += 1
+            break
+        values = ann[l].split()
+        k += 1
+        names[k] = values[0]
+        keys[k] = values[1]
+        l += 1
+    names = names[:k + 1]
+    keys = keys[:k + 1]
+
+    # read in each stream in turn until end of file
+    bnds0 = [None] * 10000
+    types0 = [None] * 10000
+    actions0 = [None] * 10000
+    nStrm1 = 0
+    while True:
+        ann[l] = ann[l].rstrip()
+        nStrm1 += 1
+        t = ann[l].split(":")
+        l += 1
+        ann[l] = ann[l].rstrip()
+        assert int(t[0][1]) == nStrm1
+        assert ann[l] == '-----------------------------'
+        l += 1
+        bnds1 = np.ones((10000, 2), dtype=int)
+        types1 = np.ones(10000, dtype=int) * -1
+        actions1 = [None] * 10000
+        k = 0
+        # start the annotations
+        while True:
+            ann[l] = ann[l].rstrip()
+            t = ann[l]
+            if not isinstance(t, str) or not t:
+                l += 1
+                break
+            t = ann[l].split()
+            type = [i for i in range(len(names)) if t[2] == names[i]]
+            type = type[0]
+            if type == None:
+                print('undefined behavior' + t[2])
+            if bnds1[k - 1, 1] != int(t[0]) - 1 and k > 0:
+                print('%d ~= %d' % (bnds1[k, 1], int(t[0]) - 1))
+            bnds1[k, :] = [int(t[0]), int(t[1])]
+            types1[k] = type
+            actions1[k] = names[type]
+            k += 1
+            l += 1
+            if l == len(ann):
+                break
+        if nStrm1 == 1:
+            nFrames = bnds1[k - 1, 1]
+        assert nFrames == bnds1[k - 1, 1]
+        bnds0[nStrm1 - 1] = bnds1[:k]
+        types0[nStrm1 - 1] = types1[:k]
+        actions0[nStrm1 - 1] = actions1[:k]
+        if l == len(ann):
+            break
+        while not ann[l]:
+            l += 1
+
+    bnds = bnds0[:nStrm1]
+    types = types0[:nStrm1]
+    actions = actions0[:nStrm1]
+
+    idx = 0
+    if len(actions[0]) < len(actions[1]):
+        idx = 1
+    type_frame = []
+    action_frame = []
+    len_bnd = []
+
+    for i in range(len(bnds[idx])):
+        numf = bnds[idx][i, 1] - bnds[idx][i, 0] + 1
+        len_bnd.append(numf)
+        action_frame.extend([actions[idx][i]] * numf)
+        type_frame.extend([types[idx][i]] * numf)
+
+    ann_dict = {
+        'keys': keys,
+        'behs': names,
+        'nstrm': nStrm1,
+        'nFrames': nFrames,
+        'behs_se': bnds,
+        'behs_dur': len_bnd,
+        'behs_bout': actions,
+        'behs_frame': action_frame
+    }
+
+    return ann_dict
+
+def parse_ann_dual(f_ann):
+    header = 'Caltech Behavior Annotator - Annotation File'
+    conf = 'Configuration file:'
+    fid = open(f_ann)
+    ann = fid.read().splitlines()
+    fid.close()
+    NFrames = []
+    # check the header
+    assert ann[0].rstrip() == header
+    assert ann[1].rstrip() == ''
+    assert ann[2].rstrip() == conf
+    # parse action list
+    l = 3
+    names = [None] * 1000
+    keys = [None] * 1000
+    types = []
+    bnds = []
+    k = -1
+
+    # get config keys and names
+    while True:
+        ann[l] = ann[l].rstrip()
+        if not isinstance(ann[l], str) or not ann[l]:
+            l += 1
+            break
+        values = ann[l].split()
+        k += 1
+        names[k] = values[0]
+        keys[k] = values[1]
+        l += 1
+    names = names[:k + 1]
+    keys = keys[:k + 1]
+
+    # read in each stream in turn until end of file
+    bnds0 = [None] * 10000
+    types0 = [None] * 10000
+    actions0 = [None] * 10000
+    nStrm1 = 0
+    while True:
+        ann[l] = ann[l].rstrip()
+        nStrm1 += 1
+        t = ann[l].split(":")
+        l += 1
+        ann[l] = ann[l].rstrip()
+        assert int(t[0][1]) == nStrm1
+        assert ann[l] == '-----------------------------'
+        l += 1
+        bnds1 = np.ones((10000, 2), dtype=int)
+        types1 = np.ones(10000, dtype=int) * -1
+        actions1 = [None] * 10000
+        k = 0
+        # start the annotations
+        while True:
+            ann[l] = ann[l].rstrip()
+            t = ann[l]
+            if not isinstance(t, str) or not t:
+                l += 1
+                break
+            t = ann[l].split()
+            type = [i for i in range(len(names)) if t[2] == names[i]]
+            type = type[0]
+            if type == None:
+                print('undefined behavior' + t[2])
+            if bnds1[k - 1, 1] != int(t[0]) - 1 and k > 0:
+                print('%d ~= %d' % (bnds1[k, 1], int(t[0]) - 1))
+            bnds1[k, :] = [int(t[0]), int(t[1])]
+            types1[k] = type
+            actions1[k] = names[type]
+            k += 1
+            l += 1
+            if l == len(ann):
+                break
+        if nStrm1 == 1:
+            nFrames = bnds1[k - 1, 1]
+        assert nFrames == bnds1[k - 1, 1]
+        bnds0[nStrm1 - 1] = bnds1[:k]
+        types0[nStrm1 - 1] = types1[:k]
+        actions0[nStrm1 - 1] = actions1[:k]
+        if l == len(ann):
+            break
+        while not ann[l]:
+            l += 1
+
+    bnds = bnds0[:nStrm1]
+    types = types0[:nStrm1]
+    actions = actions0[:nStrm1]
+
+    idx = 0
+    if len(actions[0]) < len(actions[1]):
+        idx = 1
+    type_frame = []
+    action_frame = []
+    len_bnd = []
+
+
+    for i in range(len(bnds[idx])):
+        numf = bnds[idx][i, 1] - bnds[idx][i, 0] + 1
+        len_bnd.append(numf)
+        action_frame.extend([actions[idx][i]] * numf)
+        type_frame.extend([types[idx][i]] * numf)
+
+
+    type_frame2 = []
+    action_frame2 = []
+    len_bnd2 = []
+    idx=1 if idx==0 else 0
+
+    for i in range(len(bnds[idx])):
+        numf = bnds[idx][i, 1] - bnds[idx][i, 0] + 1
+        len_bnd2.append(numf)
+        action_frame2.extend([actions[idx][i]] * numf)
+        type_frame2.extend([types[idx][i]] * numf)
+
+    ann_dict = {
+        'keys': keys,
+        'behs': names,
+        'nstrm': nStrm1,
+        'nFrames': nFrames,
+        'behs_se': bnds,
+        'behs_dur': len_bnd,
+        'behs_bout': actions,
+        'behs_frame': action_frame if 'interaction' not in action_frame else action_frame2,
+        'behs_frame2': action_frame2 if 'interaction' in action_frame2 else action_frame
+    }
+
+    return ann_dict
+
+def syncTopFront(f,num_frames,num_framesf):
+    return int(round(f / (num_framesf - 1) * (num_frames - 1))) if num_framesf > num_frames else int(round(f / (num_frames - 1) * (num_framesf - 1)))
+
+
+
+class seqIo_reader():
+    def __init__(self,fname,info=[],buildTable=True):
+        self.filename = fname
+        try:
+            self.file=open(fname,'rb')
+        except EnvironmentError as e:
+            print(os.strerror(e.errno))
+        self.header={}
+        self.seek_table=None
+        self.frames_read=-1
+        self.timestamp_length = 10
+        if info==[]:
+            self.readHeader()
+        else:
+            info.numFrames=0
+        if buildTable:
+            print("buildTable was True, so calling buildSeekTable()")
+            self.buildSeekTable(False)
+
+
+    def readHeader(self):
+        #make sure we do this at the beginning of the file
+        assert self.frames_read == -1, "Can only read header from beginning of file"
+        self.file.seek(0,0)
+        # pdb.set_trace()
+
+        # Read 1024 bytes (len of header)
+        tmp = fread(self.file,1024,np.uint8)
+        #check that the header is not all 0's
+        n=len(tmp)
+        if n<1024:raise ValueError('no header')
+        if all(tmp==0): raise ValueError('fully empty header')
+        self.file.seek(0,0)
+        #first 4 bytes stor 0XFEED next 24 store 'Norpix seq '
+        magic_number = fread(self.file,1,np.uint32)
+        name = fread(self.file,10,np.uint16)
+        name = ''.join(map(chr,name))
+        if not '{0:X}'.format(magic_number)=='FEED' or not name=='Norpix seq':raise ValueError('invalid header')
+        self.file.seek(4,1)
+        #next 8 bytes for version and header size (1024) then 512 for desc
+        version = int(fread(self.file,1,np.int32))
+        hsize =int(fread(self.file,1,np.uint32))
+        assert(hsize)==1024 ,"incorrect header size"
+        # d = self.file.read(512)
+        descr=fread(self.file,256,np.uint16)
+        # descr = ''.join(map(chr,descr))
+        # descr = ''.join(map(unichr,descr)).replace('\x00',' ')
+        descr = ''.join([chr(x) for x in descr]).replace('\x00',' ')
+        # descr = descr.encode('utf-8')
+        #read more info
+        tmp = fread(self.file,9,np.uint32)
+        assert tmp[7]==0, "incorrect origin"
+        fps = fread(self.file,1,np.float64)
+        codec = 'imageFormat' + '%03d'%tmp[5]
+        desc_format = fread(self.file,1,np.uint32)
+        padding = fread(self.file,428,np.uint8)
+        padding = ''.join(map(chr,padding))
+        #store info
+        self.header={'magicNumber':magic_number,
+                     'name':name,
+                     'seqVersion': version,
+                     'headerSize':hsize,
+                     'descr': descr,
+                     'width':int(tmp[0]),
+                     'height':int(tmp[1]),
+                     'imageBitDepth':int(tmp[2]),
+                     'imageBitDepthReal':int(tmp[3]),
+                     'imageSizeBytes':int(tmp[4]),
+                     'imageFormat':int(tmp[5]),
+                     'numFrames':int(tmp[6]),
+                     'origin':int(tmp[7]),
+                     'trueImageSize':int(tmp[8]),
+                     'fps':fps,
+                     'codec':codec,
+                     'descFormat':desc_format,
+                     'padding':padding,
+                     'nHiddenFinalFrames':0
+                     }
+        assert(self.header['imageBitDepthReal']==8)
+        # seek to end fo header
+        self.file.seek(432,1)
+        self.frames_read += 1
+
+        self.imageFormat = self.header['imageFormat']
+        if self.imageFormat in (100,200):   self.ext = 'raw'
+        elif self.imageFormat in (102,201): self.ext = 'jpg'
+        elif self.imageFormat in(0x001,0x002):  self.ext = 'png'
+        elif self.imageFormat == 101:       self.ext = 'brgb8'
+        elif self.imageFormat == 103:       self.ext = 'jbrgb'
+        else:                              raise ValueError('uknown format')
+
+        self.compressed = True if self.ext in ['jpg','jbrgb','png','brgb8'] else False
+        self.bit_depth = self.header['imageBitDepth']
+
+        # My code uses a timestamp_length of 10 bytes, old uses 8. Check if not 10
+        if self.bit_depth / 8 * (self.header['height'] * self.header['width']) + self.timestamp_length \
+                != self.header['trueImageSize']:
+            # If not 10, adjust to actual (likely 8) and print message
+            self.timestamp_length = int(self.header['trueImageSize'] \
+                                        - (self.bit_depth / 8 * (self.header['height'] * self.header['width'])))
+
+    def buildSeekTable(self,memoize=False):
+        """Build a seek table containing the offset and frame size for every frame in the video."""
+        print("in seqIo_reader.buildSeekTable()")
+        pickle_name = self.filename.strip(".seq") + ".seek"
+        if memoize:
+            if os.path.isfile(pickle_name):
+                self.seek_table = pickle.load(open(pickle_name, 'rb'))
+                return
+
+        # assert self.header['numFrames']>0
+        n=self.header['numFrames']
+        if n==0:n=1e7
+
+        seek_table = np.zeros((n)).astype(np.int64)
+        seek_table[0]=1024
+        extra = 8 # extra bytes after image data , 8 for ts then 0 or 8 empty
+        self.file.seek(1024,0)
+        #compressed case
+
+        if self.compressed:
+            i=1
+            while (True):
+                try:
+                    # size = fread(self.file,1,np.uint32)
+                    # offset = seek_table[i-1] + size +extra
+                    # seek_table[i]=offset
+                    # # seek_table[i-1,1]=size
+                    # self.file.seek(size-4+extra,1)
+
+                    size = fread(self.file, 1, np.uint32)
+                    offset = seek_table[i - 1] + size + extra
+                    # self.file.seek(size-4+extra,1)
+                    self.file.seek(offset, 0)
+                    if i == 1:
+                        if fread(self.file, 1, np.uint32) != 0:
+                            self.file.seek(-4, 1)
+                        else:
+                            extra += 8;
+                            offset += 8
+                            self.file.seek(offset, 0)
+
+                    seek_table[i] = offset
+                    # seek_table[i-1,1]=size
+                    i+=1
+                except Exception as e:
+                    break
+                    #most likely EOF
+        else:
+            #uncompressed case
+            assert (self.header['numFrames']>0)
+            frames = range(0, self.header["numFrames"])
+            offsets = [x * self.header["trueImageSize"] + 1024 for x in frames]
+            for i,offset in enumerate(offsets):
+                seek_table[i]=offset
+                # seek_table[i,1]=self.header["imageSize"]
+        if n==1e7:
+            n = np.minimum(n,i)
+            self.seek_table=seek_table[:n]
+            self.header['numFrames']=n
+        else:
+            self.seek_table=seek_table
+        if memoize:
+            pickle.dump(seek_table,open(pickle_name,'wb'))
+
+        #compute frame rate from timestamps as stored fps may be incorrect
+        # if n==1: return
+        self.getTs()
+        # ds = self.ts[1:100]-self.ts[:99]
+        # ds = ds[abs(ds-np.median(ds))<.005]
+        # if bool(np.prod(ds)): self.header['fps']=1/np.mean(ds)
+
+    def getTs(self, n=None):
+        if n==None: n=self.header['numFrames']
+        if self.compressed and self.seek_table is None:
+            self.buildSeekTable()
+
+        ts = np.zeros((n))
+        for i in range(n):
+            if not self.compressed: #uncompressed
+                self.file.seek(1024 + i*self.header['trueImageSize']+self.header['imageSizeBytes'],0)
+            else: #compressed
+                self.file.seek(self.seek_table[i],0)
+                self.file.seek(fread(self.file,1,np.uint32)-4,1)
+            # print(i)
+            ts[i]=fread(self.file,1,np.uint32)+fread(self.file,1,np.uint16)/1000.
+
+
+        self.ts=ts
+        return self.ts
+
+    def getFrame(self,index,decode=True):
+        #get frame image (I) and timestamp (ts) at which frame was recorded
+        nch = self.header['imageBitDepth']/8
+        if self.ext in ['raw','brgb8']: #read in an uncompressed image( assume imageBitDepthReal==8)
+            shape = (self.header['height'], self.header['width'])
+            self.file.seek(1024 + index*self.header['trueImageSize'],0)
+            I = fread(self.file,self.header['imageSizeBytes'],np.uint8)
+
+            if decode:
+                if nch==1:
+                    I=np.reshape(I,shape)
+                else:
+                    I=np.reshape(I,(shape,nch))
+                if nch==3:
+                    t=I[:,:,2]; I[:,:,2]=I[:,:,0]; I[:,:,1]=t
+                if self.ext=='brgb8':
+                    I= cv2.demosaicing(I, code=cv2.COLOR_BAYER_BG2BGR)
+                    # I= colour_demosaicing.demosaicing_CFA_Bayer_bilinear(I,'BGGR')
+
+        elif self.ext in ['jpg','jbrgb']:
+            self.file.seek(self.seek_table[index],0)
+            nBytes = fread(self.file,1,np.uint32)
+            data = fread(self.file,nBytes-4,np.uint8)
+            if decode:
+                I = PIL.Image.open(io.BytesIO(data))
+                if self.ext == 'jbrgb':
+                    I= cv2.demosaicing(I, code=cv2.COLOR_BAYER_BG2BGR)
+                    # I=colour_demosaicing.demosaicing_CFA_Bayer_bilinear(I,'BGGR')
+            else:
+                I = data
+
+        elif self.ext=='png':
+            self.file.seek(self.seek_table[index],0)
+            nBytes = fread(self.file,1,np.uint32)
+            I= fread(self.file,nBytes-4,np.uint8)
+            if decode:
+                I= np.array(I).transpose(range(I.shape,-1,-1))
+        else: assert(False)
+        ts = fread(self.file,1,np.uint32)+fread(self.file,1,np.uint16)/1000.
+        return np.array(I), ts
+
+    # Close the file
+    def close(self):
+        self.file.close()
+
+class seqIo_writer():
+    def __init__(self,filename,old_header):
+        self.file = open(filename,'wb')
+        self.file.seek(0,0)
+        self.header=old_header
+
+        #create space for header
+        fwrite(self.file,np.zeros(1024).astype(int),np.uint8)
+
+        assert(set(['width','height','fps','codec']).issubset(self.header.keys()))
+
+        codec = self.header['codec']
+        if   codec in ['monoraw', 'imageFormat100']:        self.frmt = 100;self.nch = 1;self.ext = 'raw'
+        elif codec in ['raw', 'imageFormat200']:            self.frmt = 200;self.nch = 3;self.ext = 'raw'
+        elif codec in ['monojpg', 'imageFormat102']:        self.frmt = 102;self.nch = 1;self.ext = 'jpg'
+        elif codec in ['jpg', 'imageFormat201']:            self.frmt = 201;self.nch = 3;self.ext = 'jpg'
+        elif codec in ['monopng', 'imageFormat001']:        self.frmt = 0x001;self.nch = 1;self.ext = 'png'
+        elif codec in ['png', 'imageFormat002']:            self.frmt = 0x002;self.nch = 3;self.ext = 'png'
+        else:                                               raise ValueError('unknown format')
+
+        self.header['imageFormat']=self.frmt
+        self.header['imageBitDepth']=8*self.nch
+        self.header['imageBitDepthReal']=8
+        nBytes = self.header['width']*self.header['height']*self.nch
+        self.header['imageSizeBytes']=nBytes
+        self.header['numFrames']=0
+        self.header['trueImageSize']=nBytes + 6 +512-np.mod(nBytes+6,512)
+
+    # Close the file
+    def close(self):
+        self.writeHeader()
+        self.file.close()
+
+    def writeHeader(self):
+        self.file.seek(0,0)
+        # first write 4 bytes to store 0XFEED, next 24 store 'Nrpix seq  '
+        fwrite(self.file,int('FEED',16),np.uint32)
+        name = np.array(['Norpix seq  ']).view(np.uint8)
+        fwrite(self.file,name, np.uint16)
+        # next 8 bytes for version (3) and header size (1024) then 512 for descr
+        fwrite(self.file,[3,1024],np.int32)
+        if not 'descr' in self.header.keys() or len(np.array([self.header['descr']]).view(np.uint8))>256: d = np.array(['No Description']).view(np.uint8)
+        else: d= np.array([self.header['descr']]).view(np.uint8)
+        d = np.concatenate((d[:np.minimum(256,len(d))],np.zeros(256-len(d)).astype(np.uint8)))
+        fwrite(self.file,d,np.uint16)
+        #write remaining info
+        vals= [self.header['width'],self.header['height'],self.header['imageBitDepth'],self.header['imageBitDepthReal'],
+               self.header['imageSizeBytes'],self.header['imageFormat'],self.header['numFrames'],0,self.header['trueImageSize']]
+        fwrite(self.file,vals,np.uint32)
+        #store frame rate nad pad with 0s
+        fwrite(self.file,self.header['fps'],np.float64)
+        fwrite(self.file,np.zeros(432),np.uint8)
+
+    def addFrame(self,I,ts=0,encode=1):
+        nCh = self.header['imageBitDepth']/8
+        ext = self.ext
+        c = self.header['numFrames']+1
+        if encode:
+            siz = [self.header['height'],self.header['width'],nCh]
+            assert(I.shape[0]==siz[0] and I.shape[1]==siz[1])
+            if len(I.shape)==3:
+                assert(I.shape[2]==siz[2] or I.shape[2]==self.nch)
+        if ext=='raw':
+            #write uncompressed image and assume imageBitDepthReal==8
+            if not encode : assert(I.size==self.header['imageSizeBytes'])
+            else:
+                if nCh==3: t=I[:,:,2]; I[:,:,2]=I[:,:,0];I[:,:,0]=t
+                if nCh==1: I=I.transpose()
+                else: I = np.transpose( np.expand_dims(I, axis=2), (2, 1, 0) )
+            # I= I.flat.view(np.uint8)
+            I= I.flat
+            fwrite(self.file,I,np.uint8)
+            pad = self.header['trueImageSize']-self.header['imageSizeBytes']-6
+        if ext =='jpg':
+            if encode:
+                #write red from to temporary jpg
+                cv2.imwrite('tmp.jpg',I, [int(cv2.IMWRITE_JPEG_QUALITY ),80])
+                # j=Image.fromarray(I.astype(np.uint8))
+                # j.save('tmp.jpg')
+                # I=Image.open('tmp.jpg')
+                fid  = open('tmp.jpg','r')
+                I = fid.read()
+                fid.close()
+                b=bytearray(I)
+                assert (b[0] == 255 and b[1] == 216 and b[-2] == 255 and b[-1] == 217); # JPG
+                os.remove('tmp.jpg')
+                I = np.array(list(b)).astype(np.uint8)
+            nbytes = len(I)+4
+            fwrite(self.file,nbytes,np.uint32)
+            # self.file.write(I)
+            fwrite(self.file,I,np.uint8)
+            pad = 10
+        if ts==0: ts = (c-1)/self.header['fps']
+        s = int(np.floor(ts))
+        ms = int(np.round(np.mod(ts,1)*1000))
+        fwrite(self.file,s,np.int32)
+        fwrite(self.file,ms,np.uint16)
+        self.header['numFrames']=c
+        if pad>0:
+            pad = np.zeros(pad).astype(np.uint8)
+            fwrite(self.file,pad,np.uint8)
+
+def seqIo_crop(fname, tname, frames):
+    """
+    Crop sub-sequence from seq file.
+
+    Frame indices are 0 indexed. frames need not be consecutive and can
+    contain duplicates. An index of -1 indicates a blank (all 0) frame. If
+    contiguous subset of frames is cropped timestamps are preserved.
+
+    USAGE
+     seqIo( fName, 'crop', tName, frames )
+
+    INPUTS
+     fName      - seq file name
+     tName      - cropped seq file name
+     frames     - frame indices (0 indexed)
+    """
+    if not isinstance(frames, np.ndarray): frames=np.array(frames)
+    sr = seqIo_reader(fname)
+    sw = seqIo_writer(tname,sr.header)
+    pad,_= sr.getFrame(0)
+    pad = np.zeros(pad.size).astype(np.uint8)
+    kp = frames>=0 & frames<sr.header['numFrames']
+    if not np.all(kp): frames = frames[kp]
+    print('%i out of bounds frames'% np.sum(~kp))
+    ordered = np.all(frames[1:]==frames[:-1]+1)
+    n= frames.size
+    k=0
+    for f in frames:
+        if f<0:
+            sw.addFrame(pad)
+            continue
+        I,ts = sr.getFrame(f)
+        k+=1
+        if ordered:
+            sw.addFrame(I,ts)
+        else:
+            sw.addFrame(I)
+    sr.close()
+    sw.close
+
+def seqIo_toImgs(fName, tDir=[], skip=1, f0=0, f1=np.inf, ext=''):
+    """
+    Extract images from seq file to target directory or array.
+
+    USAGE
+     Is = seqIo( fName, 'toImgs', [tDir], [skip], [f0], [f1], [ext] )
+
+    INPUTS
+     fName      - seq file name
+     tDir       - [] target directory (if empty extract images to array)
+     skip       - [1] skip between written frames
+     f0         - [0] first frame to write
+     f1         - [numFrames-1] last frame to write
+     ext        - [] optionally save as given type (slow, reconverts)
+
+    OUTPUTS
+     Is         - if isempty(tDir) outputs image array (else Is=[])
+    """
+    sr = seqIo_reader(fName)
+    f1 = np.minimum(f1,sr.header['numFrames']-1)
+    frames = range(f0,f1,skip)
+    n=len(frames)
+    k=0
+    #output images to array
+    if tDir==[]:
+        I,_=sr.getFrame(0)
+        d = I.shape
+        assert(len(d)==2 or len(d)==3)
+        try:
+            Is = np.zeros((I.shape+(n,))).astype(I.dtype)
+        except:
+            sr.close()
+            raise
+        for k in range(n):
+            I,ts = sr.getFrame(k)
+            if len(d)==2:
+                Is[:,:,k]=I
+            else:
+                Is[:,:,:,k]=I
+            print('saved %d' % k)
+
+        sr.close()
+    # output image directory
+    if not os.path.exists(tDir):os.makedirs(tDir)
+    if tDir.split('/')[-1]!='/':tDir+'/'
+    Is = np.array([])
+    for frame in frames:
+        f = tDir + 'I%05.' % (frame)
+        I, ts = sr.getFrame(frame)
+        if ext!='':
+            cv2.imwrite(f+ext,I)
+        else:
+            cv2.imwrite(f+sr.ext)
+        k+=1
+        print('saved %d' % frame)
+    sr.close()
+    return Is
+
+def seqIo_frImgs(fName, header=[], aviName=[], Is=[], sDir=[], name='I', ndig=5, f0=0, f1=1e6):
+    """
+    Create seq file from an array or directory of images or from an AVI file.
+
+    For info, if converting from array, only codec (e.g., 'jpg') and fps must
+    be specified while width and height and determined automatically. If
+    converting from AVI, fps is also determined automatically.
+
+    USAGE
+     seqIo( fName, 'frImgs', info, varargin )
+
+    INPUTS
+     fName      - seq file name
+     info       - defines codec, etc, see seqIo>writer
+     varargin   - additional params (struct or name/value pairs)
+      .aviName    - [] if specified create seq from avi file
+      .Is         - [] if specified create seq from image array
+      .sDir       - [] source directory
+      .skip       - [1] skip between frames
+      .name       - ['I'] base name of images
+      .nDigits    - [5] number of digits for filename index
+      .f0         - [0] first frame to read
+      .f1         - [10^6] last frame to read
+    """
+
+    if aviName!=[]: #avi movie exists
+        vc = cv2.VideoCapture(aviName)
+        if vc.isOpened(): rval = True
+        else:
+            rval = False
+            print('video not readable')
+            return
+        fps = vc.get(cv2.cv.CV_CAP_PROP_FPS)
+        NUM_FRAMES = int(vc.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
+        print(NUM_FRAMES)
+        IM_TOP_H = vc.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)
+        IM_TOP_W = vc.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)
+        header['width']=IM_TOP_W
+        header['height']=IM_TOP_H
+        header['fps']=fps
+
+        sw = seqIo_writer(fName,header)
+        print('creating seq from AVI')
+        # initialize timer
+        timer = pb.ProgressBar(widgets=['Converting ', pb.Percentage(), ' -- ',
+                                        pb.FormatLabel('Frame %(value)d'), '/',
+                                        pb.FormatLabel('%(max)d'), ' [', pb.Timer(), '] ',
+                                        pb.Bar(), ' (', pb.ETA(), ') '], maxval=NUM_FRAMES)
+        for f in range(NUM_FRAMES):
+            rval, im = vc.read()
+            if rval:
+                im= im.astype(np.uint8)
+            sw.addFrame(im)
+            timer.update(f)
+        sw.close()
+        timer.finish()
+    elif Is==[]:
+        assert(os.path.isdir(sDir))
+        sw = seqIo_writer(fName,header)
+        frmstr = '%s/%s%%0%ii.%s' % (sDir,name,ndig,header.ext)
+        for frame in range(f0,f1):
+            f = frmstr % frame
+            if not os.path.isfile(f):break
+            fid = open(f, 'r')
+            if fid<0: sw.close();  assert(False)
+            I = fid.read()
+            fid.close()
+            b = bytearray(I)
+            assert (b[0] == 255 and b[1] == 216 and b[-2] == 255 and b[-1] == 217);  # JPG
+            I = np.array(list(b)).astype(np.uint8)
+            sw.addFrame(I,0,0)
+        sw.close()
+        if frame==f0: print('No images found')
+    else:
+        nd = len(Is.shape)
+        if nd==2: nd=3
+        assert(nd<=4)
+        nFrm = Is.shape[nd-1]
+        header['height']=Is.shape[0]
+        header['width']=Is.shape[1]
+        sw =seqIo_writer(fName,header)
+        if nd==3:
+            for f in range(nFrm): sw.addFrame(Is[:,:,f])
+        if nd==4:
+            for f in range(nFrm): sw.addFrame(Is[:,:,:,f])
+        sw.close()
+
+def seqIo_convert(fName, tName, imgFun, info=[], skip=1, f0=0, f1=np.inf):
+    """
+    Convert seq file by applying imgFun(I) to each frame I.
+
+    USAGE
+     seqIo( fName, 'convert', tName, imgFun, varargin )
+
+    INPUTS
+     fName      - seq file name
+     tName      - converted seq file name
+     imgFun     - function to apply to each image
+     varargin   - additional params (struct or name/value pairs)
+      .info       - [] info for target seq file
+      .skip       - [1] skip between frames
+      .f0         - [0] first frame to read
+      .f1         - [inf] last frame to read
+    """
+    assert(fName!=tName)
+    sr = seqIo_reader(fName)
+    if info==[]: info=sr.header
+    n=sr.header['numFrames']
+    f1=np.minimum(f1,n-1)
+    I,ts=sr.getFrame(0)
+    I=imgFun(I)
+    info['width']=I.shape[1]
+    info['height']=I.shape[0]
+    sw =seqIo_writer(tName,info)
+    print('converting seq')
+    for frame in range(f0,f1,skip):
+        I, ts = sr.getFrame(frame)
+        I = imgFun(I)
+        if skip==1:
+            sw.addFrame(I,ts)
+        else:
+            sw.addFrameI
+    sw.close()
+    sr.close()
+
+def seqIo_newHeader(fName, info):
+    """
+    Replace header of seq file with provided info.
+
+    Can be used if the file fName has a corrupt header. Automatically tries
+    to compute number of frames in fName. No guarantees that it will work.
+
+    USAGE
+     seqIo( fName, 'newHeader', info )
+
+    INPUTS
+     fName      - seq file name
+     info       - info for target seq file
+    """
+    d, n = os.path.split(fName)
+    if d==[]:d='./'
+    tName=fName[:-4] + '_new'  + time.strftime("%d_%m_%Y") + fName[-4:]
+    sr = seqIo_reader(fName)
+    sw = seqIo_writer(tName,info)
+    n=sr.header['numFrames']
+    for f in range(n):
+        I,ts=sr.getFrame(f)
+        sw.addFrame(I,ts)
+    sr.close()
+    sw.close()
+
+class seqIo_dualReader():
+    """
+    seqIo_dualReader
+    Create interface sr for reading dual seq files.
+
+    Wrapper for two seq files of the same image dims and roughly the same
+    frame counts that are treated as a single reader object. getframe()
+    returns the concatentation of the two frames. For videos of different
+    frame counts, the first video serves as the "dominant" video and the
+    frame count of the second video is adjusted accordingly. Same general
+    usage as in reader, but the only supported operations are: close(),
+    getframe(), getinfo(), and seek().
+
+    USAGE
+     sr = seqIo( fNames, 'readerDual', [cache] )
+
+    INPUTS
+     fNames - two seq file names
+     cache  - [0] size of cache (see seqIo>reader)
+
+    OUTPUTS
+     sr     - interface for reading seq file
+    """
+    def __init__(self,file1,file2):
+        self.s1 = seqIo_reader(file1)
+        self.s2 = seqIo_reader(file2)
+        self.info  = self.s1.header
+        #set the display to be vertically align
+        self.info['height']=self.s1.header['height']+self.s2.header['height']
+        self.info['width']=np.maximum(self.s1.header['width'],self.s2.header['width'])
+
+        if self.s1.header['numFrames']!=self.s2.header['numFrames']:
+            print('Two videos files have different number of frames')
+            print('1st video has %d frames' % self.s1.header['numFrames'])
+            print('2nd video has %d frames' % self.s2.header['numFrames'])
+        print('first video %s is used as annotation refeence' % file1)
+
+    def getFrame(self):
+        I1,ts = self.s1.getFrame(0)
+        I2,_ = self.s2.getFrame(0)
+
+        w1 = I1.shape[1]
+        w2 = I2.shape[1]
+
+        if w1!=w2:
+            m=np.argmax(w1,w2)
+            if m==0:
+                wl = int(np.floor((w1-w2)/2.))
+                wr = w1-w2-wl
+                nd = len(I2.shape)
+                if nd==2:
+                    padl = np.zeros((I2.shape[0],wl)).astype(np.uint8)
+                    padr = np.zeros((I2.shape[0],wr)).astype(np.uint8)
+                else:
+                    padl = np.zeros((I2.shape[0],wl,I2.shape[2])).astype(np.uint8)
+                    padr = np.zeros((I2.shape[0],wr,I2.shape[2])).astype(np.uint8)
+                    I2 = np.concatenate((padl,I2,padr),axis=1)
+            else:
+                wl = int(np.floor((w2 - w1) / 2.))
+                wr = w2 - w1 - wl
+                nd = len(I2.shape)
+                if nd == 2:
+                    padl = np.zeros((I1.shape[0], wl)).astype(np.uint8)
+                    padr = np.zeros((I1.shape[0], wr)).astype(np.uint8)
+                else:
+                    padl = np.zeros((I1.shape[0], wl, I1.shape[2])).astype(np.uint8)
+                    padr = np.zeros((I1.shape[0], wr, I1.shape[2])).astype(np.uint8)
+                    I1 = np.concatenate((padl, I1, padr), axis=1)
+        I = np.hstack((I1,I2))
+        return I,ts
+
+class seqIo_extractor():
+    """
+    Create new seq files from top and fron view and syncronize them is not
+    path_vid: video path
+    vid_top: seq top video path and name
+    vid_front: seq front video path and name
+    s: start frame
+    e: end frame
+
+    """
+    def __init__(self,path_vid,vid_top,vid_front,s,e):
+        sr_top = seqIo_reader(path_vid+vid_top)
+        sr_front = seqIo_reader(path_vid+vid_front)
+        num_frames=sr_top.header['numFrames']
+        num_framesf=sr_front.header['numFrames']
+        name =os.path.dirname(video_top).split('/')[-1]
+
+        if not os.path.exists(pathvid + name + '_%06d_%06d' % (s, e)):
+            os.makedirs(pathvid + name + '_%06d_%06d' % (s, e))
+        newdir = pathvid + name + '_%06d_%06d' % (s, e)
+        video_out_top = newdir + '/' + name + '_%06d_%06d_Top_J85.seq' % (s, e)
+        video_out_front = newdir + '/' + name + '_%06d_%06d_Front_J85.seq' % (s, e)
+
+        sw_top = seqIo_writer(video_out_top, sr_top.header)
+        sw_front = seqIo_writer(video_out_front, sr_front.header)
+
+        for f in range(s - 1, e):
+            if num_framesf > num_frames:
+                I_top, ts = sr_top.getFrame(f2(f))
+                I_front, ts2 = sr_front.getFrame(f)
+            else:
+                I_top, ts = sr_top.getFrame(f)
+                I_front, ts2 = sr_front.getFrame(f2(f))
+            sw_top.addFrame(I_top, ts)
+            sw_front.addFrame(I_front, ts2)
+            print(f)
+        sw_top.close()
+        sw_front.close()
+
+    def f2(f):
+            return int(round(f / (num_framesf - 1) * (num_frames - 1))) if num_framesf > num_frames else int(round(f / (num_frames - 1) * (num_framesf - 1)))
+
+def seqIo_toVid(fName, ext='avi'):
+    """
+    seqIo_toVid
+    Create seq file to another common used format as avi or mp4.
+
+    USAGE
+     seqIo( fName, ext )
+
+    INPUTS
+     fName      - seq file name
+     ext        - video extension to convert to
+    """
+
+    assert fName[-3:]=='seq', 'Not a seq file'
+    sr = seqIo_reader(fName)
+    N  = sr.header['numFrames']
+    h = sr.header['height']
+    w = sr.header['width']
+    fps = sr.header['fps']
+
+    out = fName[:-3]+ext
+    sw = skvideo.io.FFmpegWriter(out)
+    # sw = cv2.VideoWriter(out, -1, fps, (w, h))
+    timer = pb.ProgressBar(widgets=['Converting ', pb.Percentage(), ' -- ',
+                                    pb.FormatLabel('Frame %(value)d'), '/',
+                                    pb.FormatLabel('%(max)d'), ' [', pb.Timer(), '] ',
+                                    pb.Bar(), ' (', pb.ETA(), ') '], maxval=N)
+
+    for f in range(N):
+        I, ts = sr.getFrame(f)
+        #sw.writeFrame(Image.fromarray(I))
+        sw.write(I)
+        timer.update(f)
+    timer.finish()
+    # cv2.destroyAllWindows()
+    # sw.release()
+    sw.close()
+    sr.close()
+    print(out + ' converted')
+
+
+
+
+# minimum header
+# header = {'width': IM_TOP_W,
+#           'height': IM_TOP_H,
+#           'fps': fps,
+#           'codec': 'imageFormat102'}
+# filename= '/media/cristina/MARS_data/mice_project/teresa/Mouse156_20161017_17-22-09/Mouse156_20161017_17-22-09_Top_J85.seq'
+# filename_out = filename[:-4] + '_new.seq'
+# reader = seqIo_reader(filename)
+# reader.header
+# Initialize a SEQ writer
+# writer = seqIo_writer(filename_out,reader.header)
+# I,ts = reader.getFrame(0)
+# writer.addFrame(I,ts)
+# for f in range(8):
+#     I,ts = reader.getFrame(f)
+#     print(writer.file.tell())
+#     writer.addFrame(I,ts)
+# writer.close()
+# reader.close()
+
+
+

utils/utils.py

@@ -0,0 +1,632 @@
+import os
+import io
+import pickle
+import copy
+from collections import Counter
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+import regex as re
+import numpy as np
+import pandas as pd
+from sklearn.manifold import TSNE
+from sklearn.svm import SVC
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, classification_report
+import torch
+from tqdm import tqdm
+from PIL import Image
+from transformers import AutoProcessor, AutoModel
+import streamlit as st
+from .data_loading import load_multiple_annotations, load_multiple_annotations_io
+from .data_processing import generate_label_array
+from .seqIo import seqIo_reader
+from .mp4Io import mp4Io_reader
+
+SLIP_MODEL_ID = "google/siglip-so400m-patch14-384"
+CLIP_MODEL_ID = "openai/clip-vit-base-patch32"
+
+def create_annot_fname_dict(annot_fnames: list[str])-> dict:
+    fs = re.compile(r'.*(_\d+)$')
+
+    unique_files = set()
+    for file in annot_fnames:
+        file_name = os.fsdecode(file)
+        base_name, _ = os.path.splitext(file_name)
+        if fs.match(base_name):
+            ind = len(fs.match(base_name).group(1))
+            unique_files.add(base_name[:-ind])
+        else:
+            unique_files.add(base_name)
+    
+    annot_fname_dict = {}
+    for unique_file in unique_files:
+        annot_fname_dict.update({unique_file: [file for file in annot_fnames if unique_file in file]})
+    return annot_fname_dict
+
+def create_annot_fname_dict_io(annot_fnames: list[str], annot_files: list)-> dict:
+    annot_file_dict = {}
+    for file in annot_files:
+        annot_file_dict.update({file.name : file})
+    fs = re.compile(r'.*(_\d+)$')
+
+    unique_files = set()
+    for file in annot_fnames:
+        file_name = os.fsdecode(file)
+        base_name, _ = os.path.splitext(file_name)
+        if fs.match(base_name):
+            ind = len(fs.match(base_name).group(1))
+            unique_files.add(base_name[:-ind])
+        else:
+            unique_files.add(base_name)
+    
+    annot_fname_dict = {}
+    for unique_file in unique_files:
+        annot_list = [file for file in annot_fnames if unique_file in file]
+        annot_list.sort()
+        annot_file_list = [annot_file_dict[annot_file_name] for annot_file_name in annot_list]
+        annot_fname_dict.update({unique_file: annot_file_list})
+    return annot_fname_dict
+
+def get_io_reader(uploaded_file):
+    assert uploaded_file.name[-3:]=='seq', 'Not a seq file'
+    with NamedTemporaryFile(suffix="seq", delete=False) as temp:
+        temp.write(uploaded_file.getvalue())
+        sr = seqIo_reader(temp.name)
+    return sr
+
+def load_slip_model(device):
+    return AutoModel.from_pretrained(SLIP_MODEL_ID).to(device)
+
+def load_slip_preprocessor():
+    return AutoProcessor.from_pretrained(SLIP_MODEL_ID)
+
+def load_clip_model(device):
+    return AutoModel.from_pretrained(CLIP_MODEL_ID).to(device)
+
+def load_clip_preprocessor():
+    return AutoProcessor.from_pretrained(CLIP_MODEL_ID)
+
+def encode_image(image, device, model, processor):
+    with torch.no_grad():
+        #convert_models_to_fp32(model)
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        image_features = model.get_image_features(**inputs)
+    return image_features.cpu().numpy().flatten()
+
+def generate_embeddings_stream(fnames : list[str],
+                        model = 'SLIP',
+                        downsample_rate = 4,
+                        save_csv = False)-> tuple[list, list, list]:
+    # set up model and device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    if model == 'SLIP':
+        embed_model = load_slip_model(device)
+        processor = load_slip_preprocessor()
+    elif model == 'CLIP':
+        embed_model = load_clip_model(device)
+        processor = load_clip_preprocessor()
+
+    all_video_embeddings = []
+    all_video_frames = []
+    for fname in fnames:
+        # read in file
+        is_seq = False
+        if fname[-3:] == 'seq': is_seq = True
+        
+        if is_seq:
+            sr = seqIo_reader(fname)
+        else:
+            sr = mp4Io_reader(fname)
+        N  = sr.header['numFrames']
+
+        # set up embeddings and frame arrays
+        embeddings = []
+        frames = list(range(N))[::downsample_rate]
+        print(frames)
+
+        # create progress bar
+        i = 0
+        pbar_text = lambda i: f'Creating embeddings for {fname}. {i}/{len(frames)} frames.'
+        pbar = st.progress(0, text=pbar_text(0))
+
+        # convert each frame to embeddings
+        for f in tqdm(frames):
+            img, _ = sr.getFrame(f)
+            img_arr = np.array(img)
+            if is_seq:
+                img_rgb = Image.fromarray(img_arr, 'L').convert('RGB')
+            else:
+                img_rgb = Image.fromarray(img_arr).convert('RGB')
+
+            embeddings.append(encode_image(img_rgb, device, embed_model, processor))
+
+            # update progress bar
+            i += 1
+            pbar.progress(i/len(frames), pbar_text(i))
+
+        # save csv of single file
+        if save_csv:
+            df = pd.DataFrame(embeddings)
+            df['Frame'] = frames
+
+            # save csv
+            basename = Path(fname).stem
+            df.to_csv(f'{basename}_embeddings_downsample_{downsample_rate}.csv', index=False)
+
+        all_video_embeddings.append(np.array(embeddings))
+        all_video_frames.append(frames)
+    return all_video_embeddings, all_video_frames
+
+def get_io_reader(uploaded_file):
+    if uploaded_file.name[-3:]=='seq':
+        with NamedTemporaryFile(suffix="seq", delete=False) as temp:
+            temp.write(uploaded_file.getvalue())
+            sr = seqIo_reader(temp.name)
+    else:
+        with NamedTemporaryFile(suffix="mp4", delete=False) as temp:
+            temp.write(uploaded_file.getvalue())
+            sr = mp4Io_reader(temp.name)
+    return sr
+
+def generate_embeddings_stream_io(uploaded_files : list,
+                                model = 'SLIP',
+                                downsample_rate = 4,
+                                save_csv = False)-> tuple[list, list, list]:
+    # set up model and device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    if model == 'SLIP':
+        embed_model = load_slip_model(device)
+        processor = load_slip_preprocessor()
+    elif model == 'CLIP':
+        embed_model = load_clip_model(device)
+        processor = load_clip_preprocessor()
+
+    all_video_embeddings = []
+    all_video_frames = []
+    for file in uploaded_files:
+        is_seq = False
+        if file.name[-3:] == 'seq': is_seq = True
+
+        # read in file
+        sr = get_io_reader(file)
+        N  = sr.header['numFrames']
+
+        # set up embeddings and frame arrays
+        embeddings = []
+        frames = list(range(N))[::downsample_rate]
+        print(frames)
+
+        # create progress bar
+        i = 0
+        pbar_text = lambda i: f'Creating embeddings for {file.name}. {i}/{len(frames)} frames.'
+        pbar = st.progress(0, text=pbar_text(0))
+
+        # convert each frame to embeddings
+        for f in tqdm(frames):
+            img, _ = sr.getFrame(f)
+            img_arr = np.array(img)
+            if is_seq:
+                img_rgb = Image.fromarray(img_arr, 'L').convert('RGB')
+            else:
+                img_rgb = Image.fromarray(img_arr).convert('RGB')
+
+            embeddings.append(encode_image(img_rgb, device, embed_model, processor))
+
+            # update progress bar
+            i += 1
+            pbar.progress(i/len(frames), pbar_text(i))
+
+        # save csv of single file
+        if save_csv:
+            df = pd.DataFrame(embeddings)
+            df['Frame'] = frames
+
+            # save csv
+            df.to_csv(f'embeddings_downsample_{downsample_rate}_{frames}_frames.csv', index=False)
+
+        all_video_embeddings.append(np.array(embeddings))
+        all_video_frames.append(frames)
+    return all_video_embeddings, all_video_frames
+
+def create_embeddings_csv(out: str,
+                          fnames: list[str],
+                          embeddings: list[np.ndarray],
+                          frames: list[list[int]],
+                          annotations: list[list[str]],
+                          test_fnames: None | list[str],
+                          views: None | list[str],
+                          conditions: None | list[str],
+                          downsample_rate = 4,
+                          filesystem = None):
+    """
+    Creates a .csv file containing all of the generated embeddings and provived information.
+
+    Parameters:
+    -----------
+    out : str
+        The name of the resulting file.
+    fnames : list[str]
+        Video sources for each of the embedding arrays.
+    embeddings : np.ndarray
+        The generated embeddings from the images.
+    downsample_rate : int
+        The downsample_rate used for generating the embeddings.
+    """
+    assert len(fnames) == len(embeddings)
+    assert len(embeddings) == len(frames)
+    all_embeddings = np.vstack(embeddings)
+    df = pd.DataFrame(all_embeddings)
+    
+    labels = []
+    for i, annot_fnames in enumerate(annotations):
+        _, ext = os.path.splitext(annot_fnames[0])
+        if ext == '.annot':
+            annot, _, _, sr = load_multiple_annotations(annot_fnames, filesystem=filesystem)
+            annot_labels = generate_label_array(annot, downsample_rate, len(frames[i]))
+        elif ext == '.csv':
+            if not filesystem: 
+                annot_df = pd.read_csv(annot_fnames[0], header=None)
+            else:
+                with filesystem.open(annot_fnames[0], 'r') as csv_file:
+                    annot_df = pd.read_csv(csv_file, header=None)
+            annot_labels = annot_df[0].to_list()[::downsample_rate]
+            assert len(annot_labels) == len(frames[i]), "There is a mismatch between the number of frames and number of labels. Make sure that the passed in csv file has no header."
+        else:
+            raise ValueError(f'Incompatible file for annotations used. Got a file of type "{ext}".')
+        assert len(annot_labels) == len(frames[i]), "There is a mismatch between the number of frames and number of labels. Make sure you have passed in the correct files."
+        print(annot_labels)
+        labels.append(annot_labels)
+    all_labels = np.hstack(labels)
+    print(len(all_labels))
+    df['Label'] = all_labels
+    
+    all_frames = np.hstack(frames)
+    df['Frame'] = all_frames
+    sources = [[fname for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]
+    all_sources = np.hstack(sources)
+    df['Source'] = all_sources
+
+    if test_fnames:
+        t_split = lambda x: True if x in test_fnames else False
+        test = [[t_split(fname) for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]
+    else:
+        test = [[True for _ in range(len(frames[i]))] for i, _ in enumerate(fnames)]
+    all_test = np.hstack(test)
+    df['Test'] = all_test
+
+    if views:
+        view = [[views[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    else:
+        view = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    all_view = np.hstack(view)
+    df['View'] = all_view
+    
+    if conditions:
+        condition = [[conditions[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    else:
+        condition = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    all_condition = np.hstack(condition)
+    df['Condition'] = all_condition
+    return df
+
+def create_embeddings_csv_io(out: str,
+                          fnames: list[str],
+                          embeddings: list[np.ndarray],
+                          frames: list[list[int]],
+                          annotations: list,
+                          test_fnames: None | list[str],
+                          views: None | list[str],
+                          conditions: None | list[str],
+                          downsample_rate = 4):
+    """
+    Creates a .csv file containing all of the generated embeddings and provived information.
+
+    Parameters:
+    -----------
+    out : str
+        The name of the resulting file.
+    fnames : list[str]
+        Video sources for each of the embedding arrays.
+    embeddings : np.ndarray
+        The generated embeddings from the images.
+    downsample_rate : int
+        The downsample_rate used for generating the embeddings.
+    """
+    assert len(fnames) == len(embeddings)
+    assert len(embeddings) == len(frames)
+    all_embeddings = np.vstack(embeddings)
+    df = pd.DataFrame(all_embeddings)
+    
+    labels = []
+    for i, uploaded_annots in enumerate(annotations):
+        print(i)
+        _, ext = os.path.splitext(uploaded_annots[0].name)
+        if ext == '.annot':
+            annot, _, _, sr = load_multiple_annotations_io(uploaded_annots)
+            annot_labels = generate_label_array(annot, downsample_rate, len(frames[i]))
+        elif ext == '.csv':
+            annot_df = pd.read_csv(uploaded_annots[0], header=None)
+            annot_labels = annot_df[0].to_list()[::downsample_rate]
+            assert len(annot_labels) == len(frames[i]), "There is a mismatch between the number of frames and number of labels. Make sure that the passed in csv file has no header."
+        else:
+            raise ValueError(f'Incompatible file for annotations used. Got a file of type "{ext}".')
+        assert len(annot_labels) == len(frames[i]), "There is a mismatch between the number of frames and number of labels. Make sure you have passed in the correct files."
+        print(annot_labels)
+        labels.append(annot_labels)
+    all_labels = np.hstack(labels)
+    print(len(all_labels))
+    df['Label'] = all_labels
+    
+    all_frames = np.hstack(frames)
+    df['Frame'] = all_frames
+    sources = [[fname for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]
+    all_sources = np.hstack(sources)
+    df['Source'] = all_sources
+
+    if test_fnames:
+        t_split = lambda x: True if x in test_fnames else False
+        test = [[t_split(fname) for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]
+    else:
+        test = [[True for _ in range(len(frames[i]))] for i, _ in enumerate(fnames)]
+    all_test = np.hstack(test)
+    df['Test'] = all_test
+
+    if views:
+        view = [[views[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    else:
+        view = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    all_view = np.hstack(view)
+    df['View'] = all_view
+    
+    if conditions:
+        condition = [[conditions[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    else:
+        condition = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]
+    all_condition = np.hstack(condition)
+    df['Condition'] = all_condition
+    return df
+
+def process_dataset_in_mem(embeddings_df: pd.DataFrame,
+                    specified_classes=None,
+                    classes_to_remove=None,
+                    max_class_size=None,
+                    animal_state=None,
+                    view=None,
+                    shuffle_data=False,
+                    test_videos=None):
+    """
+    Processes output generated from embeddings paired with images and behavior labels.
+
+    Parameters:
+    -----------
+    csv_path : str
+        Path to the file containing the original data. This should contain embeddings,
+        a column named `'Label'` and a column named `'Images'`.
+    specified_classes : None | list[str]
+        An optional input. Defines labels which should be kept as is in the `'Label'`
+        column and which should be changed to a default `other` label.
+    classes_to_remove : None | list[str]
+        An optional input. Drops rows from the dataframe which contain a label in the
+        list.
+    max_class_size : None | int
+        An optional input. Determines the maximum amount of rows a single label can
+        appear in for each unique label in the `'Label'` column.
+    animal_state : None | str
+        An optional input. Drops rows from the dataframe which do not contain a match
+        for `animal_state` in the text field within the `'Images'` column.
+    view : None | str
+        An optional input. Drops rows from the dataframe which do not contain a match
+        for `view` in the text field within the `'Images'` column.
+    shuffle_data : bool
+        Determines wether the dataframe should have its rows shuffled.
+    test_videos : None | list[str]
+        An optional input. Determines what rows should be in the `test` dataframe, and
+        which should be in the `train` dataframe. It drops rows from the respective
+        dataframe by keeping or dropping rows which do not contain a match for a `str`
+        in `test_videos` in the text field within the `'Images'` column, respectively.
+    
+    Returns:
+    --------
+    balanced_train_embeddings : pandas.DataFrame
+        A processed dataframe whose rows contain the embeddings for each of the images
+        at the corresponding index within `balanced_train_images`.
+    balanced_train_labels : list[str]
+        A list of labels for each of the images at the corresponing index within
+        `balanced_train_images`.
+    balanced_train_images: list[str]
+        A list of paths to images with each image at an index corresponding to a label
+        with the same index in `balanced_train_labels` and the same row index within
+        `balanced_train_embeddings`.
+    test_embeddings : pandas.DataFrame
+        A processed dataframe whose rows contain the embeddings for each of the images
+        at the corresponding index within `test_images`.
+    test_labels : list[str]
+        A list of labels for each of the images at the corresponing index within
+        `test_images`.
+    test_images : list[str]
+        A list of paths to images with each image at an index corresponding to a label
+        with the same index in `test_labels` and the same row index within
+        `test_embeddings`.
+    """
+    # Convert embeddings, labels, and images to a DataFrame for easy manipulation
+    df = copy.deepcopy(embeddings_df)
+    df_keys = [str(x) for x in df.keys()]
+    #Filter by fed or fasted
+    if 'Condition' in df_keys and animal_state:
+        df = df[df['Condition'].str.contains(animal_state, na=False)]
+
+    if 'View' in df_keys and view:
+        df = df[df['View'].str.contains(view, na=False)]
+
+    # Extract unique video names excluding the frame number
+    #unique_video_names = df['Images'].apply(lambda x: '_'.join(x.split('_')[:-1])).unique()
+    #print("\nUnique video names:\n", unique_video_names)
+
+    if classes_to_remove:
+        df = df[~df['Label'].str.contains('|'.join(classes_to_remove), na=False)]
+    elif classes_to_remove and 'all' in classes_to_remove:
+        df = df[df['Label'].str.contains('|'.join(classes_to_remove), na=False)]
+
+    # Further filter to include only specified_classes
+    if specified_classes:
+        single_match = lambda x: list(set(x.split('||')) & set(specified_classes))[0]
+        df['Label'] = df['Label'].apply(lambda x: single_match(x) if not set(x.split('||')).isdisjoint(specified_classes) else 'other')
+        specified_classes.append('other')
+
+    # Separate the DataFrame into test and training sets based on test_videos
+    if 'Test' in df_keys and test_videos:
+        test_df = df[df['Test']]
+        train_df = df[~df['Test']]
+    elif test_videos:
+        test_df = df[df['Images'].str.contains('|'.join(test_videos), na=False)]
+        train_df = df[~df['Images'].str.contains('|'.join(test_videos), na=False)]
+    else:
+        test_df = pd.DataFrame(columns=df.columns)
+        train_df = df
+    
+    # Print the number of frames in each class before balancing
+    label_counts = train_df['Label'].value_counts()
+    print("\nNumber of training frames in each class before balancing:")
+    print(label_counts)
+    
+    if max_class_size:
+        balanced_train_df = pd.concat([
+            group.sample(n=min(len(group), max_class_size), random_state=1)
+            for label, group in train_df.groupby('Label')
+        ])
+    else:
+        balanced_train_df = train_df
+
+    # Shuffle the training DataFrame
+    if shuffle_data:
+        balanced_train_df = balanced_train_df.sample(frac=1).reset_index(drop=True)
+    
+    # Convert training set back to numpy array and list
+    if not "Images" in df_keys:
+        balanced_train_embeddings = balanced_train_df.drop(columns=['Label', 'Frame', 'Source', 'Test','View','Condition']).to_numpy()
+        balanced_train_labels = balanced_train_df['Label'].tolist()
+        balanced_train_images = balanced_train_df['Frame'].tolist()
+        
+        # Convert test set back to numpy array and list
+        test_embeddings = test_df.drop(columns=['Label', 'Frame', 'Source', 'Test','View','Condition']).to_numpy()
+        test_labels = test_df['Label'].tolist()
+        test_images = test_df['Frame'].tolist()
+    else:
+        # Convert training set back to numpy array and list
+        balanced_train_embeddings = balanced_train_df.drop(columns=['Label', 'Images']).to_numpy()
+        balanced_train_labels = balanced_train_df['Label'].tolist()
+        balanced_train_images = balanced_train_df['Images'].tolist()
+        
+        # Convert test set back to numpy array and list
+        test_embeddings = test_df.drop(columns=['Label', 'Images']).to_numpy()
+        test_labels = test_df['Label'].tolist()
+        test_images = test_df['Images'].tolist()
+    
+    # Print the number of frames in each class after balancing
+    if specified_classes or max_class_size:
+        balanced_label_counts = Counter(balanced_train_labels)
+        print("\nNumber of training frames in each class after balancing:")
+        print(balanced_label_counts)
+
+    test_label_counts = test_df['Label'].value_counts()
+    # print("\nNumber of testing frames in each class:")
+    print(test_label_counts)
+    
+    return balanced_train_embeddings, balanced_train_labels, balanced_train_images, test_embeddings, test_labels, test_images
+
+def multiclass_merge_and_filter_bouts(multiclass_vector, bout_threshold, proximity_threshold):
+    # Get the unique labels in the multiclass vector (excluding zero, assuming zero is the background/no label)
+    unique_labels = np.unique(multiclass_vector)
+    unique_labels = unique_labels[unique_labels != 0]
+
+    # Initialize a vector to store the merged and filtered multiclass vector
+    merged_vector = np.zeros_like(multiclass_vector)
+
+    for label in unique_labels:
+        # Create a binary vector for the current label
+        binary_vector = (multiclass_vector == label)
+
+        # Find the start and end indices of all sequences of 1's for this label
+        starts = np.where(np.diff(np.concatenate(([0], binary_vector))) == 1)[0]
+        ends = np.where(np.diff(np.concatenate((binary_vector, [0]))) == -1)[0]
+
+        # Step 1: Merge close short bouts
+        i = 0
+        while i < len(starts) - 1:
+            # Check if the gap between the end of the current bout and the start of the next bout
+            # is within the proximity threshold
+            if starts[i + 1] - ends[i] <= proximity_threshold:
+                # Merge the two bouts by setting all elements between the start of the first
+                # and the end of the second bout to 1
+                binary_vector[ends[i]:starts[i + 1]] = 1
+                # Remove the next bout from consideration
+                starts = np.delete(starts, i + 1)
+                ends = np.delete(ends, i)
+            else:
+                i += 1
+
+        # Update the starts and ends after merging
+        starts = np.where(np.diff(np.concatenate(([0], binary_vector))) == 1)[0]
+        ends = np.where(np.diff(np.concatenate((binary_vector, [0]))) == -1)[0]
+
+        # Step 2: Remove standalone short bouts
+        for i in range(len(starts)):
+            # Check the length of the bout
+            length_of_bout = ends[i] - starts[i] + 1
+
+            # If the length is less than the threshold, set those elements to 0
+            if length_of_bout < bout_threshold:
+                binary_vector[starts[i]:ends[i] + 1] = 0
+
+        # Combine the binary vector with the merged_vector, ensuring only the current label is set
+        merged_vector[binary_vector] = label
+
+    # Return the filtered multiclass vector
+    return merged_vector
+
+def get_unique_labels(label_list: list[str]):
+    label_set = set()
+    for label in label_list:
+        individual_labels = label.split('||')
+        for individual_label in individual_labels:
+            label_set.add(individual_label)
+    return list(label_set)
+
+def get_train_test_split(train_embeds, numerical_labels, test_size=0.05, random_state=42):
+    return train_test_split(train_embeds, numerical_labels, test_size=test_size, random_state=random_state)
+
+def train_model(X_train, y_train, random_state=42):
+    # Train SVM Classifier
+    svm_clf = SVC(kernel='rbf', random_state=random_state, probability=True)
+    svm_clf.fit(X_train, y_train)
+    return svm_clf
+
+def pickle_model(model):
+    pickled = io.BytesIO()
+    pickle.dump(model, pickled)
+    return pickled
+
+def get_seq_io_reader(uploaded_file):
+    assert uploaded_file.name[-3:]=='seq', 'Not a seq file'
+    with NamedTemporaryFile(suffix="seq", delete=False) as temp:
+        temp.write(uploaded_file.getvalue())
+        sr = seqIo_reader(temp.name)
+    return sr
+
+def seq_to_arr(sr):
+    N = sr.header['numFrames']
+    images = []
+    for f in range(N):
+        I, ts = sr.getFrame(f)
+        images.append(I)
+    return np.array(images)
+
+def get_2d_embedding(embeddings: pd.DataFrame):
+    tsne = TSNE(n_jobs=4, n_components=2, random_state=42, perplexity=50)
+    embedding_2d = tsne.fit_transform(np.array(embeddings))
+    return embedding_2d
+
+
+