neuralcomputation/batik
0
1import streamlit as st2import plotly.express as px3import numpy as np4import pandas as pd5import torch6from utils.mp4Io import mp4Io_reader7from utils.seqIo import seqIo_reader8import pandas as pd9from PIL import Image10from pathlib import Path11from transformers import AutoProcessor, AutoModel12from tempfile import NamedTemporaryFile13from tqdm import tqdm14from utils.utils import create_embeddings_csv_io, process_dataset_in_mem, generate_embeddings_stream_io15from get_llava_response import get_llava_response, load_llava_checkpoint_hf16from sklearn.manifold import TSNE17from openai import OpenAI18import cv219import base6420from hdbscan import HDBSCAN, all_points_membership_vectors21import random22 23# --server.maxUploadSize 300024REPO_NAME = 'ncoria/llava-lora-vicuna-clip-5-epochs-merge'25 26def load_llava_model():27 return load_llava_checkpoint_hf(REPO_NAME)28 29def get_unique_labels(label_list: list[str]):30 label_set = set()31 for label in label_list:32 individual_labels = label.split('||')33 for individual_label in individual_labels:34 label_set.add(individual_label)35 return list(label_set)36 37SYSTEM_PROMPT = """You are a researcher studying mice interactions from videos of the inside of a resident 38intruder box where there is either just the resident mouse (the black one) or the resident and the intruder mouse (the white one).39Your 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.40Thus, you should use the visual information about the mice in the image to try to provide a detailed behavioral description of the image."""41 42@st.cache_resource43def get_io_reader(uploaded_file):44 if uploaded_file.name[-3:]=='seq':45 with NamedTemporaryFile(suffix="seq", delete=False) as temp:46 temp.write(uploaded_file.getvalue())47 sr = seqIo_reader(temp.name)48 else:49 with NamedTemporaryFile(suffix="mp4", delete=False) as temp:50 temp.write(uploaded_file.getvalue())51 sr = mp4Io_reader(temp.name)52 return sr53 54def get_image(sr, frame_no: int):55 image, _ = sr.getFrame(frame_no)56 return image57 58@st.cache_data59def get_2d_embedding(embeddings: pd.DataFrame):60 tsne = TSNE(n_jobs=4, n_components=2, random_state=42, perplexity=50)61 embedding_2d = tsne.fit_transform(np.array(embeddings))62 return embedding_2d63 64HDBSCAN_PARAMS = {65 'min_samples': 166}67 68@st.cache_data69def hdbscan_classification(umap_embeddings, embeddings_2d, cluster_range):70 max_num_clusters = -np.infty71 num_clusters = []72 min_cluster_size = np.linspace(cluster_range[0], cluster_range[1], 4)73 for min_c in min_cluster_size:74 learned_hierarchy = HDBSCAN(75 prediction_data=True, min_cluster_size=int(round(min_c * 0.01 *umap_embeddings.shape[0])),76 cluster_selection_method='leaf' ,77 **HDBSCAN_PARAMS).fit(umap_embeddings)78 num_clusters.append(len(np.unique(learned_hierarchy.labels_)))79 if num_clusters[-1] > max_num_clusters:80 max_num_clusters = num_clusters[-1]81 retained_hierarchy = learned_hierarchy82 assignments = retained_hierarchy.labels_83 assign_prob = all_points_membership_vectors(retained_hierarchy)84 soft_assignments = np.argmax(assign_prob, axis=1)85 retained_hierarchy.fit(embeddings_2d)86 return retained_hierarchy, assignments, assign_prob, soft_assignments87 88def upload_image(frame: np.ndarray):89 """returns the file ID."""90 _, encoded_image = cv2.imencode('.png', frame)91 return base64.b64encode(encoded_image.tobytes()).decode('utf-8')92 93def ask_question_with_image_gpt(file_id, system_prompt, question, api_key):94 """Asks a question about the uploaded image."""95 client = OpenAI(api_key=api_key)96 97 if file_id != None:98 response = client.chat.completions.create(99 model="gpt-4o",100 messages=[101 {"role": "system", "content": system_prompt},102 {"role": "user", "content": [103 {"type": "text", "text": question},104 {"type": "image_url", "image_url": {"url": f"data:image/jpg:base64, {file_id}"}}]105 }106 ]107 )108 else:109 response = client.chat.completions.create(110 model="gpt-4o",111 messages=[112 {"role": "system", "content": system_prompt},113 {"role": "user", "content": question}114 ]115 )116 return response.choices[0].message.content117 118def ask_question_with_image_llava(image, system_prompt, question,119 tokenizer, model, image_processor):120 outputs = get_llava_response([question],121 [image],122 system_prompt,123 tokenizer,124 model,125 image_processor,126 REPO_NAME,127 stream_output=False)128 return outputs[0]129 130def ask_summary_question(image_array, label_array, api_key):131 # load llava model132 with st.spinner("Loading LLaVA model. This can take 10 to 30 minutes. Please wait..."):133 tokenizer, model, image_processor = load_llava_model()134 135 # global variable136 system_prompt = SYSTEM_PROMPT137 138 # collect responses139 responses = []140 141 # create progress bar142 j = 0143 pbar_text = lambda j: f'Creating llava response {j}/{len(label_array)}.'144 pbar = st.progress(0, text=pbar_text(0))145 146 for i, image in enumerate(image_array):147 label = label_array[i]148 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. "149 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 "\150 "how they are positioned with respect to each other. Usually, the body parts (i.e., tail, genitals, face, body, ears, paws)"\151 "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."152 response = ask_question_with_image_llava(image, system_prompt, question,153 tokenizer, model, image_processor)154 responses.append(response)155 # update progress bar156 j += 1157 pbar.progress(j/len(label_array), pbar_text(j))158 159 system_prompt_summarize = "You are a researcher studying mice interactions from videos of the inside of a resident "\160 "intruder box where there is either just the resident mouse (the black one) or the resident and the intruder mouse (the white one). "\161 "You will be given a question about a list of descriptions from frames of these videos. "\162 "Your job is to answer the question by focusing on the behaviors of the mice and their postures "\163 "as well as any other aspects of the descriptions that may be relevant to the class label associated with them"164 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 "\165 "behavior label which captures what is described in the descriptions: \n\n"166 user_prompt_summarize = user_prompt_summarize + '\n'.join(responses)167 summary_response = ask_question_with_image_gpt(None, system_prompt_summarize, user_prompt_summarize, api_key)168 return summary_response169 170if "embeddings_df" not in st.session_state:171 st.session_state.embeddings_df = None172 173st.title('batik: behavior discovery and LLM-based interpretation')174 175api_key = st.text_input("OpenAI API Key:","")176 177st.subheader("generate or import embeddings")178 179st.text("Upload files to generate embeddings.")180with st.form('embedding_generation_settings'):181 seq_file = st.file_uploader("Choose a video file", type=['seq', 'mp4'])182 annot_files = st.file_uploader("Choose an annotation File", type=['annot','csv'], accept_multiple_files=True)183 downsample_rate = st.number_input('Downsample Rate',value=4)184 submit_embed_settings = st.form_submit_button('Create Embeddings', type='secondary')185 186st.markdown("**(Optional)** Upload embeddings.")187embeddings_csv = st.file_uploader("Choose a .csv File", type=['csv'])188 189if submit_embed_settings and seq_file is not None and annot_files is not None:190 video_embeddings, video_frames = generate_embeddings_stream_io([seq_file],191 "SLIP",192 downsample_rate,193 False)194 195 fnames = [seq_file.name]196 embeddings_df = create_embeddings_csv_io(out="file",197 fnames=fnames,198 embeddings=video_embeddings,199 frames=video_frames,200 annotations=[annot_files],201 test_fnames=None,202 views=None,203 conditions=None,204 downsample_rate=downsample_rate)205 st.session_state.embeddings_df = embeddings_df206elif embeddings_csv is not None:207 embeddings_df = pd.read_csv(embeddings_csv)208 st.session_state.embeddings_df = embeddings_df209else:210 st.text('Please upload file(s).')211 212st.divider()213st.subheader("provide video file if not yet already provided")214 215uploaded_file = st.file_uploader("Choose a video file", type=['seq', 'mp4'])216 217st.divider()218if st.session_state.embeddings_df is not None and (uploaded_file is not None or seq_file is not None):219 if seq_file is not None:220 uploaded_file = seq_file221 io_reader = get_io_reader(uploaded_file)222 print("CONVERTED SEQ")223 label_list = st.session_state.embeddings_df['Label'].to_list()224 unique_label_list = get_unique_labels(label_list)225 print(f"unique_labels: {unique_label_list}")226 #unique_label_list = ['check_genital', 'wiggle', 'lordose', 'stay', 'turn', 'top_up', 'dart', 'sniff', 'approach', 'into_male_cage']227 #unique_label_list = ['into_male_cage', 'intromission', 'male_sniff', 'mount']228 kwargs = {'embeddings_df' : st.session_state.embeddings_df, 229 'specified_classes' : unique_label_list,230 'classes_to_remove' : None,231 'max_class_size' : None,232 'animal_state' : None,233 'view' : None,234 'shuffle_data' : False,235 'test_videos' : None}236 train_embeds, train_labels, train_images, _, _, _ = process_dataset_in_mem(**kwargs)237 print("PROCESSED DATASET")238 if "Images" in st.session_state.embeddings_df.keys():239 train_images = [i for i in range(len(train_images))]240 embedding_2d = get_2d_embedding(train_embeds)241else:242 st.text('Please generate embeddings and provide video file.')243 print("GOT 2D EMBEDS")244 245if uploaded_file is not None and st.session_state.embeddings_df is not None:246 st.subheader("t-SNE Projection")247 option = st.selectbox(248 "Select Color Option",249 ("By Label", "By Time", "By Cluster")250 )251 if embedding_2d is not None:252 if option is not None:253 if option == "By Label":254 color = 'label'255 elif option == "By Time":256 color = 'frame_no'257 else:258 color = 'cluster_label'259 260 if option in ["By Label", "By Time"]:261 edf = pd.DataFrame(embedding_2d,columns=['tsne_dim_1', 'tsne_dim_2'])262 edf.insert(2,'frame_no',np.array([int(x) for x in train_images]))263 edf.insert(3, 'label', train_labels)264 fig = px.scatter(265 edf,266 x="tsne_dim_1",267 y="tsne_dim_2",268 color=color,269 hover_data=["frame_no"],270 color_discrete_sequence=px.colors.qualitative.Dark24271 )272 else:273 r, _, _, _ = hdbscan_classification(train_embeds, embedding_2d, [4, 6])274 edf = pd.DataFrame(embedding_2d,columns=['tsne_dim_1', 'tsne_dim_2'])275 edf.insert(2,'frame_no',np.array([int(x) for x in train_images]))276 edf.insert(3, 'label', train_labels)277 edf.insert(4, 'cluster_label', [str(c_id) for c_id in r.labels_.tolist()])278 fig = px.scatter(279 edf,280 x="tsne_dim_1",281 y="tsne_dim_2",282 color=color,283 hover_data=["frame_no"],284 color_discrete_sequence=px.colors.qualitative.Dark24285 )286 287 event = st.plotly_chart(fig, key="df", on_select="rerun")288 else:289 st.text("No Color Option Selected")290 else:291 st.text('No Embeddings Loaded')292 293 event_dict = event.selection294 295 if event_dict is not None:296 custom_data = []297 for point in event_dict['points']:298 data = point["customdata"][0]299 custom_data.append(int(data))300 301 if len(custom_data) > 10:302 custom_data = random.sample(custom_data, 10)303 if len(custom_data) > 1:304 col_1, col_2 = st.columns(2)305 with col_1:306 for frame_no in custom_data[::2]:307 st.image(get_image(io_reader, frame_no))308 st.caption(f"Frame {frame_no}, {train_labels[frame_no]}")309 with col_2:310 for frame_no in custom_data[1::2]:311 st.image(get_image(io_reader, frame_no))312 st.caption(f"Frame {frame_no}, {train_labels[frame_no]}")313 elif len(custom_data) == 1:314 frame_no = custom_data[0]315 st.image(get_image(io_reader, frame_no))316 st.caption(f"Frame {frame_no}, {train_labels[frame_no]}")317 else:318 st.text('No Points Selected')319 320 if len(custom_data) == 1:321 frame_no = custom_data[0]322 image = get_image(io_reader, frame_no)323 system_prompt = SYSTEM_PROMPT324 label = train_labels[frame_no]325 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. "\326 "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 "\327 "how they are positioned with respect to each other. Usually, the body parts (i.e., tail, genitals, face, body, ears, paws)" \328 "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."329 with st.spinner("Loading LLaVA model. This can take 10 to 30 minutes. Please wait..."):330 tokenizer, model, image_processor = load_llava_model()331 response = ask_question_with_image_llava(image, system_prompt, question,332 tokenizer, model, image_processor)333 st.markdown(response)334 335 elif len(custom_data) > 1:336 image_array = [get_image(io_reader, f_no) for f_no in custom_data]337 label_array = [train_labels[f_no] for f_no in custom_data]338 response = ask_summary_question(image_array, label_array, api_key)339 st.markdown(response)340