posetail 0.1.0__tar.gz

posetail-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, Allen Institute for Neural Dynamics
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

posetail-0.1.0/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: posetail
+Version: 0.1.0
+Summary: a model for tracking 2d or 3d animal pose through time
+License: BSD-3-Clause
+Project-URL: Repository, https://github.com/AllenNeuralDynamics/posetail
+Keywords: 3d point tracking,2d point tracking,multi-view tracking,pose estimation
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: aniposelib>=0.7.0
+Requires-Dist: easydict==1.13
+Requires-Dist: einops==0.7.0
+Requires-Dist: opencv-python==4.9.0.80
+Requires-Dist: toml==0.10.2
+Requires-Dist: wandb==0.19.5
+Requires-Dist: lightning<3,>=2.5.6
+Requires-Dist: rerun-sdk[notebook]<=0.22.1
+Requires-Dist: numpy==1.26
+Requires-Dist: imgaug
+Requires-Dist: imagecorruptions
+Requires-Dist: setuptools<81
+Requires-Dist: schedulefree
+Requires-Dist: decord
+Requires-Dist: tqdm
+Requires-Dist: fastapi>=0.111
+Requires-Dist: uvicorn[standard]>=0.29
+Requires-Dist: python-multipart>=0.0.9
+Requires-Dist: requests>=2.31
+Dynamic: license-file
+
+# posetail

posetail-0.1.0/README.md

@@ -0,0 +1 @@
+# posetail

posetail-0.1.0/inference_utils.py

