__author__ = "Simon Nilsson; sronilsson@gmail.com"
import os
from copy import copy, deepcopy
from typing import Dict, Optional, Union
import cv2
import numpy as np
from shapely.affinity import scale
from shapely.geometry import MultiPolygon, Polygon
try:
from typing import Literal
except:
from typing_extensions import Literal
import pandas as pd
from simba.mixins.geometry_mixin import GeometryMixin
from simba.mixins.image_mixin import ImageMixin
from simba.utils.checks import (check_float, check_if_dir_exists,
check_instance, check_int,
check_nvidea_gpu_available, check_str,
check_valid_boolean, is_img_bw)
from simba.utils.data import df_smoother, savgol_smoother
from simba.utils.enums import Formats, Methods, Options
from simba.utils.errors import FFMPEGCodecGPUError, InvalidInputError
from simba.utils.printing import SimbaTimer, stdout_success
from simba.utils.read_write import (find_all_videos_in_directory, get_fn_ext,
get_video_meta_data, remove_files,
write_df)
from simba.video_processors.video_processing import (video_bg_subtraction,
video_bg_subtraction_mp)
[docs]class BlobLocationComputer(object):
"""
Detecting and saving blob locations from video files.
.. video:: _static/img/BlobLocationComputer.webm
:width: 800
:autoplay:
:loop:
:muted:
:align: center
:param Union[str, os.PathLike] data_path: Path to a video file or a directory containing video files. The videos will be processed for blob detection.
:param Optional[bool] verbose: If True, prints progress and success messages to the console. Default is True.
:param Optional[bool] gpu: If True, GPU acceleration will be used for blob detection. Default is True.
:param Optional[int] batch_size: The number of frames to process in each batch for blob detection. Default is 2500.
:param Optional[Union[str, os.PathLike]] save_dir: Directory where the blob location data will be saved as CSV files. If None, the results will not be saved. Default is None.
:param Optional[bool] multiprocessing: If True, video background subtraction will be done using multiprocessing. Default is False.
:example:
>>> x = BlobLocationComputer(data_path=r"C:/troubleshooting/RAT_NOR/project_folder/videos/2022-06-20_NOB_DOT_4_downsampled_bg_subtracted.mp4", multiprocessing=True, gpu=True, batch_size=2000, save_dir=r"C:/blob_positions")
>>> x.run()
"""
def __init__(self,
data_path: Union[str, os.PathLike],
verbose: Optional[bool] = True,
gpu: Optional[bool] = True,
batch_size: int = 2500,
save_dir: Optional[Union[str, os.PathLike]] = None,
smoothing: Optional[str] = None,
multiprocessing: bool = False):
if os.path.isdir(data_path):
self.data_paths = find_all_videos_in_directory(directory=data_path, as_dict=True, raise_error=True).values()
elif os.path.isfile(data_path):
self.data_paths = [data_path]
else:
raise InvalidInputError(msg=f'{data_path} is not a valid directory or video file path or directory path.')
self.video_meta_data = {}
for i in self.data_paths:
self.video_meta_data[get_fn_ext(filepath=i)[1]] = get_video_meta_data(video_path=i)
if save_dir is not None:
check_if_dir_exists(in_dir=os.path.dirname(save_dir), source=self.__class__.__name__)
check_valid_boolean(value=[verbose, gpu, multiprocessing])
if gpu and not check_nvidea_gpu_available():
raise FFMPEGCodecGPUError(msg='No GPU detected.', source=self.__class__.__name__)
if smoothing is not None:
check_str(name=f'{self.__class__.__name__} smoothing', value=smoothing, options=Options.SMOOTHING_OPTIONS.value)
check_int(name=f'{self.__class__.__name__} batch_size', value=batch_size, min_value=1)
self.multiprocessing = multiprocessing
self.verbose = verbose
self.gpu = gpu
self.batch_size = batch_size
self.save_dir = save_dir
self.smoothing = smoothing
[docs] def run(self):
timer = SimbaTimer(start=True)
self.location_data = {}
for file_cnt, video_path in enumerate(self.data_paths):
video_timer = SimbaTimer(start=True)
_, video_name, ext = get_fn_ext(filepath=video_path)
temp_video_path = os.path.join(os.path.dirname(video_path), video_name + '_temp.mp4')
if not self.multiprocessing:
_ = video_bg_subtraction(video_path=video_path, verbose=self.verbose, bg_color=(0, 0, 0), fg_color=(255, 255, 255), save_path=temp_video_path)
else:
_ = video_bg_subtraction_mp(video_path=video_path, verbose=self.verbose, bg_color=(0, 0, 0), fg_color=(255, 255, 255), save_path=temp_video_path)
self.location_data[video_name] = ImageMixin.get_blob_locations(video_path=temp_video_path, gpu=self.gpu, verbose=self.verbose, batch_size=self.batch_size).astype(np.int32)
remove_files(file_paths=[temp_video_path])
video_timer.stop_timer()
print(f'Blob detection for video {video_name} ({file_cnt+1}/{len(self.data_paths)}) complete (elapsed time: {video_timer.elapsed_time_str}s)...')
timer.stop_timer()
if self.smoothing is not None:
print('Smoothing data...')
smoothened_data = {}
if self.smoothing == Methods.SAVITZKY_GOLAY.value:
for video_name, video_data in self.location_data.items():
smoothened_data[video_name] = savgol_smoother(data=video_data, fps=self.video_meta_data[video_name]['fps'], time_window=2000, source=video_name)
if self.smoothing == Methods.GAUSSIAN.value:
for video_name, video_data in self.location_data.items():
smoothened_data[video_name] = df_smoother(data=pd.DataFrame(video_data, columns=['X', 'Y']), fps=self.video_meta_data[video_name]['fps'], time_window=2000, source=video_name, method='gaussian')
self.location_data = deepcopy(smoothened_data)
del smoothened_data
if self.save_dir is not None:
for video_name, video_data in self.location_data.items():
save_path = os.path.join(self.save_dir, f'{video_name}.csv')
df = pd.DataFrame(video_data, columns=['X', 'Y'])
write_df(df=df, file_type=Formats.CSV.value, save_path=save_path)
if self.verbose:
stdout_success(f'Video blob detection complete for {len(self.data_paths)} videos, data saved at {self.save_dir}', elapsed_time=timer.elapsed_time_str)
else:
if self.verbose:
stdout_success(f'Video blob detection complete for {len(self.data_paths)} video', elapsed_time=timer.elapsed_time_str)
# x = BlobLocationComputer(data_path=r"C:\troubleshooting\RAT_NOR\project_folder\videos\2022-06-20_NOB_DOT_4_downsampled.mp4", multiprocessing=True, gpu=True, batch_size=2000, save_dir=r"C:\troubleshooting\RAT_NOR\project_folder\csv\blob_positions", smoothing='Savitzky Golay')
# x.run()
# x = BlobLocationComputer(data_path=r"C:\troubleshooting\RAT_NOR\project_folder\videos\2022-06-20_NOB_DOT_4_downsampled.mp4", multiprocessing=True, gpu=True, batch_size=2000, save_dir=r"C:\troubleshooting\RAT_NOR\project_folder\csv\blob_positions", smoothing='Savitzky Golay')
# x.run()