simba-uw-tf-dev 3.2.1__tar.gz → 3.2.2__tar.gz

{simba_uw_tf_dev-3.2.1 → simba_uw_tf_dev-3.2.2}/PKG-INFO

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

{simba_uw_tf_dev-3.2.1 → simba_uw_tf_dev-3.2.2}/Simba_UW_tf_dev.egg-info/PKG-INFO

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

{simba_uw_tf_dev-3.2.1 → simba_uw_tf_dev-3.2.2}/Simba_UW_tf_dev.egg-info/SOURCES.txt

@@ -21,6 +21,8 @@ simba/assets/.env
 simba/assets/TheGoldenLab.PNG
 simba/assets/UbuntuMono-Regular.ttf
 simba/assets/__init__.py
+simba/assets/botsort.yml
+simba/assets/bytetrack.yml
 simba/assets/dash_simba_base.css
 simba/assets/fonts/.DS_Store
 simba/assets/fonts/Playwrite ES Deco.ttf
@@ -383,7 +385,13 @@ simba/bounding_box_tools/find_boundaries.py
 simba/bounding_box_tools/visualize_boundaries.py
 simba/bounding_box_tools/yolo/__init__.py
 simba/bounding_box_tools/yolo/model.py
-simba/bounding_box_tools/yolo/visualize.py
+simba/bounding_box_tools/yolo/utils.py
+simba/bounding_box_tools/yolo/yolo_inference.py
+simba/bounding_box_tools/yolo/yolo_pose_inference.py
+simba/bounding_box_tools/yolo/yolo_pose_track_inference.py
+simba/bounding_box_tools/yolo/yolo_pose_track_visualizer.py
+simba/bounding_box_tools/yolo/yolo_pose_visualizer.py
+simba/bounding_box_tools/yolo/yolo_visualize.py
 simba/cue_light_tools/__init__.py
 simba/cue_light_tools/cue_light_analyzer.py
 simba/cue_light_tools/cue_light_clf_statistics.py

{simba_uw_tf_dev-3.2.1 → simba_uw_tf_dev-3.2.2}/setup.py

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

simba_uw_tf_dev-3.2.2/simba/assets/botsort.yml

@@ -0,0 +1,14 @@
+tracker_type: botsort
+track_high_thresh: 0.1 #Higher values make the tracker stricter (less false matches, but more ID loss).
+track_low_thresh: 0.1 #Used for low-confidence matching or future buffer merging.
+new_track_thresh: 0.80 #Too high → might miss new objects. Too low → might create too many false positives.
+track_buffer: 60 #Number of frames to keep a track alive without a match.
+match_thresh: 0.125 #Lower = stricter matching = fewer false positives, but more ID switches.
+with_reid: True
+model: auto  # Alternatively, specify the path to your ReID model
+appearance_thresh: 0.4
+gmc: False
+alpha: 0.9
+gmc_method: None
+proximity_thresh: 30
+fuse_score: 1

simba_uw_tf_dev-3.2.2/simba/assets/bytetrack.yml

@@ -0,0 +1,9 @@
+tracker_type: bytetrack
+track_thresh: 0.4         # Lowered to catch more detections
+match_thresh: 0.8
+track_high_thresh: 0.3    # Slightly more permissive
+track_low_thresh: 0.1
+track_buffer: 30
+mot20: False              # Disable stricter matching
+fuse_score: True
+new_track_thresh: 0.2     # Slightly lower to improve new ID creation

simba_uw_tf_dev-3.2.2/simba/bounding_box_tools/yolo/model.py

