simba-uw-tf-dev 5.4.5__tar.gz → 5.4.6__tar.gz

{simba_uw_tf_dev-5.4.5 → simba_uw_tf_dev-5.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: simba_uw_tf_dev
-Version: 5.4.5
+Version: 5.4.6
 Summary: Toolkit for computer classification and analysis of behaviors in experimental animals
 Home-page: https://github.com/sgoldenlab/simba
 Author: Simon Nilsson, Jia Jie Choong, Sophia Hwang

{simba_uw_tf_dev-5.4.5 → simba_uw_tf_dev-5.4.6}/README.md

@@ -1,11 +1,13 @@
 # SimBA (Simple Behavioral Analysis)
 
-[![pypi](https://img.shields.io/pypi/v/simba-uw-tf-dev.svg)]([https://pepy.tech/project/simba-uw-tf-dev](https://pypi.org/project/Simba-UW-tf-dev/))
+SimBA is a toolkit for creating supervised machine-learning classifiers of animal social and non-social behavior from pose-estimation data, without requiring a programming background.
+
+[![pypi](https://img.shields.io/pypi/v/simba-uw-tf-dev.svg)](https://pypi.org/project/Simba-UW-tf-dev/)
 [![py3.10 tests](https://github.com/sgoldenlab/simba/actions/workflows/tests_py310.yml/badge.svg)](https://github.com/sgoldenlab/simba/actions/workflows/tests_py310.yml)
 [![py3.6 tests](https://github.com/sgoldenlab/simba/actions/workflows/tests_py36.yml/badge.svg)](https://github.com/sgoldenlab/simba/actions/workflows/tests_py36.yml)
 ![Docs](https://readthedocs.org/projects/simba-uw-tf-dev/badge/?version=latest&style=flat)
 [![License: Modified BSD 3-Clause](https://img.shields.io/badge/License-Modified%20BSD%203--Clause%20(Academic%2FResearch)-blue.svg)](LICENSE)
-[![Gitter chat](https://badges.gitter.im/USER/REPO.png)](https://gitter.im/SimBA-Resource/community)
+[![Gitter chat](https://badges.gitter.im/SimBA-Resource/community.png)](https://gitter.im/SimBA-Resource/community)
 [![Download: Weights](https://img.shields.io/badge/Download-Weights-orange.svg)](https://osf.io/5t4y9/)
 [![SimBA: listserv](https://img.shields.io/static/v1?label=SimBA&message=listserv&color=blue)](https://docs.google.com/forms/d/e/1FAIpQLSfjbjae0XqNcl7GYOxmqvRsCveG-cmf4p4hBNNJ8gu5vPLHng/viewform)
 [![DOI](https://zenodo.org/badge/206670333.svg)](https://zenodo.org/badge/latestdoi/206670333)
@@ -21,12 +23,14 @@
 <img src="https://github.com/sgoldenlab/simba/blob/master/images/landing_1.gif" />
 </p>
 
+## Quickstart
 
-To install SimBA via pip, use the following command (preferably using python3.6, or python 3.10 if necessary):
 ```bash
-pip install simba-uw-tf-dev
+pip install simba-uw-tf-dev   # Python 3.6 or 3.10
+simba                         # launch the GUI
 ```
-For detailed installation instructions, click [HERE](https://github.com/sgoldenlab/simba/blob/master/docs/installation_new.md)
+
+See **[Scenario 1](docs/Scenario1.md)** for a worked example that takes raw tracking data through to a validated classifier.
 
 ## Installation ⚙️
 
@@ -34,6 +38,26 @@ For detailed installation instructions, click [HERE](https://github.com/sgoldenl
 
 - [Install SimBA using Anaconda](https://github.com/sgoldenlab/simba/blob/master/docs/anaconda_2025.md)
 
+## Scope & inputs
+
+- **Input:** pose-estimation from [DeepLabCut](https://github.com/DeepLabCut/DeepLabCut) (incl. [multi-animal](docs/Multi_animal_pose.md)), [SLEAP](https://sleap.ai), [DeepPoseKit](https://github.com/jgraving/DeepPoseKit), [DANNCE](https://github.com/spoonsso/dannce) (3D), [MARS](https://github.com/neuroethology/MARS), [FaceMap](https://github.com/MouseLand/facemap), [APT](https://github.com/kristinbranson/APT), [SuperAnimal-TopView](docs/superanimal_topview_project.md), [YOLO](https://github.com/ultralytics/ultralytics), or [blob tracking](docs/blob_track.md); user-defined pose schemes supported.
+- **Output:** per-frame behavior classifiers with standard evaluation (precision/recall, learning curves, permutation importance) and [SHAP-based explainability](docs/SHAP.md).
+- **Validation:** classifier libraries validated in mice and rats; all data, models, and annotations available on [OSF](https://osf.io/tmu6y/).
+- **Analyses** — all produce descriptive statistics, export to CSV, and can be split into time-bins:
+    - *Behavior:* bout counts, durations, latency to first event, event frequency, severity scoring.
+    - *Movement:* velocity, distance travelled, path/trajectory metrics.
+    - *Space (ROIs/zones):* time-in-zone, entries, distance moved per zone, behavior × zone; supports animal-anchored zones.
+    - *Social & directional:* which animal — or body part, or ROI — each animal is oriented toward.
+    - *Temporal structure:* behavior-sequence coupling ([FSTTC](docs/FSTTC.md)), burst detection & smoothing ([Kleinberg](docs/kleinberg_filter.md)).
+    - *Turn-key assays:* [spontaneous alternation](docs/spontaneous_alternation.md), pup retrieval, [cue-light](docs/cue_light_tutorial.md), light/dark box, freezing, circling, etc.
+    - *Unsupervised:* dimensionality reduction and clustering of behavior.
+- **Visualizations** — static plots or overlays burned onto the video, mergeable into one multi-panel video (see [Visualizations](docs/Visualizations.md)):
+    - *Tracking overlays:* pose keypoints, skeletons, animal-anchored bounding boxes, blob contours.
+    - *Spatial:* path/trajectory plots, ROI/zone overlays, location heatmaps.
+    - *Classifier:* annotated prediction videos, Gantt charts, classifier heatmaps, probability plots (incl. interactive grapher), SHAP summary plots.
+    - *Movement & directionality:* distance plots, velocity/data line plots, directionality overlays (animal→animal, →body-part, →ROI).
+    - *Assay-specific:* cue-light, light/dark box, spontaneous-alternation, circular/polar plots.
+
 ##  Documentation: Scenario tutorials
 
 To faciliate the initial use of SimBA, we provide several use scenarios. We have created these scenarios around a hypothetical experiment that take a user from initial use (completely new start) all the way through analyzing a complete experiment and then adding additional experimental datasets to an initial project.
@@ -47,32 +71,33 @@ To faciliate the initial use of SimBA, we provide several use scenarios. We have
 ### Scenario 4: [Analyzing and adding new Experimental data to a previously started project](https://github.com/sgoldenlab/simba/blob/master/docs/Scenario4_new.md)
 
 ## Tutorial 📚
+- [Analysing animal directions in SimBA](docs/directionality_between_animals.md) 🧭
 - [API](https://simba-uw-tf-dev.readthedocs.io/en/latest/index.html) 📘
-- [Visualization tools](docs/Visualizations.md) 👁️
-- [Process video using SimBA tools](docs/Tools.md) 🔨
 - [Batch pre-process video using SimBA](docs/tutorial_process_videos.md) 🏭
-- [Using DeepPoseKit in SimBA](docs/DeepPoseKit_in_SimBA.md) 📙
-- [Using DeepLabCut through SimBA](docs/Tutorial_DLC.md) 📗
-- [SimBA generic tutorial](docs/tutorial.md) 📘
-- [SimBA friendly asked questions (FAQ)](docs/FAQ.md) 📕 
-- [SimBA behavioral annotation interface](docs/label_behavior.md) 🏷️
-- [SimBA Advanced behavioral annotation interface](docs/advanced_labelling.md) 🏷️
-- [Using third-party annotation tools in SimBA](docs/third_party_annot.md) 🏷️
-- [Using user-defined ROIs in SimBA](/docs/ROI_tutorial_new.md) 🗺️
-- [Using multi-animal pose (maDLC/SLEAP/APT) in SimBA](/docs/Multi_animal_pose.md) 🐭🐭
-- [Using the SimBA data analysis and export dashboard](docs/plotly_dash.md) 📊
+- [Blob (contour) tracking in SimBA](docs/blob_track.md) 🟣
+- [Bounding boxes in SimBA](docs/anchored_rois.md)📦
+- [Compute feature subsets in SimBA](docs/feature_subsets.md) 📕
+- [Cue-light analyses in SimBA](docs/cue_light_tutorial.md)💡💡
+- [Downloading compressed data from the SimBA OSF repository](https://github.com/sgoldenlab/simba/blob/master/docs/using_OSF.md) 💾
 - [Explainable machine classifications in SimBA (SHAP)](docs/SHAP.md) 🧮
-- [Mutual exclusivity using heuristic rules in SimBA](docs/mutual_exclusivity_heuristic_rules.md) 📗
-- [Compute feature subsets in SimBA](docs/feature_subsets.md) 📕 
 - [Kleinberg markov chain classification smoothing in SimBA](docs/kleinberg_filter.md) 🔗
-- [Cue-light analyses in SimBA](docs/cue_light_tutorial.md)💡💡
-- [Bounding boxes in SimBA](docs/anchored_rois.md)📦
+- [Mutual exclusivity using heuristic rules in SimBA](docs/mutual_exclusivity_heuristic_rules.md) 📗
+- [Process video using SimBA tools](docs/Tools.md) 🔨
+- [Recommended hardware](https://github.com/sgoldenlab/simba/blob/master/misc/system_requirements.md) 🖥️
 - [Reversing the directionality of classifiers in SimBA](docs/reverse_annotations.md) ⏪
+- [SimBA Advanced behavioral annotation interface](docs/advanced_labelling.md) 🏷️
+- [SimBA behavioral annotation interface](docs/label_behavior.md) 🏷️
+- [SimBA friendly asked questions (FAQ)](docs/FAQ.md) 📕
+- [SimBA generic tutorial](docs/tutorial.md) 📘
 - [Spike-time correlation coefficients in SimBA](docs/FSTTC.md) 📔
 - [Spontaneous alternation in SimBA](/docs/spontaneous_alternation.md)🌽
-- [Analysing animal directions in SimBA](docs/directionality_between_animals.md) 🧭
-- [Recommended hardware](https://github.com/sgoldenlab/simba/blob/master/misc/system_requirements.md) 🖥️
-- [Downloading compressed data from the SimBA OSF repository](https://github.com/sgoldenlab/simba/blob/master/docs/using_OSF.md) 💾
+- [Using DeepLabCut through SimBA](docs/Tutorial_DLC.md) 📗
+- [Using DeepPoseKit in SimBA](docs/DeepPoseKit_in_SimBA.md) 📙
+- [Using multi-animal pose (maDLC/SLEAP/APT) in SimBA](/docs/Multi_animal_pose.md) 🐭🐭
+- [Using the SimBA data analysis and export dashboard](docs/plotly_dash.md) 📊
+- [Using third-party annotation tools in SimBA](docs/third_party_annot.md) 🏷️
+- [Using user-defined ROIs in SimBA](docs/roi_tutorial_new_2025.md) 🗺️
+- [Visualization tools](docs/Visualizations.md) 👁️
 
 ## Legacy documentation: General methods
 
@@ -119,75 +144,6 @@ statistics based on animals movements and locations such as:
 Furthermore, the ROI data can  be used to build potentially valuable, additional, features for random forest predictive classifiers. Such features can be used to generate a machine model that classify behaviors that depend on the spatial location of body parts in relation to the ROIs. **CAUTION**: If spatial locations are irrelevant for the behaviour being classified, then such features should *not* be included in the machine model generation as they just 
 only introduce noise.
 
-
-## Apr-16-2023: API Reference
-
-See SimBA on [readthedocs](https://simba-uw-tf-dev.readthedocs.io/) for [API reference](https://simba-uw-tf-dev.readthedocs.io/en/latest/api.html) and [example notebooks](https://simba-uw-tf-dev.readthedocs.io/en/latest/notebooks.html) when you prefer to run methods available in SimBA through the command line over the GUI. If you find that notebook examples are lacking, please reach out to us [Gitter](https://gitter.im/SimBA-Resource/community) or by opening an issue and let's chat about typing up further examples!
-
-## Apr-04-2023: SimBA version 1.55 release
-
-Although core methods remain, the SimBA code and documentation has sprawled significantly. If you are curious about the methods in the GUI, try clicking the header icons which should take you to the relevant documentation.
-
-<p align = "center">
-<img src="/images/doc_inst.png" alt=""
-	title="SimBA" />
-<p/>
-
-As always, reach out to us on [Gitter chat](https://gitter.im/SimBA-Resource/community) or open an issue if you have questions, bug reports of feature requests! 
-
-## Oct-07-2021: SimBA version 1.31 release
-
-We have released a significantly improved GUI for region-of-interest segmentation and analysis.
-[Click here to go to the new ROI documentation page.](https://github.com/sgoldenlab/simba/blob/master/docs/ROI_tutorial_new.md)
-
-## Feb-08-2021: SimBA version 1.3 release
-
-It has been nearly a year since the first public iteration of SimBA was released! We would like to thank the open-source community who have supported us and provided invaluable feedback and motivation to continue developing and supporting SimBA to where it is now. We have recently passed well over 150,000 downloads via pip install across all branches, and average between ~5000 to 10,000 weekly downloads alongside a gitter community of >100 users. We have just passed 15 citations for the SimBA preprint, which was released ~8 months ago. This would not be possible without your support. Thank you.
-
-The newest release of SimBA, v1.3, provides a significant jump in features, quality of life improvements, and bug fixes. Several are highlighted below.
-
-Please update using `pip install simba-uw-tf==1.3.7`, [this version](https://pypi.org/project/Simba-UW-tf/) has native deeplabcut and deepposekit GUI support disabled. Hence, tensorflow is not needed. Pose-estimation developers have created excellent GUIs for their pipelines, and we do a disservice to you by not supporting the most updated versions. SimBA now supports pose-estimation dataframe imports from Deeplabcut, DeepPoseKit, SLEAP, MARS and others. If you are developing a new pose-estimation method and would like it directly supported in SimBA, please let us know! 
-
-### Selected New Features
-- Easy install of SimBA via `pip` - [Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/installation_new.md)
-- Install simba using anaconda - [Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/anaconda_2025.md)
-- Introduction of SHAP for behavioral neuroscience classifier explainability and standarization- [Documentation](/docs/SHAP.md)
-- Plotly integration for immediate data visualization - [Documentation](/docs/plotly_dash.md)
-- Labelling/annotating behaviors with many third-party apps - [Documentation](/docs/third_party_annot.md)
-- Kleinberg Filter for smoothing - [Documentation](/docs/kleinberg_filter.md)
-- ROI Visualization update - [Documentation](/docs/ROI_tutorial.md)
-- User define features extraction - [Documentation](/docs/extractFeatures.md)
-- Quick line plot - Allow user to make line plots with selected bodypart and tracking data (located under **Tools**)
-- Many, many, many, many bug-fixes
-
-
-## June-12-2020: SimBA version 1.2 release
-### New Features
-- Multi-animal DLC support - [Documentation](/docs/Multi_animal_pose.md)
-- Multi-animal SLEAP support - [Documentation](/docs/Multi_animal_pose.md)
-- SimBA 'pseudo-labelling' module - [Documentation](/docs/pseudoLabel.md)
-- Easy install of SimBA via `pip` - [Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/installation_new.md)
-- Plenty of new quality-of-life features (e.g., time-bin analyzes / improved visualizations options) - [Documentation](/docs/Scenario1.md)
-- Many, many, many, many bug-fixes
-
-Please join our [Gitter chat](https://gitter.im/SimBA-Resource/community) if you have any questions, or even if you would simply like to discuss potential applications for SimBA in your work.  Please come by, stay inside, wash your hands, and check on your lab mates reguarly!
-
-## April-25-2020: SimBA pre-print manuscript release
-
-A pre-print SimBA manuscript on bioRxiv! The manuscript details the use of SimBA for generation of social predictive classifiers in rat and mouse resident-intruder protocols - please check it out using the link above. All data, pose-estimation models, and the final classifiers generated in the manuscript, can be accessed through our [OSF repository](https://osf.io/tmu6y/) and through the [Resource](https://github.com/sgoldenlab/simba#resource-) menu further down this page.
-
-## March-05-2020: SimBA version 1.1 release
-### New Features
-- Region of Interest support (ROI Module) - [Documentation](/docs/ROI_tutorial.md)
-- DeepPoseKit support (DPK Module) - [Documentation](/docs/DeepPoseKit_in_SimBA.md)
-- SimBA accepts user-defined pose-configurations (Flexible Annotation Module) - [Documentation](/docs/Pose_config.md)
-- Interactive classifer discrimination thresholding - [Documentation](/docs/validation_tutorial.md#validate-model-on-single-video)
-- Individual discrimination thresholds for classifiers - [Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/Scenario2.md#part-3-run-the-classifier-on-new-data)
-- Heatmap visualizations -[Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/Scenario2.md#part-5--visualizing-machine-predictions)
-- Multi-crop tool - [Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/Tutorial_tools.md#multi-crop-videos)
-- Post-classification validation method for false-positives - [Documentation](/docs/classifier_validation.md#classifier-validation)
-- Many, many, many bug-fixes
-
 ## What is SimBA?
 Several excellent computational frameworks exist that enable high-throughput and consistent tracking of freely moving unmarked animals. Here we introduce and distribute a pipeline that enabled users to use these pose-estimation approaches in combination with behavioral annotation and generation of supervised machine-learning behavioral predictive classifiers. We have developed this pipeline for the analysis of complex social behaviors, but have included the flexibility for users to generate predictive classifiers across other behavioral modalities with minimal effort and no specialized computational background.  
 
@@ -202,8 +158,6 @@ SimBA does not require computer science and programing experience, and SimBA is
 
 **SimBA provides several validated classifer libraries using videos filmed from above at 90° angle with pose-estimation data from 8 body parts per animal; please see our [OSF repository](https://osf.io/tmu6y/) for access to all files. SimBA now accepts any user-defined pose-estimation annotation schemes with the inclusion of the [Flexible Annotation Module in v1.1](https://github.com/sgoldenlab/simba/blob/master/docs/Pose_config.md). SimBA now supports maDLC and SLEAP for similar looking animals with the release of [maDLC/SLEAP module in v1.2](/docs/Multi_animal_pose.md).** 
 
-**Installation note:** SimBA can be installed either with TensorFlow compatability (for generating DeepLabCut, DeepPoseKit and SLEAP pose-estimation models), or without TensorFlow (for stand-alone use with classifiers and other functions). Please choose the appropriate branch for your needs, using pip install. More details are found in the [Installation Documentation](https://github.com/sgoldenlab/simba/blob/master/docs/installation_new.md).
-
 **Listserv for release information:** If you would like to receive notification for new releases of SimBA, please **[fill out this form](https://forms.gle/R47RWN4stNSJBj9D9)** and you will be added to the listserv.
 
 #### Mouse
@@ -218,39 +172,27 @@ SimBA does not require computer science and programing experience, and SimBA is
 ## Pipeline 👷
 ![](https://github.com/sgoldenlab/simba/blob/master/images/overallflow.PNG)
 
-## Resource 💾
-
-All data (classifiers etc.) is available on our [Open Science Framework repository](https://osf.io/tmu6y/). For a schematic overview of the data respository folder structure (as of March-20-2020), click [HERE](https://github.com/sgoldenlab/simba/blob/master/images/OSF_folder_structure_031820.jpg).
-
-### API reference
-- [SimBA on readthedocs](https://simba-uw-tf-dev.readthedocs.io/en/latest/index.html)
-
-### Models
-Below is a link to download trained  behavior classification models to apply it on your dataset
-- [Random forest models](https://osf.io/kwge8/) 🌲
-
-### Docker images
-- [Docker hub](https://hub.docker.com/repositories/goldenlab)
-
-
-### SimBA visualization examples
-- [YouTube playlist](https://www.youtube.com/playlist?list=PLi5Vwf0hhy1R6NDQJ3U28MOUJPfl2YWYl) 📺
-
-### Labelled images
-- [Annotated images for tracking models](https://osf.io/sr3ck/) 📷
-
-### Tracking weights
-- [DeepLabCut tracking weights](https://osf.io/sr3ck/) 🏋️
+## Resources 💾
 
-### Golden Lab webpage
-- [Sam Golden Lab UW](https://goldenneurolab.com/) 🧪🧫🐁
+- **Data, pose models & classifiers** — [OSF repository](https://osf.io/tmu6y/) 💾
+- **Trained classifiers** — [Random forest models](https://osf.io/kwge8/) 🌲
+- **Install / package** — [PyPI](https://pypi.org/project/Simba-UW-tf-dev/) 📦
+- **API reference** — [SimBA on ReadTheDocs](https://simba-uw-tf-dev.readthedocs.io/en/latest/index.html) 📘
+- **Example notebooks** — [Run SimBA from code](https://simba-uw-tf-dev.readthedocs.io/en/latest/notebooks.html) 📓
+- **Docker images** — [Docker Hub](https://hub.docker.com/repositories/goldenlab) 🐳
+- **Visualization examples** — [YouTube playlist](https://www.youtube.com/playlist?list=PLi5Vwf0hhy1R6NDQJ3U28MOUJPfl2YWYl) 📺
+- **Labelled images & tracking weights** — [DeepLabCut annotations/weights (OSF)](https://osf.io/sr3ck/) 📷
+- **Community & support** — [Gitter chat](https://gitter.im/SimBA-Resource/community) 💬
+- **Bug reports & feature requests** — [GitHub Issues](https://github.com/sgoldenlab/simba/issues) 🐛
+- **Golden Lab** — [Sam Golden Lab, UW](https://goldenneurolab.com/) 🧪
+- **Download statistics** — [Live dashboard](https://sronilsson.github.io/download_stats/) 📊
 
-### Downloads
-- [Downloads statistics dashboard](https://sronilsson.github.io/download_stats/)
+## Developer & contact 👨‍💻
 
+SimBA is developed and maintained by **[Simon Nilsson](https://github.com/sronilsson)** ([homepage](https://sronilsson.netlify.app/)). For questions, bug reports, or feature requests, reach out via [GitHub](https://github.com/sronilsson) or [open an issue](https://github.com/sgoldenlab/simba/issues).
 
 ## License 📃
-This project is licensed under the GPLv3. Note that the software is provided 'as is', without warranty of any kind, express or implied. 
+This project is licensed under the BSD 3-Clause License, modified for academic and research use only (see [LICENSE](LICENSE)). Note that the software is provided 'as is', without warranty of any kind, express or implied. 
 
 If you find **any** part of the code or data useful for your own work, please cite us. You can view and download the citation file by clicking the <kbd>Cite this repository</kbd> button at the top right of this page. Thank you 🙏!
 
@@ -265,7 +207,7 @@ If you find **any** part of the code or data useful for your own work, please ci
   		publisher = {Nature Publishing Group},
   		URL = {https://simba-uw-tf-dev.readthedocs.io/},
   		repository = {https://github.com/sgoldenlab/simba},
-    		license = {GPL-3.0}
+    		license = {BSD-3-Clause (Academic/Research Use Only)}
 			}
 ## References 📜
 
@@ -276,5 +218,4 @@ If you find **any** part of the code or data useful for your own work, please ci
 - [Jia Jie Choong](https://github.com/inoejj)
 - [Sophia Hwang](https://github.com/sophihwang26)
 
-#
-Author [Simon N](https://github.com/sronilsson), [JJ Choong](https://github.com/inoejj)
+See the [full credits page](https://simba-uw-tf-dev.readthedocs.io/en/latest/credits.html) for everyone who has contributed to SimBA.

{simba_uw_tf_dev-5.4.5 → simba_uw_tf_dev-5.4.6}/setup.py

@@ -30,7 +30,7 @@ with open(GPU_REQUIREMENTS_PATH, "r") as f:
 # Setup configuration
 setuptools.setup(
     name="simba_uw_tf_dev",
-    version="5.4.5",
+    version="5.4.6",
     author="Simon Nilsson, Jia Jie Choong, Sophia Hwang",
     author_email="sronilsson@gmail.com",
     description="Toolkit for computer classification and analysis of behaviors in experimental animals",

simba_uw_tf_dev-5.4.6/simba/data_processors/cuda/bg_subtractor.py

@@ -0,0 +1,333 @@
+__author__ = "Simon Nilsson"
+
+import math
+import os
+import queue
+import subprocess
+import tempfile
+import threading
+from typing import Optional, Tuple, Union
+
+import cv2
+import numpy as np
+from numba import cuda
+
+try:
+    import PyNvVideoCodec as nvc
+except Exception:
+    nvc = None
+try:
+    import torch
+except Exception:
+    torch = None
+
+from simba.data_processors.cuda.utils import _cuda_luminance_pixel_to_grey
+from simba.utils.checks import (check_file_exist_and_readable,
+                                check_if_dir_exists, check_if_valid_img,
+                                check_if_valid_rgb_tuple, check_int, check_str,
+                                check_valid_boolean)
+from simba.utils.enums import Formats
+from simba.utils.errors import FFMPEGCodecGPUError, SimBAPAckageVersionError
+from simba.utils.printing import SimbaTimer, stdout_information, stdout_success
+from simba.utils.read_write import (get_fn_ext, get_video_meta_data,
+                                    is_video_color,
+                                    read_img_batch_from_video_gpu)
+
+THREADS_PER_BLOCK = (32, 32, 1)
+NVENC_CODECS = ('h264', 'hevc', 'av1')
+
+
+@cuda.jit()
+def _bg_subtraction_kernel(imgs, avg_img, results, is_clr, fg_clr, threshold, bg_clr):
+    """Device: uint8 background subtraction, 3-channel in/out (BGR). Writes bg colour
+    explicitly so the results buffer needs no host pre-fill and can be reused."""
+    x, y, n = cuda.grid(3)
+    if n < 0 or n > (imgs.shape[0] - 1):
+        return
+    if y < 0 or y > (imgs.shape[1] - 1):
+        return
+    if x < 0 or x > (imgs.shape[2] - 1):
+        return
+    if is_clr[0] == 1:
+        b1, g1, r1 = imgs[n][y][x][0], imgs[n][y][x][1], imgs[n][y][x][2]
+        b2, g2, r2 = avg_img[y][x][0], avg_img[y][x][1], avg_img[y][x][2]
+        b_diff = b1 - b2 if b1 > b2 else b2 - b1
+        g_diff = g1 - g2 if g1 > g2 else g2 - g1
+        r_diff = r1 - r2 if r1 > r2 else r2 - r1
+        grey_diff = _cuda_luminance_pixel_to_grey(b_diff, g_diff, r_diff)
+        if grey_diff > threshold[0]:
+            if fg_clr[0] != -1:
+                results[n][y][x][0] = fg_clr[0]
+                results[n][y][x][1] = fg_clr[1]
+                results[n][y][x][2] = fg_clr[2]
+            else:
+                results[n][y][x][0] = b1
+                results[n][y][x][1] = g1
+                results[n][y][x][2] = r1
+        else:
+            results[n][y][x][0] = bg_clr[0]
+            results[n][y][x][1] = bg_clr[1]
+            results[n][y][x][2] = bg_clr[2]
+    else:
+        v1, v2 = imgs[n][y][x][0], avg_img[y][x][0]
+        grey_diff = v1 - v2 if v1 > v2 else v2 - v1
+        out = (v1 if fg_clr[0] != -1 else 255) if grey_diff > threshold[0] else 0
+        results[n][y][x][0] = out
+        results[n][y][x][1] = out
+        results[n][y][x][2] = out
+
+
+@cuda.jit()
+def _bg_subtraction_argb_kernel(imgs, avg_img, results, is_clr, fg_clr, threshold, bg_clr):
+    """Device: RGB(NVDEC) in -> ARGB out for NVENC. R/B are swapped in the luminance
+    call so the mask matches the CPU (bgr24) path. NVENC 'ARGB' wants byte order
+    [B, G, R, A] (little-endian ARGB), verified by per-slot calibration."""
+    x, y, n = cuda.grid(3)
+    if n < 0 or n > (imgs.shape[0] - 1):
+        return
+    if y < 0 or y > (imgs.shape[1] - 1):
+        return
+    if x < 0 or x > (imgs.shape[2] - 1):
+        return
+    if is_clr[0] == 1:
+        r1, g1, b1 = imgs[n][y][x][0], imgs[n][y][x][1], imgs[n][y][x][2]
+        r2, g2, b2 = avg_img[y][x][0], avg_img[y][x][1], avg_img[y][x][2]
+        r_diff = r1 - r2 if r1 > r2 else r2 - r1
+        g_diff = g1 - g2 if g1 > g2 else g2 - g1
+        b_diff = b1 - b2 if b1 > b2 else b2 - b1
+        grey_diff = _cuda_luminance_pixel_to_grey(b_diff, g_diff, r_diff)
+        if grey_diff > threshold[0]:
+            if fg_clr[0] != -1:
+                out_r, out_g, out_b = fg_clr[0], fg_clr[1], fg_clr[2]
+            else:
+                out_r, out_g, out_b = r1, g1, b1
+        else:
+            out_r, out_g, out_b = bg_clr[0], bg_clr[1], bg_clr[2]
+    else:
+        v1, v2 = imgs[n][y][x][0], avg_img[y][x][0]
+        out_r = out_g = out_b = (v1 if fg_clr[0] != -1 else 255) if (v1 - v2 if v1 > v2 else v2 - v1) > threshold[0] else 0
+    results[n][y][x][0] = out_b
+    results[n][y][x][1] = out_g
+    results[n][y][x][2] = out_r
+    results[n][y][x][3] = 255
+
+
+class BackgroundSubtractorCUDA(object):
+    """
+    Remove background from a video using GPU acceleration.
+
+    .. video:: _static/img/video_bg_subtraction.webm
+       :width: 800
+       :autoplay:
+       :loop:
+       :muted:
+       :align: center
+
+    Uses a fully on-GPU pipeline when PyNvVideoCodec is available: frames are decoded on
+    the NVDEC engine straight to device memory, background-subtracted by a CUDA kernel,
+    and re-encoded on the NVENC engine - no host round-trips and no CPU video encoding.
+    When PyNvVideoCodec is not installed, transparently falls back to an ffmpeg-decode +
+    CUDA-kernel + cv2-encode pipeline whose three stages are overlapped across threads.
+
+    .. seealso::
+       Reference CPU implementation: :func:`simba.video_processors.video_processing.video_bg_subtraction`.
+       To create an ``avg_frm``, see :func:`simba.video_processors.video_processing.create_average_frm`.
+
+    .. note::
+       The NVENC back-end encodes H.264 (YUV420), so its output pixels are perceptually
+       equivalent but not bit-identical to the CPU (mpeg4) back-end - both are lossy. The
+       computed foreground/background classification is the same for both.
+
+    .. csv-table:: Expected runtimes by back-end (RTX 4070, 1280x1024, 1802 frames)
+       :header: BACK-END, TIME (S), FPS, SPEEDUP
+       :widths: 45, 18, 18, 18
+       :align: center
+
+       "nvenc (NVDEC -> CUDA -> NVENC)", 2.47, 731, 20x
+       "cv2 fallback (ffmpeg + CUDA + cv2)", 10.87, 166, 4.6x
+       "original bg_subtraction_cuda", 50.08, 36, 1x
+
+    .. csv-table:: NVENC back-end: FPS by image size
+       :file: ../../../docs/tables/bg_subtraction_cuda_nvenc.csv
+       :widths: 40, 30, 30
+       :align: center
+       :header-rows: 1
+
+    .. note::
+       ``batch_size`` does not affect the NVENC back-end, which decodes and encodes
+       frame-by-frame - measured FPS at 1280x1024 is flat across batch sizes
+       (1 -> 729, 100 -> 735, 500 -> 736). ``batch_size`` only tunes the cv2 fallback.
+
+    :param Union[str, os.PathLike] video_path: Path to the video to remove the background from.
+    :param np.ndarray avg_frm: Average frame of the video (BGR), e.g. from :func:`~simba.video_processors.video_processing.create_average_frm`.
+    :param Optional[Union[str, os.PathLike]] save_path: Output path. If None, saved next to the input with a ``_bg_removed`` suffix.
+    :param Optional[Tuple[int, int, int]] bg_clr: Background fill colour. Default (0, 0, 0).
+    :param Optional[Tuple[int, int, int]] fg_clr: Foreground colour. If None, original pixel colours are kept.
+    :param Optional[int] batch_size: Frames processed per batch on the fallback (cv2) back-end. Default 500.
+    :param Optional[int] threshold: Difference threshold (0-255) above which a pixel is foreground. Default 50.
+    :param Optional[bool] use_nvenc: Force the back-end. None (default) = auto (NVENC if available, else cv2). True raises if PyNvVideoCodec is missing.
+    :param str codec: NVENC codec when using the NVENC back-end ('h264', 'hevc', or 'av1'). Default 'h264'.
+    :param int queue_size: Pipeline queue depth for the cv2 fallback back-end. Default 3.
+    :param bool verbose: If True, print progress. Default True.
+
+    :example:
+
+    >>> from simba.video_processors.video_processing import create_average_frm
+    >>> video_path = "/mnt/c/troubleshooting/mitra/project_folder/videos/clipped/592_MA147_Gq_CNO_0515.mp4"
+    >>> avg_frm = create_average_frm(video_path=video_path)
+    >>> subtractor = BackgroundSubtractorCUDA(video_path=video_path, avg_frm=avg_frm, fg_clr=(255, 255, 255))
+    >>> subtractor.run()
+    """
+
+    def __init__(self,
+                 video_path: Union[str, os.PathLike],
+                 avg_frm: np.ndarray,
+                 save_path: Optional[Union[str, os.PathLike]] = None,
+                 bg_clr: Optional[Tuple[int, int, int]] = (0, 0, 0),
+                 fg_clr: Optional[Tuple[int, int, int]] = None,
+                 batch_size: Optional[int] = 500,
+                 threshold: Optional[int] = 50,
+                 use_nvenc: Optional[bool] = None,
+                 codec: str = 'h264',
+                 queue_size: int = 3,
+                 verbose: bool = True):
+
+        check_file_exist_and_readable(file_path=video_path)
+        check_if_valid_img(data=avg_frm, source=self.__class__.__name__)
+        check_if_valid_rgb_tuple(data=bg_clr)
+        check_int(name=f'{self.__class__.__name__} batch_size', value=batch_size, min_value=1)
+        check_int(name=f'{self.__class__.__name__} threshold', value=threshold, min_value=0, max_value=255)
+        check_int(name=f'{self.__class__.__name__} queue_size', value=queue_size, min_value=1)
+        check_str(name=f'{self.__class__.__name__} codec', value=codec, options=NVENC_CODECS)
+        check_valid_boolean(value=verbose, source=f'{self.__class__.__name__} verbose', raise_error=True)
+        if use_nvenc is not None:
+            check_valid_boolean(value=use_nvenc, source=f'{self.__class__.__name__} use_nvenc', raise_error=True)
+        if fg_clr is not None:
+            check_if_valid_rgb_tuple(data=fg_clr)
+        if save_path is not None:
+            check_if_dir_exists(in_dir=os.path.dirname(str(save_path)), source=f'{self.__class__.__name__} save_path')
+
+        if use_nvenc is None or use_nvenc:
+            missing = [name for name, mod in (('PyNvVideoCodec', nvc), ('torch', torch)) if mod is None]
+            nvenc_ok = not missing
+            if use_nvenc and not nvenc_ok:
+                raise SimBAPAckageVersionError(msg=f'use_nvenc=True but the NVENC back-end requires {" and ".join(missing)}, which could not be found in this environment.', source=self.__class__.__name__)
+            use_nvenc = nvenc_ok
+
+        self.video_path = video_path
+        self.batch_size, self.threshold = batch_size, threshold
+        self.use_nvenc, self.codec, self.queue_size, self.verbose = use_nvenc, codec, queue_size, verbose
+        self.video_meta = get_video_meta_data(video_path=video_path)
+        self.w, self.h = self.video_meta['width'], self.video_meta['height']
+        if save_path is None:
+            in_dir, video_name, _ = get_fn_ext(filepath=video_path)
+            save_path = os.path.join(in_dir, f'{video_name}_bg_removed.mp4')
+        self.save_path = save_path
+
+        avg = cv2.resize(avg_frm, (self.w, self.h)).astype(np.uint8)
+        if use_nvenc:
+            avg = avg[..., ::-1].copy()
+        is_color = np.array([1]) if is_video_color(video_path) else np.array([0])
+        fg = np.array(fg_clr).astype(np.int32) if fg_clr is not None else np.array([-1]).astype(np.int32)
+        self.y_dev = cuda.to_device(avg)
+        self.fg_dev = cuda.to_device(fg)
+        self.is_clr_dev = cuda.to_device(is_color)
+        self.thr_dev = cuda.to_device(np.array([threshold]).astype(np.int32))
+        self.bg_dev = cuda.to_device(np.array(bg_clr).astype(np.int32))
+
+    def run(self) -> Union[str, os.PathLike]:
+        """Run background subtraction and write the output video. Returns the save path."""
+        timer = SimbaTimer(start=True)
+        self._run_nvenc() if self.use_nvenc else self._run_cv2()
+        timer.stop_timer()
+        if self.verbose:
+            stdout_success(msg=f'Video saved at {self.save_path} (back-end: {"nvenc" if self.use_nvenc else "cv2"})', elapsed_time=timer.elapsed_time_str)
+        return self.save_path
+
+    def _run_nvenc(self) -> None:
+        """Fully on-GPU back-end: NVDEC decode -> CUDA kernel -> NVENC encode -> mux."""
+        argb_t = torch.empty((1, self.h, self.w, 4), dtype=torch.uint8, device='cuda:0')
+        argb_dev = cuda.as_cuda_array(argb_t)
+        decoder = nvc.SimpleDecoder(str(self.video_path), use_device_memory=True, output_color_type=nvc.OutputColorType.RGB)
+        encoder = nvc.CreateEncoder(self.w, self.h, 'ARGB', False, codec=self.codec)
+        grid = (math.ceil(self.w / 32), math.ceil(self.h / 32), 1)
+        raw_path = os.path.join(tempfile.gettempdir(), f'_bgsub_{os.getpid()}.{self.codec}')
+        n_dec = len(decoder)
+        with open(raw_path, 'wb') as raw:
+            for i in range(n_dec):
+                t = torch.from_dlpack(decoder[i]).unsqueeze(0)
+                x_dev = cuda.as_cuda_array(t)
+                _bg_subtraction_argb_kernel[grid, THREADS_PER_BLOCK](x_dev, self.y_dev, argb_dev, self.is_clr_dev, self.fg_dev, self.thr_dev, self.bg_dev)
+                bitstream = encoder.Encode(argb_t[0])
+                if bitstream:
+                    raw.write(bytes(bitstream))
+                if self.verbose and (i + 1) % 500 == 0:
+                    stdout_information(msg=f'Encoded {i + 1}/{n_dec} frames...', source=self.__class__.__name__)
+            tail = encoder.EndEncode()
+            if tail:
+                raw.write(bytes(tail))
+        command = f'ffmpeg -y -loglevel error -framerate {self.video_meta["fps"]} -i "{raw_path}" -c copy "{self.save_path}"'
+        subprocess.call(command, shell=True)
+        try:
+            os.remove(raw_path)
+        except OSError:
+            pass
+
+    def _run_cv2(self) -> None:
+        """Fallback back-end: ffmpeg GPU decode + CUDA kernel + cv2 encode, with the
+        decode/compute/encode stages overlapped across a thread pipeline."""
+        n = self.video_meta['frame_count']
+        batch_cnt = int(max(1, np.ceil(n / self.batch_size)))
+        frm_batches = np.array_split(np.arange(0, n), batch_cnt)
+        writer = cv2.VideoWriter(str(self.save_path), cv2.VideoWriter_fourcc(*Formats.MP4_CODEC.value), self.video_meta['fps'], (self.w, self.h))
+        results_dev = cuda.device_array((self.batch_size, self.h, self.w, 3), dtype=np.uint8)
+        grid_x, grid_y = math.ceil(self.w / 32), math.ceil(self.h / 32)
+        decode_q: queue.Queue = queue.Queue(maxsize=self.queue_size)
+        encode_q: queue.Queue = queue.Queue(maxsize=self.queue_size)
+        errors = []
+
+        def _decoder():
+            try:
+                for fb in frm_batches:
+                    imgs = read_img_batch_from_video_gpu(video_path=self.video_path, start_frm=fb[0], end_frm=fb[-1], out_format='array')
+                    decode_q.put(np.ascontiguousarray(imgs))
+            except Exception as e:
+                errors.append(('decoder', e))
+            finally:
+                decode_q.put(None)
+
+        def _encoder():
+            try:
+                while True:
+                    res = encode_q.get()
+                    if res is None:
+                        break
+                    for img in res:
+                        writer.write(img)
+            except Exception as e:
+                errors.append(('encoder', e))
+
+        dec_t = threading.Thread(target=_decoder, daemon=True)
+        enc_t = threading.Thread(target=_encoder, daemon=True)
+        dec_t.start()
+        enc_t.start()
+        processed = 0
+        while True:
+            imgs = decode_q.get()
+            if imgs is None:
+                break
+            batch_n = imgs.shape[0]
+            x_dev = cuda.to_device(imgs)
+            bpg = (grid_x, grid_y, math.ceil(batch_n / THREADS_PER_BLOCK[2]))
+            _bg_subtraction_kernel[bpg, THREADS_PER_BLOCK](x_dev, self.y_dev, results_dev, self.is_clr_dev, self.fg_dev, self.thr_dev, self.bg_dev)
+            encode_q.put(results_dev[:batch_n].copy_to_host())
+            processed += 1
+            if self.verbose:
+                stdout_information(msg=f'Processed batch {processed}/{len(frm_batches)}...', source=self.__class__.__name__)
+        encode_q.put(None)
+        dec_t.join()
+        enc_t.join()
+        writer.release()
+        if errors:
+            stage, exc = errors[0]
+            raise FFMPEGCodecGPUError(msg=f'{stage} thread failed: {exc}', source=self.__class__.__name__)