yolo

sgoldenlab · Nov 16, 2024 · 2af2ad6 · 2af2ad6
1 parent 491c061
commit 2af2ad6
Show file tree

Hide file tree

Showing 8 changed files with 553 additions and 47 deletions.
diff --git a/docs/_static/custom.js b/docs/_static/custom.js
@@ -0,0 +1,4 @@
+window.dataLayer = window.dataLayer || [];
+function gtag() { dataLayer.push(arguments); }
+gtag('js', new Date());
+gtag('config', 'G-PEKR9R5J47');
diff --git a/docs/conf.py b/docs/conf.py
@@ -38,9 +38,7 @@
 mathjax_path = "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/MathJax.js?config=TeX-AMS-MML_HTMLorMML"
 
 
-intersphinx_mapping = {
-  'python': ('https://docs.python.org/3', None),
-  }
+intersphinx_mapping = {'python': ('https://docs.python.org/3', None)}
 html_favicon = "../simba/assets/icons/readthedocs_logo.png"
 html_logo = "../simba/assets/icons/readthedocs_logo.png"
 latex_engine = 'xelatex'
@@ -65,7 +63,11 @@
 
 html_theme = 'sphinx_rtd_theme'
 html_static_path = ['_static']
-html_css_files = [
-    'css/simba_theme.css',  # Include your existing CSS file
-    'custom.css',  # Include your additional CSS file
-]
+html_css_files = ['css/simba_theme.css',  # Include your existing CSS file
+                  'custom.css']  # Include your additional CSS file
+
+
+html_js_files = [
+    "https://www.googletagmanager.com/gtag/js?id=G-PEKR9R5J47",
+    "custom.js"
+]
diff --git a/docs/nb/yolo_ex_2.ipynb b/docs/nb/yolo_ex_2.ipynb
diff --git a/setup.py b/setup.py
@@ -29,7 +29,7 @@
 # Setup configuration
 setuptools.setup(
     name="Simba-UW-tf-dev",
-    version="2.3.4",
+    version="2.3.5",
     author="Simon Nilsson, Jia Jie Choong, Sophia Hwang",
     author_email="[email protected]",
     description="Toolkit for computer classification and analysis of behaviors in experimental animals",

diff --git a/simba/bounding_box_tools/yolo/model.py b/simba/bounding_box_tools/yolo/model.py
@@ -59,6 +59,7 @@ def inference_yolo(weights_path: Union[str, os.PathLike],
                    batch_size: Optional[int] = 4,
                    smoothing_method: Optional[Literal['savitzky-golay', 'bartlett', 'blackman', 'boxcar', 'cosine', 'gaussian', 'hamming', 'exponential']] = None,
                    smoothing_time_window: Optional[int] = None,
+                   interpolate: Optional[bool] = False,
                    stream: Optional[bool] = True) -> Union[None, Dict[str, pd.DataFrame]]:
 
     """
@@ -75,14 +76,15 @@ def inference_yolo(weights_path: Union[str, os.PathLike],
     :param Optional[bool] gpu: If True, performs inference on the GPU. Defaults to False.
     :param Optional[bool] stream: If True, iterate over a generater. Default to True. Recommended on longer videos.
     :param Optional[bool] batch_size: Number of frames to process in parallel. Default to 4.
-
+    :param Optional[bool] interpolate: If True, interpolates missing bounding boxes using ``nearest`` method.
 
     :example:
     >>> inference_yolo(weights_path=r"/mnt/c/troubleshooting/coco_data/mdl/train8/weights/best.pt", video_path=r"/mnt/c/troubleshooting/mitra/project_folder/videos/FRR_gq_Saline_0624.mp4", save_dir=r"/mnt/c/troubleshooting/coco_data/mdl/results", verbose=True, gpu=True)
     """
     COORD_COLS = ['X1', 'Y1', 'X2', 'Y2', 'X3', 'Y3', 'X4', 'Y4']
     OUT_COLS = ['FRAME', 'CLASS_ID', 'CLASS_NAME', 'CONFIDENCE', 'X1', 'Y1', 'X2', 'Y2', 'X3', 'Y3', 'X4', 'Y4']
     SMOOTHING_METHODS = ('savitzky-golay', 'bartlett', 'blackman', 'boxcar', 'cosine', 'gaussian', 'hamming', 'exponential')
+    NEAREST = 'nearest'
 
     if isinstance(video_path, list):
         check_valid_lst(data=video_path, source=f'{inference_yolo.__name__} video_path', valid_dtypes=(str, np.str_,), min_len=1)
@@ -92,7 +94,7 @@ def inference_yolo(weights_path: Union[str, os.PathLike],
     for i in video_path:
         _ = get_video_meta_data(video_path=i)
     check_file_exist_and_readable(file_path=weights_path)
-    check_valid_boolean(value=[half_precision, gpu, verbose, stream], source=inference_yolo.__name__)
+    check_valid_boolean(value=[half_precision, gpu, verbose, stream, interpolate], source=inference_yolo.__name__)
     check_int(name=f'{inference_yolo.__name__} batch_size', value=batch_size, min_value=1)
     if save_dir is not None:
         check_if_dir_exists(in_dir=save_dir, source=f'{inference_yolo.__name__} save_dir')
@@ -105,28 +107,34 @@ def inference_yolo(weights_path: Union[str, os.PathLike],
     if gpu:
         model.export(format='engine')
         model.to('cuda')
+    class_dict = model.names
     results = {}
     for path in video_path:
         _, video_name, _ = get_fn_ext(filepath=path)
         video_meta_data = get_video_meta_data(video_path=path)
         video_out = []
         video_predictions = model.predict(source=path, half=half_precision, batch=batch_size, stream=stream)
         for frm_cnt, video_prediction in enumerate(video_predictions):
-            class_names = video_prediction.names
             if video_prediction.obb is not None:
                 boxes = np.array(video_prediction.obb.data.cpu()).astype(np.float32)
             else:
                 boxes = np.array(video_prediction.boxes.data.cpu()).astype(np.float32)
-            classes = np.unique(boxes[:, -1])
-            for c in classes:
-                cls_data = boxes[np.argwhere(boxes[:, -1] == c)].reshape(-1, boxes.shape[1])
-                cls_data = cls_data[np.argmax(boxes[:, -2].flatten())]
-                if video_prediction.obb is not None:
-                    box = yolo_obb_data_to_bounding_box(center_x=cls_data[0], center_y=cls_data[1], width=cls_data[2], height=cls_data[3], angle=cls_data[4]).flatten()
+            for c in list(class_dict.keys()):
+                cls_data = boxes[np.argwhere(boxes[:, -1] == c)]
+                if cls_data.shape[0] == 0:
+                    video_out.append(np.array([frm_cnt, c, class_dict[c], -1, -1, -1, -1, -1, -1, -1, -1, -1]))
                 else:
-                    box = np.array([cls_data[0], cls_data[1], cls_data[2], cls_data[1], cls_data[2], cls_data[3], cls_data[0], cls_data[3]]).astype(np.int32)
-                video_out.append([frm_cnt, cls_data[-1], class_names[cls_data[-1]], cls_data[-2]] + list(box))
+                    cls_data = cls_data.reshape(-1, boxes.shape[1])[np.argmax(boxes[:, -2].flatten())]
+                    if video_prediction.obb is not None:
+                        box = yolo_obb_data_to_bounding_box(center_x=cls_data[0], center_y=cls_data[1], width=cls_data[2], height=cls_data[3], angle=cls_data[4]).flatten()
+                    else:
+                        box = np.array([cls_data[0], cls_data[1], cls_data[2], cls_data[1], cls_data[2], cls_data[3], cls_data[0], cls_data[3]]).astype(np.int32)
+                    video_out.append([frm_cnt, cls_data[-1], class_dict[cls_data[-1]], cls_data[-2]] + list(box))
         results[video_name] = pd.DataFrame(video_out, columns=OUT_COLS)
+        if interpolate:
+            for cord_col in COORD_COLS:
+                results[video_name][cord_col] = results[video_name][cord_col].astype(np.int32).replace(to_replace=-1, value=np.nan)
+                results[video_name][cord_col] = results[video_name][cord_col].interpolate(method=NEAREST, axis=0).ffill().bfill()
         if smoothing_method is not None:
             if smoothing_method != 'savitzky-golay':
                 smoothened = df_smoother(data=results[video_name][COORD_COLS], fps=video_meta_data['fps'], time_window=smoothing_time_window, source=inference_yolo.__name__, method=smoothing_method)

diff --git a/simba/mixins/geometry_mixin.py b/simba/mixins/geometry_mixin.py
@@ -1596,18 +1596,20 @@ def multiframe_shape_distance(self,
         pool.terminate()
         return results
 
-    def multiframe_minimum_rotated_rectangle(
-        self,
-        shapes: List[Polygon],
-        video_name: Optional[str] = None,
-        verbose: Optional[bool] = False,
-        animal_name: Optional[bool] = None,
-        core_cnt: int = -1,
-    ) -> List[Polygon]:
+    def multiframe_minimum_rotated_rectangle(self,
+                                             shapes: List[Polygon],
+                                             video_name: Optional[str] = None,
+                                             verbose: Optional[bool] = False,
+                                             animal_name: Optional[bool] = None,
+                                             core_cnt: int = -1) -> List[Polygon]:
+
         """
-        Compute the minimum rotated rectangle for each Polygon in a list using mutiprocessing.
+        Compute the minimum rotated rectangle for each Polygon in a list using multiprocessing.
 
         :param List[Polygon] shapes: List of Polygons.
+        :param Optional[str] video_name: Optional video name to print (if verbose is True).
+        :param Optional[str] animal_name: Optional animal name to print (if verbose is True).
+        :param Optional[bool] verbose: If True, prints progress.
         :param core_cnt: Number of CPU cores to use for parallel processing. Default is -1, which uses all available cores.
         """
 
@@ -1645,9 +1647,8 @@ def multiframe_minimum_rotated_rectangle(
                 results.append(result)
 
         timer.stop_timer()
-        stdout_success(
-            msg="Rotated rectangles complete.", elapsed_time=timer.elapsed_time_str
-        )
+        if verbose:
+            stdout_success(msg="Rotated rectangles complete.", elapsed_time=timer.elapsed_time_str)
         pool.join()
         pool.terminate()
         return results
@@ -3772,7 +3773,7 @@ def geometries_to_exterior_keypoints(geometries: List[Polygon], core_cnt: Option
                 pool.imap(GeometryMixin._geometries_to_exterior_keypoints_helper, geometries, chunksize=1)):
                 results.append(mp_return)
         results = [i for xs in results for i in xs]
-        return np.array(results).astype(np.int32)
+        return np.ascontiguousarray(np.array(results)).astype(np.int32)
 
     @staticmethod
     @njit("(int32[:, :, :],)", parallel=True)

diff --git a/simba/mixins/timeseries_features_mixin.py b/simba/mixins/timeseries_features_mixin.py
@@ -5,12 +5,10 @@
 
 import numpy as np
 import pandas as pd
-from numba import (boolean, float32, float64, int64, jit, njit, prange, typed,
-                   types)
+from numba import (boolean, float32, float64, int64, jit, njit, prange, typed, types)
 from numba.typed import Dict, List
 from numpy.lib.stride_tricks import as_strided
-from statsmodels.tsa.stattools import (adfuller, grangercausalitytests, kpss,
-                                       zivot_andrews)
+from statsmodels.tsa.stattools import (adfuller, grangercausalitytests, kpss, zivot_andrews)
 
 from simba.utils.errors import InvalidInputError
 
@@ -22,11 +20,10 @@
 import typing
 from typing import Optional, Tuple, get_type_hints
 
-from simba.utils.checks import (check_float, check_instance, check_int,
-                                check_str, check_that_column_exist,
-                                check_valid_array, check_valid_lst)
+from simba.utils.checks import (check_float, check_instance, check_int, check_str, check_that_column_exist, check_valid_array, check_valid_lst)
 from simba.utils.enums import Formats
 from simba.utils.read_write import find_core_cnt
+from simba.mixins.statistics_mixin import Statistics
 
 
 class TimeseriesFeatureMixin(object):
@@ -2357,4 +2354,55 @@ def sliding_path_aspect_ratio(x: np.ndarray,
             else:
                 results[r - 1] = (w / h) * px_per_mm
 
-        return results
+        return results
+
+    @staticmethod
+    def radial_eccentricity(x: np.ndarray, reference_point: np.ndarray):
+        """
+        Compute the radial eccentricity of a set of points relative to a reference point.
+
+        Radial eccentricity quantifies the degree of elongation in the spatial distribution
+        of points. The value ranges between 0 and 1, where: - 0 indicates a perfectly circular distribution. - Values approaching 1 indicate a highly elongated or linear distribution.
+
+        :param np.ndarray x: 2-dimensional numpy array representing the input data with shape (n, m), where n is the number of frames and m is the coordinates.
+        :param np.ndarray data: A 1D array of shape (n_dimensions,) representing the reference point with  respect to which the radial eccentricity is calculated.
+
+        :example:
+        >>> points = np.random.randint(0, 1000, (100000, 2))
+        >>> reference_point = np.mean(points, axis=0)
+        >>> TimeseriesFeatureMixin.radial_eccentricity(x=points, reference_point=reference_point)
+        """
+
+        check_valid_array(data=x, source=f"{TimeseriesFeatureMixin.radial_eccentricity.__name__} x", accepted_ndims=(2,), accepted_axis_1_shape=[2,], accepted_dtypes=Formats.NUMERIC_DTYPES.value)
+        check_valid_array(data=reference_point, source=f"{TimeseriesFeatureMixin.radial_eccentricity.__name__} reference_point", accepted_ndims=(1,), accepted_dtypes=Formats.NUMERIC_DTYPES.value)
+        centered_points = x - reference_point
+        cov_matrix = Statistics.cov_matrix(data=centered_points.astype(np.float32))
+        eigenvalues, _ = np.linalg.eig(cov_matrix)
+        eigenvalues = np.sort(eigenvalues)[::-1]
+        return np.sqrt(1 - eigenvalues[1] / eigenvalues[0])
+
+
+    @staticmethod
+    def radial_dispersion_index(x: np.ndarray, reference_point: np.ndarray) -> float:
+        """
+        Compute the Radial Dispersion Index (RDI) for a set of points relative to a reference point.
+
+        The RDI quantifies the variability in radial distances of points from the reference point, normalized by the mean radial distance.
+        For example, the radial dispersion from an ROI center.
+
+        :param np.ndarray x: 2-dimensional numpy array representing the input data with shape (n, m), where n is the number of frames and m is the coordinates.
+        :param np.ndarray reference_point: A 1D array of shape (n_dimensions,) representing the reference point with  respect to which the radial dispertion index is calculated.
+        :rtype: float
+
+        :example:
+        >>> points = np.random.randint(0, 1000, (100000, 2))
+        >>> reference_point = np.mean(points, axis=0)
+        >>> TimeseriesFeatureMixin.radial_dispersion_index(x=points, reference_point=reference_point)
+        """
+
+        check_valid_array(data=x, source=f"{TimeseriesFeatureMixin.radial_dispersion_index.__name__} x", accepted_ndims=(2,), accepted_axis_1_shape=[2,], accepted_dtypes=Formats.NUMERIC_DTYPES.value)
+        check_valid_array(data=reference_point, source=f"{TimeseriesFeatureMixin.radial_dispersion_index.__name__} reference_point", accepted_ndims=(1,), accepted_dtypes=Formats.NUMERIC_DTYPES.value)
+        radial_distances = np.linalg.norm(x - reference_point, axis=1)
+        return np.std(radial_distances) / np.mean(radial_distances)
+
+