@@ -0,0 +1,166 @@
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+from typing import Dict, List, Optional, Tuple, Union
+
+try:
+    from typing import Literal
+except:
+    from typing_extensions import Literal
+
+import numpy as np
+import pandas as pd
+import torch
+from ultralytics import YOLO
+
+from simba.data_processors.cuda.utils import _is_cuda_available
+from simba.third_party_label_appenders.converters import \
+    yolo_obb_data_to_bounding_box
+from simba.utils.checks import (check_file_exist_and_readable, check_float,
+                                check_if_dir_exists, check_int, check_str,
+                                check_valid_array, check_valid_boolean,
+                                check_valid_lst, check_valid_tuple, get_fn_ext)
+from simba.utils.data import df_smoother, savgol_smoother
+from simba.utils.enums import Formats
+from simba.utils.errors import InvalidInputError, SimBAGPUError
+from simba.utils.printing import SimbaTimer, stdout_success
+from simba.utils.read_write import find_core_cnt, get_video_meta_data
+
+
+def inference_yolo_pose_tracks(weights_path: Union[str, os.PathLike],
+                               video_path: Union[Union[str, os.PathLike], List[Union[str, os.PathLike]]],
+                               keypoint_names: Tuple[str, ...],
+                               botsort_config_path: Union[str, os.PathLike],
+                               verbose: Optional[bool] = False,
+                               save_dir: Optional[Union[str, os.PathLike]] = None,
+                               device: Union[Literal['cpu'], int] = 0,
+                               format: Optional[str] = None,
+                               batch_size: Optional[int] = 4,
+                               torch_threads: int = 8,
+                               half_precision: bool = True,
+                               stream: bool = False,
+                               threshold: float = 0.7,
+                               max_tracks: Optional[int] = 2):
+
+    OUT_COLS = ['FRAME', 'CLASS_ID', 'CLASS_NAME', 'CONFIDENCE', 'X1', 'Y1', 'X2', 'Y2', 'X3', 'Y3', 'X4', 'Y4']
+    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)
+    elif isinstance(video_path, str):
+        check_file_exist_and_readable(file_path=video_path)
+        video_path = [video_path]
+    for i in video_path:
+        _ = get_video_meta_data(video_path=i)
+    check_file_exist_and_readable(file_path=weights_path)
+    check_file_exist_and_readable(file_path=botsort_config_path)
+    check_valid_boolean(value=verbose, source=inference_yolo.__name__)
+    check_int(name=f'{inference_yolo.__name__} batch_size', value=batch_size, min_value=1)
+    check_float(name=f'{inference_yolo.__name__} threshold', value=threshold, min_value=10e-6, max_value=1.0)
+    check_valid_tuple(x=keypoint_names, source=f'{inference_yolo.__name__} keypoint_names', min_integer=1, valid_dtypes=(str,))
+    if max_tracks is not None:
+        check_int(name=f'{inference_yolo.__name__} max_tracks', value=max_tracks, min_value=1)
+    if save_dir is not None:
+        check_if_dir_exists(in_dir=save_dir, source=f'{inference_yolo.__name__} save_dir')
+    keypoint_col_names = [f'{i}_{s}'.upper() for i in keypoint_names for s in ['x', 'y', 'p']]
+    OUT_COLS.extend(keypoint_col_names)
+    torch.set_num_threads(torch_threads)
+    model = load_yolo_model(weights_path=weights_path, device=device, format=format)
+    results = {}
+    class_dict = model.names
+    timer = SimbaTimer(start=True)
+    for path in video_path:
+        _, video_name, _ = get_fn_ext(filepath=path)
+        _ = get_video_meta_data(video_path=path)
+        video_out = []
+        predictions = model.track(source=path, stream=False, tracker=botsort_config_path)
+        print(predictions)
+        #video_predictions = model.predict(source=path, half=half_precision, batch=batch_size, stream=stream)
+
+
+
+# VIDEO_PATH = "/mnt/d/netholabs/yolo_videos/2025-05-28_19-46-56.mp4"
+# BOTSORT_PATH = "/mnt/c/projects/simba/simba/simba/assets/botsort.yml"
+# inference_yolo_pose_tracks(weights_path=r"/mnt/d/netholabs/yolo_mdls_1/train/weights/best.pt",
+#                            video_path=VIDEO_PATH,
+#                            save_dir=r"/mnt/d/netholabs/yolo_test/results",
+#                            verbose=True,
+#                            device=0,
+#                            format='onnx',
+#                            keypoint_names=('nose', 'ear_left', 'ear_right', 'lateral_left', 'center', 'lateral_right', 'tail_base'),
+#                            batch_size=32,
+#                            botsort_config_path=BOTSORT_PATH)
+
+
+#
+#
+#
+# #fit_yolo(initial_weights=r"/mnt/d/yolo_weights/yolo11n-pose.pt", model_yaml=r"/mnt/d/netholabs/yolo_data_1/map.yaml", save_path=r"/mnt/d/netholabs/yolo_mdls_1", batch=32, epochs=100)
+#
+# #video_path = "/mnt/d/netholabs/videos/2025-04-17_17-09-28.h264"
+# video_path = "/mnt/d/netholabs/videos_/2025-05-27_20-59-48.mp4"
+# video_path = "/mnt/d/netholabs/yolo_videos/input/mp4_20250606083508/2025-05-28_19-50-23.mp4"
+# #video_path = "/mnt/c/Users/sroni/Downloads/2025-05-27_00-32-26.mp4"
+#
+#
+# #
+# inference_yolo_pose(weights_path=r"/mnt/d/netholabs/yolo_mdls_1/train/weights/best.pt",
+#                     video_path=video_path,
+#                     save_dir=r"/mnt/d/netholabs/yolo_test/results",
+#                     verbose=True,
+#                     device=0,
+#                     format='onnx',
+#                     keypoint_names=('nose', 'ear_left', 'ear_right', 'lateral_left', 'center', 'lateral_right', 'tail_base'),
+#                     batch_size=32)
+#
+#
+# #
+
+
+
+# fit_yolo(initial_weights=r"/mnt/d/yolo_weights/yolov8n.pt",
+#          model_yaml=r"/mnt/d/netholabs/imgs/yolo_train_test_val/map.yaml", save_path=r"/mnt/d/netholabs/imgs/yolo_mdl", batch=32, epochs=100)
+
+# fit_yolo(initial_weights=r"/mnt/c/troubleshooting/coco_data/weights/yolov8n-obb.pt",
+#          model_yaml=r"/mnt/c/troubleshooting/coco_data/model.yaml",
+#          project_path=r"/mnt/c/troubleshooting/coco_data/mdl",
+#          batch=16, epochs=100)
+
+
+
+# fit_yolo(initial_weights=r"/mnt/c/troubleshooting/coco_data/weights/yolov8n.pt",
+#          model_yaml=r"/mnt/d/netholabs/yolo_test/yolo_train/map.yaml",
+#          save_path=r"/mnt/d/netholabs/yolo_test/yolo_mdl", epochs=150, batch=24)
+
+
+
+# inference_yolo(weights_path=r"/mnt/d/netholabs/yolo_test/yolo_mdl/train10/weights/best.pt",
+#                video_path=r"/mnt/d/netholabs/out/2025-04-17_17-05-07.mp4",
+#                save_dir=r"/mnt/d/netholabs/yolo_test/results",
+#                verbose=True,
+#                gpu=False)
+
+# inference_yolo(weights=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)
+#
+#
+
+
+# r = inference_yolo(weights=r"C:\troubleshooting\coco_data\mdl\train\weights\best.pt",
+#                video_path=r"C:\troubleshooting\mitra\project_folder\videos\clipped\501_MA142_Gi_CNO_0514_clipped.mp4",
+#                batch=100)
+
+
+#fit_yolo(initial_weights=r"C:\troubleshooting\coco_data\weights\yolov8n-obb.pt", data=r"C:\troubleshooting\coco_data\model.yaml", project_path=r"C:\troubleshooting\coco_data\mdl", batch=16)
+
+
+# """
+# initial_weights=r"C:\Users\sroni\Downloads\yolov8n.pt", data=r"C:\troubleshooting\coco_data\model.yaml", epochs=30, project_path=r"C:\troubleshooting\coco_data\mdl", batch=16)
+# """
+#
+#
+#
+
+
+

