neuralcomputation/batik
0
1import os2import io3import pickle4import copy5from collections import Counter6from pathlib import Path7from tempfile import NamedTemporaryFile8import regex as re9import numpy as np10import pandas as pd11from sklearn.manifold import TSNE12from sklearn.svm import SVC13from sklearn.model_selection import train_test_split14from sklearn.metrics import accuracy_score, classification_report15import torch16from tqdm import tqdm17from PIL import Image18from transformers import AutoProcessor, AutoModel19import streamlit as st20from .data_loading import load_multiple_annotations, load_multiple_annotations_io21from .data_processing import generate_label_array22from .seqIo import seqIo_reader23from .mp4Io import mp4Io_reader24 25SLIP_MODEL_ID = "google/siglip-so400m-patch14-384"26CLIP_MODEL_ID = "openai/clip-vit-base-patch32"27 28def create_annot_fname_dict(annot_fnames: list[str])-> dict:29 fs = re.compile(r'.*(_\d+)$')30 31 unique_files = set()32 for file in annot_fnames:33 file_name = os.fsdecode(file)34 base_name, _ = os.path.splitext(file_name)35 if fs.match(base_name):36 ind = len(fs.match(base_name).group(1))37 unique_files.add(base_name[:-ind])38 else:39 unique_files.add(base_name)40 41 annot_fname_dict = {}42 for unique_file in unique_files:43 annot_fname_dict.update({unique_file: [file for file in annot_fnames if unique_file in file]})44 return annot_fname_dict45 46def create_annot_fname_dict_io(annot_fnames: list[str], annot_files: list)-> dict:47 annot_file_dict = {}48 for file in annot_files:49 annot_file_dict.update({file.name : file})50 fs = re.compile(r'.*(_\d+)$')51 52 unique_files = set()53 for file in annot_fnames:54 file_name = os.fsdecode(file)55 base_name, _ = os.path.splitext(file_name)56 if fs.match(base_name):57 ind = len(fs.match(base_name).group(1))58 unique_files.add(base_name[:-ind])59 else:60 unique_files.add(base_name)61 62 annot_fname_dict = {}63 for unique_file in unique_files:64 annot_list = [file for file in annot_fnames if unique_file in file]65 annot_list.sort()66 annot_file_list = [annot_file_dict[annot_file_name] for annot_file_name in annot_list]67 annot_fname_dict.update({unique_file: annot_file_list})68 return annot_fname_dict69 70def get_io_reader(uploaded_file):71 assert uploaded_file.name[-3:]=='seq', 'Not a seq file'72 with NamedTemporaryFile(suffix="seq", delete=False) as temp:73 temp.write(uploaded_file.getvalue())74 sr = seqIo_reader(temp.name)75 return sr76 77def load_slip_model(device):78 return AutoModel.from_pretrained(SLIP_MODEL_ID).to(device)79 80def load_slip_preprocessor():81 return AutoProcessor.from_pretrained(SLIP_MODEL_ID)82 83def load_clip_model(device):84 return AutoModel.from_pretrained(CLIP_MODEL_ID).to(device)85 86def load_clip_preprocessor():87 return AutoProcessor.from_pretrained(CLIP_MODEL_ID)88 89def encode_image(image, device, model, processor):90 with torch.no_grad():91 #convert_models_to_fp32(model)92 inputs = processor(images=image, return_tensors="pt").to(device)93 image_features = model.get_image_features(**inputs)94 return image_features.cpu().numpy().flatten()95 96def generate_embeddings_stream(fnames : list[str],97 model = 'SLIP',98 downsample_rate = 4,99 save_csv = False)-> tuple[list, list, list]:100 # set up model and device101 device = "cuda" if torch.cuda.is_available() else "cpu"102 os.environ['CUDA_VISIBLE_DEVICES'] = '0'103 if model == 'SLIP':104 embed_model = load_slip_model(device)105 processor = load_slip_preprocessor()106 elif model == 'CLIP':107 embed_model = load_clip_model(device)108 processor = load_clip_preprocessor()109 110 all_video_embeddings = []111 all_video_frames = []112 for fname in fnames:113 # read in file114 is_seq = False115 if fname[-3:] == 'seq': is_seq = True116 117 if is_seq:118 sr = seqIo_reader(fname)119 else:120 sr = mp4Io_reader(fname)121 N = sr.header['numFrames']122 123 # set up embeddings and frame arrays124 embeddings = []125 frames = list(range(N))[::downsample_rate]126 print(frames)127 128 # create progress bar129 i = 0130 pbar_text = lambda i: f'Creating embeddings for {fname}. {i}/{len(frames)} frames.'131 pbar = st.progress(0, text=pbar_text(0))132 133 # convert each frame to embeddings134 for f in tqdm(frames):135 img, _ = sr.getFrame(f)136 img_arr = np.array(img)137 if is_seq:138 img_rgb = Image.fromarray(img_arr, 'L').convert('RGB')139 else:140 img_rgb = Image.fromarray(img_arr).convert('RGB')141 142 embeddings.append(encode_image(img_rgb, device, embed_model, processor))143 144 # update progress bar145 i += 1146 pbar.progress(i/len(frames), pbar_text(i))147 148 # save csv of single file149 if save_csv:150 df = pd.DataFrame(embeddings)151 df['Frame'] = frames152 153 # save csv154 basename = Path(fname).stem155 df.to_csv(f'{basename}_embeddings_downsample_{downsample_rate}.csv', index=False)156 157 all_video_embeddings.append(np.array(embeddings))158 all_video_frames.append(frames)159 return all_video_embeddings, all_video_frames160 161def get_io_reader(uploaded_file):162 if uploaded_file.name[-3:]=='seq':163 with NamedTemporaryFile(suffix="seq", delete=False) as temp:164 temp.write(uploaded_file.getvalue())165 sr = seqIo_reader(temp.name)166 else:167 with NamedTemporaryFile(suffix="mp4", delete=False) as temp:168 temp.write(uploaded_file.getvalue())169 sr = mp4Io_reader(temp.name)170 return sr171 172def generate_embeddings_stream_io(uploaded_files : list,173 model = 'SLIP',174 downsample_rate = 4,175 save_csv = False)-> tuple[list, list, list]:176 # set up model and device177 device = "cuda" if torch.cuda.is_available() else "cpu"178 os.environ['CUDA_VISIBLE_DEVICES'] = '0'179 with st.spinner('Loading multimodal model...'):180 if model == 'SLIP':181 embed_model = load_slip_model(device)182 processor = load_slip_preprocessor()183 elif model == 'CLIP':184 embed_model = load_clip_model(device)185 processor = load_clip_preprocessor()186 187 all_video_embeddings = []188 all_video_frames = []189 for file in uploaded_files:190 is_seq = False191 if file.name[-3:] == 'seq': is_seq = True192 193 # read in file194 sr = get_io_reader(file)195 N = sr.header['numFrames']196 197 # set up embeddings and frame arrays198 embeddings = []199 frames = list(range(N))[::downsample_rate]200 print(frames)201 202 # create progress bar203 i = 0204 pbar_text = lambda i: f'Creating embeddings for {file.name}. {i}/{len(frames)} frames.'205 pbar = st.progress(0, text=pbar_text(0))206 207 # convert each frame to embeddings208 for f in tqdm(frames):209 img, _ = sr.getFrame(f)210 img_arr = np.array(img)211 if is_seq:212 img_rgb = Image.fromarray(img_arr, 'L').convert('RGB')213 else:214 img_rgb = Image.fromarray(img_arr).convert('RGB')215 216 embeddings.append(encode_image(img_rgb, device, embed_model, processor))217 218 # update progress bar219 i += 1220 pbar.progress(i/len(frames), pbar_text(i))221 222 # save csv of single file223 if save_csv:224 df = pd.DataFrame(embeddings)225 df['Frame'] = frames226 227 # save csv228 df.to_csv(f'embeddings_downsample_{downsample_rate}_{N}_frames.csv', index=False)229 230 all_video_embeddings.append(np.array(embeddings))231 all_video_frames.append(frames)232 return all_video_embeddings, all_video_frames233 234def create_embeddings_csv(out: str,235 fnames: list[str],236 embeddings: list[np.ndarray],237 frames: list[list[int]],238 annotations: list[list[str]],239 test_fnames: None | list[str],240 views: None | list[str],241 conditions: None | list[str],242 downsample_rate = 4,243 filesystem = None):244 """245 Creates a .csv file containing all of the generated embeddings and provived information.246 247 Parameters:248 -----------249 out : str250 The name of the resulting file.251 fnames : list[str]252 Video sources for each of the embedding arrays.253 embeddings : np.ndarray254 The generated embeddings from the images.255 downsample_rate : int256 The downsample_rate used for generating the embeddings.257 """258 assert len(fnames) == len(embeddings)259 assert len(embeddings) == len(frames)260 all_embeddings = np.vstack(embeddings)261 df = pd.DataFrame(all_embeddings)262 263 labels = []264 for i, annot_fnames in enumerate(annotations):265 _, ext = os.path.splitext(annot_fnames[0])266 if ext == '.annot':267 annot, _, _, sr = load_multiple_annotations(annot_fnames, filesystem=filesystem)268 annot_labels = generate_label_array(annot, downsample_rate, len(frames[i]))269 elif ext == '.csv':270 if not filesystem: 271 annot_df = pd.read_csv(annot_fnames[0], header=None)272 else:273 with filesystem.open(annot_fnames[0], 'r') as csv_file:274 annot_df = pd.read_csv(csv_file, header=None)275 annot_labels = annot_df[0].to_list()[::downsample_rate]276 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."277 else:278 raise ValueError(f'Incompatible file for annotations used. Got a file of type "{ext}".')279 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."280 print(annot_labels)281 labels.append(annot_labels)282 all_labels = np.hstack(labels)283 print(len(all_labels))284 df['Label'] = all_labels285 286 all_frames = np.hstack(frames)287 df['Frame'] = all_frames288 sources = [[fname for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]289 all_sources = np.hstack(sources)290 df['Source'] = all_sources291 292 if test_fnames:293 t_split = lambda x: True if x in test_fnames else False294 test = [[t_split(fname) for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]295 else:296 test = [[True for _ in range(len(frames[i]))] for i, _ in enumerate(fnames)]297 all_test = np.hstack(test)298 df['Test'] = all_test299 300 if views:301 view = [[views[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]302 else:303 view = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]304 all_view = np.hstack(view)305 df['View'] = all_view306 307 if conditions:308 condition = [[conditions[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]309 else:310 condition = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]311 all_condition = np.hstack(condition)312 df['Condition'] = all_condition313 return df314 315def create_embeddings_csv_io(out: str,316 fnames: list[str],317 embeddings: list[np.ndarray],318 frames: list[list[int]],319 annotations: list,320 test_fnames: None | list[str],321 views: None | list[str],322 conditions: None | list[str],323 downsample_rate = 4):324 """325 Creates a .csv file containing all of the generated embeddings and provived information.326 327 Parameters:328 -----------329 out : str330 The name of the resulting file.331 fnames : list[str]332 Video sources for each of the embedding arrays.333 embeddings : np.ndarray334 The generated embeddings from the images.335 downsample_rate : int336 The downsample_rate used for generating the embeddings.337 """338 assert len(fnames) == len(embeddings)339 assert len(embeddings) == len(frames)340 all_embeddings = np.vstack(embeddings)341 df = pd.DataFrame(all_embeddings)342 343 labels = []344 for i, uploaded_annots in enumerate(annotations):345 print(i)346 _, ext = os.path.splitext(uploaded_annots[0].name)347 if ext == '.annot':348 annot, _, _, sr = load_multiple_annotations_io(uploaded_annots)349 annot_labels = generate_label_array(annot, downsample_rate, len(frames[i]))350 elif ext == '.csv':351 annot_df = pd.read_csv(uploaded_annots[0], header=None)352 annot_labels = annot_df[0].to_list()[::downsample_rate]353 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."354 else:355 raise ValueError(f'Incompatible file for annotations used. Got a file of type "{ext}".')356 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."357 print(annot_labels)358 labels.append(annot_labels)359 all_labels = np.hstack(labels)360 print(len(all_labels))361 df['Label'] = all_labels362 363 all_frames = np.hstack(frames)364 df['Frame'] = all_frames365 sources = [[fname for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]366 all_sources = np.hstack(sources)367 df['Source'] = all_sources368 369 if test_fnames:370 t_split = lambda x: True if x in test_fnames else False371 test = [[t_split(fname) for _ in range(len(frames[i]))] for i, fname in enumerate(fnames)]372 else:373 test = [[True for _ in range(len(frames[i]))] for i, _ in enumerate(fnames)]374 all_test = np.hstack(test)375 df['Test'] = all_test376 377 if views:378 view = [[views[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]379 else:380 view = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]381 all_view = np.hstack(view)382 df['View'] = all_view383 384 if conditions:385 condition = [[conditions[i] for _ in range(len(frames[i]))] for i in range(len(fnames))]386 else:387 condition = [[None for _ in range(len(frames[i]))] for i in range(len(fnames))]388 all_condition = np.hstack(condition)389 df['Condition'] = all_condition390 return df391 392def process_dataset_in_mem(embeddings_df: pd.DataFrame,393 specified_classes=None,394 classes_to_remove=None,395 max_class_size=None,396 animal_state=None,397 view=None,398 shuffle_data=False,399 test_videos=None):400 """401 Processes output generated from embeddings paired with images and behavior labels.402 403 Parameters:404 -----------405 csv_path : str406 Path to the file containing the original data. This should contain embeddings,407 a column named `'Label'` and a column named `'Images'`.408 specified_classes : None | list[str]409 An optional input. Defines labels which should be kept as is in the `'Label'`410 column and which should be changed to a default `other` label.411 classes_to_remove : None | list[str]412 An optional input. Drops rows from the dataframe which contain a label in the413 list.414 max_class_size : None | int415 An optional input. Determines the maximum amount of rows a single label can416 appear in for each unique label in the `'Label'` column.417 animal_state : None | str418 An optional input. Drops rows from the dataframe which do not contain a match419 for `animal_state` in the text field within the `'Images'` column.420 view : None | str421 An optional input. Drops rows from the dataframe which do not contain a match422 for `view` in the text field within the `'Images'` column.423 shuffle_data : bool424 Determines wether the dataframe should have its rows shuffled.425 test_videos : None | list[str]426 An optional input. Determines what rows should be in the `test` dataframe, and427 which should be in the `train` dataframe. It drops rows from the respective428 dataframe by keeping or dropping rows which do not contain a match for a `str`429 in `test_videos` in the text field within the `'Images'` column, respectively.430 431 Returns:432 --------433 balanced_train_embeddings : pandas.DataFrame434 A processed dataframe whose rows contain the embeddings for each of the images435 at the corresponding index within `balanced_train_images`.436 balanced_train_labels : list[str]437 A list of labels for each of the images at the corresponing index within438 `balanced_train_images`.439 balanced_train_images: list[str]440 A list of paths to images with each image at an index corresponding to a label441 with the same index in `balanced_train_labels` and the same row index within442 `balanced_train_embeddings`.443 test_embeddings : pandas.DataFrame444 A processed dataframe whose rows contain the embeddings for each of the images445 at the corresponding index within `test_images`.446 test_labels : list[str]447 A list of labels for each of the images at the corresponing index within448 `test_images`.449 test_images : list[str]450 A list of paths to images with each image at an index corresponding to a label451 with the same index in `test_labels` and the same row index within452 `test_embeddings`.453 """454 # Convert embeddings, labels, and images to a DataFrame for easy manipulation455 df = copy.deepcopy(embeddings_df)456 df_keys = [str(x) for x in df.keys()]457 #Filter by fed or fasted458 if 'Condition' in df_keys and animal_state:459 df = df[df['Condition'].str.contains(animal_state, na=False)]460 461 if 'View' in df_keys and view:462 df = df[df['View'].str.contains(view, na=False)]463 464 # Extract unique video names excluding the frame number465 #unique_video_names = df['Images'].apply(lambda x: '_'.join(x.split('_')[:-1])).unique()466 #print("\nUnique video names:\n", unique_video_names)467 468 if classes_to_remove:469 df = df[~df['Label'].str.contains('|'.join(classes_to_remove), na=False)]470 elif classes_to_remove and 'all' in classes_to_remove:471 df = df[df['Label'].str.contains('|'.join(classes_to_remove), na=False)]472 473 # Further filter to include only specified_classes474 if specified_classes:475 single_match = lambda x: list(set(x.split('||')) & set(specified_classes))[0]476 df['Label'] = df['Label'].apply(lambda x: single_match(x) if not set(x.split('||')).isdisjoint(specified_classes) else 'other')477 specified_classes.append('other')478 479 # Separate the DataFrame into test and training sets based on test_videos480 if 'Test' in df_keys and test_videos:481 test_df = df[df['Test']]482 train_df = df[~df['Test']]483 elif test_videos:484 test_df = df[df['Images'].str.contains('|'.join(test_videos), na=False)]485 train_df = df[~df['Images'].str.contains('|'.join(test_videos), na=False)]486 else:487 test_df = pd.DataFrame(columns=df.columns)488 train_df = df489 490 # Print the number of frames in each class before balancing491 label_counts = train_df['Label'].value_counts()492 print("\nNumber of training frames in each class before balancing:")493 print(label_counts)494 495 if max_class_size:496 balanced_train_df = pd.concat([497 group.sample(n=min(len(group), max_class_size), random_state=1)498 for label, group in train_df.groupby('Label')499 ])500 else:501 balanced_train_df = train_df502 503 # Shuffle the training DataFrame504 if shuffle_data:505 balanced_train_df = balanced_train_df.sample(frac=1).reset_index(drop=True)506 507 # Convert training set back to numpy array and list508 if not "Images" in df_keys:509 balanced_train_embeddings = balanced_train_df.drop(columns=['Label', 'Frame', 'Source', 'Test','View','Condition']).to_numpy()510 balanced_train_labels = balanced_train_df['Label'].tolist()511 balanced_train_images = balanced_train_df['Frame'].tolist()512 513 # Convert test set back to numpy array and list514 test_embeddings = test_df.drop(columns=['Label', 'Frame', 'Source', 'Test','View','Condition']).to_numpy()515 test_labels = test_df['Label'].tolist()516 test_images = test_df['Frame'].tolist()517 else:518 # Convert training set back to numpy array and list519 balanced_train_embeddings = balanced_train_df.drop(columns=['Label', 'Images']).to_numpy()520 balanced_train_labels = balanced_train_df['Label'].tolist()521 balanced_train_images = balanced_train_df['Images'].tolist()522 523 # Convert test set back to numpy array and list524 if 'Test' in test_df:525 test_embeddings = test_df.drop(columns=['Label', 'Images', 'Test']).to_numpy()526 else:527 test_embeddings = test_df.drop(columns=['Label', 'Images']).to_numpy()528 529 test_labels = test_df['Label'].tolist()530 test_images = test_df['Images'].tolist()531 532 # Print the number of frames in each class after balancing533 if specified_classes or max_class_size:534 balanced_label_counts = Counter(balanced_train_labels)535 print("\nNumber of training frames in each class after balancing:")536 print(balanced_label_counts)537 538 test_label_counts = test_df['Label'].value_counts()539 # print("\nNumber of testing frames in each class:")540 print(test_label_counts)541 542 return balanced_train_embeddings, balanced_train_labels, balanced_train_images, test_embeddings, test_labels, test_images543 544def multiclass_merge_and_filter_bouts(multiclass_vector, bout_threshold, proximity_threshold):545 # Get the unique labels in the multiclass vector (excluding zero, assuming zero is the background/no label)546 unique_labels = np.unique(multiclass_vector)547 unique_labels = unique_labels[unique_labels != 0]548 549 # Initialize a vector to store the merged and filtered multiclass vector550 merged_vector = np.zeros_like(multiclass_vector)551 552 for label in unique_labels:553 # Create a binary vector for the current label554 binary_vector = (multiclass_vector == label)555 556 # Find the start and end indices of all sequences of 1's for this label557 starts = np.where(np.diff(np.concatenate(([0], binary_vector))) == 1)[0]558 ends = np.where(np.diff(np.concatenate((binary_vector, [0]))) == -1)[0]559 560 # Step 1: Merge close short bouts561 i = 0562 while i < len(starts) - 1:563 # Check if the gap between the end of the current bout and the start of the next bout564 # is within the proximity threshold565 if starts[i + 1] - ends[i] <= proximity_threshold:566 # Merge the two bouts by setting all elements between the start of the first567 # and the end of the second bout to 1568 binary_vector[ends[i]:starts[i + 1]] = 1569 # Remove the next bout from consideration570 starts = np.delete(starts, i + 1)571 ends = np.delete(ends, i)572 else:573 i += 1574 575 # Update the starts and ends after merging576 starts = np.where(np.diff(np.concatenate(([0], binary_vector))) == 1)[0]577 ends = np.where(np.diff(np.concatenate((binary_vector, [0]))) == -1)[0]578 579 # Step 2: Remove standalone short bouts580 for i in range(len(starts)):581 # Check the length of the bout582 length_of_bout = ends[i] - starts[i] + 1583 584 # If the length is less than the threshold, set those elements to 0585 if length_of_bout < bout_threshold:586 binary_vector[starts[i]:ends[i] + 1] = 0587 588 # Combine the binary vector with the merged_vector, ensuring only the current label is set589 merged_vector[binary_vector] = label590 591 # Return the filtered multiclass vector592 return merged_vector593 594def get_unique_labels(label_list: list[str]):595 label_set = set()596 for label in label_list:597 individual_labels = label.split('||')598 for individual_label in individual_labels:599 label_set.add(individual_label)600 return list(label_set)601 602def get_train_test_split(train_embeds, numerical_labels, test_size=0.05, random_state=42):603 return train_test_split(train_embeds, numerical_labels, test_size=test_size, random_state=random_state)604 605def train_model(X_train, y_train, random_state=42):606 # Train SVM Classifier607 svm_clf = SVC(kernel='rbf', random_state=random_state, probability=True)608 svm_clf.fit(X_train, y_train)609 return svm_clf610 611def pickle_model(model):612 pickled = io.BytesIO()613 pickle.dump(model, pickled)614 return pickled615 616def get_seq_io_reader(uploaded_file):617 assert uploaded_file.name[-3:]=='seq', 'Not a seq file'618 with NamedTemporaryFile(suffix="seq", delete=False) as temp:619 temp.write(uploaded_file.getvalue())620 sr = seqIo_reader(temp.name)621 return sr622 623def seq_to_arr(sr):624 N = sr.header['numFrames']625 images = []626 for f in range(N):627 I, ts = sr.getFrame(f)628 images.append(I)629 return np.array(images)630 631def get_2d_embedding(embeddings: pd.DataFrame):632 tsne = TSNE(n_jobs=4, n_components=2, random_state=42, perplexity=50)633 embedding_2d = tsne.fit_transform(np.array(embeddings))634 return embedding_2d635 636 637 638