Trainable implementation of full autoencoder.

configs/data_configs/alm_side.json

@@ -0,0 +1,8 @@
+{
+"proportional_h_coord_top": 0.21666666666666667,
+"target_h": 120,
+"target_w": 112,
+"extension": ".png",
+"trial_pattern": "^CW44_20240522153954_side_trial_"
+
+}

configs/model_configs/alm_default.json

@@ -0,0 +1,24 @@
+{
+"num_blocks": 2,
+"in_channels_0": 1,
+"out_channels_0": 16,
+"in_channels_1": 16,
+"out_channels_1": 32,
+"kernel_preconv": 3,
+"stride_preconv": 1,
+"pool_size_preconv_0": 2,
+"pool_size_preconv_1": 4,
+"use_batch_norm_preconv": false,
+"kernel_residual": 3,
+"stride_residual": 1,
+"use_batch_norm_residual": false,
+"pool_size_residual_0": null,
+"pool_size_residual_1": 4,
+"n_layers_residual": 3,
+"out_conv": 288,
+"out_linear": 128,
+"embed_size": 16,
+"use_batch_norm_linear": false,
+"image_height": 120,
+"image_width": 112
+}

configs/train_configs/alm_default.json

@@ -0,0 +1,17 @@
+{
+"scheduler":"linear",
+"start_factor":0.1,
+"end_factor":1,
+"warmup_steps":10,
+"max_epochs":500,
+"weight_decay":0,
+"l2_weight":1,
+"learning_rate":0,
+"batch_size":10,
+"num_workers":2,
+"subsample_rate":10,
+"subsample_offset":0,
+"max_epochs":500,
+"accelerator":"gpu"
+"fast_dev_run":false
+}

configs/train_configs/alm_default_dev.json

@@ -0,0 +1,17 @@
+{
+"scheduler":"linear",
+"start_factor":0.1,
+"end_factor":1,
+"warmup_steps":10,
+"max_epochs":500,
+"weight_decay":0,
+"l2_weight":1,
+"learning_rate":0,
+"batch_size":10,
+"num_workers":2,
+"subsample_rate":10,
+"subsample_offset":0,
+"max_epochs":100,
+"accelerator":"cpu",
+"fast_dev_run":true
+}

notes/preprocessing.md

@@ -0,0 +1,13 @@
+# Preprocessing description
+
+## Default cropping 
+
+[precrop](../figures/CW44_naive_pre_crop_img.png)
+
+[crop](../figures/CW44_naive_default_crop_img.png)
+
+## Antialiasing 
+
+Differences greater than 1e-3. Due to aliasing I'm pretty sure. 
+
+[alias](../figures/preprocessing_difference_alias.png)

requirements_cpu.txt

@@ -1,4 +1,10 @@
-pip install torch  # should specify a version later
-pip install numpy
-pip install matplotlib
-pip install pytest
+torch  # should specify a version later
+torchvision
+joblib
+pytorch_lightning
+numpy
+matplotlib
+pytest
+opencv-python
+fire
+scikit-image

scripts/eval_single_session.py