simba_uw_tf_dev-3.2.2/simba/bounding_box_tools/yolo/utils.py

@@ -0,0 +1,210 @@
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+from typing import Optional, Union
+import torch
+try:
+    from typing import Literal
+except:
+    from typing_extensions import Literal
+
+import numpy as np
+from ultralytics import YOLO
+
+from simba.utils.checks import (check_file_exist_and_readable, check_float, check_int, check_str, check_valid_boolean, check_valid_array, check_if_dir_exists, check_instance, check_if_valid_img)
+from simba.utils.errors import SimBAGPUError, InvalidInputError
+from simba.data_processors.cuda.utils import _is_cuda_available
+from simba.utils.enums import Formats
+from simba.utils.read_write import find_core_cnt, get_video_meta_data
+
+def check_valid_device(device: Union[Literal['cpu'], int]):
+    if isinstance(device, str):
+        check_str(name=f'{load_yolo_model.__name__} format', value=device.lower(), options=['cpu'], raise_error=True)
+    else:
+        check_int(name=f'{load_yolo_model.__name__} device', value=device, min_value=0, raise_error=True)
+        gpu_available, gpus = _is_cuda_available()
+        if not gpu_available:
+            raise SimBAGPUError(msg=f'No GPU detected but device {device} passed', source=load_yolo_model.__name__)
+        if device not in list(gpus.keys()):
+            raise SimBAGPUError(msg=f'Unaccepted GPU device {device} passed. Accepted: {list(gpus.keys())}', source=load_yolo_model.__name__)
+
+
+
+def fit_yolo(weights_path: Union[str, os.PathLike],
+             model_yaml: Union[str, os.PathLike],
+             save_path: Union[str, os.PathLike],
+             epochs: int = 25,
+             batch: Union[int, float] = 16,
+             plots: bool = True,
+             imgsz: int = 640,
+             format: Optional[str] = None,
+             device:  Union[Literal['cpu'], int] = 0,
+             verbose: bool = True,
+             workers: int = 8):
+    """
+    Trains a YOLO model using specified initial weights and a configuration YAML file.
+
+    .. note::
+       `Download initial weights <https://huggingface.co/Ultralytics>`__.
+       `Example model_yaml <https://github.com/sgoldenlab/simba/blob/master/misc/ex_yolo_model.yaml>`__.
+
+    :param initial_weights: Path to the pre-trained YOLO model weights (usually a `.pt` file). Example weights can be found [here](https://huggingface.co/Ultralytics).
+    :param model_yaml: YAML file containing paths to the training, validation, and testing datasets and the object class mappings. Example YAML file can be found [here](https://github.com/sgoldenlab/simba/blob/master/misc/ex_yolo_model.yaml).
+    :param save_path: Directory path where the trained model, logs, and results will be saved.
+    :param epochs: Number of epochs to train the model. Default is 5.
+    :param batch: Batch size for training. Default is 16.
+    :return: None. The trained model and associated training logs are saved in the specified `project_path`.
+
+    :example:
+    >>> fit_yolo(initial_weights=r"C:\troubleshooting\coco_data\weights\yolov8n-obb.pt", data=r"C:\troubleshooting\coco_data\model.yaml", save_path=r"C:\troubleshooting\coco_data\mdl", batch=16)
+
+    """
+    VALID_FORMATS = ["onnx", "engine", "torchscript", "onnxsimplify", "coreml", "openvino", "pb", "tf", "tflite"]
+
+    if not _is_cuda_available()[0]:
+        raise SimBAGPUError(msg='No GPU detected.', source=fit_yolo.__name__)
+    check_file_exist_and_readable(file_path=weights_path)
+    check_file_exist_and_readable(file_path=model_yaml)
+    check_valid_boolean(value=verbose, source=f'{fit_yolo.__name__} verbose', raise_error=True)
+    check_valid_boolean(value=plots, source=f'{fit_yolo.__name__} plots', raise_error=True)
+    check_if_dir_exists(in_dir=save_path)
+    if format is not None: check_str(name=f'{fit_yolo.__name__} format', value=format.lower(), options=VALID_FORMATS, raise_error=True)
+    check_int(name=f'{fit_yolo.__name__} epochs', value=epochs, min_value=1)
+    check_int(name=f'{fit_yolo.__name__} imgsz', value=imgsz, min_value=1)
+    check_int(name=f'{fit_yolo.__name__} workers', value=workers, min_value=-1, unaccepted_vals=[0], max_value=find_core_cnt()[0])
+    if workers == -1: workers = find_core_cnt()[0]
+    check_valid_device(device=device)
+    model = load_yolo_model(weights_path=weights_path, verbose=verbose, format=format, device=device)
+    model.train(data=model_yaml, epochs=epochs, project=save_path, batch=batch, plots=plots, imgsz=imgsz, workers=workers)
+
+
+def load_yolo_model(weights_path: Union[str, os.PathLike],
+                    verbose: bool = True,
+                    format: Optional[str] = None,
+                    device: Union[Literal['cpu'], int] = 0):
+
+    """
+    Load a YOLO model.
+
+    :param Union[str, os.PathLike] weights_path: Path to model weights (.pt, .engine, etc).
+    :param bool verbose: Whether to print loading info.
+    :param Optional[str] format: Export format, one of VALID_FORMATS or None to skip export.
+    :param Union[Literal['cpu'], int]  device: Device to load model on. 'cpu', int GPU index.
+
+    :example:
+    >>> load_yolo_model(weights_path=r"/mnt/c/troubleshooting/coco_data/mdl/train8/weights/best.pt", format="onnx", device=0)
+    """
+
+    VALID_FORMATS = ["onnx", "engine", "torchscript", "onnxsimplify", "coreml", "openvino", "pb", "tf", "tflite"]
+    check_file_exist_and_readable(file_path=weights_path)
+    check_valid_boolean(value=verbose, source=f'{load_yolo_model.__name__} verbose', raise_error=True)
+    if format is not None: check_str(name=f'{load_yolo_model.__name__} format', value=format.lower(), options=VALID_FORMATS, raise_error=True)
+    check_valid_device(device=device)
+    model = YOLO(weights_path, verbose=verbose)
+    model.to(device=device)
+    if format is not None:
+        model.export(format=format)
+
+    return model
+
+def _get_undetected_obs(frm_id: int, class_id: int, class_name: str, value_cnt: int) -> np.ndarray:
+    """Helper to get a missing observation of specified class and frame """
+    check_int(name=f'{_get_undetected_obs.__name__} frm_id', value=frm_id, min_value=0)
+    check_int(name=f'{_get_undetected_obs.__name__} class_id', value=class_id, min_value=0)
+    check_str(name=f'{_get_undetected_obs.__name__} class_name', value=class_name, raise_error=True)
+    check_int(name=f'{_get_undetected_obs.__name__} value_cnt', value=value_cnt, min_value=0)
+    return np.array([frm_id, class_id, class_name] + [-1] * value_cnt)
+
+
+def filter_yolo_keypoint_data(bbox_data: np.ndarray,
+                              keypoint_data: np.ndarray,
+                              class_id: Optional[int] = None,
+                              confidence: Optional[float] = None,
+                              class_idx: Optional[int] = None,
+                              confidence_idx: Optional[int] = None):
+    """
+    Helper to filters YOLO bounding box and keypoint data based on class ID and/or confidence threshold.
+
+    :param np.ndarray bbox_data: A 2D array of shape (N, M) representing YOLO bounding box data, where each row corresponds to one detection and contains class and confidence values.
+    :param np.ndarray bbox_data: A 3D array of shape (N, 2, 3) representing keypoints for each detection, where K is the number of keypoints per detection.
+    :param Optional[int] class_id: Target class ID to filter detections. Defaults to None.
+    :param Optional[float] confidence: Minimum confidence threshold to keep detections. Must be in [0, 1]. Defaults to None.
+    :param int confidence_idx: Index in `bbox_data` where confidence value is stored. Defaults to 5.
+    :param int class_idx: Index in `bbox_data` where class ID is stored. Defaults to 6.
+    """
+
+    if class_id is None and confidence is None:
+        raise InvalidInputError(msg='Provide at least one filter condition')
+
+    check_valid_array(data=bbox_data, source=filter_yolo_keypoint_data.__name__, accepted_ndims=(2,), accepted_dtypes=Formats.NUMERIC_DTYPES.value, min_axis_0=1)
+    check_valid_array(data=keypoint_data, source=filter_yolo_keypoint_data.__name__, accepted_ndims=(3,), accepted_dtypes=Formats.NUMERIC_DTYPES.value, accepted_axis_0_shape=[bbox_data.shape[0],])
+    class_id is not None and check_int(name=filter_yolo_keypoint_data.__name__, value=class_id, min_value=0, raise_error=True)
+    confidence is not None and check_float(name=filter_yolo_keypoint_data.__name__, value=confidence, min_value=0, max_value=1.0, raise_error=True)
+    class_id is not None and check_int(name=f'{filter_yolo_keypoint_data.__name__} class_idx', value=class_idx, min_value=-bbox_data.shape[1], max_value=bbox_data.shape[1])
+    confidence is not None and check_int(name=f'{filter_yolo_keypoint_data.__name__} confidence_idx', value=confidence_idx, min_value=-bbox_data.shape[1], max_value=bbox_data.shape[1])
+
+    if class_id is not None:
+        cls_idx = np.argwhere(bbox_data[:, class_idx] == class_id).flatten()
+        bbox_data, keypoint_data = bbox_data[cls_idx], keypoint_data[cls_idx]
+    if confidence is not None:
+        cls_idx = np.argwhere(bbox_data[:, confidence_idx] >= confidence).flatten()
+        bbox_data, keypoint_data = bbox_data[cls_idx], keypoint_data[cls_idx]
+
+    return bbox_data, keypoint_data
+
+
+def yolo_predict(model: YOLO,
+                 source: Union[str, os.PathLike, np.ndarray],
+                 half: bool = False,
+                 batch_size: Optional[int] = 4,
+                 stream: bool = False,
+                 imgsz: int = 640,
+                 device:  Union[Literal['cpu'], int] = 0,
+                 threshold: float = 0.25,
+                 max_detections: int = 300,
+                 verbose: bool = True):
+
+    """
+    Produce YOLO predictions.
+
+    :param Union[str, os.PathLike] model: Loaded ultralytics.YOLO model. Returned by :func:`~simba.bounding_box_tools.yolo.model.load_yolo_model`.
+    :param Union[str, os.PathLike, np.ndarray] source: Path to video, video stream, directory, image, or image as loaded array.
+    :param bool half: Whether to use half precision (FP16) for inference to speed up processing.
+    :param bool stream: If True, return a generator that yields results one by one. Useful for stream or large videos.
+    :param int imgsz: Size to resize input images to (square dimension). Must be positive integer.
+    :param Optional[int] batch_size: If stream is False, then the number of images to process in each batch.
+    :param Union[Literal['cpu'], int] device: Device identifier for inference. 'cpu' to force CPU inference. E.g., integer index of the GPU device (e.g., 0 for 'cuda:0').
+    :param float threshold: Confidence threshold for filtering predictions. Only detections with confidence >= threshold are returned. Must be between 0.0 and 1.0.
+    :param int max_detections: Maximum number of detections per image/frame to return.
+    :param bool verbose: If True, print inference progress and summary information.
+    :returns: YOLO results or generator of YOLO results.
+    """
+
+    check_valid_boolean(value=half, source=f'{yolo_predict.__name__} half', raise_error=True)
+    check_valid_boolean(value=stream, source=f'{yolo_predict.__name__} stream', raise_error=True)
+    check_valid_boolean(value=verbose, source=f'{yolo_predict.__name__} verbose', raise_error=True)
+    check_int(name=f'{yolo_predict.__name__} imgsz', value=imgsz, min_value=1, raise_error=False)
+    check_int(name=f'{yolo_predict.__name__} max_detections', value=max_detections, min_value=1, raise_error=False)
+    check_float(name=f'{yolo_predict.__name__} threshold', value=threshold, min_value=0.0, max_value=1.0, raise_error=False)
+    if not stream:
+        v, _ = check_int(name=f'{yolo_predict.__name__} batch_size', value=batch_size, min_value=1, raise_error=False)
+        if not v: raise InvalidInputError(msg=f'If not stream, pass valid batch_size integer. Got {batch_size}', source=yolo_predict.__name__)
+    check_instance(source=f'{yolo_predict.__name__} source', instance=source, accepted_types=(str, os.PathLike, np.ndarray), raise_error=True)
+    if isinstance(source, (str, os.PathLike,)):
+        _ = get_video_meta_data(video_path=source)
+    else:
+        check_if_valid_img(data=source, source=f'{yolo_predict.__name__} source', raise_error=True)
+    check_instance(source=f'{yolo_predict.__name__} model', instance=model, accepted_types=(YOLO,), raise_error=True)
+    check_valid_device(device=device)
+    return model.predict(source=source,
+                         half=half,
+                         batch=batch_size,
+                         stream=stream,
+                         device=device,
+                         conf=threshold,
+                         max_det=max_detections,
+                         verbose=verbose)
+
+
+#fit_yolo(weights_path="/mnt/d/yolo_weights/yolo11n-pose.pt", model_yaml="/mnt/d/ares/data/termite_2/yolo/map.yaml", save_path="/mnt/d/ares/data/termite_2/yolo/mdl", batch=32, epochs=1)
+#fit_yolo(weights_path="/mnt/d/yolo_weights/yolo11n-pose.pt", model_yaml="/mnt/d/ares/data/ant/yolo/map.yaml", save_path="/mnt/d/ares/data/ant/yolo/mdl", batch=18, epochs=100, workers=12)