@@ -0,0 +1,256 @@
+import os
+import cv2 
+import glob
+
+import torch
+
+import numpy as np
+
+from collections import defaultdict
+
+from posetail.datasets.utils import get_dirs
+from train_utils import *
+
+
+
+def get_checkpoint(wandb_prefix, run_id, checkpoint = None):
+
+    if checkpoint is not None: 
+        checkpoint_fmt = str(checkpoint).zfill(8)
+        checkpoint_path = os.path.join(
+            wandb_prefix, run_id, 'files', 'checkpoints', 
+            f'checkpoint_{checkpoint_fmt}.pth')
+        
+    else:
+        checkpoints = sorted(glob.glob(
+            os.path.join(wandb_prefix, run_id, 'files', 'checkpoints', '*.pth')))
+        checkpoint_path = checkpoints[-1]
+
+    return checkpoint_path
+
+
+def load_predictions(data_path, device):
+
+    data = np.load(data_path) 
+
+    coords_pred = torch.from_numpy(data['coords_pred']).to(device)
+    vis_pred = torch.from_numpy(data['vis_pred']).to(device)
+    conf_pred = torch.from_numpy(data['conf_pred']).to(device)
+
+    coords_true = torch.from_numpy(data['coords_true']).to(device)
+    vis_true = torch.from_numpy(data['vis_true']).to(device)
+
+    fnums = torch.from_numpy(data['fnums']).to(device)
+    video_path = ''.join(data['video_path'])
+
+    return coords_pred, vis_pred, conf_pred, coords_true, vis_true, fnums, video_path
+
+
+def combine_predictions(prefix): 
+
+    # traverse results for a particular dataset
+    for session in get_dirs(prefix): 
+
+        session_path = os.path.join(prefix, session)
+
+        for trial in get_dirs(session_path): 
+
+            trial_path = os.path.join(session_path, trial)
+
+            # skip if there are no npz prediction files
+            prediction_paths = sorted(glob.glob(os.path.join(trial_path, 'predictions', 'predictions_*.npz')))
+            if len(prediction_paths) == 0:
+                print(f'skipping... no prediction paths found at {trial_path}')
+                continue
+
+            data = [np.load(p) for p in prediction_paths] 
+
+            # extract metadata 
+            keys_to_exclude = ['coords_pred', 'vis_pred', 'conf_pred',
+                               'coords_true', 'vis_true', 'fnums']
+
+            sample_info = {k: data[0][k] for k in data[0].keys() if k not in keys_to_exclude}
+
+            # combine predictions from each consecutive time period
+            coords_pred = np.concatenate([d['coords_pred'] for d in data], axis = 0)
+
+            vis_pred = np.concatenate([d['vis_pred'] for d in data], axis = 0)
+            conf_pred = np.concatenate([d['conf_pred']for d in data], axis = 0)
+            coords_true = np.concatenate([d['coords_true'] for d in data], axis = 0)
+            vis_true = np.concatenate([d['vis_true'] for d in data], axis = 0)
+            fnums = np.concatenate([d['fnums'] for d in data], axis = 0)
+
+            results = {
+                'coords_pred': coords_pred, 
+                'vis_pred': vis_pred, 
+                'conf_pred': conf_pred,
+                'coords_true': coords_true,
+                'vis_true': vis_true,
+                'fnums': fnums, 
+            }
+
+            results.update(sample_info)
+
+            # save combined data 
+            predictions_fname = os.path.join(prefix, session, trial, f'predictions.npz')
+            np.savez(predictions_fname, **results)
+            print(f'predictions saved to {predictions_fname}')
+
+
+def predict_on_dataset_3d(model, dataloader, outpath, device, 
+                          max_kpts = 1000, debug_ix = None):
+
+    torch.set_float32_matmul_precision('high')
+    model.eval()
+
+    for j, batch in enumerate(dataloader):
+
+        if debug_ix and j == debug_ix: 
+            break
+
+        views = [view.to(device) for view in batch.views]
+        coords = batch.coords.to(device)
+        vis = batch.vis
+        fnums = batch.fnums.cpu().numpy()
+        cgroup = batch.cgroup 
+        sample_info = batch.sample_info
+        
+        # fallback if visibilities are not provided
+        if vis is None: 
+            vis = get_vis_true(coords)
+
+        if cgroup: 
+            cgroup = [dict_to_device(cam_dict, device) for cam_dict in cgroup]
+        
+        # can do multiple passes if there are a lot of keypoints to predict 
+        # (helps reduce memory)
+        n_passes = np.ceil(coords.shape[2] / max_kpts).astype(int)
+        coords_pred = []
+        vis_pred = []
+        conf_pred = []
+        coords_true = []
+        vis_true = []
+
+        for i in range(n_passes): 
+
+            coords_subset = coords[:, :, i * max_kpts : i * max_kpts + max_kpts, :]
+            vis_subset = vis[:, :, i * max_kpts : i * max_kpts + max_kpts, :]
+
+             # mask NaNs, don't want to pass in coords that are NaN in the first frame
+            coords_first = coords_subset[:, 0, :, :]  # B, N, R
+            valid_mask = ~torch.isnan(coords_first).any(dim = -1)  # B, N
+            coords_valid = coords_first[:, valid_mask[0], :]  # B, n, R
+
+            # get model predictions given coords in the first frame
+            with torch.no_grad():
+
+                outputs = model(
+                    views = views, 
+                    coords = coords_valid, 
+                    camera_group = cgroup
+                )
+
+                # populate valid predictions in side of coords
+                B, S, N, R = coords_subset.shape
+                output_coords_valid = torch.full((B, S, N, R), float('nan'),
+                    device = coords_subset.device, dtype = coords_subset.dtype)
+                output_vis_valid = torch.full((B, S, N, 1), float('nan'), 
+                    device = coords_subset.device, dtype = coords_subset.dtype)
+                output_conf_valid = torch.full((B, S, N, 1), float('nan'),
+                    device = coords_subset.device, dtype = coords_subset.dtype)
+
+                output_coords_valid[:, :, valid_mask[0], :] = outputs['coords_pred']
+                output_vis_valid[:, :, valid_mask[0], :] = outputs['vis_pred']
+                output_conf_valid[:, :, valid_mask[0], :] = outputs['conf_pred']
+
+                print('true', coords_subset[:, :, 0, :])
+                print('pred', outputs['coords_pred'][:, :, 0, :])
+
+                coords_pred.append(torch.squeeze(outputs['coords_pred'], dim = 0).cpu().numpy())
+                vis_pred.append(torch.squeeze(outputs['vis_pred'], dim = 0).cpu().numpy())
+                conf_pred.append(torch.squeeze(outputs['conf_pred'], dim = 0).cpu().numpy())
+                coords_true.append(torch.squeeze(coords_subset, dim = 0).cpu().numpy())
+                vis_true.append(torch.squeeze(vis_subset, dim = 0).cpu().numpy())
+
+        results = {
+            'coords_pred': np.concatenate(coords_pred, axis = 1), 
+            'vis_pred': np.concatenate(vis_pred, axis = 1), 
+            'conf_pred': np.concatenate(conf_pred, axis = 1),
+            'coords_true': np.concatenate(coords_true, axis = 1),
+            'vis_true': np.concatenate(vis_true, axis = 1),
+        }
+        results.update({'fnums': fnums})
+
+        keys_to_exclude = ['fnums']
+        if sample_info.subject_ids is not None: 
+            results.update({'subject_ids': sample_info.subject_ids})
+        else: 
+            keys_to_exclude.append('subject_ids')
+
+        results.update({k: sample_info[k] for k in sample_info.keys() if k not in keys_to_exclude})
+
+        # save predictions
+        start_ix = str(results['start_ix']).zfill(8)
+        predictions_outpath = os.path.join(
+            outpath, results['session'], results['trial'], 'predictions')
+        os.makedirs(predictions_outpath, exist_ok = True)
+        predictions_fname = os.path.join(predictions_outpath, f'predictions_{start_ix}.npz')
+        np.savez(predictions_fname, **results)
+        print(f'predictions saved to {predictions_fname}')
+
+    return outpath
+
+
+def generate_video_2d(video_path, results_path, outpath, run_id, scale, device): 
+    # NOTE: deprecated for now
+    # TODO: get camera group and project coords
+    (coords_pred, vis_pred, conf_pred, coords_true, 
+    vis_true, fnums, video_path) = load_predictions(results_path, device)
+
+    coords_true = coords_true.cpu().numpy().astype(int)
+    coords_true[..., 0] = coords_true[..., 0] * scale[0]
+    coords_true[..., 1] = coords_true[..., 1] * scale[1]
+
+    coords_pred = coords_pred.cpu().numpy().astype(int)
+    coords_pred[..., 0] = coords_pred[..., 0] * scale[0]
+    coords_pred[..., 1] = coords_pred[..., 1] * scale[1]
+
+    cap = cv2.VideoCapture(video_path)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = 30.0 # cap.get(cv2.CAP_PROP_FPS)
+
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    runid = run_id.split('-')[-1]
+    video_name = os.path.splitext(os.path.basename(video_path))[0]
+    video_outpath = os.path.join(outpath, f'{video_name}_{runid}.mp4')
+    out = cv2.VideoWriter(video_outpath, fourcc, fps, (frame_width, frame_height))
+
+    i = 0
+    j = 0
+    ret = True
+
+    while ret:
+
+        ret, frame = cap.read() 
+
+        if not ret: 
+            break
+        
+        if i not in fnums:
+            out.write(frame)
+
+        else:
+            for coord_true, coord_pred in zip(coords_true[j], coords_pred[j]):
+                cv2.circle(frame, tuple(coord_true), 5, (0, 255, 0), -1)
+                cv2.circle(frame, tuple(coord_pred), 5, (0, 0, 255), -1)
+
+            out.write(frame)
+            j += 1
+
+        i += 1
+
+    cap.release()
+    out.release()
+
+    return video_outpath