@@ -0,0 +1,98 @@
+"""
+Evaluate a single session autoencoder. 
+
+Assumes that we are given a path to a directory `preds/{modeltype}/date/time/`. Will use the configuration parameters stored there 
+"""
+import os
+import datetime
+from tqdm import tqdm
+import json
+import numpy as np
+import fire
+import torch
+import pytorch_lightning as pl
+from behavioral_autoencoder.module import SingleSessionModule
+from behavioral_autoencoder.dataloading import SessionFramesDataModule
+from behavioral_autoencoder.dataset import CropResizeProportion
+
+here = os.path.join(os.path.abspath(os.path.dirname(__file__)))
+
+def eval_trialwise(model,test_dataset,mean_image,path):
+    """
+    """
+    for i, image_sequence in tqdm(enumerate(test_dataset)):
+        reconstructs,latents = model(image_sequence[None,:].cuda())
+        reconstructs_centered = reconstructs + mean_image.cuda()
+        folder = test_dataset.trial_folders[i]
+        savepath = os.path.join(path,folder)
+        try:
+            os.mkdir(savepath)
+        except FileExistsError:    
+            pass
+        np.save(os.path.join(savepath,"reconstruct.npy"),reconstructs_centered.cpu().detach().numpy())
+        np.save(os.path.join(savepath,"latents.npy"),latents.cpu().detach().numpy())
+
+def main(data_path,data_config_path,eval_config_path):
+    """
+    get a model path, and use it to load in a given model. 
+    """
+    saved_checkpoint_path = os.path.join(".","models","single_session_autoencoder","03-07-25","18_05_06","epoch=99-step=33100.ckpt")
+    data_dir = os.path.join("home","ubuntu","Data","CW35","2023_12_15","Frames")
+    metadata_dir = os.path.join(".","preds","single_session_autoencoder","03-07-25","18_05_06")
+    video_fps = 400
+    delay_start_time = 2.5+1.3 ## pre-sample and sample time intervals. 
+    delay_end_time = 2.5+1.3+3 ## delay is 3 seconds.
+    subsample = 10
+    eval_batch_size=1
+
+    ## Load in data related stuff
+
+    with open(data_config_path,"r") as f:
+        data_process_config = json.load(f)
+
+    with open(eval_config_path,"r") as f: 
+        eval_config = json.load(f)
+
+    alm_cropping = CropResizeProportion(data_config_path)
+    data_config = {
+            "data_path":data_path,
+            "transform":alm_cropping,
+            "extension":data_process_config["extension"],
+            "trial_pattern":data_process_config["trial_pattern"],
+            "frame_subset":[f"frame_{i:06d}.png" for i in np.arange(int(delay_start_time*video_fps),int(delay_end_time*video_fps),subsample)]
+            }
+
+    date = datetime.datetime.now().strftime("%m-%d-%y")
+    time = datetime.datetime.now().strftime("%H_%M_%S")
+    datestamp_eval = os.path.join(here,"eval",date)
+    timestamp_eval = os.path.join(here,"eval",date,time)
+    for path in [datestamp_eval,timestamp_eval]:
+        try:
+            os.mkdir(path)
+        except FileExistsError:    
+            pass
+
+    sfdm = SessionFramesDataModule(
+            data_config,
+            eval_config["batch_size"],
+            eval_config["num_workers"],
+            eval_config["subsample_rate"],
+            eval_config["subsample_offset"],
+            eval_config["val_subsample_rate"],
+            eval_config["val_subsample_offset"]
+            )
+
+    model = SingleSessionModule.load_from_checkpoint(
+            checkpoint_path=saved_checkpoint_path
+            )
+
+    sfdm.setup("test")
+
+    eval_trialwise(model,sfdm.dataset,sfdm.mean_image,path)
+
+
+    import pdb; pdb.set_trace()
+
+if __name__ == "__main__":
+    fire.Fire(main)
+

scripts/train_single_session.py