simba_uw_tf_dev-3.2.2/simba/bounding_box_tools/yolo/yolo_inference.py

@@ -0,0 +1,161 @@
+
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+from typing import Dict, List, Optional, Tuple, Union
+
+try:
+    from typing import Literal
+except:
+    from typing_extensions import Literal
+
+import numpy as np
+import pandas as pd
+import torch
+
+from simba.third_party_label_appenders.converters import yolo_obb_data_to_bounding_box
+from simba.utils.checks import (check_file_exist_and_readable, check_float, check_if_dir_exists, check_int, check_str, check_valid_boolean, check_valid_lst, get_fn_ext, check_valid_tuple, check_valid_array)
+from simba.utils.data import df_smoother, savgol_smoother
+from simba.utils.printing import SimbaTimer, stdout_success
+from simba.utils.read_write import find_core_cnt, get_video_meta_data
+from simba.bounding_box_tools.yolo.utils import load_yolo_model, yolo_predict, check_valid_device, _get_undetected_obs
+
+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'
+
+class YoloInference():
+    """
+    Performs object detection inference on a video using a YOLO model.
+
+    This class performs YOLO-based object detection on one or more video files. It supports GPU acceleration,
+    batch processing, streaming, and optional result saving. The model returns bounding box coordinates and
+    class confidence scores for each frame. Results can be smoothed or interpolated to handle detection gaps.
+
+    .. seealso::
+       To perform bounding box and keypoint (pose) detection, see :func:`~simba.bounding_box_tools.yolo.yolo_pose_inference.YOLOPoseInference`
+
+    :param Union[str, os.PathLike] weights_path: Path to the YOLO model weights file.
+    :param Union[str, os.PathLike] or List[Union[str, os.PathLike]] video_path: Path(s) to the input video file(s) for performing inference.
+    :param Optional[bool] verbose: If True, outputs progress information and timing. Defaults to False.
+    :param Optional[Union[str, os.PathLike]] save_dir: Directory to save the inference results (e.g., as CSV). If None, results are returned. Defaults to None.
+    :param Optional[bool] half_precision: If True, uses half precision (fp16) for inference. Reduces memory usage and speeds up inference on supported devices. Defaults to True.
+    :param Union[Literal['cpu'], int] device: Device to use for inference. Use 'cpu' for CPU or GPU index (e.g., 0 for CUDA:0). Defaults to 0.
+    :param Optional[int] batch_size: Number of frames to process in parallel. Defaults to 4.
+    :param int core_cnt: Number of CPU cores to use for preprocessing or parallel work. Defaults to 8.
+    :param float threshold: Confidence threshold for object detections. Detections below this value are discarded. Defaults to 0.25.
+    :param int max_detections: Maximum number of detections per frame. Defaults to 300.
+    :param Optional[Literal['savitzky-golay', 'bartlett', 'blackman', 'boxcar', 'cosine', 'gaussian', 'hamming', 'exponential']] smoothing_method: Optional smoothing method to apply to bounding box sequences over time to reduce jitter.
+    :param Optional[int] smoothing_time_window: Number of frames to consider for temporal smoothing. Effective only if smoothing_method is specified.
+    :param bool interpolate: If True, interpolates missing bounding boxes using the 'nearest' method. Useful for filling detection gaps.
+    :param int imgsz: Input image size for inference. Must be a square value. Defaults to 640.
+    :param Optional[bool] stream: If True, performs inference using a generator that yields predictions iteratively, reducing memory usage. Defaults to True.
+    :return: If `save_dir` is None, returns a dictionary of results keyed by video file names. Otherwise, results are saved to `save_dir`.
+    :rtype: Union[None, Dict[str, pd.DataFrame]]
+
+    :example:
+    >>> video_path = "/mnt/d/netholabs/yolo_videos/input/mp4_20250606083508/2025-05-28_19-50-23.mp4"
+    >>> i = YoloInference(weights_path=r"/mnt/c/troubleshooting/coco_data/mdl/train8/weights/best.pt", video_path=video_path, save_dir=r"/mnt/c/troubleshooting/coco_data/mdl/results", verbose=True, gpu=True)
+    >>> i.run()
+    """
+    def __init__(self,
+                 weights_path: Union[str, os.PathLike],
+                 video_path: Union[Union[str, os.PathLike], List[Union[str, os.PathLike]]],
+                 verbose: Optional[bool] = False,
+                 save_dir: Optional[Union[str, os.PathLike]] = None,
+                 half_precision: Optional[bool] = True,
+                 device: Union[Literal['cpu'], int] = 0,
+                 batch_size: Optional[int] = 4,
+                 core_cnt: int = 8,
+                 threshold: float = 0.25,
+                 max_detections: int = 300,
+                 smoothing_method: Optional[Literal['savitzky-golay', 'bartlett', 'blackman', 'boxcar', 'cosine', 'gaussian', 'hamming', 'exponential']] = None,
+                 smoothing_time_window: Optional[int] = None,
+                 interpolate: bool = False,
+                 imgsz: int = 640,
+                 stream: Optional[bool] = True) -> Union[None, Dict[str, pd.DataFrame]]:
+
+        if isinstance(video_path, list):
+            check_valid_lst(data=video_path, source=f'{self.__class__.__name__} video_path', valid_dtypes=(str, np.str_,), min_len=1)
+        elif isinstance(video_path, str):
+            check_file_exist_and_readable(file_path=video_path)
+            video_path = [video_path]
+        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, verbose, stream, interpolate], source=self.__class__.__name__)
+        check_int(name=f'{self.__class__.__name__} batch_size', value=batch_size, min_value=1)
+        check_int(name=f'{self.__class__.__name__} imgsz', value=imgsz, min_value=1)
+        check_float(name=f'{self.__class__.__name__} threshold', value=threshold, min_value=0.0, max_value=1.0)
+        check_int(name=f'{self.__class__.__name__} max_detections', value=max_detections, min_value=1)
+        check_valid_device(device=device)
+        if save_dir is not None:
+            check_if_dir_exists(in_dir=save_dir, source=f'{self.__class__.__name__} save_dir')
+        if smoothing_method is not None:
+            check_str(name=f'{self.__class__.__name__} smoothing', value=smoothing_method, options=SMOOTHING_METHODS)
+            check_float(name=f'{self.__class__.__name__} smoothing_time_window', value=smoothing_time_window, min_value=10e-6)
+        torch.set_num_threads(core_cnt)
+        self.model = load_yolo_model(weights_path=weights_path, verbose=verbose, device=device)
+        self.video_path, self.half_precision, self.stream, self.batch_size = video_path, half_precision, stream, batch_size
+        self.interpolate, self.smoothing_method, self.smoothing_time_window = interpolate, smoothing_method, smoothing_time_window
+        self.save_dir, self.verbose, self.imgsz, self.core_cnt, self.device = save_dir, verbose, imgsz, core_cnt,device
+        self.threshold, self.max_detections = threshold, max_detections
+
+    def run(self):
+        class_dict = self.model.names
+        results = {}
+        timer = SimbaTimer(start=True)
+        for path in self.video_path:
+            _, video_name, _ = get_fn_ext(filepath=path)
+            video_meta_data = get_video_meta_data(video_path=path)
+            video_out = []
+            video_predictions = yolo_predict(model=self.model, source=path, half=self.half_precision, batch_size=self.batch_size, stream=self.stream, imgsz=self.imgsz, device=self.device, threshold=self.threshold, max_detections=self.max_detections, verbose=self.verbose)
+            for frm_cnt, video_prediction in enumerate(video_predictions):
+                boxes = video_prediction.obb.data if video_prediction.obb is not None else video_prediction.boxes.data
+                boxes = boxes.cpu().numpy().astype(np.float32)
+                detected_classes = np.unique(boxes[:, -1]).astype(int) if boxes.size > 0 else []
+                for class_id, class_name in class_dict.items():
+                    if class_id not in detected_classes:
+                        video_out.append(_get_undetected_obs(frm_id=frm_cnt, class_id=class_id, class_name=class_name, value_cnt=9))
+                        continue
+                    cls_data = boxes[np.argwhere(boxes[:, -1] == class_id)]
+                    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 self.interpolate:
+                for cord_col in COORD_COLS:
+                    results[video_name][cord_col] = results[video_name][cord_col].astype(np.float32).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 self.smoothing_method is not None:
+                if self.smoothing_method != 'savitzky-golay':
+                    smoothened = df_smoother(data=results[video_name][COORD_COLS], fps=video_meta_data['fps'], time_window=self.smoothing_time_window, source=self.__class__.__name__, method=self.smoothing_method)
+                else:
+                    smoothened = savgol_smoother(data=results[video_name][COORD_COLS], fps=video_meta_data['fps'], time_window=self.smoothing_time_window, source=self.__class__.__name__)
+                results[video_name].update(smoothened)
+        timer.stop_timer()
+        if not self.save_dir:
+            if self.verbose:
+                stdout_success(f'YOLO results created', elapsed_time=timer.elapsed_time_str)
+            return results
+        else:
+            for k, v in results.items():
+                save_path = os.path.join(self.save_dir, f'{k}.csv')
+                v.to_csv(save_path)
+            if self.verbose:
+                stdout_success(f'YOLO results saved in {self.save_dir} directory', elapsed_time=timer.elapsed_time_str)
+
+#video_path = r"/mnt/c/troubleshooting/mitra/project_folder/videos/501_MA142_Gi_CNO_0521.mp4"
+#video_path = "/mnt/d/netholabs/yolo_videos/input/mp4_20250606083508/2025-05-28_19-50-23.mp4"
+# video_path = "/mnt/d/netholabs/yolo_videos/2025-05-28_19-50-23.mp4"
+# weights_path = "/mnt/d/netholabs/yolo_mdls_1/train/weights/best.pt"
+# save_dir = '/mnt/d/netholabs/yolo_videos_out'
+# i = YoloInference(weights_path=weights_path,
+#                   video_path=video_path,
+#                   save_dir=save_dir,
+#                   verbose=True)
+# i.run()