diff --git a/.gitignore b/.gitignore
index e8f12a4..df7e3cf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,73 +1,92 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# Virtual environments
-venv/
-ENV/
-env/
-.venv/
-.env/
-pip-wheel-metadata/
-*.log
-
-# Jupyter Notebook checkpoints
-.ipynb_checkpoints
-
-# VSCode or PyCharm
-.vscode/
-.idea/
-
-# Pytest
-.cache/
-.pytest_cache/
-
-# MyPy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-
-# Docker
-*.pid
-*.tar
-*.log
-*.sqlite
-*.db
-
-# System files
-.DS_Store
-Thumbs.db
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual environments
+venv/
+ENV/
+env/
+.venv/
+.env/
+pip-wheel-metadata/
+*.log
+
+# Jupyter Notebook checkpoints
+.ipynb_checkpoints
+
+# VSCode or PyCharm
+.vscode/
+.idea/
+
+# Pytest
+.cache/
+.pytest_cache/
+
+# MyPy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+
+# Docker
+*.pid
+*.tar
+*.log
+*.sqlite
+*.db
+
+# System files
+.DS_Store
+Thumbs.db
+
+# Videos and Images
+*.jpg
+*.png
+*.mkv
+*.mp4
+
+# Data Folders
+heatseek/counting/clips
+heatseek/counting/detections_413_full_thermal_vid
+heatseek/counting/detections_414_full_thermal_vid
+heatseek/counting/detections
+heatseek/counting/detections_120b
+heatseek/counting/detections_30b
+heatseek/counting/detections_60b
+heatseek/counting/detections_synthetic
+heatseek/counting/detections_test
+heatseek/outputs
+heatseek/counting/verify_detections.ipynb
diff --git a/LICENSE b/LICENSE
index 0014fed..a4b1cce 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,21 +1,21 @@
-MIT License
-
-Copyright (c) 2025 SDZWA
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+MIT License
+
+Copyright (c) 2025 SDZWA
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
index 9d74e17..b25f2eb 100644
--- a/README.md
+++ b/README.md
@@ -1,179 +1,179 @@
-# HeatSeek
-
-**HeatSeek** is a command-line tool for downloading thermal datasets, preprocessing video, training YOLO models, and tracking objects in thermal imagery. The current focus of the package is on thermal imagery of bats.
-
-
- 
- 
-
-
-## Installation
-
-First, clone the repository:
-
-```bash
-git clone ...
-cd heatseek
-```
-It is recommended to build a conda environment for correct dependencies:
-
-```bash
-conda env create -f environment.yml
-```
-Then install the package:
-
-```bash
-pip install . --no-build-isolation
-```
-
-Alternatively, for development:
-
-```bash
-pip install -e . --no-build-isolation
-```
-
----
-
-## Usage
-
-### Download Weights if Necassary
-
-```bash
-heatseek train \
- --data-yaml path/to/data.yaml \
- --weights yolov8s.pt \
- --epochs 400 \
- --batch 16 \
- --imgsz 640
-```
-### Download dataset from RoboFlow
-
-```bash
-heatseek download --api-key API_KEY --workspace WORKSPACE --project PROJECT_NAME --version VERSION_NUM --nc CLASS_NUM --names CLASS_NAME
-```
-
-### Train YOLO Model
-
-```bash
-heatseek train \
- --data-yaml path/to/data.yaml \
- --weights yolov11s.pt \
- --epochs 400 \
- --batch 16 \
- --imgsz 640
-```
-
----
-
-### Preprocess Thermal Video
-
-```bash
-heatseek preprocess \
- --input raw_video.mp4 \
- --output bg_reduced.mp4
-```
-
-Removes background using thermal contrast thresholding.
-
----
-
-### Track Objects
-
-```bash
-heatseek track \
- --input video.mp4 \
- --output tracks.json \
- --weights yolov8.pt
-```
-
-Runs YOLO-based detection and object tracking.
-
----
-
-### Preprocess + Track in One Step
-
-```bash
-heatseek pretrack \
- --input raw.mp4 \
- --preprocessed bg_reduced.mp4 \
- --output tracks.mp4 \
- --weights yolov8.pt
-```
-
-Runs background reduction then tracking in one go.
-
----
-
----
-
-### Density Annotator
-
-```bash
-heatseek annotate \
- --input raw.mp4 /path/to/data/folder\
-```
-
-Click to annotate points for density data. Press 's' to save and 'n' to skip
-
-
- 
-
-
-
-Creates annotations of density data to train density count estimator models.
-
----
-
-
-## Dependencies
-
-Ensure you have the following installed via:
-
-
-```bash
-pip install -r requirements.txt
-```
----
-
-## Project Structure
-
-```
-heatseek/
-├── cli.py # Entry point
-├── data_utils.py # Dataset download and YAML writing
-├── train.py # Training logic
-├── preprocess.py # Background reduction
-├── detect_track.py # Detection and tracking
-```
-
----
-## Acknowledgements
-
-Example images and datasets:
-https://universe.roboflow.com/goingbatty/
-
----
-## License
-
-MIT License
-
-Copyright (c) 2025 SDZWA
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
+# HeatSeek
+
+**HeatSeek** is a command-line tool for downloading thermal datasets, preprocessing video, training YOLO models, and tracking objects in thermal imagery. The current focus of the package is on thermal imagery of bats.
+
+
+
+
+
+
+## Installation
+
+First, clone the repository:
+
+```bash
+git clone ...
+cd heatseek
+```
+It is recommended to build a conda environment for correct dependencies:
+
+```bash
+conda env create -f environment.yml
+```
+Then install the package:
+
+```bash
+pip install . --no-build-isolation
+```
+
+Alternatively, for development:
+
+```bash
+pip install -e . --no-build-isolation
+```
+
+---
+
+## Usage
+
+### Download Weights if Necassary
+
+```bash
+heatseek train \
+ --data-yaml path/to/data.yaml \
+ --weights yolov8s.pt \
+ --epochs 400 \
+ --batch 16 \
+ --imgsz 640
+```
+### Download dataset from RoboFlow
+
+```bash
+heatseek download --api-key API_KEY --workspace WORKSPACE --project PROJECT_NAME --version VERSION_NUM --nc CLASS_NUM --names CLASS_NAME
+```
+
+### Train YOLO Model
+
+```bash
+heatseek train \
+ --data-yaml path/to/data.yaml \
+ --weights yolov11s.pt \
+ --epochs 400 \
+ --batch 16 \
+ --imgsz 640
+```
+
+---
+
+### Preprocess Thermal Video
+
+```bash
+heatseek preprocess \
+ --input raw_video.mp4 \
+ --output bg_reduced.mp4
+```
+
+Removes background using thermal contrast thresholding.
+
+---
+
+### Track Objects
+
+```bash
+heatseek track \
+ --input video.mp4 \
+ --output tracks.json \
+ --weights yolov8.pt
+```
+
+Runs YOLO-based detection and object tracking.
+
+---
+
+### Preprocess + Track in One Step
+
+```bash
+heatseek pretrack \
+ --input raw.mp4 \
+ --preprocessed bg_reduced.mp4 \
+ --output tracks.mp4 \
+ --weights yolov8.pt
+```
+
+Runs background reduction then tracking in one go.
+
+---
+
+---
+
+### Density Annotator
+
+```bash
+heatseek annotate \
+ --input raw.mp4 /path/to/data/folder\
+```
+
+Click to annotate points for density data. Press 's' to save and 'n' to skip
+
+
+
+
+
+
+Creates annotations of density data to train density count estimator models.
+
+---
+
+
+## Dependencies
+
+Ensure you have the following installed via:
+
+
+```bash
+pip install -r requirements.txt
+```
+---
+
+## Project Structure
+
+```
+heatseek/
+├── cli.py # Entry point
+├── data_utils.py # Dataset download and YAML writing
+├── train.py # Training logic
+├── preprocess.py # Background reduction
+├── detect_track.py # Detection and tracking
+```
+
+---
+## Acknowledgements
+
+Example images and datasets:
+https://universe.roboflow.com/goingbatty/
+
+---
+## License
+
+MIT License
+
+Copyright (c) 2025 SDZWA
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
diff --git a/environment.yml b/environment.yml
index 40859ef..c57b1a2 100644
--- a/environment.yml
+++ b/environment.yml
@@ -1,11 +1,11 @@
-name: heatseek-env
-channels:
- - defaults
- - conda-forge
- - pytorch
-
-dependencies:
- - python=3.9
- - pip
- - pip:
- - -r requirements.txt
+name: heatseek-env
+channels:
+ - defaults
+ - conda-forge
+ - pytorch
+
+dependencies:
+ - python=3.9
+ - pip
+ - pip:
+ - -r requirements.txt
diff --git a/examples/boson_record.py b/examples/boson_record.py
new file mode 100644
index 0000000..9e1f4b5
--- /dev/null
+++ b/examples/boson_record.py
@@ -0,0 +1,81 @@
+"""Example script to record thermal imagery using FLIR Boson
+Radiometric camera.
+
+This script demonstrates how to use the 'BosonCapture' class from
+the 'heatseek' package to record both radiometric data and a normalized
+viewable video file from a connected FLIR Boson camera.
+
+Usage:
+ python boson_record.py \
+ --serial_port path/to/serial/port \
+ --video_port path/to/video/port \
+ --raw_output_path path/to/output.npy \
+ --video_output_path path/to/output.mp4 \
+ --bosonsdk_path path/to/BosonSDK/SDK_USER_PERMISSIONS \
+ --recording_time 60
+
+If no arguments are provided, default filenames are used and script records
+for 10s.
+"""
+
+import argparse
+from time import sleep
+from heatseek.boson_capture import BosonCapture
+
+
+def main():
+ """Run exmaple thermal recording using BosonCapture
+
+ Parses arguments, initializes camera, starts recording for specified
+ duration, then stops and releases camera.
+
+ Command-line args:
+ --serial_port (str, opt): Path to serial port.
+ Default is \dev\ttyACM0.
+ --video_port (str, opt): Path to video port.
+ Default is \dev\video0.
+ --raw_output_path (str, opt): Filepath to save
+ radiometric output (.npy)
+ --video_output_path (str, opt): Filepath to save
+ normalized video (.mp4)
+ --bosonsdk_path (str, opt) Filepath to boson sdk folder
+ --recording_time (int, opt): Duration of recording in seconds.
+ Default is 10s.
+
+ Returns:
+ None
+ """
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--serial_port',
+ type=str,
+ help='path to serial port for radiometric data. Default is /dev/ttyACM0.')
+ parser.add_argument('--video_port',
+ type=str,
+ help='path to video port. Default is /dev/video0')
+ parser.add_argument('--raw_output_path',
+ type=str,
+ help='filepath to save raw radiometric output')
+ parser.add_argument('--video_output_path',
+ type=str,
+ help='filepath to save normalized video output')
+ parser.add_argument('--bosonsdk_path',
+ type=str,
+ help='filepath to boson sdk')
+ parser.add_argument('--recording_time',
+ type=int,
+ help='time in s to record')
+ args = parser.parse_args()
+
+ camera = BosonCapture(serial_port=args.serial_port,
+ video_port=args.video_port,
+ sdkpath=args.bosonsdk_path)
+ camera.start_recording(raw=args.raw_output_path,
+ norm=args.video_output_path)
+ sleep(args.recording_time or 10)
+ camera.stop_recording()
+ camera.release_camera()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py
new file mode 100644
index 0000000..b959fe5
--- /dev/null
+++ b/heatseek/boson_capture.py
@@ -0,0 +1,258 @@
+"""Implementation of Capture abstract class for FLIR Boson
+Radiometric Thermal camera
+
+This module provides the 'Boson_Capture' class, which implements the standard
+camera interface defined by the 'Capture' abstract base class. It handles
+initializing the BosonSDK, capturing single frames, recording radiometric
+frames, and creating a viewable mp4 video with a color map.
+"""
+
+import threading
+import sys
+import os
+from time import sleep
+from importlib import import_module
+import numpy as np
+import cv2
+from heatseek.capture import Capture
+
+
+class BosonCapture(Capture):
+ """Camera interface for FLIR Boson Radiometric Thermal camera
+
+ Implements 'Capture' abstract base class for Boson Radiometric thermal
+ camera.
+
+ Attributes:
+ serial_port (str): path to serial port for radiometric data
+ video_port (str): path to video port for noramlized video data
+ cam_api (module): imported FLIR Boson SDK API modulde
+ enums (module): imported FLIR Boson Enum definitions
+ camera (pyClient): camera object created using boson SDK
+ isRecording (bool): True when recording session in progress
+ raw_data_fpath (str): file path to save raw radiometric data
+ viewable_video_fpath (str): file path to save normalized thermal video
+ raw_mm (np.memmap): memory mapped array for storing raw
+ radiometric frames
+ recording_thread (threading.Thread): thread running recording loop
+ height (int): frame pixel height
+ width (int): frame pixel width
+ """
+
+ def __init__(self, serial_port=None, video_port=None, sdkpath=None):
+ """Initializes boson camera interface
+
+ Loads Boson SDK, connects to camera and configures radiometric
+ parameters
+
+ Args:
+ serial_port (str, opt): path to serial port.
+ Defaults to \dev\ttyACM0.
+ video_port (str, opt): path to video port.
+ Defaults to \dev\video0.
+ skdpath (str, opt): path to FLIR Boson SDK folder.
+ Defaults to ~/BosonSDK/SDK_USER_PERMISSIONS
+ """
+
+ self.serial_port = serial_port or '/dev/ttyACM0'
+ self.video_port = video_port or '/dev/video0'
+ self.sdkpath = sdkpath or '~/BosonSDK/SDK_USER_PERMISSIONS'
+ self.recording = False
+ self.height = 256
+ self.width = 320
+ self.raw_data_fpath = None
+ self.viewable_video_fpath = None
+ self.recording_thread = None
+ self.raw_mm = None
+
+ # import Boson SDK
+ path = os.path.expanduser(self.sdkpath)
+ assert os.path.exists(path), 'SDK Path does not exist'
+ sys.path.append(path)
+ self.cam_api = import_module(
+ 'SDK_USER_PERMISSIONS.ClientFiles_Python.Client_API'
+ )
+ self.enums = import_module(
+ 'SDK_USER_PERMISSIONS.ClientFiles_Python.EnumTypes'
+ )
+
+ # create camera object & configure radiometric parameters
+ self.camera = self.cam_api.pyClient(manualport=self.serial_port)
+ self.setup()
+
+ def setup(self):
+ """Configure radiometric parameters for capture
+
+ Sets High Gain Mode, enables TLinear mode, sets IR16 mode,
+ configures window transmissivity, refreshes LUT, & performs FFC
+ """
+
+ print('Setting up Radiometry...')
+
+ # Set Radiometric Parameters
+ self._configure(self.camera.bosonSetGainMode,
+ self.enums.FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN,
+ description='High Gain Mode')
+ self._configure(self.camera.TLinearSetControl,
+ self.enums.FLR_ENABLE_E.FLR_ENABLE,
+ description='TLinear Mode')
+ self._configure(self.camera.sysctrlSetUsbVideoIR16Mode,
+ self.enums.FLR_SYSCTRL_USBIR16_MODE_E.FLR_SYSCTRL_USBIR16_MODE_TLINEAR,
+ description='IR16 Mode')
+ self._configure(self.camera.radiometrySetTransmissionWindow,
+ 100,
+ description='Window Transimission')
+ self._configure(self.camera.TLinearRefreshLUT,
+ self.enums.FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN,
+ description='LUT Refresh')
+ self._configure(self.camera.bosonRunFFC,
+ description='Flat Field Correction')
+
+ def _configure(self, func, *args, delay=1, success_code=0, description=''):
+ """run camera paramter configuration function until it succeeds
+
+ Args:
+ func: function to update camera parameter
+ *args: args for func
+ delay (int): delay before trying again
+ success_code (int): success code returned from func
+ description (str): name of setting
+ """
+
+ print(f'\nSetting {description}')
+ sleep(0.5)
+ result = func(*args)
+ while result != success_code:
+ sleep(delay)
+ result = func(*args)
+ print('Success!')
+
+ def take_image(self):
+ """Capture single image frame from camera.
+
+ Currently placeholder method - will be implemented to return ndarray
+ """
+
+ print('Taking Image- PLACEHOLDER')
+
+ def start_recording(self, raw=None, norm=None):
+ """Begin thread for continuous recording
+
+ Args:
+ raw (str, opt): filepath to save raw radiometric data.
+ Defaults to output.npy
+ norm (str, opt): filepath to save normalized mp4 video.
+ Defaults to output.mp4
+ """
+
+ if self.recording:
+ print('Recording in Progress!')
+ return
+
+ self.raw_data_fpath = raw or 'output.npy'
+ self.viewable_video_fpath = norm or 'output.mp4'
+
+ self.recording = True
+ self.recording_thread = threading.Thread(
+ target=self._record_loop, daemon=True
+ )
+ self.recording_thread.start()
+ print('Staring Recording...')
+
+ def _record_loop(self):
+ """Internal method: recording loop to continuously capture frames
+
+ Saves raw radiometric frames and writes a viewable mp4 video.
+ Stops recording when self.recording is False or preallocated
+ memory is full
+ """
+
+ # video settings
+ cap = cv2.VideoCapture(self.video_port, cv2.CAP_V4L2)
+ cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
+ cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
+ cap.set(cv2.CAP_PROP_CONVERT_RGB, 0)
+ cap.set(cv2.CAP_PROP_FOURCC,
+ cv2.VideoWriter_fourcc('Y', '1', '6', ' '))
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ writer = cv2.VideoWriter(self.viewable_video_fpath,
+ fourcc,
+ 60,
+ (self.width, self.height))
+
+ # raw video setup
+ self.raw_mm = np.memmap(self.raw_data_fpath,
+ dtype=np.uint16,
+ mode='w+',
+ shape=(50000, self.height, self.width))
+ frame_index = 0
+
+ try:
+ while self.recording:
+ ret, frame = cap.read()
+ if not ret:
+ print('Frame grab failed. Stopping Recording')
+ break
+
+ # Viewable frame
+ frame_8bit = cv2.normalize(frame, None, 0, 255,
+ norm_type=cv2.NORM_MINMAX).astype(np.uint8)
+ frame_color = cv2.applyColorMap(frame_8bit,
+ cv2.COLORMAP_INFERNO).astype(np.uint8)
+ writer.write(frame_color)
+
+ # Raw Video
+ if frame_index > self.raw_mm.shape[0]:
+ print('Reached Preallocated Size, Stopping Recording')
+ break
+ self.raw_mm[frame_index] = frame
+ frame_index += 1
+ self.raw_mm.flush()
+
+ finally:
+ self._finalize_recording()
+
+ def stop_recording(self):
+ """Stop ongoing recording session.
+
+ Sets isRecording flag to False, waits for recording thread to finish,
+ and cleans up resources.
+ """
+
+ if not self.recording:
+ print('No recording in progress')
+ return
+
+ print('Stopping Recording...')
+ self.recording = False
+
+ if self.recording_thread is not None:
+ self.recording_thread.join()
+ self.recording_thread = None
+
+ def _finalize_recording(self):
+ """Finalize recording.
+
+ Flushes and deletes memory-mapped raw frame array, closes video writer,
+ and prints file locations.
+ """
+
+ self.recording = False
+
+ if self.raw_mm is not None:
+ self.raw_mm.flush()
+ del self.raw_mm
+ self.raw_mm = None
+
+ print('Recording Successfully Completed')
+ print(f'Radiometric data saved: {self.raw_data_fpath}')
+ print(f'Viewable Video: {self.viewable_video_fpath}')
+
+ def release_camera(self):
+ """Release camera resources.
+
+ Closes camera connection via boson SDK
+ """
+
+ self.camera.Close()
+ print('Release Camera')
diff --git a/heatseek/capture.py b/heatseek/capture.py
new file mode 100644
index 0000000..de53e7d
--- /dev/null
+++ b/heatseek/capture.py
@@ -0,0 +1,70 @@
+"""
+Deifnes an abstract base class for thermal camera interfaces.
+
+This modele provides a generic 'Capture' abstract base class that
+defines a consistent API for different camera hardware implementations.
+Concrete subclasses ('Boson_Capture') can inherit from 'Capture' and
+implement all abstract methods to support their specific SDKs.
+"""
+
+from abc import ABC, abstractmethod
+
+
+class Capture(ABC):
+ """Abstract base class for thermal camera interfaces
+
+ Provides standard API for intializing, taking images, and
+ managing recordings across different camera types.
+ """
+
+ def __init__(self):
+ """Initialize capture interface
+
+ This method should handle hardware initalization and
+ SDK loading
+
+ Args:
+ camera_id (int, opt): Device identifier for camera.
+ Defaults to 0.
+ """
+
+ @abstractmethod
+ def setup(self):
+ """Prepare camera for capture
+
+ This method shoudl handle any configuration of camera settings
+ that needs to happen before data collection
+ """
+
+ @abstractmethod
+ def take_image(self):
+ """Capture and return a single frame
+
+ This method should capture and return a single frame as a np.array
+
+ Returns:
+ Any (np.array): captured image
+ """
+
+ @abstractmethod
+ def start_recording(self):
+ """Begin video recording
+
+ This method should capture frames continuously and save to disk
+ """
+
+ @abstractmethod
+ def stop_recording(self):
+ """Stop ongoing recording
+
+ This method should stop any current video recordings and ensure
+ all collected data is saved properly to disk
+ """
+
+ @abstractmethod
+ def release_camera(self):
+ """Release camera hardware and close
+
+ This method should handle any releasing and/or shutting down of
+ camera hardware
+ """
diff --git a/heatseek/cli.py b/heatseek/cli.py
index f31a2a3..5ffaccc 100644
--- a/heatseek/cli.py
+++ b/heatseek/cli.py
@@ -1,126 +1,126 @@
-#!/usr/bin/env python3
-# heatseek/cli.py
-
-import argparse
-import requests
-from .data_utils import download_dataset, write_data_yaml
-from .train import train
-from .preprocess import reduce_background
-from .detect_track import detect_and_track
-from .density_annotator import annotate_folder
-def main():
- parser = argparse.ArgumentParser(prog="heatseek")
- subs = parser.add_subparsers(dest="cmd", required=True)
-
- # getweights
- getw = subs.add_parser("getweights", help="Download pretrained YOLO weights")
- getw.add_argument(
- "--url",
- default="https://sandiegozoo.box.com/shared/static/0edcehha4y8yli9h0bzettitftbd1jdl.pt",
- help="URL to download weights from",
- )
- getw.add_argument(
- "--output", default="yolo_model.pt", help="Where to save the downloaded weights"
- )
-
- # download (Roboflow)
- dl = subs.add_parser("download", help="Fetch Roboflow dataset and write data.yaml")
- dl.add_argument("--api-key", required=True, help="Roboflow API key")
- dl.add_argument("--workspace", required=True, help="Roboflow workspace name")
- dl.add_argument("--project", required=True, help="Roboflow project name")
- dl.add_argument("--version", type=int, required=True, help="Roboflow version number")
- dl.add_argument("--nc", type=int, default=3, help="Number of classes")
- dl.add_argument(
- "--names",
- nargs="+",
- default=["hot_bat", "cold_bat", "other"],
- help="List of class names",
- )
-
- # train
- tr = subs.add_parser("train", help="Train YOLO on your dataset")
- tr.add_argument("--data-yaml", required=True, help="Path to your data.yaml")
- tr.add_argument("--weights", default="yolo11s.pt", help="Pretrained weights file")
- tr.add_argument("--epochs", type=int, default=400, help="Number of epochs")
- tr.add_argument("--batch", type=int, default=16, help="Batch size")
- tr.add_argument("--imgsz", type=int, default=640, help="Image size (pixels)")
- tr.add_argument(
- "--device", type=int, default=0, help="CUDA device index (uses CPU if none)"
- )
-
- # preprocess
- pre = subs.add_parser("preprocess", help="Reduce background in a video")
- pre.add_argument("--input", required=True, help="Input video path")
- pre.add_argument("--output", required=True, help="Output video path")
- pre.add_argument(
- "--config",
- default="heatseek/config/preproc_config.yaml",
- help="YAML config for optical‐flow thresholds",
- )
-
- # track
- track = subs.add_parser("track", help="Detect + track objects in a video")
- track.add_argument("--input", required=True, help="Input video path")
- track.add_argument("--output", required=True, help="Output path for results")
- track.add_argument("--weights", required=True, help="Weights for detection")
-
- # annotator
- annot = subs.add_parser("annotate", help="Annotate a video with detections")
- annot.add_argument("--folder", required=True, help="Input folder path") #TODO add save folder path
-
- # pretrack
- pretrack = subs.add_parser("pretrack", help="Preprocess (bg-reduce) then detect+track")
- pretrack.add_argument("--input", required=True, help="Input video path")
- pretrack.add_argument(
- "--preprocessed", required=True, help="Where to dump preprocessed video"
- )
- pretrack.add_argument("--output", required=True, help="Output path for tracking")
- pretrack.add_argument("--weights", required=True, help="Weights for detection")
- pretrack.add_argument(
- "--config",
- default="heatseek/config/preproc_config.yaml",
- help="YAML config for optical-flow thresholds",
- )
-
- args = parser.parse_args()
-
- if args.cmd == "getweights":
- print(f"Downloading weights from {args.url}…")
- resp = requests.get(args.url)
- resp.raise_for_status()
- with open(args.output, "wb") as f:
- f.write(resp.content)
- print(f"Weights saved to {args.output}")
-
- elif args.cmd == "download":
- ds_dir = download_dataset(
- args.api_key, args.workspace, args.project, args.version
- )
- yaml_path = write_data_yaml(ds_dir, args.nc, args.names)
- print("→ data.yaml written at", yaml_path)
-
- elif args.cmd == "train":
- train(
- data_yaml=args.data_yaml,
- weights=args.weights,
- epochs=args.epochs,
- batch=args.batch,
- imgsz=args.imgsz,
- device_idx=args.device,
- )
-
- elif args.cmd == "preprocess":
- reduce_background(args.input, args.output, args.config)
-
- elif args.cmd == "track":
- detect_and_track(args.input, args.output, args.weights)
-
- elif args.cmd == "pretrack":
- reduce_background(args.input, args.preprocessed, args.config)
- detect_and_track(args.preprocessed, args.output, args.weights)
- elif args.cmd == "annotate":
- annotate_folder(args.folder, "annotations.json")
-
-
-if __name__ == "__main__":
- main()
+#!/usr/bin/env python3
+# heatseek/cli.py
+
+import argparse
+import requests
+from .data_utils import download_dataset, write_data_yaml
+from .train import train
+from .preprocess import reduce_background
+from .detect_track import detect_and_track
+from .density_annotator import annotate_folder
+def main():
+ parser = argparse.ArgumentParser(prog="heatseek")
+ subs = parser.add_subparsers(dest="cmd", required=True)
+
+ # getweights
+ getw = subs.add_parser("getweights", help="Download pretrained YOLO weights")
+ getw.add_argument(
+ "--url",
+ default="https://sandiegozoo.box.com/shared/static/0edcehha4y8yli9h0bzettitftbd1jdl.pt",
+ help="URL to download weights from",
+ )
+ getw.add_argument(
+ "--output", default="yolo_model.pt", help="Where to save the downloaded weights"
+ )
+
+ # download (Roboflow)
+ dl = subs.add_parser("download", help="Fetch Roboflow dataset and write data.yaml")
+ dl.add_argument("--api-key", required=True, help="Roboflow API key")
+ dl.add_argument("--workspace", required=True, help="Roboflow workspace name")
+ dl.add_argument("--project", required=True, help="Roboflow project name")
+ dl.add_argument("--version", type=int, required=True, help="Roboflow version number")
+ dl.add_argument("--nc", type=int, default=3, help="Number of classes")
+ dl.add_argument(
+ "--names",
+ nargs="+",
+ default=["hot_bat", "cold_bat", "other"],
+ help="List of class names",
+ )
+
+ # train
+ tr = subs.add_parser("train", help="Train YOLO on your dataset")
+ tr.add_argument("--data-yaml", required=True, help="Path to your data.yaml")
+ tr.add_argument("--weights", default="yolo11s.pt", help="Pretrained weights file")
+ tr.add_argument("--epochs", type=int, default=400, help="Number of epochs")
+ tr.add_argument("--batch", type=int, default=16, help="Batch size")
+ tr.add_argument("--imgsz", type=int, default=640, help="Image size (pixels)")
+ tr.add_argument(
+ "--device", type=int, default=0, help="CUDA device index (uses CPU if none)"
+ )
+
+ # preprocess
+ pre = subs.add_parser("preprocess", help="Reduce background in a video")
+ pre.add_argument("--input", required=True, help="Input video path")
+ pre.add_argument("--output", required=True, help="Output video path")
+ pre.add_argument(
+ "--config",
+ default="heatseek/config/preproc_config.yaml",
+ help="YAML config for optical‐flow thresholds",
+ )
+
+ # track
+ track = subs.add_parser("track", help="Detect + track objects in a video")
+ track.add_argument("--input", required=True, help="Input video path")
+ track.add_argument("--output", required=True, help="Output path for results")
+ track.add_argument("--weights", required=True, help="Weights for detection")
+
+ # annotator
+ annot = subs.add_parser("annotate", help="Annotate a video with detections")
+ annot.add_argument("--folder", required=True, help="Input folder path") #TODO add save folder path
+
+ # pretrack
+ pretrack = subs.add_parser("pretrack", help="Preprocess (bg-reduce) then detect+track")
+ pretrack.add_argument("--input", required=True, help="Input video path")
+ pretrack.add_argument(
+ "--preprocessed", required=True, help="Where to dump preprocessed video"
+ )
+ pretrack.add_argument("--output", required=True, help="Output path for tracking")
+ pretrack.add_argument("--weights", required=True, help="Weights for detection")
+ pretrack.add_argument(
+ "--config",
+ default="heatseek/config/preproc_config.yaml",
+ help="YAML config for optical-flow thresholds",
+ )
+
+ args = parser.parse_args()
+
+ if args.cmd == "getweights":
+ print(f"Downloading weights from {args.url}…")
+ resp = requests.get(args.url)
+ resp.raise_for_status()
+ with open(args.output, "wb") as f:
+ f.write(resp.content)
+ print(f"Weights saved to {args.output}")
+
+ elif args.cmd == "download":
+ ds_dir = download_dataset(
+ args.api_key, args.workspace, args.project, args.version
+ )
+ yaml_path = write_data_yaml(ds_dir, args.nc, args.names)
+ print("→ data.yaml written at", yaml_path)
+
+ elif args.cmd == "train":
+ train(
+ data_yaml=args.data_yaml,
+ weights=args.weights,
+ epochs=args.epochs,
+ batch=args.batch,
+ imgsz=args.imgsz,
+ device_idx=args.device,
+ )
+
+ elif args.cmd == "preprocess":
+ reduce_background(args.input, args.output, args.config)
+
+ elif args.cmd == "track":
+ detect_and_track(args.input, args.output, args.weights)
+
+ elif args.cmd == "pretrack":
+ reduce_background(args.input, args.preprocessed, args.config)
+ detect_and_track(args.preprocessed, args.output, args.weights)
+ elif args.cmd == "annotate":
+ annotate_folder(args.folder, "annotations.json")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/heatseek/config/preproc_config.yaml b/heatseek/config/preproc_config.yaml
index ea0127f..f0ccf20 100644
--- a/heatseek/config/preproc_config.yaml
+++ b/heatseek/config/preproc_config.yaml
@@ -1,9 +1,9 @@
-# preproc_config.yaml
-motion_thresh: 2.75
-flow_params:
- pyr_scale: 0.5
- levels: 3
- winsize: 4
- iterations: 3
- poly_n: 3
- poly_sigma: 1.2
+# preproc_config.yaml
+motion_thresh: 2.75
+flow_params:
+ pyr_scale: 0.5
+ levels: 3
+ winsize: 4
+ iterations: 3
+ poly_n: 3
+ poly_sigma: 1.2
diff --git a/heatseek/counting/CountLine.py b/heatseek/counting/CountLine.py
new file mode 100644
index 0000000..f94d927
--- /dev/null
+++ b/heatseek/counting/CountLine.py
@@ -0,0 +1,81 @@
+import numpy as np
+
+class CountLine():
+ '''
+ Counts everytime a bat crosses a line defined in space (y = constant or x = constant)
+ '''
+
+ def __init__(self, line_value, line_dim=1, total_frames=None):
+ '''
+ line_value: the value that defines the position of the
+ line the bats are crossing.
+ line_dim: whether the line is horizontal or verical
+ (0 for vertical, 1 for horizontal)
+ total_frames: total_frames in video
+
+ '''
+ self.line_value = line_value
+ self.line_dim = line_dim
+ self.total_frames = total_frames
+ if total_frames:
+ # How many bats are crossing the line in each frame
+ self.num_crossing = np.zeros(total_frames)
+ self.forward = np.zeros(total_frames)
+ self.backward = np.zeros(total_frames)
+ self.bat_ids_crossed = []
+ # What frames did crosses occur
+ self.frame_cross = []
+
+ def _crossing_frame(self, track, forward):
+ """First frame index where the bat is on the far side of the line."""
+ coords = track['track'][:, self.line_dim]
+ beyond = coords <= self.line_value if forward else coords >= self.line_value
+ return int(np.argmax(beyond)) + track['first_frame']
+
+ def is_crossing(self, track, track_ind):
+ '''
+ Checks if given track is crossing the line.
+
+ track: a bat track
+ track_ind: number track in list of tracks
+
+ returns 1 is forward crossing
+ -1 if backward crossing
+ 0 if no crossing
+ '''
+
+ # Crossing line going away
+ if track['track'][0, self.line_dim] >= self.line_value:
+ # last frame was below line
+ if track['track'][-1, self.line_dim] <= self.line_value:
+ # frame_num = (np.argmin(track['track'][-1, self.line_dim] <= self.line_value)
+ # + track['first_frame'])
+ frame_num = self._crossing_frame(track, forward=True)
+ # this frame above or on line
+ # So bat has crossed line
+ if self.total_frames:
+ self.num_crossing[frame_num] += 1
+ self.forward[frame_num] += 1
+ self.bat_ids_crossed.append(track_ind)
+ self.frame_cross.append(frame_num)
+
+ return (1, frame_num)
+
+
+ # Crossing line coming back
+ if track['track'][0, self.line_dim] <= self.line_value:
+ # last frame was above line
+ if track['track'][-1, self.line_dim] >= self.line_value:
+ # this frame below or on line
+ # So bat has crossed line coming back
+ # frame_num = (np.argmin(track['track'][-1, self.line_dim] >= self.line_value)
+ # + track['first_frame'])
+ frame_num = self._crossing_frame(track, forward=False)
+ if self.total_frames:
+ self.num_crossing[frame_num] -= 1
+ self.backward[frame_num] += 1
+ self.bat_ids_crossed.append(-track_ind)
+ self.frame_cross.append(frame_num)
+ return (-1, frame_num)
+ return (0, None)
+
\ No newline at end of file
diff --git a/heatseek/counting/README.md b/heatseek/counting/README.md
new file mode 100644
index 0000000..dd319b7
--- /dev/null
+++ b/heatseek/counting/README.md
@@ -0,0 +1,128 @@
+# Counting Pipeline
+This pipeline consists of 3 stages:
+
+1. Segmentation
+2. Raw Tracking
+3. Counting-Line Tracking
+
+## Segmentation
+
+This script is the **first stage** of the counting pipeline. It turns a raw thermal video into a set of per-frame blob detections that every downstream stage (tracking, line-crossing counting) consumes.
+
+### What it does
+
+For each frame of the input video it:
+
+1. **Builds a foreground mask.** Two modes are supported, selected by the `radiometric` flag:
+ - **Radiometric** (FLIR Boson): the 8-bit pixel is mapped back to a temperature in °C, and any pixel warmer than `median(background) + threshold` is flagged. Bats read *hotter* (brighter) than the scene.
+ - **Non-radiometric** (generic thermal): any pixel *darker* than `median(background) - threshold` is flagged. Bats are assumed to read *darker* than the scene. This is the active path for the config below (`radiometric: False`).
+2. **Extracts blobs** from the mask via `bat_functions.get_blob_info` — center, area, contour, and bounding rect per detection.
+3. **Previews the mask** on the first processed frame and prompts for confirmation before committing to the full run, so a bad threshold is caught early.
+
+After the pass it:
+
+- **Runs a density sanity check.** `max_bats` computes the per-frame pixel displacement of a bat, `d = (speed × f) / (depth × fps)`, where `f = focal_length / pixel_pitch` is the focal length in pixels. From this it derives a maximum trackable count, `n_max = (width × height) / (16 · d²)`. If the 99th-percentile per-frame detection count exceeds `n_max`, the swarm is too dense to associate reliably and the run aborts. (Further information about this math provided [here](https://drive.google.com/uc?export=download&id=1JrdbtFeGjzSSB6oRCtV8dOb5-wmQ3uFg))
+- **Saves detections** to the output folder as `centers.npy`, `size.npy`, `rects.npy`, and chunked `contours-compressed-*.npy`.
+- **Writes an overlay video** with contours and centers drawn on the original frames for visual QA.
+- **Emits tracker thresholds.** It logs a recommended minimum distance threshold of `ceil(d)` and a maximum of `≤ 2 × min`, which parameterize the downstream tracking stage.
+
+### Why it matters
+
+Detection quality is the ceiling on everything that follows. A missed or split blob here cannot be recovered by the tracker. The density check also acts as a gate. It prevents feeding a swarm into the tracker when the geometry makes correct frame-to-frame association impossible, and it hands the tracker the distance thresholds derived from the same physics.
+
+### Config parameters
+
+| Parameter | Units / Type | Description |
+|---|---|---|
+| `video_path` | path | Input thermal video to process. |
+| `output_folder` | path | Directory where detections and the overlay video are written. |
+| `threshold` | intensity or °C | Mask sensitivity. In **radiometric** mode, degrees Celsius above background. In **non-radiometric** mode, intensity units below background. |
+| `min_val_radiometric` | cK | Lower bound of the radiometric scale used to map 8-bit pixels back to temperature. **Radiometric mode only.** |
+| `max_val_radiometric` | cK | Upper bound of the radiometric scale. **Radiometric mode only.** |
+| `radiometric` | bool | Selects the mask function: `True` → FLIR Boson temperature path, `False` → generic intensity path. |
+| `speed` | m/s | Bat flight speed; use an upper-bound estimate for the safest density check. |
+| `focal_length` | m | Camera focal length. |
+| `pixel_pitch` | m/px | Physical spacing between adjacent pixel centers (typically given in µm; convert to meters if necessary). |
+| `depth` | m | Distance from the camera to the swarm. |
+
+## Raw Tracking
+
+This script is the **second stage** of the counting pipeline. It consumes the per-frame detections produced by the previous stage (`centers.npy`, `size.npy`, `contours-compressed-*.npy`) and links them into persistent, identity-bearing tracks (one trajectory per bat) that the downstream line-crossing counter can use to count each individual exactly once.
+
+### What it does
+
+1. **Splits the video into overlapping segments.** `build_camera_dicts` divides the full frame range into `num_video_segments` chunks. Adjacent chunks overlap by `segment_overlap_seconds × fps` frames so a bat crossing a segment boundary can still be recovered on both sides. The final segment always runs to the end of the video. Segments whose output already exists are skipped, so interrupted runs resume without reprocessing.
+2. **Tracks each segment in parallel.** `run_tracking` distributes the segments across a multiprocessing pool; each worker calls `kbf.find_tracks`, which associates detections frame-to-frame under the distance thresholds below. A track may coast (be interpolated forward without a detection) for up to `max_unseen_frames` frames before it is terminated, so a brief missed detection doesn't fragment one bat into two tracks. The first contours file is dropped to correct an off-by-one alignment against the centers/sizes arrays.
+3. **Merges segments into one track list.** `combine_tracks` stitches the per-segment files into a single `raw_tracks.npy`, keeping only tracks that started before each segment's overlap region so a bat seen in two overlapping segments isn't counted twice. This step is idempotent. It won't overwrite an existing merge.
+4. **Writes an overlay video** (`rt_overlay_video.mp4`) with a distinctly colored dot and heading arrow per track. Coasting frames (where the tracker guessed the position, i.e. `pos_index` is `nan`) are drawn faded and hollow so real detections and interpolated ones are visually distinguishable.
+
+### Why it matters
+
+Detections on their own have no identity. The same bat in two consecutive frames is just two unrelated blobs. This stage resolves those blobs into single trajectories, which is exactly what counting requires (count the track, not the blob). The coasting mechanism prevents count inflation from momentary detection dropouts, and the overlapping-segment scheme lets long videos be tracked in parallel without losing tracks at chunk boundaries.
+
+### Config parameters
+
+| Parameter | Units / Type | Description |
+|---|---|---|
+| `video_path` | path | Source thermal video. Used to read `fps` and to render the overlay; detections are read from `output_folder`. |
+| `output_folder` | path | Directory holding the detection files from stage 1 and where segment tracks, `raw_tracks.npy`, and the overlay are written. |
+| `max_distance_threshold` | px | Maximum association distance. A detection farther than this from a track is not linked to it — effectively a ceiling on per-frame displacement. |
+| `min_distance_threshold` | px | Minimum association distance used by the tracking logic; detections closer than this may be merged or filtered. |
+| `max_distance_threshold_noise` | px | Distance threshold applied when handling noise / spurious detections during association. |
+| `max_unseen_frames` | frames | How many consecutive frames a track may coast without a matching detection before it is terminated. |
+| `num_video_segments` | int | Number of overlapping segments the video is split into for parallel tracking. |
+| `segment_overlap_seconds` | seconds | Overlap between adjacent segments, converted to frames via `fps`, so boundary-crossing tracks are recoverable. |
+
+## Crossing & Counting
+
+This script is the **final stage** of the pipeline. It takes the trajectories from raw tracking (`raw_tracks.npy`) and turns them into an actual directional count by measuring how tracks cross a counting line through the middle of the frame.
+
+### What it does
+
+1. **Filters and selects crossing tracks.** `save_crossing_tracks_from_raw_tracks` loads the raw tracks, drops anything shorter than two points (`threshold_short_tracks`) to remove spurious fragments, then keeps only the tracks that actually cross the midline (`measure_crossing_bats`). The result is written to `crossing_tracks.npy`.
+2. **Counts crossings and their direction.** `find_crossing_frames` walks each track and emits an event at every frame where the track's position changes sign relative to the counting line, i.e. every geometric crossing. Direction is signed: `+1` = increasing along the counting axis (down / right), `-1` = decreasing (up / left).
+3. **Visualizes the result.** `visualize_crossing_tracks` plots a sample of crossing tracks colored by net direction alongside a histogram of crossings over time, and prints per-direction totals and the net count.
+4. **Renders a counting overlay** (`ct_overlay.mp4`) showing the counting line, per-track dots and heading arrows, a white flash ring on the exact crossing frame, and a live cumulative **OUT / IN / NET** tally burned into each frame.
+
+Counting axis is chosen at runtime by exactly one flag:
+
+- `--count_out` — horizontal counting line at `height / 2`; counts vertical (up/down) crossings. Upward crossings tally as **OUT**, downward as **IN**.
+- `--count_across` — vertical counting line at `width / 2`; counts horizontal (left/right) crossings.
+
+### Why it matters
+
+This is where the pipeline produces its deliverable: a number. Tracks by themselves describe motion. This stage reduces them to counted, direction-resolved line crossings. Filtering short tracks first keeps detector noise out of that number, and the signed counting gives net flow (e.g. how many bats left the roost) rather than just a raw total.
+
+### Config parameters
+
+| Parameter | Units / Type | Description |
+|---|---|---|
+| `input_video_path` | path | Source thermal video. Used to read `fps`, `frame_width`, and `frame_height`, and to render the counting overlay. |
+| `raw_tracks_file` | path | The `raw_tracks.npy` produced by the tracking stage. Its parent directory is where `crossing_tracks.npy` is written. |
+
+**Runtime flags** (not in the config; pass exactly one on the command line):
+
+| Flag | Description |
+|---|---|
+| `--count_out` | Count vertical crossings over a horizontal midline (OUT = upward, IN = downward). |
+| `--count_across` | Count horizontal crossings over a vertical midline. |
+
+## Running the Pipeline
+
+1) Perform inference on the frames of the video to generate detection centroids, contours, and boundaries for blobs using `video_inference.py`
+ ```bash
+ cd path/to/working/dir
+ python -m video_inference --config path/to/video_inference_config.yaml
+ ```
+
+2) Convert the detection data into raw tracks using `detections_to_tracks.py`
+ ```bash
+ cd path/to/working/dir
+ python -m detections_to_tracks --config path/to/detections_to_tracks_config.yaml
+ ```
+
+3) Generate a file and visualization of tracks that cross the midline using `crossing_tracks.py`
+ ```bash
+ cd path/to/working/dir
+ python -m crossing_tracks --config path/to/crossing_tracks_config.yaml
+ ```
\ No newline at end of file
diff --git a/heatseek/counting/__init__.py b/heatseek/counting/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/heatseek/counting/bat_functions.py b/heatseek/counting/bat_functions.py
new file mode 100644
index 0000000..d8efa36
--- /dev/null
+++ b/heatseek/counting/bat_functions.py
@@ -0,0 +1,351 @@
+import os
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from scipy import signal
+from scipy.optimize import linear_sum_assignment
+from heatseek.counting.CountLine import CountLine
+import heatseek.counting.koger_tracking as ktf
+
+def get_blob_info(binary_image, background=None, size_threshold=0):
+
+ '''
+ Get contours from binary image. Then find center and average radius of each contour
+
+ binary_image: 2D image
+ background: 2D array used to see locally how dark the background is
+ size_threshold: radius above which blob is considered real
+ '''
+
+ contours, hierarchy = cv2.findContours(binary_image.astype(np.uint8).copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+ centers = []
+ # Size of bounding rectangles
+ sizes = []
+ areas = []
+ # angle of bounding rectangle
+ angles = []
+ rects = []
+ good_contours = []
+ contours = [np.squeeze(contour) for contour in contours]
+
+ for contour_ind, contour in enumerate(contours):
+
+
+
+ if len(contour.shape) > 1:
+
+ rect = cv2.minAreaRect(contour)
+
+ if background is not None:
+ darkness = background[int(rect[0][1]), int(rect[0][0])]
+ if darkness < 30:
+ dark_size_threshold = size_threshold + 22
+ elif darkness < 50:
+ dark_size_threshold = size_threshold + 15
+ elif darkness < 80:
+ dark_size_threshold = size_threshold + 10
+ elif darkness < 100:
+ dark_size_threshold = size_threshold + 5
+ # elif darkness < 130:
+ # dark_size_threshold = size_threshold + 3
+ else:
+ dark_size_threshold = size_threshold
+ else:
+ dark_size_threshold = 0 # just used in if statement
+
+ area = rect[1][0] * rect[1][1]
+
+ if (area >= dark_size_threshold) or background is None:
+ centers.append(rect[0])
+ sizes.append(rect[1])
+ angles.append(rect[2])
+ good_contours.append(contour)
+ areas.append(area)
+ rects.append(rect)
+ if centers:
+ centers = np.stack(centers, 0)
+ sizes = np.stack(sizes, 0)
+ else:
+ centers = np.zeros((0,2))
+
+ return (centers, np.array(areas), good_contours, angles, sizes, rects)
+
+
+def add_all_points_as_new_tracks(raw_track_list, positions, contours,
+ sizes, current_frame_ind, noise):
+ """ When there are no active tracks, add all new points to new tracks.
+
+ Args:
+ raw_track_list (list): list of tracks
+ positions (numpy array): p x 2
+ contours (list): p contours
+ current_frame_ind (int): current frame index
+ noise: how much noise to add to tracks initially
+ """
+
+ for ind, (position, contour, size) in enumerate(zip(positions, contours, sizes)):
+ raw_track_list.append(
+ ktf.create_new_track(first_frame=current_frame_ind,
+ first_position=position, pos_index=ind,
+ noise=noise, contour=contour, size=size
+ )
+ )
+
+ return raw_track_list
+
+
+
+def find_tracks(first_frame_ind, positions,
+ contours_files=None, contours_list=None,
+ sizes_list=None, max_frame=None, verbose=True,
+ tracks_file=None, max_distance_threshold=None, min_distance_threshold=None,
+ max_distance_threshold_noise=None, max_unseen_time=None):
+ """ Take in positions of all individuals in frames and find tracks.
+
+ Args:
+ first_frame_ind (int): index of first frame of these tracks
+ positions (list): n x 2 for each frame
+ contours_files (list): list of files for contour info from each frame
+ contours_list: already loaded list of contours, only used if contours_file
+ is None
+ sizes_list (list): sizes info from each frame
+
+ return list of all tracks found
+ """
+
+ raw_track_list = []
+ min_new_track_distance = 1
+ min_distance_big = 30
+
+ #try to connect points to the next frame
+ if max_frame is None:
+ max_frame = len(positions) #positions = centers
+
+ contours_file_ind = 0
+ previous_contours_seen = 0
+ if contours_files:
+ contours_list = np.load(contours_files[contours_file_ind], allow_pickle=True)
+ while first_frame_ind >= previous_contours_seen + len(contours_list):
+ contours_file_ind += 1
+ previous_contours_seen += len(contours_list)
+ contours_list = np.load(contours_files[contours_file_ind], allow_pickle=True)
+ print(f'using {contours_files[contours_file_ind]}')
+ elif not contours_list:
+ print("Needs contour_files or contour_list")
+ return
+
+
+ contours_ind = first_frame_ind - previous_contours_seen - 1
+
+
+ for frame_ind in range(first_frame_ind, max_frame):
+ contours_ind += 1
+
+ if contours_files:
+ if contours_ind >= len(contours_list):
+ # load next file
+ try:
+ contours_file_ind += 1
+ contours_list = np.load(contours_files[contours_file_ind], allow_pickle=True)
+ contours_ind = 0
+ except:
+ if tracks_file:
+ tracks_file_error = os.path.splitext(tracks_file)[0] + f'-error-{frame_ind}.npy'
+ print(tracks_file_error)
+ np.save(tracks_file_error, np.array(raw_track_list, dtype=object))
+ #get tracks that are still active (have been seen within the specified time)
+ active_list = ktf.calculate_active_list(raw_track_list, max_unseen_time, frame_ind)
+
+ if verbose:
+ if frame_ind % 10000 == 0:
+ print('frame {} processed.'.format(frame_ind))
+ if tracks_file:
+ np.save(tracks_file, np.array(raw_track_list, dtype=object))
+ if len(active_list) == 0:
+ #No existing tracks to connect to
+ #Every point in next frame must start a new track
+ raw_track_list = add_all_points_as_new_tracks(
+ raw_track_list, positions[frame_ind], contours_list[contours_ind],
+ sizes_list[frame_ind], frame_ind, noise=1
+ )
+ continue
+
+ # Make sure there are new points to add
+ new_positions = None
+ row_ind = None
+ col_ind = None
+ new_sizes = None
+ new_position_indexes = None
+ distance = None
+ contours = None
+ if len(positions[frame_ind]) != 0:
+
+ #positions from the next step
+ new_positions = positions[frame_ind]
+ contours = [np.copy(contour) for contour in contours_list[contours_ind]]
+ new_sizes = sizes_list[frame_ind]
+
+ raw_track_list = ktf.calculate_max_distance(
+ raw_track_list, active_list, max_distance_threshold,
+ max_distance_threshold_noise, min_distance_threshold,
+ use_size=True, min_distance_big=min_distance_big
+ )
+
+ distance = ktf.calculate_distances(
+ new_positions, raw_track_list, active_list
+ )
+
+ max_distance = ktf.create_max_distance_array(
+ distance, raw_track_list, active_list
+ )
+
+ assert distance.shape[1] == len(new_positions)
+ assert distance.shape[1] == len(contours)
+ assert distance.shape[1] == len(new_sizes)
+
+ # Some new points could be too far away from every existing track
+ raw_track_list, distance, new_positions, new_position_indexes, new_sizes, contours = ktf.process_points_without_tracks(
+ distance, max_distance, raw_track_list, new_positions, contours,
+ frame_ind, new_sizes
+ )
+
+
+ if distance.shape[1] > 0:
+ # There are new points can be assigned to existing tracks
+ #connect the dots from one frame to the next
+
+ row_ind, col_ind = linear_sum_assignment(np.log(distance + 1))
+
+ # In casese where there are fewer new points than existing tracks
+ # some tracks won't get new point. Just assign them to
+ # the closest point
+ row_ind, col_ind = ktf.filter_tracks_without_new_points(
+ raw_track_list, distance, row_ind, col_ind, active_list, frame_ind
+ )
+ # Check if tracks with big bats got assigned to small points which are
+ # probably noise
+ row_ind, col_ind = ktf.fix_tracks_with_small_points(
+ raw_track_list, distance, row_ind, col_ind, active_list, new_sizes, frame_ind)
+ # see if points got assigned to tracks that are farther
+ # than max_threshold_distance
+ # This happens when the closer track gets assigned
+ # to a differnt point
+ row_ind, col_ind = ktf.filter_bad_assigns(raw_track_list, active_list, distance, max_distance,
+ row_ind, col_ind
+ )
+
+
+ raw_track_list = ktf.update_tracks(raw_track_list, active_list, frame_ind,
+ row_ind, col_ind, new_positions,
+ new_position_indexes, new_sizes, contours,
+ distance, min_new_track_distance)
+ raw_track_list = ktf.remove_noisy_tracks(raw_track_list)
+ raw_track_list = ktf.finalize_tracks(raw_track_list)
+ if tracks_file:
+ np.save(tracks_file, np.array(raw_track_list, dtype=object))
+ print('{} final save.'.format(os.path.basename(os.path.dirname(tracks_file))))
+ return raw_track_list
+
+def threshold_short_tracks(raw_track_list, min_length_threshold=2):
+ """Only return tracks that are longer than min_length_threshold."""
+
+ track_list = []
+ for track_num, track in enumerate(raw_track_list):
+ if isinstance(track['track'], list):
+ track['track'] = np.array(track['track'])
+ track_length = track['track'].shape[0]
+ if track_length >= min_length_threshold:
+ track_list.append(track)
+ return track_list
+
+def get_rects(track):
+ """ Fit rotated bounding rectangles to each contour in track.
+
+ track: track dict with 'contour' key linked to list of cv2 contours
+ """
+ rects = []
+ for contour in track['contour']:
+ if len(contour.shape) > 1:
+ rect = cv2.minAreaRect(contour)
+ rects.append(rect[1])
+ else:
+ rects.append((np.nan, np.nan))
+
+ return np.array(rects)
+
+def get_wingspan(track):
+ """ Estimate wingspan in pixels from average of peak sizes of longest
+ rectangle edges.
+ """
+
+ if not 'rects' in track.keys():
+ track['rects'] = get_rects(track)
+
+ max_edge = np.nanmax(track['rects'], 1)
+ max_edge = max_edge[~np.isnan(max_edge)]
+ peaks = signal.find_peaks(max_edge)[0]
+ if len(peaks) != 0:
+ mean_wing = np.nanmean(max_edge[peaks])
+ else:
+ mean_wing = np.nanmean(max_edge)
+
+ return mean_wing
+
+def measure_crossing_bats(track_list, frame_height=None, frame_width=None,
+ count_across=False, count_out=True, num_frames=None,
+ with_rects=True, ):
+
+ """ Find and quantify all tracks that cross middle line.
+
+ track_list: list of track dicts
+ frame_height: height of frame in pixels
+ frame_width: width of frame in pixels
+ count_across: count horizontal tracks
+ count_out: count vertical tracks
+ num_frames: number of frames in observation
+ with_rects: if True calculate rects if not already
+ in track and estimate wingspan and body size
+
+ """
+ if count_across:
+ assert frame_width, "If vertical must specify frame width."
+ across_line = CountLine(int(frame_width/2), line_dim=0, total_frames=num_frames)
+ if count_out:
+ assert frame_height, "If horizontal must specify frame height."
+ out_line = CountLine(int(frame_height/2), line_dim=1, total_frames=num_frames)
+
+ crossing_track_list = []
+
+ for track_ind, track in enumerate(track_list):
+ out_result = None
+ across_result = None
+
+ if count_out:
+ out_result, out_frame_num = out_line.is_crossing(track, track_ind)
+
+ if count_across:
+ across_result, across_frame_num = across_line.is_crossing(track, track_ind)
+
+ if out_result or across_result:
+ crossing_track_list.append(track)
+
+ if count_out:
+ crossing_track_list[-1]['crossed'] = out_result
+ crossing_track_list[-1]['crossed_frame'] = out_frame_num
+ if count_across:
+ crossing_track_list[-1]['across_crossed'] = across_result
+ crossing_track_list[-1]['across_crossed_frame'] = across_frame_num
+
+ track[id] = track_ind
+
+ if with_rects:
+ if not 'rects' in track.keys():
+ track['rects'] = get_rects(track)
+
+ crossing_track_list[-1]['mean_wing'] = get_wingspan(track)
+
+
+ return crossing_track_list
+
diff --git a/heatseek/counting/config/crossing_tracks_config.yaml b/heatseek/counting/config/crossing_tracks_config.yaml
new file mode 100644
index 0000000..1cc7cee
--- /dev/null
+++ b/heatseek/counting/config/crossing_tracks_config.yaml
@@ -0,0 +1,2 @@
+'input_video_path' : ''
+'raw_tracks_file' : detections/raw_tracks.npy
\ No newline at end of file
diff --git a/heatseek/counting/config/detections_to_tracks_config.yaml b/heatseek/counting/config/detections_to_tracks_config.yaml
new file mode 100644
index 0000000..97e64a6
--- /dev/null
+++ b/heatseek/counting/config/detections_to_tracks_config.yaml
@@ -0,0 +1,8 @@
+'video_path' : ''
+'output_folder': detections
+'max_distance_threshold' : 14
+'min_distance_threshold' : 8
+'max_distance_threshold_noise' : 10
+'max_unseen_frames' : 1
+'num_video_segments' : 10
+'segment_overlap_seconds' : 10
\ No newline at end of file
diff --git a/heatseek/counting/config/video_inference_config.yaml b/heatseek/counting/config/video_inference_config.yaml
new file mode 100644
index 0000000..5d9640a
--- /dev/null
+++ b/heatseek/counting/config/video_inference_config.yaml
@@ -0,0 +1,10 @@
+'video_path' : ''
+'output_folder': detections
+'threshold' : 10
+'min_val_radiometric' : 28000
+'max_val_radiometric' : 32000
+'radiometric' : False
+'speed' : 9
+'focal_length' : 0.019
+'pixel_pitch' : 0.000017
+'depth' : 45
\ No newline at end of file
diff --git a/heatseek/counting/crossing_tracks.py b/heatseek/counting/crossing_tracks.py
new file mode 100644
index 0000000..aacc428
--- /dev/null
+++ b/heatseek/counting/crossing_tracks.py
@@ -0,0 +1,274 @@
+from heatseek.counting.bat_functions import threshold_short_tracks, measure_crossing_bats
+import numpy as np
+import matplotlib.pyplot as plt
+import argparse
+import cv2
+from heatseek.counting.detections_to_tracks import make_palette, build_frame_lookup, _fade
+from collections import Counter
+from pathlib import Path
+import os
+import yaml
+
+def save_crossing_tracks_from_raw_tracks(file, out_file, count_across, count_out, frame_height=176, frame_width=176, verbose=True):
+ """ Save all crossing tracks after preproccesing all tracks in observation.
+
+ Parameters:
+ file: full path to a numpy file that is a list of track objects
+ out_file: full path where list of crossing tracks should be saved
+ count_across: count horizontal tracks
+ count_out: count vertical tracks
+
+ Nothing is returned
+ """
+
+ raw_track_list = np.load(file, allow_pickle=True)
+
+ if verbose:
+ print(f"{len(raw_track_list)} raw tracks in observation.")
+
+ # Get rid of tracks less than two points long
+ tracks_list = threshold_short_tracks(raw_track_list, min_length_threshold=2)
+ # Get list of tracks that cross the mid line
+ crossing_tracks_list = measure_crossing_bats(tracks_list, frame_height=frame_height, frame_width=frame_width,
+ count_across=count_across, count_out=count_out)
+
+ if verbose:
+ print(f"{len(crossing_tracks_list)} tracks crossing counting line",
+ "in observation.")
+
+ np.save(out_file, np.array(crossing_tracks_list, dtype=object))
+
+def visualize_crossing_tracks(crossing_tracks_file, count_out, count_across, frame_height=176, frame_width=176, num_tracks=100):
+ if count_out == count_across:
+ raise ValueError("Provide exactly one of count_out=True or count_across=True.")
+ if count_across and frame_width is None:
+ raise ValueError("frame_width required when count_across=True.")
+
+ midline_y = frame_height // 2
+ midline_x = frame_width // 2
+
+ crossing_tracks = np.load(crossing_tracks_file, allow_pickle=True)
+
+ print(f'Total crossing tracks: {len(crossing_tracks)}')
+ if len(crossing_tracks) == 0:
+ print('No crossing tracks found.')
+ return
+
+ print(f'Track keys: {crossing_tracks[0].keys()}')
+
+ if count_out:
+ line_dim, line_value = 1, midline_y
+ direction_labels = 'blue=upward, red=downward'
+ pos_label, neg_label = 'down-to-up', 'up-to-down'
+ else:
+ line_dim, line_value = 0, midline_x
+ direction_labels = 'blue=leftward, red=rightward'
+ pos_label, neg_label = 'right-to-left', 'left-to-right'
+
+ events = [] # (frame, direction, track_ind); direction +1 = down/right
+ for ti, tr in enumerate(crossing_tracks):
+ for f, d in find_crossing_frames(tr, line_value, line_dim):
+ events.append((f, d, ti))
+
+ net_by_track = Counter()
+ for f, d, ti in events:
+ net_by_track[ti] += d
+
+ num_tracks = min(num_tracks, len(crossing_tracks))
+ fig, axes = plt.subplots(1, 2, figsize=(15, 6))
+
+ ax = axes[0]
+ for track_ind in np.linspace(0, len(crossing_tracks), num_tracks, endpoint=False, dtype=int):
+ track = crossing_tracks[track_ind]
+ # net > 0 means net downward motion across the line
+ color = 'red' if net_by_track[track_ind] > 0 else 'blue'
+ ax.plot(track['track'][:, 0], track['track'][:, 1], color=color, alpha=0.3, linewidth=1.5)
+
+ if count_out:
+ ax.axhline(y=midline_y, color='green', linewidth=2, linestyle='--', label='midline')
+ ax.invert_yaxis()
+ else:
+ ax.axvline(x=midline_x, color='green', linewidth=2, linestyle='--', label='midline')
+
+ ax.set_title(f'Sample of {num_tracks} crossing tracks\n{direction_labels}')
+ ax.legend()
+
+ ax = axes[1]
+ crossing_frames = [f for f, _, _ in events]
+ ax.hist(crossing_frames, bins=100)
+ ax.set_xlabel('Frame number')
+ ax.set_ylabel('Number of crossings')
+ ax.set_title('Bat crossings over time')
+
+ pos = sum(1 for _, d, _ in events if d > 0) # down / right
+ neg = sum(1 for _, d, _ in events if d < 0) # up / left
+
+ print('Crossing frames:')
+ print(f'{pos_label}: {neg}')
+ print(f'{neg_label}: {pos}')
+ print(f'total events: {len(events)} net: {neg - pos}')
+
+ for f, d, ti in sorted(events):
+ print(f' frame {f}: {pos_label if d < 0 else neg_label} (track {ti})')
+
+ plt.tight_layout()
+ plt.show()
+
+def find_crossing_frames(track, line_value, line_dim):
+ """Yield (abs_frame, direction). direction = +1 increasing along line_dim
+ (down / right), -1 decreasing (up / left)."""
+ coord = np.asarray(track['track'], dtype=float)[:, line_dim] - line_value
+ first = int(track['first_frame'])
+ events = []
+
+ for i in range(1, len(coord)):
+ if coord[i - 1] * coord[i] < 0:
+ events.append((first + i, 1 if coord[i] > 0 else -1))
+
+ return events
+
+
+def create_crossing_overlay_video(input_video_path, output_video_path, crossing_tracks,
+ count_out=True, count_across=False,
+ frame_height=176, frame_width=176,
+ lookback=4, min_arrow_len=1.5, arrow_scale=2.5,
+ dot_radius=4):
+ if count_out == count_across:
+ raise ValueError("Set exactly one of count_out / count_across.")
+
+ if count_out:
+ line_dim, line_value = 1, int(frame_height // 2) # horizontal line, y
+ else:
+ line_dim, line_value = 0, int(frame_width // 2) # vertical line, x
+
+ palette = make_palette(len(crossing_tracks))
+ per_frame = build_frame_lookup(crossing_tracks, lookback=lookback)
+
+ crossing_events = set() # (frame, track_id) — for the flash ring
+ pos_per_frame = Counter() # +1: increasing along line_dim (down / right)
+ neg_per_frame = Counter() # -1: decreasing (up / left)
+
+ for track_id, tr in enumerate(crossing_tracks):
+ for f, d in find_crossing_frames(tr, line_value, line_dim):
+ crossing_events.add((f, track_id))
+ if d > 0:
+ pos_per_frame[f] += 1
+ else:
+ neg_per_frame[f] += 1
+
+ cum_pos = 0
+ cum_neg = 0
+
+ cap = cv2.VideoCapture(input_video_path)
+ fps = cap.get(cv2.CAP_PROP_FPS) or 60.0
+ out = None
+ frame_idx = 0
+
+ while True:
+ ret, frame = cap.read()
+
+ if not ret:
+ break
+
+ if frame.dtype != np.uint8:
+ frame = cv2.normalize(frame, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+
+ if frame.ndim == 2:
+ frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
+
+ frame = np.ascontiguousarray(frame)
+
+ if out is None:
+ h, w = frame.shape[:2]
+
+ if (count_out and h != frame_height) or (count_across and w != frame_width):
+ print(f"WARNING: video is {w}x{h} but counting used "
+ f"frame_width={frame_width}, frame_height={frame_height}. "
+ f"Line and track coords will be misaligned unless these match.")
+
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h))
+
+ if not out.isOpened():
+ raise RuntimeError(f"VideoWriter failed to open for {output_video_path}")
+
+ if count_out:
+ cv2.line(frame, (0, line_value), (w, line_value), (255, 255, 255), 1)
+ else:
+ cv2.line(frame, (line_value, 0), (line_value, h), (255, 255, 255), 1)
+
+ for (x, y, track_id, is_coasting, dx, dy) in per_frame.get(frame_idx, []):
+ color = palette[track_id]
+ cx, cy = int(round(x)), int(round(y))
+
+ if is_coasting:
+ draw_color = _fade(color)
+ cv2.circle(frame, (cx, cy), dot_radius, draw_color, 1)
+ arrow_thick = 1
+ else:
+ draw_color = color
+ cv2.circle(frame, (cx, cy), dot_radius, draw_color, -1)
+ arrow_thick = 2
+
+ mag = (dx * dx + dy * dy) ** 0.5
+
+ if mag >= min_arrow_len:
+ ex = int(round(x + dx * arrow_scale))
+ ey = int(round(y + dy * arrow_scale))
+ cv2.arrowedLine(frame, (cx, cy), (ex, ey), draw_color, arrow_thick, tipLength=0.35)
+
+ if (frame_idx, track_id) in crossing_events:
+ cv2.circle(frame, (cx, cy), dot_radius + 5, (255, 255, 255), 2) # crossing flash
+
+ cum_pos += pos_per_frame.get(frame_idx, 0)
+ cum_neg += neg_per_frame.get(frame_idx, 0)
+
+ lines = [f"OUT: {cum_neg}", f"IN: {cum_pos}", f"NET: {cum_neg - cum_pos}"]
+
+ for i, txt in enumerate(lines):
+ y = 15 + i * 14
+ cv2.putText(frame, txt, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0,0,0), 3, cv2.LINE_AA)
+ cv2.putText(frame, txt, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255,255,255), 1, cv2.LINE_AA)
+
+ out.write(frame)
+ frame_idx += 1
+
+ cap.release()
+
+ if out is not None:
+ out.release()
+
+ print(f'Crossing overlay saved to {output_video_path} '
+ f'({len(crossing_tracks)} crossing tracks)')
+
+def main():
+ parser = argparse.ArgumentParser(description='Get bat crossing tracks across the video')
+ parser.add_argument('--config', help='Path to Video Inference Config File')
+ parser.add_argument('--count_across', action='store_true', help='Count horizontal crossings across vertical midline')
+ parser.add_argument('--count_out', action='store_true', help='Count vertical crossings across horizontal midline')
+ args = parser.parse_args()
+ config_path = Path(args.config)
+
+ if not config_path.exists():
+ raise FileNotFoundError(f'Config file not found: {config_path}')
+
+ with open(config_path, 'r') as f:
+ config = yaml.safe_load(f)
+
+ cap = cv2.VideoCapture(config['input_video_path'])
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ frame_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
+ frame_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
+ output_video_file = os.path.join(os.path.dirname(config['input_video_path']), 'ct_overlay.mp4')
+ crossing_tracks_file = os.path.join(os.path.dirname(config['raw_tracks_file']), 'crossing_tracks.npy')
+ save_crossing_tracks_from_raw_tracks(file=config['raw_tracks_file'], out_file=crossing_tracks_file,
+ count_across=args.count_across, count_out=args.count_out, frame_height=frame_height,
+ frame_width=frame_width)
+ visualize_crossing_tracks(crossing_tracks_file=crossing_tracks_file, count_out=args.count_out, count_across=args.count_across,
+ frame_height=frame_height, frame_width=frame_width)
+ crossing_tracks = np.load(crossing_tracks_file, allow_pickle=True)
+ create_crossing_overlay_video(input_video_path=config['input_video_path'], output_video_path=output_video_file,
+ crossing_tracks=crossing_tracks, count_out=args.count_out, count_across=args.count_across,
+ frame_height=frame_height, frame_width=frame_width)
+if __name__ == '__main__':
+ main()
diff --git a/heatseek/counting/detections_to_tracks.py b/heatseek/counting/detections_to_tracks.py
new file mode 100644
index 0000000..4fed421
--- /dev/null
+++ b/heatseek/counting/detections_to_tracks.py
@@ -0,0 +1,369 @@
+import numpy as np
+import os
+import glob
+import heatseek.counting.bat_functions as kbf
+from multiprocessing import Pool
+import argparse
+from collections import defaultdict
+import cv2
+from pathlib import Path
+import yaml
+
+def track(camera_dict, max_distance_threshold, min_distance_threshold, max_distance_threshold_noise, max_unseen_time):
+ """
+ Run multi-object tracking on precomputed detection data for a single camera.
+
+ Loads saved contour, center, and size data from a camera's output folder,
+ then runs the tracking algorithm (kbf.find_tracks) over a specified frame
+ range. The first contours file is skipped to avoid an off-by-one alignment
+ issue with the centers/sizes arrays. Results are saved to a .npy file in
+ the camera folder.
+
+ If no contours files are found, a warning is printed and tracking is skipped.
+
+ Parameters:
+ camera_dict (dict): Configuration dict with the following keys:
+ - 'camera_folder' (str): Path to the directory containing
+ precomputed detection files (contours, centers, sizes).
+ - 'first_frame' (int): Index of the first frame to track.
+ - 'max_frame' (int): Index of the last frame to track (inclusive).
+ max_distance_threshold (float): Maximum distance threshold for track
+ association. Tracks with detections farther apart than this threshold
+ will not be linked together.
+ min_distance_threshold (float): Minimum distance threshold for track
+ association. Tracks with detections closer than this threshold may be
+ merged or filtered out, depending on the tracking algorithm's logic.
+
+ Output:
+ Writes a raw tracks file to:
+ /first_frame__max_val__raw_tracks.npy
+
+ Returns:
+ None
+ """
+ camera_folder = camera_dict['camera_folder']
+ first_frame = camera_dict['first_frame']
+ max_frame = camera_dict['max_frame']
+ print(f"begun: frames {first_frame} to {max_frame}")
+
+ contours_files = sorted(
+ glob.glob(os.path.join(camera_folder, 'contours-compressed-*.npy'))
+ )
+
+ if contours_files:
+ contours_files = contours_files[1:]
+ centers = np.load(os.path.join(camera_folder, 'centers.npy'), allow_pickle=True)
+ sizes = np.load(os.path.join(camera_folder, 'size.npy'), allow_pickle=True)
+ tracks_file = os.path.join(camera_folder, f'first_frame_{first_frame}_max_val_{max_frame}_raw_tracks.npy')
+ raw_tracks = kbf.find_tracks(first_frame, centers, contours_files=contours_files,
+ sizes_list=sizes, tracks_file=tracks_file,
+ max_frame=max_frame, max_distance_threshold=max_distance_threshold,
+ min_distance_threshold=min_distance_threshold,
+ max_distance_threshold_noise=max_distance_threshold_noise, max_unseen_time=max_unseen_time)
+ else:
+ print("Missing contour files.")
+
+def build_camera_dicts(output_folder, num_groups=10, fps=30, overlap_seconds=10):
+ """
+ Divide a video's detection data into overlapping frame groups and return
+ a list of tracking job descriptors for any groups not yet processed.
+
+ Loads the precomputed centers array to determine total frame count, then
+ splits the full range into num_groups evenly spaced segments. Adjacent
+ segments overlap by overlap_seconds * fps frames so that tracks crossing
+ a segment boundary can still be recovered. The final segment always runs
+ to the end of the video (max_frame=None).
+
+ Segments whose raw tracks output file already exists are skipped, allowing
+ interrupted runs to resume without reprocessing completed chunks.
+
+ Parameters:
+ output_folder (str): Path to the camera output directory containing
+ 'centers.npy' and where raw tracks files will be written.
+ num_groups (int): Number of frame segments to divide the video into.
+ Defaults to 10.
+ fps (int): Frames per second of the source video, used to convert
+ overlap_seconds to a frame count. Defaults to 60.
+ overlap_seconds (int or float): Seconds of overlap between adjacent
+ segments. Defaults to 10.
+
+ Returns:
+ list[dict]: One dict per unprocessed segment, each containing:
+ - 'camera_folder' (str): Same as output_folder.
+ - 'first_frame' (int): First frame index for this segment
+ (overlap-adjusted for all segments after the first).
+ - 'max_frame' (int or None): Exclusive end frame index,
+ or None for the final segment.
+ """
+ centers_file = os.path.join(output_folder, 'centers.npy')
+ centers = np.load(centers_file, allow_pickle=True)
+
+ if num_groups <= 1 or len(centers) <= fps * overlap_seconds:
+ return [{'camera_folder': output_folder,
+ 'first_frame': 0,
+ 'max_frame': None}]
+
+
+ overlap_frames = int(fps * overlap_seconds)
+ camera_dicts = []
+
+ max_vals = np.linspace(0, len(centers), num_groups, dtype=int)[1:].tolist()
+ max_vals[-1] = None
+ min_vals = np.linspace(0, len(centers), num_groups, dtype=int)[:-1]
+ min_vals[1:] = min_vals[1:] - overlap_frames
+
+ for min_val, max_val in zip(min_vals, max_vals):
+ min_val = int(np.max([min_val, 0]))
+ camera_dict = {'camera_folder': output_folder,
+ 'first_frame': min_val,
+ 'max_frame': max_val}
+ tracks_basename = f'first_frame_{min_val}_max_val_{max_val}_raw_tracks.npy'
+ tracks_file = os.path.join(output_folder, tracks_basename)
+ if not os.path.exists(tracks_file):
+ camera_dicts.append(camera_dict)
+ print(tracks_file)
+
+ return camera_dicts
+
+def combine_overlapping_tracks(output_folder, first_group=0, last_group=None, save=False):
+ """
+ Merge per-segment raw track files into a single deduplicated track list.
+
+ Loads all 'first_frame_*.npy' track files from output_folder, sorts them
+ by their start frame, and stitches them together by retaining only tracks
+ that began before the overlap region of each segment. This prevents tracks
+ detected in the overlapping frames of adjacent segments from being counted
+ twice. For the final segment, all remaining tracks are included regardless
+ of start frame.
+
+ Any tracks whose 'track', 'pos_index', or 'size' fields are plain lists
+ are converted to numpy arrays before being appended.
+
+ Parameters:
+ output_folder (str): Path to the directory containing per-segment
+ 'first_frame_*.npy' track files.
+ first_group (int): Index of the first segment file to include.
+ Defaults to 0 (all segments).
+ last_group (int or None): Exclusive index of the last segment file to
+ include. Defaults to None (all segments through the end).
+ save (bool): If True, writes the merged track list to
+ /raw_tracks.npy and prints a confirmation message.
+ Defaults to False.
+
+ Returns:
+ None. Results are printed to stdout and optionally saved to disk.
+ """
+ track_files = glob.glob(os.path.join(output_folder, 'first_frame*.npy'))
+ track_files = sorted(track_files, key=lambda f: int(os.path.basename(f).split('_')[2]))
+
+ track_groups = []
+ for file in track_files:
+ track_groups.append(np.load(file, allow_pickle=True))
+
+ for track_file in track_files:
+ print(os.path.basename(track_file))
+
+ first_overlap_frames = [int(os.path.basename(f).split('_')[2]) for f in track_files[1:]]
+ first_overlap_frames.append(None)
+ print(first_overlap_frames)
+
+ all_tracks = []
+
+ for group_ind, track_group in enumerate(track_groups[first_group:last_group]):
+ if group_ind >= len(track_groups) - 1:
+ for track in track_group:
+ if type(track['track']) == list:
+ track['track'] = np.stack(track['track'])
+ track['pos_index'] = np.stack(track['pos_index'])
+ if 'size' in track:
+ track['size'] = np.stack(track['size'])
+ all_tracks.append(track)
+ break
+
+ for track_ind, track in enumerate(track_group):
+ if track['first_frame'] < first_overlap_frames[first_group + group_ind]:
+ all_tracks.append(track)
+
+ all_tracks_file = os.path.join(output_folder, 'raw_tracks.npy')
+
+ if save:
+ np.save(all_tracks_file, all_tracks)
+ print(f'Saved {len(all_tracks)} tracks to {all_tracks_file}')
+
+def combine_tracks(output_folder):
+ """
+ Merge overlapping track segments into a single file, skipping if already done.
+
+ Calls combine_overlapping_tracks() only if raw_tracks.npy does not already
+ exist in output_folder, making this safe to call multiple times without
+ overwriting a completed merge.
+
+ Parameters:
+ output_folder (str): Path to the directory containing per-segment
+ 'first_frame_*.npy' track files and where 'raw_tracks.npy'
+ will be written.
+
+ Returns:
+ None
+ """
+ if not os.path.exists(os.path.join(output_folder, 'raw_tracks.npy')):
+ combine_overlapping_tracks(output_folder, save=True)
+
+def run_tracking(output_folder, max_distance_threshold, min_distance_threshold, processes=5, max_distance_threshold_noise=10,
+ max_unseen_frames=1, num_groups=10, fps=30, overlap_seconds=10):
+ """
+ Run multi-object tracking across all unprocessed frame segments in parallel.
+
+ Builds the list of pending tracking jobs via build_camera_dicts(), then
+ distributes them across a multiprocessing pool. Each worker calls track()
+ on one segment dict. Segments whose output file already exists are
+ automatically skipped by build_camera_dicts().
+
+ Parameters:
+ output_folder (str): Path to the camera output directory containing
+ precomputed detection files and where raw track files will be saved.
+ processes (int): Number of parallel worker processes. Defaults to 5.
+
+ Returns:
+ None
+ """
+ camera_dicts = build_camera_dicts(output_folder, num_groups=num_groups, fps=fps, overlap_seconds=overlap_seconds)
+ print(f'Tracking {len(camera_dicts)} groups')
+ args = [(d, max_distance_threshold, min_distance_threshold, max_distance_threshold_noise, max_unseen_frames) for d in camera_dicts]
+
+ with Pool(processes=processes) as pool:
+ pool.starmap(track, args)
+
+def make_palette(n):
+ """One visually distinct BGR color per track, evenly spaced in hue."""
+ colors = []
+
+ for i in range(max(n, 1)):
+ hue = int(179 * i / max(n, 1))
+ hsv = np.uint8([[[hue, 220, 255]]])
+ bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0][0]
+ colors.append((int(bgr[0]), int(bgr[1]), int(bgr[2])))
+
+ return colors
+
+
+def build_frame_lookup(all_tracks, lookback=4):
+ """Map absolute frame index -> list of per-object draw records.
+
+ Each record: (x, y, track_id, is_coasting, dx, dy)
+ - position is track['track'][i], coordinate order (x, y)
+ - frame for row i is track['first_frame'] + i (arrays are dense)
+ - is_coasting: True when pos_index[i] is nan (tracker guessed this frame)
+ - (dx, dy): instantaneous heading from row (i - lookback) to row i
+ """
+ per_frame = defaultdict(list)
+
+ for track_id, tr in enumerate(all_tracks):
+ positions = np.asarray(tr['track'], dtype=float) # (N, 2), (x, y)
+ pos_index = np.asarray(tr['pos_index'], dtype=float) # nan on coasting frames
+ first = int(tr['first_frame'])
+ n = positions.shape[0]
+
+ for i in range(n):
+ x, y = positions[i]
+ is_coasting = bool(np.isnan(pos_index[i]))
+ j = max(0, i - lookback)
+ dx = positions[i, 0] - positions[j, 0]
+ dy = positions[i, 1] - positions[j, 1]
+ per_frame[first + i].append((x, y, track_id, is_coasting, dx, dy))
+
+ return per_frame
+
+
+def _fade(color, factor=0.45):
+ """Blend a BGR color toward mid-gray to de-emphasize coasting frames."""
+ gray = 110
+ return tuple(int(c * factor + gray * (1 - factor)) for c in color)
+
+def create_overlay_video(input_video_path, output_video_path, all_tracks,
+ lookback=4, min_arrow_len=1.5, arrow_scale=2.5,
+ dot_radius=4, min_track_length=1):
+ if min_track_length > 1:
+ all_tracks = [t for t in all_tracks
+ if np.asarray(t['track']).shape[0] >= min_track_length]
+
+ palette = make_palette(len(all_tracks))
+ per_frame = build_frame_lookup(all_tracks, lookback=lookback)
+
+ cap = cv2.VideoCapture(input_video_path)
+ fps = cap.get(cv2.CAP_PROP_FPS) or 60.0
+ out = None
+ frame_idx = 0
+
+ while True:
+ ret, frame = cap.read()
+ if not ret:
+ break
+
+ if frame.dtype != np.uint8:
+ frame = cv2.normalize(frame, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+ if frame.ndim == 2:
+ frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
+ frame = np.ascontiguousarray(frame)
+
+ if out is None:
+ h, w = frame.shape[:2]
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h))
+ if not out.isOpened():
+ raise RuntimeError(f"VideoWriter failed to open for {output_video_path}")
+
+ for (x, y, track_id, is_coasting, dx, dy) in per_frame.get(frame_idx, []):
+ color = palette[track_id]
+ cx, cy = int(round(x)), int(round(y))
+ if is_coasting:
+ draw_color = _fade(color)
+ cv2.circle(frame, (cx, cy), dot_radius, draw_color, 1)
+ arrow_thick = 1
+ else:
+ draw_color = color
+ cv2.circle(frame, (cx, cy), dot_radius, draw_color, -1)
+ arrow_thick = 2
+
+ mag = (dx * dx + dy * dy) ** 0.5
+ if mag >= min_arrow_len:
+ ex = int(round(x + dx * arrow_scale))
+ ey = int(round(y + dy * arrow_scale))
+ cv2.arrowedLine(frame, (cx, cy), (ex, ey), draw_color,
+ arrow_thick, tipLength=0.35)
+
+ out.write(frame)
+ frame_idx += 1
+
+ cap.release()
+ if out is not None:
+ out.release()
+ print(f'Overlay video saved to {output_video_path} '
+ f'({len(all_tracks)} tracks, lookback={lookback})')
+
+def main():
+ parser = argparse.ArgumentParser(description='Get centers and contours from detections from model inference')
+ parser.add_argument('--config', help='Path to Raw Tracking Config File')
+ args = parser.parse_args()
+ config_path = Path(args.config)
+
+ if not config_path.exists():
+ raise FileNotFoundError(f'Config file not found: {config_path}')
+
+ with open(config_path, 'r') as f:
+ config = yaml.safe_load(f)
+
+ cap = cv2.VideoCapture(config['video_path'])
+ fps = cap.get(cv2.CAP_PROP_FPS)
+
+ run_tracking(output_folder=config['output_folder'], max_distance_threshold=config['max_distance_threshold'],
+ min_distance_threshold=config['min_distance_threshold'],
+ max_distance_threshold_noise=config['max_distance_threshold_noise'],
+ max_unseen_frames=config['max_unseen_frames'], num_groups=config['num_video_segments'], fps=fps,
+ overlap_seconds=config['segment_overlap_seconds'])
+ combine_tracks(config['output_folder'])
+ raw_tracks = np.load(os.path.join(config['output_folder'], 'raw_tracks.npy'), allow_pickle=True)
+ create_overlay_video(config['video_path'], os.path.join(config['output_folder'], 'rt_overlay_video.mp4'), raw_tracks, lookback=1)
+
+if __name__ == '__main__':
+ main()
diff --git a/heatseek/counting/koger_tracking.py b/heatseek/counting/koger_tracking.py
new file mode 100644
index 0000000..c766391
--- /dev/null
+++ b/heatseek/counting/koger_tracking.py
@@ -0,0 +1,782 @@
+import numpy as np
+
+def create_new_track(first_frame, first_position, pos_index, class_label=None,
+ head_position=None, noise=10, contour=None, size=None,
+ debug=False):
+ """ Create the dictionary which discribes a new track.
+
+ Args:
+ first_frame: first frame track appears
+ last_frame: frame in which it was last seen
+ pos_index: detection index in frame
+ class_label: if multiclass detection
+ noise: Sometimes a false point will momentarily pop up.
+ Treating this noise as a real track causes problems because it
+ constrains the search space of a real nearby tracks. Instead
+ we only treat a track as a real track if it's alreadybeen added
+ to.
+ debug: Add list to track where track creation events can be recorded
+ """
+ new_track = {'track': [first_position],
+ 'first_frame': first_frame,
+ 'last_frame': first_frame,
+ 'pos_index': [pos_index],
+ 'noise': noise
+ }
+
+ if class_label is not None:
+ new_track['class'] = [class_label]
+ if contour is not None:
+ new_track['contour'] = [contour]
+ if size is not None:
+ new_track['size'] = [size]
+ if debug is not None:
+ new_track['debug'] = ['{} created'.format(first_frame)]
+
+ return new_track
+
+def update_unmatched_track(track, debug=False):
+ ''' Update track for next frame when there is no new point to add.
+
+ Basically add a bunch of nans and point stays put. If
+ track has positive noise value then noise value increases.
+
+ Args:
+ tracks: track dictionary
+ '''
+
+ track['track'].append(track['track'][-1])
+ # track['track'].append(predicted_position(track))
+ track['pos_index'].append(np.nan)
+ if 'size' in track:
+ track['size'].append(np.nan)
+ if 'rects' in track:
+ track['rects'].append(np.array([[np.nan, np.nan]]))
+ if 'angles' in track:
+ track['angles'].append(np.nan)
+ if 'contour' in track:
+ track['contour'].append(np.array([np.nan]))
+ if track['noise'] > 0:
+ # isn't a confirmed real track yet
+ track['noise'] += 1
+ if 'class_label' in track:
+ track['class_label'].append(np.nan)
+ if debug:
+ track['debug'].append('Track wasn\'t matched.')
+
+ return track
+
+def update_matched_track(track, frame_index, new_position,
+ new_position_index, size=np.nan,
+ rect=[np.nan, np.nan], angle=np.nan,
+ contour=np.array(np.nan), class_label=np.nan):
+ ''' Update track for next frame when there is a new point to add.
+
+ Args:
+ track: track dictionary
+ frame_index (int): current observation frame
+ new_position (np array): new position to add to track
+ new_position_index (int): raw index of new point
+ size (int): size of object being added
+ rect (object): cv2 rect object
+ angle (float): angle of point being added
+ contour (object): cv2 contour
+ class_label(int): detected class label
+ '''
+
+
+ track['track'].append(new_position)
+ track['pos_index'].append(new_position_index)
+ track['last_frame'] = frame_index
+
+ if track['noise'] > 0:
+ # Maybe this should go to 0 after found
+ # track['noise'] -= 1
+ track['noise'] = 0
+ if 'size' in track:
+ track['size'].append(size)
+ if 'rects' in track:
+ track['rects'].append(rect)
+ if 'angles' in track:
+ track['angles'].append(angle)
+ if 'contour' in track:
+ track['contour'].append(contour)
+ if 'class_label' in track:
+ track['class_label'].append(class_label)
+
+ return track
+
+def _get_track_lengths(frame_ind, track_list, active_list):
+ """ Calculate the current length of every active track.
+
+ Args:
+ frame_ind (num): current frame index
+ track_list (list): list of all tracks
+ active_list (list): list of all track indexes that are active
+
+ return array of lengths of each track
+ """
+
+ track_lengths = np.zeros(len(active_list))
+ for active_num, track_num in enumerate(active_list):
+ track_length = frame_ind - track_list[track_num]['first_frame']
+ track_lengths[active_num] = track_length
+
+ return track_lengths
+
+def _get_track_sizes(frame_ind, track_list, active_list):
+ """ Calculate the current size of every active track.
+
+ Args:
+ frame_ind (num): current frame index
+ track_list (list): list of all tracks
+ active_list (list): list of all track indexes that are active
+
+ return array of lengths of each track
+ """
+
+ track_sizes = np.zeros(len(active_list))
+ for active_num, track_num in enumerate(active_list):
+ track_size = track_list[track_num]['size'][-1]
+ track_sizes[active_num] = track_size
+
+ return track_sizes
+
+
+def filter_tracks_without_new_points(track_list, distance, row_ind,
+ col_ind, active_list, frame_ind,
+ debug=False, use_size=True):
+ """ Deal with instances where some tracks don't have new points.
+
+ This happens when there isn't a new point close enough to existing tracks
+ or when there are fewer new points than existing tracks. When it is the
+ latter, the longer track takes presedence.
+
+ Args:
+ track_list (list): all tracks
+ distance (np array): distances for every old point new point pair
+ row_ind (np array): from linear sum assignment
+ col_ind (np array): from linear sum assignment
+ active_list (list): list of active track indexes
+ frame_ind (int): current frame index
+ """
+
+ row_ind_full = np.arange(len(active_list))
+ col_ind_full = np.zeros(len(active_list), dtype=int)
+ duplicates = []
+ to_delete = [] # less competive track for the same point, or nothing close
+
+ for r_ind in row_ind_full:
+ if r_ind in row_ind:
+ # This track has been paired to a new point
+ col_ind_full[r_ind] = col_ind[np.where(row_ind == r_ind)]
+ if debug:
+ # add debug info
+ track_list[active_list[r_ind]]['debug'].append(
+ '{} assigned normal point'.format(frame_ind))
+ else:
+ # This track wasn't assigned to a new point with the linear sum assignment
+ if np.min(distance[r_ind]) < track_list[active_list[r_ind]]['max_distance']:
+ # There is a new point within this tracks assignment range
+ duplicates.append(np.argmin(distance[r_ind]))
+ col_ind_full[r_ind] = duplicates[-1]
+ if debug:
+ # add debug info
+ track_list[active_list[r_ind]]['debug'].append(
+ '{} wasn\'t assigned point but one is near'.format(frame_ind))
+ else:
+ # This track wasn't assigned a new point and their isn't one close by
+ to_delete.append(r_ind)
+ if debug:
+ # add debug info
+ track_list[active_list[r_ind]]['debug'].append(
+ '{} wasn\'t assigned point and none near'.format(frame_ind))
+
+ track_lengths = _get_track_lengths(frame_ind, track_list, active_list)
+ if use_size:
+ track_sizes = _get_track_sizes(frame_ind, track_list, active_list)
+
+ for duplicate in duplicates:
+ competing_tracks = np.squeeze(np.argwhere(col_ind_full == duplicate))
+ longest_track = np.max(track_lengths[competing_tracks])
+ if np.sum(track_lengths[competing_tracks]==longest_track) > 1:
+ if use_size:
+ dominant_track_ind = np.argmax(track_sizes[competing_tracks])
+ else:
+ # Just takes the first track with max length
+ dominant_track_ind = np.argmax(track_lengths[competing_tracks])
+ else:
+ dominant_track_ind = np.argmax(track_lengths[competing_tracks])
+
+ # Tracks that want the same point but are shorter
+ # (Following 3 lines remove all but dominant track_ind)
+ to_delete.extend(competing_tracks[:dominant_track_ind])
+ if dominant_track_ind < len(competing_tracks):
+ to_delete.extend(competing_tracks[dominant_track_ind+1:])
+ if to_delete:
+ to_delete_a = np.array(to_delete)
+ col_ind_full = np.delete(col_ind_full, to_delete_a)
+ row_ind_full = np.delete(row_ind_full, to_delete_a)
+ if debug:
+ # add debug info
+ for ind in to_delete_a:
+ track_list[active_list[ind]]['debug'].append(
+ '{} no new point given'.format(frame_ind))
+
+
+ return row_ind_full, col_ind_full
+
+def fix_tracks_with_small_points(track_list, distance, row_ind,
+ col_ind, active_list, size_list, frame_ind,
+ debug=False):
+ """ Big bat tracks should connect to big points.
+
+ Sometimes noise pops up next to a bat and is used instead of the next bat
+ point. Usually there is an obvious size difference.
+
+ Args:
+ track_list (list): all tracks
+ distance (np array): distances for every old point new point pair
+ row_ind (np array): from linear sum assignment
+ col_ind (np array): from linear sum assignment
+ active_list (list): list of active track indexes
+ size_list (list): sizes of all new possible points
+ frame_ind (int): current frame index
+ """
+
+ row_ind_full = np.arange(len(active_list))
+ col_ind_full = np.zeros(len(active_list), dtype=int)
+ duplicates = []
+ to_delete = [] # less competive track for the same point, or nothing close
+
+ for r_ind in row_ind_full:
+ if r_ind in row_ind:
+ # This track has been paired to a new point
+ col_ind_full[r_ind] = col_ind[np.where(row_ind == r_ind)]
+ track = track_list[active_list[r_ind]]
+ # if new point is less than 20% of last point probably something different
+ min_new_size = track['size'][-1] / 5
+ if min_new_size < 2:
+ # old point is already small, don't worry about it
+ continue
+ new_size = size_list[col_ind[np.where(row_ind == r_ind)]]
+ if new_size < min_new_size:
+ potential_new_inds = np.argwhere(size_list > min_new_size)
+ # take the closest new point that is a reasonable size
+ if np.any(potential_new_inds):
+ new_ind = np.argmin(distance[r_ind, potential_new_inds])
+ if distance[r_ind, potential_new_inds[new_ind]] < track['max_distance']:
+ duplicates.append(potential_new_inds[new_ind])
+ col_ind_full[r_ind] = duplicates[-1]
+ if debug:
+ # add debug info
+ track_list[active_list[r_ind]]['debug'].append(
+ '{} assigned too small point'.format(frame_ind))
+
+ track_lengths = _get_track_lengths(frame_ind, track_list, active_list)
+ track_sizes = _get_track_sizes(frame_ind, track_list, active_list)
+
+ for duplicate in duplicates:
+ competing_tracks = np.squeeze(np.argwhere(col_ind_full == duplicate))
+ if not competing_tracks.shape:
+ # There is only one track after this point
+ continue
+ longest_track = np.max(track_lengths[competing_tracks])
+ if np.sum(track_lengths[competing_tracks]==longest_track) > 1:
+ dominant_track_ind = np.argmax(track_sizes[competing_tracks])
+ else:
+ dominant_track_ind = np.argmax(track_lengths[competing_tracks])
+
+ # Tracks that want the same point but are shorter
+ to_delete.extend(competing_tracks[:dominant_track_ind])
+ if dominant_track_ind < len(competing_tracks):
+ to_delete.extend(competing_tracks[dominant_track_ind+1:])
+ if to_delete:
+ to_delete_a = np.array(to_delete)
+ col_ind_full = np.delete(col_ind_full, to_delete_a)
+ row_ind_full = np.delete(row_ind_full, to_delete_a)
+ if debug:
+ # add debug info
+ for ind in to_delete_a:
+ track_list[active_list[ind]]['debug'].append(
+ '{} no new point given after too small'.format(frame_ind))
+
+
+ return row_ind_full, col_ind_full
+
+
+def create_max_distance_array(distance, track_list, active_list):
+ """Create array that contains max acceptable distance between all pairs of points.
+
+ Args:
+ distance (np array): distance between all new and old points
+ track_list (list): list of all tracks
+ active_list (list): indexes of all active tracks
+
+ return array of same size as distance
+ """
+
+ max_distance = np.zeros_like(distance)
+
+ for active_num, track_num in enumerate(active_list):
+ max_distance[active_num, :] = track_list[track_num]['max_distance']
+
+ return max_distance
+
+def filter_bad_assigns(track_list, active_list, distance, max_distance, row_ind,
+ col_ind, double_assign=False, debug=False):
+ """ Deal with instances where point is assigned to track that is too far.
+
+ Args:
+ track_list (list): list of all tracks
+ active_list (list): inds of all currently active tracks
+ distance (np array): distance between all pairs of old and new points
+ max_distance (np array): max allowed distance between every new and old point
+ row_ind (np array): row index for each active track
+ col_ind (np array): col_index for reach active track
+ double_asign (boolean): all two tracks assigned to same point
+ """
+
+ bad_assign = distance[row_ind, col_ind] > max_distance[:, 0][row_ind]
+
+ if np.any(bad_assign):
+ bad_assign_points = np.where(bad_assign)[0]
+
+ # Assign multiple tracks to nearby points,
+ # in cases where track got assigned to somewhere far away
+ # because closer track got assigned to point first
+ # this case could come up when two animals get too close
+ # so they merge to one point
+ if double_assign:
+ col_ind[bad_assign_points] = np.argmin(
+ distance[row_ind[bad_assign_points],:], 1)
+
+ # There may be some tracks that just don't have any new points near by.
+ # Filter those out
+ not_valid_assign = distance[row_ind, col_ind] > max_distance[:, 0][row_ind]
+ if np.any(not_valid_assign):
+ if debug:
+ for r_ind in np.argwhere(not_valid_assign):
+ # print('test', np.argwhere(not_valid_assign), r_ind, not_valid_assign)
+ track_list[active_list[r_ind[0]]]['debug'].append('Distance is too far to next point.')
+ valid_assign = distance[row_ind, col_ind] <= max_distance[:, 0][row_ind]
+ col_ind = col_ind[valid_assign]
+ row_ind = row_ind[valid_assign]
+
+ return row_ind, col_ind
+
+
+def process_points_without_tracks(distance, max_distance, track_list,
+ new_positions, contours=None, frame_ind=None,
+ sizes=None, noise=1):
+ """ Find all points that are too far away from existing tracks and create new tracks.
+
+ Args:
+ distance (np array): distance between every new and old point
+ max_disatnce (np array): max allowed distance between every new and old point
+ track_list (list): list of all tracks
+ new_positions (np array): posible new positions in next frame
+ contours (list): list of all contours in frame
+ frame_ind (int): current frame number
+ sizes (np array): all sizes of detections in currrent frame
+ noise: noise value for new tracks
+ """
+
+ is_max_distance = distance > max_distance
+ # New point is too far away from every existing track
+ new_track = np.all(is_max_distance, 0)
+ new_track_ind = None
+ new_position_indexes = np.arange(new_positions.shape[0])
+ if np.any(new_track):
+ new_track_ind = np.where(new_track)[0]
+ for ind in new_track_ind:
+ if contours is not None:
+ contour = contours[ind]
+ else:
+ contour = None
+ if sizes is not None:
+ size = sizes[ind]
+ else:
+ size = None
+ track_list.append(
+ create_new_track(first_frame=frame_ind,
+ first_position=new_positions[ind],
+ pos_index=ind, noise=noise,
+ contour=contour, size=size
+ )
+ )
+
+ # Get rid of new points that are too far away and were
+ # just added as new tracks
+ distance = np.delete(distance, new_track_ind, 1)
+ new_positions = np.delete(new_positions, new_track_ind, 0)
+ new_position_indexes = np.delete(new_position_indexes, new_track_ind)
+ if sizes is not None:
+ sizes = np.delete(sizes, new_track_ind)
+ if contours is not None:
+ for track_ind in new_track_ind[::-1]:
+ contours.pop(track_ind)
+ if (sizes is not None) and (contours is not None):
+ return track_list, distance, new_positions, new_position_indexes, sizes, contours
+ elif (sizes is not None) or (contours is not None):
+ print("Warning sort out return statement for this case if you want to use.")
+ return None
+ else:
+ return track_list, distance, new_positions, new_position_indexes,
+
+def finalize_track(track):
+ """Call when track is finished to turn track lists in into numpy arrays.
+
+ When tracks are being created positions and position indexes are added to list
+ for efficieny since final size of array is unknown while track is still being
+ created.
+
+ Args:
+ track: track dict as defined in create_new_track
+
+ Return track dict with 'track' and 'pos_index' items as arrays
+
+ """
+ track['track'] = np.stack(track['track'])
+ track['pos_index'] = np.stack(track['pos_index'])
+ if 'size' in track:
+ track['size'] = np.stack(track['size'])
+
+ return track
+
+def finalize_tracks(track_list):
+ """ Convert tracks to array and get rid of extra points at end.
+
+ Args:
+ track_list (list): list of all tracks
+
+ return modified track list"""
+ for track_ind, track in enumerate(track_list):
+ track = finalize_track(track)
+ #number of extra points at the end of track that were added hoping
+ #that the point would reapear nearby. Since the tracking is now
+ # finished. We can now get rid of these extra points tacked on to the end
+ old_shape = track['track'].shape[0]
+ last_real_index = track['last_frame'] - track['first_frame'] + 1
+ track['track'] = track['track'][:last_real_index]
+ track['pos_index'] = track['pos_index'][:last_real_index]
+ if 'size' in track:
+ track['size'] = track['size'][:last_real_index]
+ if 'contour' in track:
+ track['contour'] = track['contour'][:last_real_index]
+ track_list[track_ind] = track
+ return track_list
+
+# def predicted_position(track):
+# """Constant-velocity forward prediction from the last two real positions."""
+# t = track['track']
+# if len(t) < 2:
+# return np.asarray(t[-1], dtype=float)
+# v = np.asarray(t[-1], dtype=float) - np.asarray(t[-2], dtype=float)
+# return np.asarray(t[-1], dtype=float) + v
+
+# def predicted_position(track, k=3):
+# """Constant-velocity prediction using mean velocity over the last k steps."""
+# t = np.asarray(track['track'], dtype=float)
+# if len(t) < 2:
+# return t[-1]
+# n = min(k, len(t) - 1)
+# v = (t[-1] - t[-1 - n]) / n # averaged over n frames — jitter divides by n
+# return t[-1] + v
+
+#returns an array of shape (len(active_list), positions1.shape[0])
+#row is distance from every new point to last point in row's active list
+def calculate_distances(new_positions, track_list, active_list):
+ """ Calculate the distance between every new position and every active track.
+ Distance between
+
+ Args:
+ new_positions (numpy array): p x 2
+ track_list (list): list of all tracks
+ active_list (list): index of tracks that could still be added to
+
+ return 2d array of distance betwen every combination of points
+
+ """
+ #positions from last step
+ old_positions = [track_list[track_num]['track'][-1] for track_num in active_list]
+ # old_positions = [predicted_position(track_list[n]) for n in active_list]
+ old_positions = np.stack(old_positions)
+
+ x_diff = (np.expand_dims(new_positions[:, 1], 0)
+ - np.expand_dims(old_positions[:, 1], 1)
+ )
+
+ y_diff = (np.expand_dims(new_positions[:, 0], 0)
+ - np.expand_dims(old_positions[:, 0], 1)
+ )
+ return np.sqrt(x_diff ** 2 + y_diff ** 2)
+
+def calculate_active_list(track_list, max_unseen_time, frame_num, debug=False):
+ active_list = []
+ for track_num in range(len(track_list)):
+ if frame_num - track_list[track_num]['last_frame'] <= max_unseen_time:
+ active_list.append(track_num)
+ else:
+ if debug:
+ track_list[track_num]['debug'].append('{} no longer active'.format(frame_num))
+ return active_list
+
+def calculate_max_distance(track_list, active_list, max_distance_threshold,
+ max_distance_noise, min_distance, use_size=False,
+ size_dict=None, min_distance_big=None):
+
+ """ Calculate the max distance to search for new points for each track.
+
+ The max distance is determined by the minimum of a fixed upper threshold
+ or .45 x the distance to the closest neighbor. However, established tracks
+ defined by those that have a noise value of 0 or below are considered
+ possible neighbors. This means new tracks don't restrict existing tracks
+ search area. A minimum distance also sets a floor on the seach distance such
+ that very close neighbors don't limit all possible connections.
+
+ Args:
+ track_list: list of all tracks
+ active_list: list of tracks that could be added to
+ max_distance_threshold: upper distance threshold
+ max_distance_noise: upper distance threshold tracks with noise values
+ above 0
+ min_distance: lower distance threshold
+ use_size: if objects have assosiated size and want to use that for
+ additional rules
+ size_dict: information about point sizes
+ min_distance_big: min distance for large sized points.
+ """
+
+ # only check distances to established tracks defined by a noise value of
+ # 0 or below
+ #original code
+ # positions0 = [track_list[active_list[0]]['track'][-1]]
+ # if len(active_list) > 1:
+ # for track_num in active_list[1:]:
+ # if track_list[track_num]['noise'] <= 0:
+ # positions0.append(track_list[track_num]['track'][-1])
+ # positions0 = np.stack(positions0)
+
+ #changed (hopefully more robust) code
+ positions0 = [track_list[t]['track'][-1]
+ for t in active_list
+ if track_list[t]['noise'] <= 0]
+
+ if positions0:
+ positions0 = np.stack(positions0)
+ distance = calculate_distances(positions0, track_list, active_list)
+ distance[np.where(distance == 0)] = float("inf")
+ closest_neighbor = np.clip(np.min(distance, 1) * 0.45,
+ min_distance, max_distance_threshold)
+ else:
+ # no confirmed neighbors -> no crowding constraint
+ closest_neighbor = np.full(len(active_list), float(max_distance_threshold))
+
+ # distance = calculate_distances(positions0, track_list, active_list)
+ # # closest point will be itself, so make zero distance
+ # # bigger than other distances
+ # distance[np.where(distance == 0)] = float("inf")
+ # HYPER PARAMETER
+ # # Don't connect to points that are closer to other points
+ # closest_neighbor = np.min(distance, 1) * .45
+ # # Even if neighbors are all far away, have a max threshold to look for new points
+ # closest_neighbor[np.where(closest_neighbor > max_distance_threshold)] = max_distance_threshold
+ # # However, even if neighbors are very close, should be able to connect
+ # # to points within radius of min distance
+ # closest_neighbor[np.where(closest_neighbor < min_distance)] = min_distance
+ for active_ind, track_num in enumerate(active_list):
+ track_list[track_num]['max_distance'] = closest_neighbor[active_ind]
+ if track_list[track_num]['noise'] > 0:
+ if closest_neighbor[active_ind] > max_distance_noise:
+ track_list[track_num]['max_distance'] = max_distance_noise
+
+ if use_size:
+ # Only is objects have related size (added for bats)
+ size = track_list[track_num]['size'][-1]
+ max_distance = track_list[track_num]['max_distance']
+ # if size < 30:
+ # max_distance = np.min([15, max_distance])
+ # # max_distance = np.min([45, max_distance])
+ # elif size < 120:
+ # max_distance = np.min([20, max_distance])
+ # #use the below conditions if these sizes dont work out
+ # # if size < 10:
+ # # max_distance = np.min([15, max_distance])
+ # # elif size < 40:
+ # # max_distance = np.min([20, max_distance])
+ # elif not np.isnan(size):
+ # if min_distance_big:
+ # # Even is points near by, give room to look around
+ # max_distance = np.max([min_distance_big, max_distance])
+
+ max_distance = np.max([min_distance, max_distance])
+ track_list[track_num]['max_distance'] = max_distance
+ return track_list
+
+def add_interpolated_points(frame_ind, track, new_position):
+ """ Replace points since last seen with interpolated estimates.
+
+ Args:
+ frame_ind (int): current frame index in observation
+ track (track obj): the track that is being modified
+ new_position (np array): new point being added
+
+ return updated track
+ """
+
+ gap_distance = (new_position - track['track'][-1])
+
+ missed_steps = frame_ind - track['last_frame'] - 1
+ step_distance = gap_distance / (missed_steps + 1)
+ for step in range(missed_steps):
+ track['track'][-step - 1] = (track['track'][-step - 1]
+ + (missed_steps - step) * step_distance)
+
+ return track
+
+def remove_noisy_tracks(track_list, max_noise=2):
+ """ Delete tracks that have noise values that are too high.
+
+ Args:
+ track_list (list): list of all tracks
+ max_noise: remove track if track has this noise value or higher
+
+ return modified track list
+ """
+
+ # Traverse the list in reverse order so if there are multiple tracks that
+ # need to be removed the indexing doesn't get messed up
+ for track_num in range(len(track_list) - 1, -1, -1):
+ if track_list[track_num]['noise'] >= max_noise:
+ del track_list[track_num]
+
+ return track_list
+
+def create_tracks_for_leftover_points(track_list, col_inds, frame_ind,
+ min_new_track_distance, distance,
+ new_positions, new_position_indexes,
+ contours=None, sizes=None, noise=1):
+
+ """ There can be new points that are close enough to existsing
+ tracks to prevent them from being added in the beginning that don't
+ end up being connected to existing tracks. This are added now.
+
+ Args:
+ track_list (list): list of all tracks
+ col_inds (np array): assossiates new points to columns in distance
+ frame_ind (int): current frame number in observation
+ min_new_track_distance (int): closest new point can be to existing tracks
+ distance (np array): distance between all existing tracks and new points
+ new_positions (np.array): location of all new points n x 2
+ new_position_indexes (np array): raw index of each new point
+ contours (list): list of all contours
+ sizes (list): all sizes in frame
+ noise: noise value for new tracks
+ """
+
+
+ # There are possible new points
+ for pos_ind in range(new_positions.shape[0]):
+ if pos_ind in col_inds:
+ # This point was already added to an existing track
+ continue
+ # Only add points that aren't too close to existing tracks
+ # This just a conservative choice in case of an object
+ # being detected twice and creating two tracks that cause trouble
+ # for each other
+ if np.min(distance[:, pos_ind]) > min_new_track_distance:
+ # This new point isn't too close to existing tracks
+ if contours is not None:
+ contour = contours[pos_ind]
+ else:
+ contour = None
+ if sizes is not None:
+ size = sizes[pos_ind]
+ else:
+ size = None
+ track_list.append(
+ create_new_track(first_frame=frame_ind,
+ first_position=new_positions[pos_ind],
+ pos_index=new_position_indexes[pos_ind],
+ noise=noise,
+ contour=contour,
+ size=size
+ )
+ )
+ return track_list
+
+
+
+def update_tracks(track_list, active_list, frame_index, row_inds, col_inds,
+ new_positions, new_position_indexes, new_sizes=None,
+ new_contours=None, distance=None, min_new_track_distance=None,
+ debug=False, new_track_noise=1):
+ """ Update tracks depending on if they have new points or not.
+
+ Args:
+ track_list (list): list of all tracks
+ active_list (list): list of indexes of active tracks
+ frame_index (int): current frame index
+ row_inds (np array): links active tracks to distance array
+ col_inds (np array): links new points to distance array
+ new_positions (np array): positions of new points n x 2
+ new_position_indexes (np array): raw indexes of new points
+ new_sizes (list): sizes of new points
+ new_contours (list): list of new contours
+ distance (np array): distance between all tracks and new points
+ min_new_track_distance (int): how close are new points allowed
+ to be to old points to start new track
+
+ return updated track list
+ """
+
+ active_list = np.array(active_list)
+ for track_num, track in enumerate(track_list):
+ if track_num in active_list:
+ if row_inds is None:
+ track_list[track_num] = update_unmatched_track(track)
+ continue
+ if track_num in active_list[row_inds] and len(new_positions) != 0:
+ row_count = np.where(track_num == active_list[row_inds])[0]
+ new_position = new_positions[col_inds[row_count[0]]]
+ if track['last_frame'] != frame_index - 1:
+ # This is a refound track, linearly interpolate
+ # from when last seen
+ track = add_interpolated_points(frame_index, track, new_position)
+ new_position_index = new_position_indexes[col_inds[row_count[0]]]
+ if new_sizes is not None:
+ new_size = new_sizes[col_inds[row_count][0]]
+ else:
+ new_size = None
+ if new_contours is not None:
+ new_contour = new_contours[col_inds[row_count][0]]
+ else:
+ new_contour = None
+ track_list[track_num] = update_matched_track(
+ track, frame_index, new_position, new_position_index,
+ size=new_size, contour=new_contour
+ )
+ if debug:
+ track_list[track_num]['debug'].append(
+ '{} row_ind: {}, col_ind: {} pos_ind: {} dist: {}'.format(
+ frame_index, row_inds[row_count[0]], col_inds[row_count[0]],
+ new_position_index, distance[row_inds[row_count[0]],
+ col_inds[row_count[0]]]
+ )
+ )
+ else:
+ track_list[track_num] = update_unmatched_track(track)
+
+
+ if distance is not None:
+ if distance.shape[0] < distance.shape[1]:
+ # Add new tracks for new points that weren't added to existing tracks
+ # but weren't far enough away before to aleady get a new track
+ track_list = create_tracks_for_leftover_points(
+ track_list, col_inds, frame_index, min_new_track_distance,
+ distance, new_positions, new_position_indexes, new_contours,
+ new_sizes, noise=new_track_noise
+ )
+
+ return track_list
\ No newline at end of file
diff --git a/heatseek/counting/video_inference.py b/heatseek/counting/video_inference.py
new file mode 100644
index 0000000..fcd5c8c
--- /dev/null
+++ b/heatseek/counting/video_inference.py
@@ -0,0 +1,225 @@
+import os
+import cv2
+import logging
+import numpy as np
+import argparse
+import heatseek.counting.bat_functions as bat_functions
+import time
+import matplotlib.pyplot as plt
+import math
+from pathlib import Path
+import yaml
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+SHIFT = 4 # 1/16-pixel precision
+FACTOR = 1 << SHIFT # 16
+
+def iter_frames(video_file, process_every_n_frames=1):
+ """Yield frames one at a time without storing them all in memory."""
+ video_name = os.path.splitext(os.path.basename(video_file))[0]
+ cap = cv2.VideoCapture(video_file)
+ fps = cap.get(cv2.CAP_PROP_FPS) or 1.0
+ frame_count = 0
+
+ while True:
+ ret, frame = cap.read()
+ if not ret:
+ break
+ if frame_count % process_every_n_frames == 0:
+ yield {
+ 'frame_idx': frame_count,
+ 'timestamp': frame_count / fps,
+ 'image': frame,
+ }
+ frame_count += 1
+
+ cap.release()
+ logging.info(f'{video_name}: done. {frame_count} total frames.')
+
+# USE ONLY FOR FLIR BOSON
+def get_mask_radiometric(image, min_val=28000, max_val=32000, threshold=5.0):
+ gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+ k100 = (gray.astype(np.float32) / 255.0) * (max_val - min_val) + min_val
+ temp_celsius = (k100 / 100.0) - 273.15
+ background_temp = np.median(temp_celsius)
+ mask = (temp_celsius > (background_temp + threshold)).astype(np.uint8) * 255
+ return mask
+
+# USE FOR NON-RADIOMETRIC THERMAL CAMS
+def get_mask_nonradiometric(image, threshold=10):
+ if len(image.shape) == 3:
+ gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+ else:
+ gray = image.astype(np.float32)
+ background = np.median(gray)
+ mask = (gray < (background - threshold)).astype(np.uint8) * 255
+ return mask
+
+def run_inference(video_file, min_val=28000, max_val=32000,
+ threshold=5.0, process_every_n_frames=1, radiometric=False):
+ """Process every frame, storing only detection metadata — not images."""
+ centers_list = []
+ contours_list = []
+ sizes_list = []
+ rects_list = []
+ start_time = time.perf_counter()
+ for record in iter_frames(video_file, process_every_n_frames):
+
+ if radiometric:
+ mask = get_mask_radiometric(record['image'], min_val, max_val, threshold)
+ else:
+ mask = get_mask_nonradiometric(record['image'], threshold=threshold)
+
+ if record['frame_idx'] == process_every_n_frames:
+ plt.imshow(mask, cmap='gray')
+ plt.title(f'Mask preview (threshold={threshold})')
+ plt.axis('off')
+ plt.show()
+
+ if input('Mask look okay? [y/n]: ').strip().lower() != 'y':
+ raise SystemExit('Aborted: mask unsatisfactory.')
+
+ centers, areas, contours, _, _, rects = bat_functions.get_blob_info(mask)
+
+ centers_list.append(centers)
+ sizes_list.append(areas)
+ contours_list.append(contours)
+ rects_list.append(rects)
+
+ if len(centers_list) % 1000 == 0:
+ logging.info(f'Processed {len(centers_list)} frames.')
+
+ end_time = time.perf_counter()
+ logging.info(f'Inference completed in {end_time - start_time:.2f} seconds.')
+ logging.info(f'Processed {len(centers_list)} frames.')
+
+ return centers_list, contours_list, sizes_list, rects_list
+
+def max_bats(speed, focal_len, pixel_pitch, depth, fps, width, height):
+ """
+ Find the maximum number of bats that can fit in a frame
+
+ Params:
+ - speed [m/s] - bat speed (use upper limit for safest evaluation)
+ - focal_len [m] - focal length of camera
+ - pixel_pitch [m/px] - physical distance between the centers of each pixel; this is typically measured in micrometers and will need to be converted to m
+ - depth [m] - distance between camera and swarm (can be estimated via Google Maps)
+ - fps - video's fps
+ - width - width of frame
+ - height - height of frame
+
+ Returns:
+ - n_max - maximum number of bats/frame
+ """
+ area = width * height
+ f = focal_len / pixel_pitch
+ d = (speed * f) / (depth * fps)
+ n_max = area / (16 * d**2)
+ return d, n_max
+
+def bats_within_limit(speed, focal_len, pixel_pitch, depth, fps, width, height, centers_list):
+ d, n_max_bats = max_bats(speed, focal_len, pixel_pitch, depth, fps, width, height)
+ n_bats = np.array([len(c) for c in centers_list])
+ n_peak = np.percentile(n_bats, 99) # or n_bats.max()
+ return n_peak <= n_max_bats
+
+
+def save_detections(centers_list, contours_list, sizes_list, rects_list,
+ output_folder, num_contour_files=15):
+ """Save per-frame detections to disk."""
+ os.makedirs(output_folder, exist_ok=True)
+ file_num = 0
+ new_contours = []
+
+ for frame_ind, cs in enumerate(contours_list):
+ if frame_ind % int(len(contours_list) / num_contour_files) == 0:
+ file_name = f'contours-compressed-{file_num:02d}.npy'
+ np.save(os.path.join(output_folder, file_name), np.array(new_contours, dtype=object))
+ new_contours = []
+ file_num += 1
+
+ new_contours.append([])
+
+ for c in cs:
+ cc = np.squeeze(cv2.approxPolyDP(c, 0.1, closed=True))
+ new_contours[-1].append(cc)
+
+ file_name = f'contours-compressed-{file_num:02d}.npy'
+ np.save(os.path.join(output_folder, file_name), np.array(new_contours, dtype=object))
+ np.save(os.path.join(output_folder, 'size.npy'), np.array(sizes_list, dtype=object))
+ np.save(os.path.join(output_folder, 'rects.npy'), np.array(rects_list, dtype=object))
+ np.save(os.path.join(output_folder, 'centers.npy'), np.array(centers_list, dtype=object))
+ logging.info(f'Saved detections for {len(centers_list)} frames to {output_folder}')
+
+def create_overlay_video(input_video_path, output_video_path, centers_list, contours_list):
+ cap = cv2.VideoCapture(input_video_path)
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+
+ for frame_idx, record in enumerate(iter_frames(input_video_path)):
+ frame = record['image']
+ centers = centers_list[frame_idx]
+ contours = contours_list[frame_idx]
+
+ # Draw contours and centers on the frame
+ for contour in contours:
+ cv2.drawContours(frame, [contour], -1, (0, 255, 0), 2)
+ for center in centers:
+ cx = int(round(float(center[0]) * FACTOR))
+ cy = int(round(float(center[1]) * FACTOR))
+ cv2.circle(frame, (cx, cy), 5 * FACTOR, (0, 0, 255), -1, shift=SHIFT)
+
+ out.write(frame)
+
+ cap.release()
+ out.release()
+ logging.info(f'Overlay video saved to {output_video_path}')
+
+def main():
+ parser = argparse.ArgumentParser(description='Process a thermal video and detect objects.')
+ parser.add_argument('--config', help='Path to Video Inference Config File')
+ args = parser.parse_args()
+ config_path = Path(args.config)
+
+ if not config_path.exists():
+ raise FileNotFoundError(f'Config file not found: {config_path}')
+
+ with open(config_path, 'r') as f:
+ config = yaml.safe_load(f)
+
+ centers_list, contours_list, sizes_list, rects_list = run_inference(video_file=config['video_path'],
+ min_val=config['min_val_radiometric'],
+ max_val=config['max_val_radiometric'],
+ threshold=config['threshold'], process_every_n_frames=1,
+ radiometric=config['radiometric'])
+ cap = cv2.VideoCapture(config['video_path'])
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
+ height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
+ limit_bool = bats_within_limit(speed=config['speed'], focal_len=config['focal_length'], pixel_pitch=config['pixel_pitch'],
+ depth=config['depth'], fps=fps, width=width, height=height, centers_list=centers_list)
+
+ if not limit_bool:
+ raise SystemExit('Too many bats in frame to process')
+
+ sum = 0
+
+ for i in range(len(centers_list)):
+ sum += len(centers_list[i])
+
+ logging.info(f'Average detections per frame: {sum / len(centers_list)}')
+
+ save_detections(centers_list, contours_list, sizes_list, rects_list, config['output_folder'])
+ create_overlay_video(config['video_path'], os.path.join(config['output_folder'], 'detections_overlay_video.mp4'),
+ centers_list, contours_list)
+ min_dist_threshold , _ = max_bats(speed=config['speed'], focal_len=config['focal_length'], pixel_pitch=config['pixel_pitch'],
+ depth=config['depth'], fps=fps, width=width, height=height)
+ logging.info(f'Set Minimum Distance Threshold for known tracks to {math.ceil(min_dist_threshold)} when generating raw tracks')
+ logging.info(f'Set Maximum Distance threshold to <= 2 * minimum distance threshold')
+
+if __name__ == '__main__':
+ main()
diff --git a/heatseek/data_utils.py b/heatseek/data_utils.py
index 2d64217..1aea404 100644
--- a/heatseek/data_utils.py
+++ b/heatseek/data_utils.py
@@ -1,32 +1,32 @@
-# heatseek/data_utils.py
-import os
-import yaml
-from roboflow import Roboflow
-
-
-def download_dataset(api_key: str,
- workspace: str,
- project: str,
- version: int,
- save_format: str = "yolov11") -> str:
- """Downloads from Roboflow and returns the local dataset folder path."""
- rf = Roboflow(api_key=api_key)
- proj = rf.workspace(workspace).project(project)
- ds = proj.version(version).download(save_format)
- return ds.location
-
-
-def write_data_yaml(dataset_dir: str,
- nc: int,
- names: list[str]):
- """Overwrites /data.yaml with our train/val/names spec."""
- data = {
- "train": os.path.join(dataset_dir, "train", "images"),
- "val": os.path.join(dataset_dir, "valid", "images"),
- "nc": nc,
- "names": names
- }
- path = os.path.join(dataset_dir, "data.yaml")
- with open(path, "w") as f:
- yaml.dump(data, f)
- return path
+# heatseek/data_utils.py
+import os
+import yaml
+from roboflow import Roboflow
+
+
+def download_dataset(api_key: str,
+ workspace: str,
+ project: str,
+ version: int,
+ save_format: str = "yolov11") -> str:
+ """Downloads from Roboflow and returns the local dataset folder path."""
+ rf = Roboflow(api_key=api_key)
+ proj = rf.workspace(workspace).project(project)
+ ds = proj.version(version).download(save_format)
+ return ds.location
+
+
+def write_data_yaml(dataset_dir: str,
+ nc: int,
+ names: list[str]):
+ """Overwrites /data.yaml with our train/val/names spec."""
+ data = {
+ "train": os.path.join(dataset_dir, "train", "images"),
+ "val": os.path.join(dataset_dir, "valid", "images"),
+ "nc": nc,
+ "names": names
+ }
+ path = os.path.join(dataset_dir, "data.yaml")
+ with open(path, "w") as f:
+ yaml.dump(data, f)
+ return path
diff --git a/heatseek/density_annotator.py b/heatseek/density_annotator.py
index 5e2f4de..1a2c9f5 100644
--- a/heatseek/density_annotator.py
+++ b/heatseek/density_annotator.py
@@ -1,51 +1,51 @@
-import cv2
-import os
-import json
-import argparse
-
-annotations = {}
-current_points = []
-#TODO add save folder path
-def click_event(event, x, y, flags, param):
- global current_points
- if event == cv2.EVENT_LBUTTONDOWN:
- current_points.append((x, y))
- cv2.circle(param, (x, y), 3, (0, 0, 255), -1)
- cv2.imshow("Annotator", param)
-
-def annotate_image(image_path):
- global current_points
- current_points = []
- img = cv2.imread(image_path)
- if img is None:
- print(f"Failed to load image: {image_path}")
- return None
-
- display_img = img.copy()
- cv2.imshow("Annotator", display_img)
- cv2.setMouseCallback("Annotator", click_event, display_img)
-
- print(f"Annotating {image_path} — Press 's' to save, 'n' to skip.")
- while True:
- key = cv2.waitKey(0)
- if key == ord('s'):
- return current_points
- elif key == ord('n'):
- return None
- elif key == 27: # ESC key
- print("Exiting.")
- exit()
-
-def annotate_folder(image_dir, output_json):
- image_files = sorted([f for f in os.listdir(image_dir) if f.lower().endswith((".jpg", ".png"))])
-
- for fname in image_files:
- img_path = os.path.join(image_dir, fname)
- pts = annotate_image(img_path)
- if pts is not None:
- annotations[fname] = pts
-
- cv2.destroyAllWindows()
- with open(output_json, "w") as f:
- json.dump(annotations, f, indent=2)
- print(f"Annotations saved to {output_json}")
+import cv2
+import os
+import json
+import argparse
+
+annotations = {}
+current_points = []
+#TODO add save folder path
+def click_event(event, x, y, flags, param):
+ global current_points
+ if event == cv2.EVENT_LBUTTONDOWN:
+ current_points.append((x, y))
+ cv2.circle(param, (x, y), 3, (0, 0, 255), -1)
+ cv2.imshow("Annotator", param)
+
+def annotate_image(image_path):
+ global current_points
+ current_points = []
+ img = cv2.imread(image_path)
+ if img is None:
+ print(f"Failed to load image: {image_path}")
+ return None
+
+ display_img = img.copy()
+ cv2.imshow("Annotator", display_img)
+ cv2.setMouseCallback("Annotator", click_event, display_img)
+
+ print(f"Annotating {image_path} — Press 's' to save, 'n' to skip.")
+ while True:
+ key = cv2.waitKey(0)
+ if key == ord('s'):
+ return current_points
+ elif key == ord('n'):
+ return None
+ elif key == 27: # ESC key
+ print("Exiting.")
+ exit()
+
+def annotate_folder(image_dir, output_json):
+ image_files = sorted([f for f in os.listdir(image_dir) if f.lower().endswith((".jpg", ".png"))])
+
+ for fname in image_files:
+ img_path = os.path.join(image_dir, fname)
+ pts = annotate_image(img_path)
+ if pts is not None:
+ annotations[fname] = pts
+
+ cv2.destroyAllWindows()
+ with open(output_json, "w") as f:
+ json.dump(annotations, f, indent=2)
+ print(f"Annotations saved to {output_json}")
diff --git a/heatseek/detect_track.py b/heatseek/detect_track.py
index 6f2255e..83edc45 100644
--- a/heatseek/detect_track.py
+++ b/heatseek/detect_track.py
@@ -1,38 +1,38 @@
-# heatseek/detect_track.py
-import cv2
-import torch
-from ultralytics import YOLO
-
-
-def detect_and_track(in_path: str,
- out_path: str,
- weights: str,
- tracker_cfg: str = "botsort.yaml",
- device_idx: int = 0):
- device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
- model = YOLO(weights).to(device)
-
- cap = cv2.VideoCapture(in_path)
- fps = cap.get(cv2.CAP_PROP_FPS)
- w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
- h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
- fourcc = cv2.VideoWriter_fourcc(*"mp4v")
- out = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
-
- seen_ids = set()
- while True:
- ret, frame = cap.read()
- if not ret:
- break
-
- results = model.track(frame, persist=True, tracker=tracker_cfg)
- if hasattr(results[0], "boxes") and results[0].boxes.id is not None:
- for obj_id in results[0].boxes.id.cpu().tolist():
- seen_ids.add(int(obj_id))
-
- anno = results[0].plot()
- out.write(anno)
-
- cap.release()
- out.release()
- print(f"[detect_track] done. unique IDs seen: {len(seen_ids)}")
+# heatseek/detect_track.py
+import cv2
+import torch
+from ultralytics import YOLO
+
+
+def detect_and_track(in_path: str,
+ out_path: str,
+ weights: str,
+ tracker_cfg: str = "botsort.yaml",
+ device_idx: int = 0):
+ device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
+ model = YOLO(weights).to(device)
+
+ cap = cv2.VideoCapture(in_path)
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+ out = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
+
+ seen_ids = set()
+ while True:
+ ret, frame = cap.read()
+ if not ret:
+ break
+
+ results = model.track(frame, persist=True, tracker=tracker_cfg)
+ if hasattr(results[0], "boxes") and results[0].boxes.id is not None:
+ for obj_id in results[0].boxes.id.cpu().tolist():
+ seen_ids.add(int(obj_id))
+
+ anno = results[0].plot()
+ out.write(anno)
+
+ cap.release()
+ out.release()
+ print(f"[detect_track] done. unique IDs seen: {len(seen_ids)}")
diff --git a/heatseek/preprocess.py b/heatseek/preprocess.py
index 1a7cdbb..a0e0f6b 100644
--- a/heatseek/preprocess.py
+++ b/heatseek/preprocess.py
@@ -1,86 +1,86 @@
-import cv2
-import numpy as np
-from tqdm import tqdm
-import yaml
-
-
-def reduce_background(in_path: str, out_path: str, yaml_path: str = "heatseek/config/preproc_config.yaml"):
- """
- Optical flow background reduction using parameters from a YAML config.
-
- Args:
- in_path (str): Path to input video file.
- out_path (str): Path to save processed video.
- yaml_path (str): Path to YAML configuration file containing motion_thresh and flow_params.
- """
- with open(yaml_path, 'r') as f:
- config = yaml.safe_load(f)
-
- motion_thresh = config.get('motion_thresh', 1.0)
- flow_cfg = config.get('flow_params', {})
- pyr_scale = flow_cfg.get('pyr_scale', 0.5)
- levels = flow_cfg.get('levels', 3)
- winsize = flow_cfg.get('winsize', 7)
- iterations = flow_cfg.get('iterations', 3)
- poly_n = flow_cfg.get('poly_n', 5)
- poly_sigma = flow_cfg.get('poly_sigma', 1.2)
- flags = flow_cfg.get('flags', 0)
- cap = cv2.VideoCapture(in_path)
- if not cap.isOpened():
- raise ValueError(f"[ERROR] Cannot open video file: {in_path}")
-
- fps = cap.get(cv2.CAP_PROP_FPS)
- width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
- height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
- fourcc = cv2.VideoWriter_fourcc(*"mp4v")
- out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
-
- ret, frame = cap.read()
- if not ret:
- cap.release()
- raise ValueError(f"[ERROR] Could not read first frame from {in_path}")
-
- prev_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
- total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
- for _ in tqdm(range(total_frames), desc="bg-reduce"):
- ret, frame = cap.read()
- if not ret:
- break
-
- gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
- flow = cv2.calcOpticalFlowFarneback(
- prev_gray, gray, None,
- pyr_scale, levels, winsize,
- iterations, poly_n, poly_sigma, flags
- )
- mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
-
- mask = (mag > motion_thresh).astype(np.uint8) * 255
- #kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
- #mask = cv2.erode(mask, kernel, iterations=1)
- dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
- mask = (dist > 1.5).astype(np.uint8) * 255
- mask_color = cv2.merge([mask, mask, mask])
- masked_pixels = np.count_nonzero(mask)
- total_pixels = mask.shape[0] * mask.shape[1]
- mask_ratio = masked_pixels / total_pixels
- #print(f"Mask covers {mask_ratio*100:.1f}% of the frame")
-
- radius = 3.5 # px
- area_per_bat = np.pi * (radius**2) # ~50.3 px²
- est_bats = masked_pixels / area_per_bat
- #print(f"≈{est_bats:.1f} ")
-
- # if you want an integer:
- num_bats = int(round(est_bats))
- #print(f"Estimated bat count: {num_bats}")
- fg = cv2.bitwise_and(frame, mask_color)
- out.write(fg)
-
- prev_gray = gray
-
- cap.release()
- out.release()
-
+import cv2
+import numpy as np
+from tqdm import tqdm
+import yaml
+
+
+def reduce_background(in_path: str, out_path: str, yaml_path: str = "heatseek/config/preproc_config.yaml"):
+ """
+ Optical flow background reduction using parameters from a YAML config.
+
+ Args:
+ in_path (str): Path to input video file.
+ out_path (str): Path to save processed video.
+ yaml_path (str): Path to YAML configuration file containing motion_thresh and flow_params.
+ """
+ with open(yaml_path, 'r') as f:
+ config = yaml.safe_load(f)
+
+ motion_thresh = config.get('motion_thresh', 1.0)
+ flow_cfg = config.get('flow_params', {})
+ pyr_scale = flow_cfg.get('pyr_scale', 0.5)
+ levels = flow_cfg.get('levels', 3)
+ winsize = flow_cfg.get('winsize', 7)
+ iterations = flow_cfg.get('iterations', 3)
+ poly_n = flow_cfg.get('poly_n', 5)
+ poly_sigma = flow_cfg.get('poly_sigma', 1.2)
+ flags = flow_cfg.get('flags', 0)
+ cap = cv2.VideoCapture(in_path)
+ if not cap.isOpened():
+ raise ValueError(f"[ERROR] Cannot open video file: {in_path}")
+
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+ out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+
+ ret, frame = cap.read()
+ if not ret:
+ cap.release()
+ raise ValueError(f"[ERROR] Could not read first frame from {in_path}")
+
+ prev_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+ total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+ for _ in tqdm(range(total_frames), desc="bg-reduce"):
+ ret, frame = cap.read()
+ if not ret:
+ break
+
+ gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+ flow = cv2.calcOpticalFlowFarneback(
+ prev_gray, gray, None,
+ pyr_scale, levels, winsize,
+ iterations, poly_n, poly_sigma, flags
+ )
+ mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
+
+ mask = (mag > motion_thresh).astype(np.uint8) * 255
+ #kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
+ #mask = cv2.erode(mask, kernel, iterations=1)
+ dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
+ mask = (dist > 1.5).astype(np.uint8) * 255
+ mask_color = cv2.merge([mask, mask, mask])
+ masked_pixels = np.count_nonzero(mask)
+ total_pixels = mask.shape[0] * mask.shape[1]
+ mask_ratio = masked_pixels / total_pixels
+ #print(f"Mask covers {mask_ratio*100:.1f}% of the frame")
+
+ radius = 3.5 # px
+ area_per_bat = np.pi * (radius**2) # ~50.3 px²
+ est_bats = masked_pixels / area_per_bat
+ #print(f"≈{est_bats:.1f} ")
+
+ # if you want an integer:
+ num_bats = int(round(est_bats))
+ #print(f"Estimated bat count: {num_bats}")
+ fg = cv2.bitwise_and(frame, mask_color)
+ out.write(fg)
+
+ prev_gray = gray
+
+ cap.release()
+ out.release()
+
print(f"[video_preproc] saved → {out_path}")
\ No newline at end of file
diff --git a/heatseek/train.py b/heatseek/train.py
index 57536cc..8a16f4a 100644
--- a/heatseek/train.py
+++ b/heatseek/train.py
@@ -1,40 +1,40 @@
-# heatseek/train.py
-import torch
-from ultralytics import YOLO
-
-def train(data_yaml: str,
- weights: str = "yolo11s.pt",
- epochs: int = 400,
- batch: int = 16,
- imgsz: int = 640,
- device_idx: int = 0):
- """
- Train a YOLO model using ultralytics API.
-
- Args:
- data_yaml (str): Path to data.yaml file.
- weights (str): Pre-trained weights file path.
- epochs (int): Number of training epochs.
- batch (int): Batch size.
- imgsz (int): Image size.
- device_idx (int): GPU device index; uses CPU if CUDA unavailable.
-
- Returns:
- model: The trained YOLO model instance.
- """
- device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
- print(f"[train] using device {device}")
- model = YOLO(weights)
- # ultralytics YOLO handles device internally; explicit .to() not required but can be used
- model.train(
- data=data_yaml,
- epochs=epochs,
- batch=batch,
- imgsz=imgsz,
- device=device,
- scale=0.5,
- mosaic=1.0,
- mixup=0.0,
- copy_paste=0.1,
- )
- return model
+# heatseek/train.py
+import torch
+from ultralytics import YOLO
+
+def train(data_yaml: str,
+ weights: str = "yolo11s.pt",
+ epochs: int = 400,
+ batch: int = 16,
+ imgsz: int = 640,
+ device_idx: int = 0):
+ """
+ Train a YOLO model using ultralytics API.
+
+ Args:
+ data_yaml (str): Path to data.yaml file.
+ weights (str): Pre-trained weights file path.
+ epochs (int): Number of training epochs.
+ batch (int): Batch size.
+ imgsz (int): Image size.
+ device_idx (int): GPU device index; uses CPU if CUDA unavailable.
+
+ Returns:
+ model: The trained YOLO model instance.
+ """
+ device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
+ print(f"[train] using device {device}")
+ model = YOLO(weights)
+ # ultralytics YOLO handles device internally; explicit .to() not required but can be used
+ model.train(
+ data=data_yaml,
+ epochs=epochs,
+ batch=batch,
+ imgsz=imgsz,
+ device=device,
+ scale=0.5,
+ mosaic=1.0,
+ mixup=0.0,
+ copy_paste=0.1,
+ )
+ return model
diff --git a/ideas.txt b/ideas.txt
index c9603bf..0f60024 100644
--- a/ideas.txt
+++ b/ideas.txt
@@ -1,9 +1,9 @@
-TODOs:
-
-Implement Gaussian Kernel to get Density estimates.
-Implement U-Net to get counts
-Try different settings; it might be easier to begin with the owl dataset.
-
--- Each dot needs to be a gaussian kernel then all densities need to be normalized that the sum equals to the true count
-Then train a regression model to estimate the count based on the density map:
-https://arxiv.org/abs/1907.02724
+TODOs:
+
+Implement Gaussian Kernel to get Density estimates.
+Implement U-Net to get counts
+Try different settings; it might be easier to begin with the owl dataset.
+
+-- Each dot needs to be a gaussian kernel then all densities need to be normalized that the sum equals to the true count
+Then train a regression model to estimate the count based on the density map:
+https://arxiv.org/abs/1907.02724
diff --git a/requirements.txt b/requirements.txt
index 1a50ebd..702dff2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,11 +1,11 @@
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.6.0+cu118
-torchvision==0.21+cu118
-torchaudio==2.6.0+cu118
-numpy
-opencv-python
-PyYAML
-roboflow
-setuptools
-tqdm
-ultralytics
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.6.0+cu118
+torchvision==0.21+cu118
+torchaudio==2.6.0+cu118
+numpy
+opencv-python
+PyYAML
+roboflow
+setuptools
+tqdm
+ultralytics
diff --git a/setup.py b/setup.py
index 7da7fd6..bf49b72 100644
--- a/setup.py
+++ b/setup.py
@@ -1,22 +1,22 @@
-# setup.py
-from setuptools import setup, find_packages
-
-setup(
- name="heatseek",
- version="0.1.0",
- author="Mani Amani",
- packages=find_packages(),
- install_requires=[
- "roboflow>=1.0",
- "ultralytics>=8.0",
- "opencv-python",
- "torch",
- "tqdm",
- "pyyaml",
- ],
- entry_points={
- "console_scripts": [
- "heatseek= heatseek.cli:main",
- ],
- },
-)
+# setup.py
+from setuptools import setup, find_packages
+
+setup(
+ name="heatseek",
+ version="0.1.0",
+ author="Mani Amani",
+ packages=find_packages(),
+ install_requires=[
+ "roboflow>=1.0",
+ "ultralytics>=8.0",
+ "opencv-python",
+ "torch",
+ "tqdm",
+ "pyyaml",
+ ],
+ entry_points={
+ "console_scripts": [
+ "heatseek= heatseek.cli:main",
+ ],
+ },
+)