@@ -0,0 +1,121 @@
+"""Train a single session autoencoder on provided data. 
+
+Saves checkpoints for the corresponding model, outputs to tensorboard, and finally dumps all predictions and latents into a save directory. 
+
+"""
+import os
+import fire
+import json
+import joblib
+import datetime
+import pytorch_lightning as pl
+from behavioral_autoencoder.module import SingleSessionModule
+from behavioral_autoencoder.dataset import CropResizeProportion
+from behavioral_autoencoder.dataloading import SessionFramesDataModule
+from behavioral_autoencoder.eval import get_all_predicts_latents,get_dl_predicts_latents
+from pytorch_lightning.callbacks import ModelCheckpoint,LearningRateMonitor
+from pytorch_lightning.loggers import TensorBoardLogger
+
+here = os.path.join(os.path.abspath(os.path.dirname(__file__)))
+
+def main(model_config_path, train_config_path, data_path, data_config_path):
+    """This main function takes as input four paths. These paths indicate the model configuration parameters, training configuration parameters, path to the data directory, and cropping configuration, respectively. By default we assume that we are training a single session autoencoder.  
+    """
+    print("\n=== Starting Single Session Autoencoder Training ===")
+
+    ## Model setup 
+    print("\nLoading configurations...")
+    with open(model_config_path,"r") as f:
+        model_config = json.load(f)
+    with open(train_config_path,"r") as f:
+        train_config = json.load(f)
+    model_name = "single_session_autoencoder"
+
+    print(f"Model config: {model_config}")
+    print(f"Training config: {train_config}")
+
+    hparams = {
+            "model":model_name,
+            "model_config":model_config,
+            "train_config":train_config
+            }
+
+    print("\nInitializing model...")
+    ssm = SingleSessionModule(hparams)
+
+    ## Data setup 
+    with open(data_config_path,"r") as f:
+        data_process_config = json.load(f)
+    print("\nSetting up data...")
+    alm_cropping = CropResizeProportion(data_config_path)
+    data_config = {
+            "data_path":data_path,
+            "transform":alm_cropping,
+            "data_config_path":data_config_path,
+            "extension":data_process_config["extension"],
+            "trial_pattern":data_process_config["trial_pattern"]
+            }
+    print(f"Data config: {data_config}")
+
+    print("Initializing data module...")
+    sfdm = SessionFramesDataModule(
+            data_config,
+            train_config["batch_size"],
+            train_config["num_workers"],
+            train_config["subsample_rate"],
+            train_config["subsample_offset"],
+            train_config["val_subsample_rate"],
+            train_config["val_subsample_offset"]
+            )
+
+    import pdb; pdb.set_trace()
+    ## Set up logging and trainer
+    print("\nSetting up logging and checkpoints...")
+    date=datetime.datetime.now().strftime("%m-%d-%y")
+    time=datetime.datetime.now().strftime("%H_%M_%S")
+    timestamp_model = os.path.join(here,"models",model_name,date,time)
+    timestamp_pred = os.path.join(here,"preds",model_name,date,time)
+    print(f"Model will be saved to: {timestamp_model}")
+    print(f"Predictions will be saved to: {timestamp_pred}")
+
+    logger = TensorBoardLogger("tb_logs",name="test_single_session_auto",log_graph=True)
+    checkpoint = ModelCheckpoint(monitor="mse/val", mode="min", save_last=True, dirpath=timestamp_model)
+    lr_monitor = LearningRateMonitor(logging_interval='epoch')
+
+    print("\nInitializing trainer...")
+    trainer = pl.Trainer(
+        fast_dev_run=train_config["fast_dev_run"],
+        max_epochs=train_config["max_epochs"],
+        accelerator=train_config["accelerator"],
+        enable_checkpointing=True,
+        callbacks=[checkpoint,lr_monitor],
+        log_every_n_steps=1,
+        logger=logger,
+        enable_progress_bar=True,
+    )
+
+    ## Fit the model
+    print(f"\nStarting training for {train_config['max_epochs']} epochs...")
+    trainer.fit(ssm,sfdm)
+    print("Training completed!")
+
+    ## Get out predictions 
+    print("\nGenerating predictions and latents...")
+    preds,latents = get_dl_predicts_latents(ssm,sfdm.val_dataloader(),sfdm.mean_image,train_config["batch_size"],train_config["num_workers"])
+
+    ## Save out all relevant metadata
+    print("\nSaving results...")
+    os.makedirs(timestamp_pred, exist_ok=True)
+    joblib.dump(preds,os.path.join(timestamp_pred,"preds"))
+    joblib.dump(latents,os.path.join(timestamp_pred,"latents"))
+    with open(os.path.join(timestamp_pred,"model_config"),"w") as f:
+        json.dump(hparams, f)
+    with open(os.path.join(timestamp_pred,"data_config"),"w") as f:
+        data_config["data_path"] = data_path
+        json.dump(data_config, f)
+
+    print("\n=== Training Complete ===")
+    print(f"Results saved to: {timestamp_pred}")
+
+if __name__ == "__main__":
+    fire.Fire(main)

src/behavioral_autoencoder/README.md

@@ -0,0 +1,7 @@
+# README 
+This package is organized as follows: 
+- `networks.py` contains actual descriptions of network architectures which can be used as components of an autoencoder model. 
+- `module.py` contains the logic which specifies which combination of network architectures correspond to what kind of models. 
+- `metrics.py` contains custom build evaluation metrics. 
+- `dataloading.py` contains code for the dataloaders with preprocessing. 
+- `dataset.py` contains code to specify how datasets should be structured.

src/behavioral_autoencoder/data_utils.py

@@ -64,7 +64,7 @@ def check_exists(frame_dir: Path) -> bool:
         return response == "y"
     return True
 
-def main(video_dir,frame_dir,video_suffix=".avi"):
+def main(video_dir,frame_dir,video_suffix=".avi",match_str=None):
     """Given a directory of videos, write to a different directory with the following structure:
         1. one subdirectory per video file, as well as a metadata file `annotations.txt`.
         2. within each subdirectory, pngs per individual frames.
@@ -76,12 +76,15 @@ def main(video_dir,frame_dir,video_suffix=".avi"):
             video_dir: directory containing video files. 
             frame_dir: directory to write frames to. 
             video_suffix (default=".avi"): suffix of video files to consider. 
+            match_str: string to find within the video names to write only a subset. 
     """
     video_dir = Path(video_dir)
     frame_dir = Path(frame_dir)
     # 1. Get a directory which contains video files. Store video file names.  
     video_files = os.listdir(video_dir)
     video_files = [f for f in video_files if f.endswith(video_suffix)]
+    if match_str is not None:
+        video_files = [f for f in video_files if match_str in f]
     # 2. Check that the directory we care about exists. If it doesn't create. If it does, ask user.  
     video_files_write = []
     for video_file in video_files:
@@ -99,7 +102,6 @@ def main(video_dir,frame_dir,video_suffix=".avi"):
         with open(frame_dir / "annotations.txt", "a") as f:
             f.write(f"{video_dir_name} {first} {last} 0\n")
 
-def npy_to_frames()
 
 if __name__ == "__main__":
     fire.Fire(main)