From 81e40f31c19804da6df41323482f8dddc6c716e9 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 29 Sep 2025 13:35:41 -0700 Subject: [PATCH 01/30] add script for recording video with boson sdk --- boson_scripts/record.py | 122 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 boson_scripts/record.py diff --git a/boson_scripts/record.py b/boson_scripts/record.py new file mode 100644 index 0000000..8f59720 --- /dev/null +++ b/boson_scripts/record.py @@ -0,0 +1,122 @@ +import cv2 +import numpy as np +from SDK_USER_PERMISSIONS import * +from time import sleep + +# SETUP +def setup(func, *args, delay=1, success_code=0, description=''): + '''run camera setup fucntion until it succeeds''' + + result = func(*args) + while result != success_code: + sleep(delay) + result = func(*args) + print(f'{description} Set') + + +if __name__ == "__main__": + + myCam = CamAPI.pyClient(manualport="/dev/ttyACM0") #Boson COM port on windows, check device manager + + # Set Radiometric Parameters + setup(myCam.bosonSetGainMode, + FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN, + description='High Gain Mode') + setup(myCam.TLinearSetControl, + FLR_ENABLE_E.FLR_ENABLE, + description='TLinear Mode') + setup(myCam.sysctrlSetUsbVideoIR16Mode, + FLR_SYSCTRL_USBIR16_MODE_E.FLR_SYSCTRL_USBIR16_MODE_TLINEAR, + description='IR16 Mode') + setup(myCam.radiometrySetTransmissionWindow, + 100, + description='Window Transimission') + setup(myCam.TLinearRefreshLUT, + FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN, + description='LUT Refresh') + setup(myCam.bosonRunFFC, + description='Flat Field Correction') + + +''' +success = 3 +while success != 0: + success = myCam.bosonSetGainMode(FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN) #Low gain mode also available if tem> + sleep(1) +print('Gain Mode Set') + +success = 3 +while success != 0: + success = myCam.TLinearSetControl(FLR_ENABLE_E.FLR_ENABLE) + sleep(1) +print('TLinear Mode Set') + +success = 3 +while success != 0: + success = myCam.sysctrlSetUsbVideoIR16Mode(FLR_SYSCTRL_USBIR16_MODE_E.FLR_SYSCTRL_USBIR16_MODE_TLINEAR) + sleep(1) +print(f'IR16 Mode Mode') + +success = 3 +while success != 0: + success = myCam.radiometrySetTransmissionWindow(100) #100% transmission (no window in front of Boson) + sleep(1) +print(f'Window Mode Set') + +success = 3 +while success != 0: + success = myCam.TLinearRefreshLUT(FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN) #necessary after setting radiometry> + sleep(1) +print(f'LUT Mode Set') + +success = 3 +while success != 0: + success = myCam.bosonRunFFC() + sleep(1) +print(f'FFC Complete') + + + +cap = cv2.VideoCapture(0, cv2.CAP_V4L2) +cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 256) +cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320) +cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) # don't auto-convert to RGB +cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y','1', '6', ' ')) + + +#cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('Y','1','6',' ')) +#fourcc = cv2.VideoWriter_fourcc(*'Y16 ') +#out = cv2.VideoWriter('output.mp4', fourcc, 20.0, (w, h)) + + +while True: + ret, frame = cap.read() + min_val, max_val = np.min(frame), np.max(frame) +# print(f'frame size: {frame.shape}') +# print('Raw Values: ') +# print(f'Min {min_val} | Max {max_val}') +# print(f'Unique Values: {len(np.unique(frame))}') +# convert to 8bit for viewing + if frame.dtype == np.uint16: + print(f'\nConverting to 8bit for viewing') + unique_vals = np.unique(frame) + frame_8bit = ((frame-min_val)/(max_val-min_val) *255).astype(np.uint8) + +# Center Temp + center_raw = frame[int(frame.shape[0]/2), int(frame.shape[1]/2)] + center_c = round((center_raw/100) - 273, 1) + center_f = round(center_c * 9/5 + 32, 1) + +# print('Converted Temp Values: ') + print(f'Center temp: Temp C - {center_c} | Temp F - {center_f}') + + cv2.imshow('Cam', frame_8bit) + if cv2.waitKey(1)==ord('q'): + break + +cv2.destroyAllWindows() + +cap.release() +''' +myCam.Close() + From 1c935f1b6f80ce3263b67dfb63fdda57dac1f4da Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 29 Sep 2025 13:38:03 -0700 Subject: [PATCH 02/30] Create README.md --- boson_scripts/README.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 boson_scripts/README.md diff --git a/boson_scripts/README.md b/boson_scripts/README.md new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/boson_scripts/README.md @@ -0,0 +1 @@ + From b536eec971cbf6f39bc74292335376bc504ba261 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 29 Sep 2025 13:40:52 -0700 Subject: [PATCH 03/30] Update README.md --- boson_scripts/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/boson_scripts/README.md b/boson_scripts/README.md index 8b13789..226ff60 100644 --- a/boson_scripts/README.md +++ b/boson_scripts/README.md @@ -1 +1,2 @@ - +Download Boson 3.0 IDD & SDK: https://oem.flir.com/products/boson/?vertical=lwir&segment=oem +Modifications to make compatible with Raspberry Pi ARM architecture: https://flir.custhelp.com/app/answers/detail/a_id/6471/~/compiling-the-boson-sdk-for-arm64-linux-platforms From 3a53134c8fececa1e5c647dc5264e94a012cf244 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Tue, 30 Sep 2025 07:51:52 -0900 Subject: [PATCH 04/30] Update recording script --- boson_scripts/record.py | 135 ++++++++++++++++++---------------------- 1 file changed, 59 insertions(+), 76 deletions(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index 8f59720..d9c6dae 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -2,8 +2,9 @@ import numpy as np from SDK_USER_PERMISSIONS import * from time import sleep +import argparse -# SETUP +# setup camera parameters def setup(func, *args, delay=1, success_code=0, description=''): '''run camera setup fucntion until it succeeds''' @@ -13,11 +14,30 @@ def setup(func, *args, delay=1, success_code=0, description=''): result = func(*args) print(f'{description} Set') +def get_center_temp(frame): + '''get frame center temp in C and F''' + + center_raw = frame[int(frame.shape[0]/2), int(frame.shape[1]/2)] + center_c = round((center_raw/100) - 273, 1) + center_f = round(center_c * 9/5 + 32, 1) + + return center_c, center_f if __name__ == "__main__": + # Connect to Camera myCam = CamAPI.pyClient(manualport="/dev/ttyACM0") #Boson COM port on windows, check device manager + # User set parameters + parser = argparse.ArgumentParser() + parser.add_argument('--raw_output_path', + default='output.npy', + help='filepath to save raw radiometric output') + parser.add_argument('--video_output_path', + default='output.mp4', + help='filepath to save normalized video output') + args = parser.parse_args() + # Set Radiometric Parameters setup(myCam.bosonSetGainMode, FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN, @@ -37,86 +57,49 @@ def setup(func, *args, delay=1, success_code=0, description=''): setup(myCam.bosonRunFFC, description='Flat Field Correction') + #open camera + cap = cv2.VideoCapture(0, cv2.CAP_V4L2) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 256) + cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320) + cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) # don't auto-convert to RGB + cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y','1', '6', ' ')) + -''' -success = 3 -while success != 0: - success = myCam.bosonSetGainMode(FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN) #Low gain mode also available if tem> - sleep(1) -print('Gain Mode Set') - -success = 3 -while success != 0: - success = myCam.TLinearSetControl(FLR_ENABLE_E.FLR_ENABLE) - sleep(1) -print('TLinear Mode Set') - -success = 3 -while success != 0: - success = myCam.sysctrlSetUsbVideoIR16Mode(FLR_SYSCTRL_USBIR16_MODE_E.FLR_SYSCTRL_USBIR16_MODE_TLINEAR) - sleep(1) -print(f'IR16 Mode Mode') - -success = 3 -while success != 0: - success = myCam.radiometrySetTransmissionWindow(100) #100% transmission (no window in front of Boson) - sleep(1) -print(f'Window Mode Set') - -success = 3 -while success != 0: - success = myCam.TLinearRefreshLUT(FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN) #necessary after setting radiometry> - sleep(1) -print(f'LUT Mode Set') - -success = 3 -while success != 0: - success = myCam.bosonRunFFC() - sleep(1) -print(f'FFC Complete') - - - -cap = cv2.VideoCapture(0, cv2.CAP_V4L2) -cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 256) -cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320) -cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) # don't auto-convert to RGB -cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y','1', '6', ' ')) - - -#cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('Y','1','6',' ')) -#fourcc = cv2.VideoWriter_fourcc(*'Y16 ') -#out = cv2.VideoWriter('output.mp4', fourcc, 20.0, (w, h)) - - -while True: - ret, frame = cap.read() - min_val, max_val = np.min(frame), np.max(frame) -# print(f'frame size: {frame.shape}') -# print('Raw Values: ') -# print(f'Min {min_val} | Max {max_val}') -# print(f'Unique Values: {len(np.unique(frame))}') -# convert to 8bit for viewing - if frame.dtype == np.uint16: - print(f'\nConverting to 8bit for viewing') - unique_vals = np.unique(frame) + # output video settings + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + writer = cv2.VideoWriter(args.video_output_path, fourcc, 30, (320, 256)) + + raw_frames = [] + print('recording...') + + while True: + ret, frame = cap.read() + if not ret: + print('frame grab failed, stopping') + break + + raw_frames.append(frame) + + # Viewable Window + min_val, max_val = np.min(frame), np.max(frame) frame_8bit = ((frame-min_val)/(max_val-min_val) *255).astype(np.uint8) + frame_color = cv2.applyColorMap(frame_8bit, cv2.COLORMAP_INFERNO) + writer.write(frame_color) -# Center Temp - center_raw = frame[int(frame.shape[0]/2), int(frame.shape[1]/2)] - center_c = round((center_raw/100) - 273, 1) - center_f = round(center_c * 9/5 + 32, 1) + # Center Temp + temp_c, temp_f = get_center_temp(frame) + print(f'Center temp: Temp C - {temp_c} | Temp F - {temp_f}') -# print('Converted Temp Values: ') - print(f'Center temp: Temp C - {center_c} | Temp F - {center_f}') + cv2.imshow('color', frame_color) + if cv2.waitKey(1)==ord('q'): + print(f'recording saved at {args.video_output_path}') + break - cv2.imshow('Cam', frame_8bit) - if cv2.waitKey(1)==ord('q'): - break + raw_frames = np.stack(raw_frames, axis=0) + np.save(args.raw_output_path, raw_frames) -cv2.destroyAllWindows() + cap.release() + writer.release() -cap.release() -''' -myCam.Close() + myCam.Close() From 6a962dff29330e6d925212c81760910375ed5712 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Tue, 30 Sep 2025 11:52:33 -0900 Subject: [PATCH 05/30] flake8 and pylint --- boson_scripts/record.py | 71 +++++++++++++++++++++++++++++++---------- 1 file changed, 55 insertions(+), 16 deletions(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index d9c6dae..e480cc8 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -1,12 +1,35 @@ -import cv2 -import numpy as np -from SDK_USER_PERMISSIONS import * +''' +This script records radiometric data from the FLIR Boson thermal camera + +Usage: + $ python record.py \ + --raw_output_path path/to/save/raw/data.npy \ + --video_output_path path/to/save/video/data.mp4 \ + --bosonsdk_path path/to/boson/sdk/files +''' + + from time import sleep import argparse +import sys +import os +import numpy as np +import cv2 # setup camera parameters def setup(func, *args, delay=1, success_code=0, description=''): - '''run camera setup fucntion until it succeeds''' + '''run camera setup fucntion 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 + + Returns: + None + ''' result = func(*args) while result != success_code: @@ -14,19 +37,26 @@ def setup(func, *args, delay=1, success_code=0, description=''): result = func(*args) print(f'{description} Set') -def get_center_temp(frame): - '''get frame center temp in C and F''' - center_raw = frame[int(frame.shape[0]/2), int(frame.shape[1]/2)] +def get_center_temp(frame_16bit): + '''get frame center temp in C and F + + Args: + frame_16bit: raw radiometric 16bit frame + + Returns: + center_c: center temperature in Celsius + center_f: center temperature in Fahrenheit + ''' + + center_raw = frame_16bit[int(frame_16bit.shape[0]/2), int(frame_16bit.shape[1]/2)] center_c = round((center_raw/100) - 273, 1) center_f = round(center_c * 9/5 + 32, 1) return center_c, center_f -if __name__ == "__main__": - # Connect to Camera - myCam = CamAPI.pyClient(manualport="/dev/ttyACM0") #Boson COM port on windows, check device manager +if __name__ == "__main__": # User set parameters parser = argparse.ArgumentParser() @@ -36,8 +66,19 @@ def get_center_temp(frame): parser.add_argument('--video_output_path', default='output.mp4', help='filepath to save normalized video output') + parser.add_argument('--bosonsdk_path', + default='~/BosonSDK/SDK_USER_PERMISSIONS', + help='filepath for Boson SDK') args = parser.parse_args() + # import Boson SDK + sdk_path = os.path.expanduser(args.bosonsdk_path) + sys.path.append(sdk_path) + from SDK_USER_PERMISSIONS import * + + # Connect to Camera + myCam = CamAPI.pyClient(manualport="/dev/ttyACM0") + # Set Radiometric Parameters setup(myCam.bosonSetGainMode, FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN, @@ -57,13 +98,12 @@ def get_center_temp(frame): setup(myCam.bosonRunFFC, description='Flat Field Correction') - #open camera + # open camera cap = cv2.VideoCapture(0, cv2.CAP_V4L2) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 256) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320) cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) # don't auto-convert to RGB - cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y','1', '6', ' ')) - + cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y', '1', '6', ' ')) # output video settings fourcc = cv2.VideoWriter_fourcc(*'mp4v') @@ -82,7 +122,7 @@ def get_center_temp(frame): # Viewable Window min_val, max_val = np.min(frame), np.max(frame) - frame_8bit = ((frame-min_val)/(max_val-min_val) *255).astype(np.uint8) + frame_8bit = ((frame - min_val)/(max_val - min_val) * 255).astype(np.uint8) frame_color = cv2.applyColorMap(frame_8bit, cv2.COLORMAP_INFERNO) writer.write(frame_color) @@ -91,7 +131,7 @@ def get_center_temp(frame): print(f'Center temp: Temp C - {temp_c} | Temp F - {temp_f}') cv2.imshow('color', frame_color) - if cv2.waitKey(1)==ord('q'): + if cv2.waitKey(1) == ord('q'): print(f'recording saved at {args.video_output_path}') break @@ -102,4 +142,3 @@ def get_center_temp(frame): writer.release() myCam.Close() - From 69f0061d457851f21c20b32179ef66c97f8cc19b Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Tue, 7 Oct 2025 09:28:57 -0900 Subject: [PATCH 06/30] Update 16bit to 8bit normalization method --- boson_scripts/record.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index e480cc8..0365321 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -122,15 +122,16 @@ def get_center_temp(frame_16bit): # Viewable Window min_val, max_val = np.min(frame), np.max(frame) - frame_8bit = ((frame - min_val)/(max_val - min_val) * 255).astype(np.uint8) + 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) writer.write(frame_color) + cv2.imshow('color', frame_color) - # Center Temp + # Display Center Temp temp_c, temp_f = get_center_temp(frame) print(f'Center temp: Temp C - {temp_c} | Temp F - {temp_f}') - cv2.imshow('color', frame_color) + # Stop Recording if cv2.waitKey(1) == ord('q'): print(f'recording saved at {args.video_output_path}') break From 7bd5adc1009fbcee357008d0717d152c57b46037 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 20 Oct 2025 09:18:57 -0900 Subject: [PATCH 07/30] Make setup delay longer --- boson_scripts/record.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index 0365321..6c3503b 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -17,7 +17,7 @@ import cv2 # setup camera parameters -def setup(func, *args, delay=1, success_code=0, description=''): +def setup(func, *args, delay=2, success_code=0, description=''): '''run camera setup fucntion until it succeeds Args: From 8495b635b04c1299beda67f0baa4ab1bbc8d2ba7 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 20 Oct 2025 09:24:16 -0900 Subject: [PATCH 08/30] Make compatible with headless setup --- boson_scripts/record.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index 6c3503b..44f8dc0 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -108,6 +108,7 @@ def get_center_temp(frame_16bit): # output video settings fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter(args.video_output_path, fourcc, 30, (320, 256)) + headless = not os.environ.get("DISPLAY") raw_frames = [] print('recording...') @@ -125,7 +126,8 @@ def get_center_temp(frame_16bit): 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) writer.write(frame_color) - cv2.imshow('color', frame_color) + if not headless: + cv2.imshow('color', frame_color) # Display Center Temp temp_c, temp_f = get_center_temp(frame) @@ -141,5 +143,7 @@ def get_center_temp(frame_16bit): cap.release() writer.release() + if not headless: + cv2.destroyAllWindows() myCam.Close() From 9c1f06d67fd072edfa964e800fc07ce67a46f027 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 20 Oct 2025 09:38:02 -0900 Subject: [PATCH 09/30] Fix stop record method --- boson_scripts/record.py | 77 +++++++++++++++++++++++------------------ 1 file changed, 44 insertions(+), 33 deletions(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index 44f8dc0..2b9bf34 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -16,6 +16,7 @@ import numpy as np import cv2 + # setup camera parameters def setup(func, *args, delay=2, success_code=0, description=''): '''run camera setup fucntion until it succeeds @@ -49,7 +50,8 @@ def get_center_temp(frame_16bit): center_f: center temperature in Fahrenheit ''' - center_raw = frame_16bit[int(frame_16bit.shape[0]/2), int(frame_16bit.shape[1]/2)] + center_raw = frame_16bit[int(frame_16bit.shape[0]/2), + int(frame_16bit.shape[1]/2)] center_c = round((center_raw/100) - 273, 1) center_f = round(center_c * 9/5 + 32, 1) @@ -113,37 +115,46 @@ def get_center_temp(frame_16bit): raw_frames = [] print('recording...') - while True: - ret, frame = cap.read() - if not ret: - print('frame grab failed, stopping') - break - - raw_frames.append(frame) - - # Viewable Window - min_val, max_val = np.min(frame), np.max(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) - writer.write(frame_color) + try: + while True: + ret, frame = cap.read() + if not ret: + print('frame grab failed, stopping') + break + + raw_frames.append(frame) + + # Viewable Window + min_val, max_val = np.min(frame), np.max(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) + writer.write(frame_color) + if not headless: + cv2.imshow('color', frame_color) + + # Display Center Temp + temp_c, temp_f = get_center_temp(frame) + print(f'Center temp: Temp C - {temp_c} | Temp F - {temp_f}') + + # Stop Recording + #if cv2.waitKey(1) == ord('q'): + # print(f'recording saved at {args.video_output_path}') + # break + + except KeyboardInterrupt: + print('Stopping Recording...') + + finally: + raw_frames = np.stack(raw_frames, axis=0) + np.save(args.raw_output_path, raw_frames) + + print('Recordings Saved') + print(f'Radiometric data: {args.raw_output_path}') + print(f'Normalized video: {args.video_output_path}') + + cap.release() + writer.release() if not headless: - cv2.imshow('color', frame_color) - - # Display Center Temp - temp_c, temp_f = get_center_temp(frame) - print(f'Center temp: Temp C - {temp_c} | Temp F - {temp_f}') - - # Stop Recording - if cv2.waitKey(1) == ord('q'): - print(f'recording saved at {args.video_output_path}') - break - - raw_frames = np.stack(raw_frames, axis=0) - np.save(args.raw_output_path, raw_frames) - - cap.release() - writer.release() - if not headless: - cv2.destroyAllWindows() + cv2.destroyAllWindows() - myCam.Close() + myCam.Close() From d386808a91bd6d50acadaeec4f632992bc9aa77d Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 20 Oct 2025 09:53:11 -0900 Subject: [PATCH 10/30] Fix status statement formatting --- boson_scripts/record.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index 2b9bf34..48cba73 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -32,11 +32,12 @@ def setup(func, *args, delay=2, success_code=0, description=''): None ''' + print(f'\nSetting {description}...') result = func(*args) while result != success_code: sleep(delay) result = func(*args) - print(f'{description} Set') + print(f'Success!') def get_center_temp(frame_16bit): @@ -113,7 +114,7 @@ def get_center_temp(frame_16bit): headless = not os.environ.get("DISPLAY") raw_frames = [] - print('recording...') + print('\n\nRecording Started...') try: while True: @@ -134,23 +135,20 @@ def get_center_temp(frame_16bit): # Display Center Temp temp_c, temp_f = get_center_temp(frame) - print(f'Center temp: Temp C - {temp_c} | Temp F - {temp_f}') + sys.stdout.write(f'\rCenter temp: Temp C - {temp_c} | Temp F - {temp_f}') + sys.stdout.flush() - # Stop Recording - #if cv2.waitKey(1) == ord('q'): - # print(f'recording saved at {args.video_output_path}') - # break except KeyboardInterrupt: - print('Stopping Recording...') + print('\n\nRecording Stopped...') finally: raw_frames = np.stack(raw_frames, axis=0) np.save(args.raw_output_path, raw_frames) - print('Recordings Saved') + print('\nRecordings Saved') print(f'Radiometric data: {args.raw_output_path}') - print(f'Normalized video: {args.video_output_path}') + print(f'Normalized video: {args.video_output_path}\n\n') cap.release() writer.release() From e9071011e2ad55c945bd1d6c3ccafbc78f7f35cd Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 20 Oct 2025 09:54:57 -0900 Subject: [PATCH 11/30] Fix flake8 --- boson_scripts/record.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index 48cba73..e61606e 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -37,7 +37,7 @@ def setup(func, *args, delay=2, success_code=0, description=''): while result != success_code: sleep(delay) result = func(*args) - print(f'Success!') + print('Success!') def get_center_temp(frame_16bit): @@ -138,7 +138,6 @@ def get_center_temp(frame_16bit): sys.stdout.write(f'\rCenter temp: Temp C - {temp_c} | Temp F - {temp_f}') sys.stdout.flush() - except KeyboardInterrupt: print('\n\nRecording Stopped...') From 2d7fb57f1172bff5e76a65feb145a9a37f556037 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 22 Oct 2025 13:00:47 -0900 Subject: [PATCH 12/30] Fix live view error --- boson_scripts/record.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/boson_scripts/record.py b/boson_scripts/record.py index e61606e..ef9a095 100644 --- a/boson_scripts/record.py +++ b/boson_scripts/record.py @@ -33,6 +33,7 @@ def setup(func, *args, delay=2, success_code=0, description=''): ''' print(f'\nSetting {description}...') + sleep(.5) result = func(*args) while result != success_code: sleep(delay) @@ -112,7 +113,6 @@ def get_center_temp(frame_16bit): fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter(args.video_output_path, fourcc, 30, (320, 256)) headless = not os.environ.get("DISPLAY") - raw_frames = [] print('\n\nRecording Started...') @@ -132,6 +132,7 @@ def get_center_temp(frame_16bit): writer.write(frame_color) if not headless: cv2.imshow('color', frame_color) + cv2.waitKey(1) # Display Center Temp temp_c, temp_f = get_center_temp(frame) From c63e46b182d25994bf7bd0b7d48f37d21f788df2 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 5 Nov 2025 09:59:43 -1000 Subject: [PATCH 13/30] Add abstract class to handle different cameras --- heatseek/capture.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 heatseek/capture.py diff --git a/heatseek/capture.py b/heatseek/capture.py new file mode 100644 index 0000000..442157f --- /dev/null +++ b/heatseek/capture.py @@ -0,0 +1,25 @@ +from abc import ABC, abstractmethod + +class Capture(ABC): + def __init__(self, camera_id=0): + self.camera_id = camera_id + + @abstractmethod + def setup(self): + pass + + @abstractmethod + def take_image(self): + pass + + @abstractmethod + def start_recording(self): + pass + + @abstractmethod + def stop_recording(self): + pass + + @abstractmethod + def release_camera(self): + pass From c7a5181ac46e4897d3ef04d2370d16fdba7a6a9f Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 5 Nov 2025 12:32:19 -1000 Subject: [PATCH 14/30] Add boson_capture file --- heatseek/boson_capture.py | 84 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 heatseek/boson_capture.py diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py new file mode 100644 index 0000000..083fbcb --- /dev/null +++ b/heatseek/boson_capture.py @@ -0,0 +1,84 @@ +from capture import Capture +from time import sleep +from importlib import import_module +import sys +import os +import numpy as np +import cv2 + + +class Boson_Capture(Capture): + def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): + super().__init__(camera_id) + + # import Boson SDK + path = os.path.expanduser(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 + port = f'/dev/ttyACM{self.camera_id}' + self.camera = self.cam_api.pyClient(manualport=port) + self.setup() + + def setup(self): + 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 + + Returns: + None + ''' + + result = func(*args) + while result != success_code: + sleep(delay) + result = func(*args) + print(f'{description} Set') + + def take_image(self): + print('Taking Image') + + + def start_recording(self): + print('Starting Recording...') + + def stop_recording(self): + print('Stopping Recording...') + + def release_camera(self): + self.camera.Close() + print('Release Camera') + +boson_capture = Boson_Capture() +boson_capture.release_camera() + From af995da2267b4f1750ebac1c6ddba5ca7b5300ce Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 5 Nov 2025 13:59:03 -1000 Subject: [PATCH 15/30] Add start recording loop --- heatseek/boson_capture.py | 43 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index 083fbcb..57469ae 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -10,6 +10,9 @@ class Boson_Capture(Capture): def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): super().__init__(camera_id) + self.isRecording = False + self.height = 256 + self.width = 320 # import Boson SDK path = os.path.expanduser(sdkpath) @@ -68,10 +71,44 @@ def configure(self, func, *args, delay=1, success_code=0, description=''): def take_image(self): print('Taking Image') - - def start_recording(self): + def start_recording(self, raw='output.npy', norm='output.mp4'): print('Starting Recording...') + # video settings + self.cap = cv2.VideoCapture(self.camera_id, 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(norm, fourcc, 30, (self.width, self.width)) + + # raw video setup + self.raw_mm = np.memmap(raw, dtype=np.uint16, mode='w+', shape=(50000, self.height, self.width)) + self.frame_index = 0 + + self.isRecording = True + + while self.isRecording: + ret, frame = cap.read() + if not ret: + print('frame grab failed, stopping') + break + + # Viewable frame + min_val, max_val = np.min(frame), np.max(frame) + frame_8bit = ((frame - min_val)/max_val - min_val) * 255).astype(np.uint8) + frame_color = cv2.applyColorMap(frame_8bit, cv2.COLORMAP_INFERNO) + writer.write(frame_color) + + # Raw Video + if self.frame_index > self.raw_mm.shape[0]: + print('reached preallocated size, stopping') + break + self.raw_mm[self.frame_index] = frame + self.frame_index += 1 + self.raw_mm.flush() + def stop_recording(self): print('Stopping Recording...') @@ -79,6 +116,8 @@ def release_camera(self): self.camera.Close() print('Release Camera') + +# TEST CODE boson_capture = Boson_Capture() boson_capture.release_camera() From 539ab3d351b4427274e627382723110604069afa Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Thu, 6 Nov 2025 08:55:29 -1000 Subject: [PATCH 16/30] Complete start and stop recording methods --- heatseek/boson_capture.py | 97 ++++++++++++++++++++++++++++----------- 1 file changed, 69 insertions(+), 28 deletions(-) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index 57469ae..fa74222 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -1,6 +1,7 @@ from capture import Capture from time import sleep from importlib import import_module +import threading import sys import os import numpy as np @@ -72,45 +73,82 @@ def take_image(self): print('Taking Image') def start_recording(self, raw='output.npy', norm='output.mp4'): - print('Starting Recording...') + if self.isRecording: + print('Recording in Progress!') + return + + self.raw_data_fpath = raw + self.viewable_video_fpath = norm + + self.isRecording = True + self.recording_thread = threading.Thread(target=self._record_loop, daemon=True) + self.recording_thread.start() + print('Staring Recording...') + + def _record_loop(self): # video settings self.cap = cv2.VideoCapture(self.camera_id, 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', ' ')) + self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height) + self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width) + self.cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) + self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y', '1', '6', ' ')) fourcc = cv2.VideoWriter_fourcc(*'mp4v') - writer = cv2.VideoWriter(norm, fourcc, 30, (self.width, self.width)) + writer = cv2.VideoWriter(self.viewable_video_fpath, fourcc, 30, (self.width, self.height)) + # raw video setup - self.raw_mm = np.memmap(raw, dtype=np.uint16, mode='w+', shape=(50000, self.height, self.width)) + self.raw_mm = np.memmap(self.raw_data_fpath, dtype=np.uint16, mode='w+', shape=(50000, self.height, self.width)) self.frame_index = 0 - self.isRecording = True - - while self.isRecording: - ret, frame = cap.read() - if not ret: - print('frame grab failed, stopping') - break - - # Viewable frame - min_val, max_val = np.min(frame), np.max(frame) - frame_8bit = ((frame - min_val)/max_val - min_val) * 255).astype(np.uint8) - frame_color = cv2.applyColorMap(frame_8bit, cv2.COLORMAP_INFERNO) - writer.write(frame_color) - - # Raw Video - if self.frame_index > self.raw_mm.shape[0]: - print('reached preallocated size, stopping') - break - self.raw_mm[self.frame_index] = frame - self.frame_index += 1 - self.raw_mm.flush() + try: + while self.isRecording: + ret, frame = self.cap.read() + if not ret: + print('frame grab failed, stopping') + break + + # Viewable frame + min_val, max_val = np.min(frame), np.max(frame) + #frame_8bit = ((frame - min_val)/(max_val - min_val) * 255).astype(np.uint8) + 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 self.frame_index > self.raw_mm.shape[0]: + print('reached preallocated size, stopping') + break + self.raw_mm[self.frame_index] = frame + self.frame_index += 1 + self.raw_mm.flush() + + finally: + self._finalize_recording() def stop_recording(self): + if not self.isRecording: + print('No recording in progress') + return + print('Stopping Recording...') + self.isRecording = False + + if self.recording_thread is not None: + self.recording_thread.join() + self.record_thread = None + + def _finalize_recording(self): + self.isRecording = 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('Radiometric data saved: {self.raw_data_fpath}') + print('Viewable Video: {self.viewable_video_fpath}') def release_camera(self): self.camera.Close() @@ -119,5 +157,8 @@ def release_camera(self): # TEST CODE boson_capture = Boson_Capture() +boson_capture.start_recording() +sleep(10) +boson_capture.stop_recording() boson_capture.release_camera() From 479902144501ecaeb6dcbead2dbdf169c4e06d33 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Thu, 6 Nov 2025 09:23:58 -1000 Subject: [PATCH 17/30] Document & pylint capture.py --- heatseek/capture.py | 56 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 51 insertions(+), 5 deletions(-) diff --git a/heatseek/capture.py b/heatseek/capture.py index 442157f..19fa73b 100644 --- a/heatseek/capture.py +++ b/heatseek/capture.py @@ -1,25 +1,71 @@ +''' +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, camera_id=0): + '''Initialize capture interface + + This method should handle hardware initalization and + SDK loading + + Args: + camera_id (int, opt): Device identifier for camera. + Defaults to 0. + ''' self.camera_id = camera_id @abstractmethod def setup(self): - pass + '''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): - pass + '''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): - pass + '''Begin video recording + + This method should capture frames continuously and save to disk + ''' @abstractmethod def stop_recording(self): - pass + '''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): - pass + '''Release camera hardware and close + + This method should handle any releasing and/or shutting down of + camera hardware + ''' From 5c4107e3edf33c6e6019f38301704e2026d42dfb Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Thu, 6 Nov 2025 12:46:06 -1000 Subject: [PATCH 18/30] pylint and flake8 boson_capture.py --- heatseek/boson_capture.py | 164 -------------------------------------- 1 file changed, 164 deletions(-) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index fa74222..e69de29 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -1,164 +0,0 @@ -from capture import Capture -from time import sleep -from importlib import import_module -import threading -import sys -import os -import numpy as np -import cv2 - - -class Boson_Capture(Capture): - def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): - super().__init__(camera_id) - self.isRecording = False - self.height = 256 - self.width = 320 - - # import Boson SDK - path = os.path.expanduser(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 - port = f'/dev/ttyACM{self.camera_id}' - self.camera = self.cam_api.pyClient(manualport=port) - self.setup() - - def setup(self): - 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 - - Returns: - None - ''' - - result = func(*args) - while result != success_code: - sleep(delay) - result = func(*args) - print(f'{description} Set') - - def take_image(self): - print('Taking Image') - - def start_recording(self, raw='output.npy', norm='output.mp4'): - if self.isRecording: - print('Recording in Progress!') - return - - self.raw_data_fpath = raw - self.viewable_video_fpath = norm - - self.isRecording = True - self.recording_thread = threading.Thread(target=self._record_loop, daemon=True) - self.recording_thread.start() - print('Staring Recording...') - - def _record_loop(self): - - # video settings - self.cap = cv2.VideoCapture(self.camera_id, cv2.CAP_V4L2) - self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height) - self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width) - self.cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) - self.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, 30, (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)) - self.frame_index = 0 - - try: - while self.isRecording: - ret, frame = self.cap.read() - if not ret: - print('frame grab failed, stopping') - break - - # Viewable frame - min_val, max_val = np.min(frame), np.max(frame) - #frame_8bit = ((frame - min_val)/(max_val - min_val) * 255).astype(np.uint8) - 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 self.frame_index > self.raw_mm.shape[0]: - print('reached preallocated size, stopping') - break - self.raw_mm[self.frame_index] = frame - self.frame_index += 1 - self.raw_mm.flush() - - finally: - self._finalize_recording() - - def stop_recording(self): - if not self.isRecording: - print('No recording in progress') - return - - print('Stopping Recording...') - self.isRecording = False - - if self.recording_thread is not None: - self.recording_thread.join() - self.record_thread = None - - def _finalize_recording(self): - self.isRecording = 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('Radiometric data saved: {self.raw_data_fpath}') - print('Viewable Video: {self.viewable_video_fpath}') - - def release_camera(self): - self.camera.Close() - print('Release Camera') - - -# TEST CODE -boson_capture = Boson_Capture() -boson_capture.start_recording() -sleep(10) -boson_capture.stop_recording() -boson_capture.release_camera() - From 5c67bb81e9387ee2170049aed70bc244c94e1da1 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Thu, 6 Nov 2025 12:50:48 -1000 Subject: [PATCH 19/30] Fix bad commit --- heatseek/boson_capture.py | 265 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 265 insertions(+) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index e69de29..cb5ddbc 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -0,0 +1,265 @@ +''' +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 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: + camera_id (int): used to select the appropriate USB port + 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, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): + '''Initializes boson camera interface + + Loads Boson SDK, connects to camera and configures radiometric + parameters + + Args: + camera_id (int, opt): USB port index for camera. Defaults to 0 + skdpath (str, opt): path to FLIR Boson SDK folder. + Defaults to ~/BosonSDK/SDK_USER_PERMISSIONS + ''' + super().__init__(camera_id) + + self.recording = False + self.height = 256 + self.width = 320 + self.raw_data_fpath = None + self.viewable_video_fpath = None + self.recording_thread = None + self.cap = None + self.raw_mm = None + self.frame_index = None + + # import Boson SDK + path = os.path.expanduser(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 + port = f'/dev/ttyACM{self.camera_id}' + self.camera = self.cam_api.pyClient(manualport=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='output.npy', norm='output.mp4'): + '''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 + self.viewable_video_fpath = norm + + 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 + self.cap = cv2.VideoCapture(self.camera_id, cv2.CAP_V4L2) + self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height) + self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width) + self.cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) + self.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, + 30, + (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 = self.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 self.frame_index > self.raw_mm.shape[0]: + print('Reached Preallocated Size, Stopping Recording') + break + self.raw_mm[self.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') + + +# TEST CODE +boson_capture = BosonCapture() +boson_capture.start_recording() +sleep(10) +boson_capture.stop_recording() +boson_capture.release_camera() From 136b21b728aaf77bf9a9e1b2e772aa192617d59e Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 10 Nov 2025 11:45:05 -1000 Subject: [PATCH 20/30] Fix recording frame rate and final pylint --- heatseek/boson_capture.py | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index cb5ddbc..aef6336 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -1,5 +1,4 @@ -''' -Implementation of Capture abstract class for FLIR Boson +'''Implementation of Capture abstract class for FLIR Boson Radiometric Thermal camera This module provides the 'Boson_Capture' class, which implements the standard @@ -18,7 +17,6 @@ from capture import Capture - class BosonCapture(Capture): '''Camera interface for FLIR Boson Radiometric Thermal camera @@ -59,9 +57,9 @@ def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): self.raw_data_fpath = None self.viewable_video_fpath = None self.recording_thread = None - self.cap = None +# self.cap = None self.raw_mm = None - self.frame_index = None +# self.frame_index = None # import Boson SDK path = os.path.expanduser(sdkpath) @@ -167,16 +165,16 @@ def _record_loop(self): ''' # video settings - self.cap = cv2.VideoCapture(self.camera_id, cv2.CAP_V4L2) - self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height) - self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width) - self.cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) - self.cap.set(cv2.CAP_PROP_FOURCC, + cap = cv2.VideoCapture(self.camera_id, 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, - 30, + 60, (self.width, self.height)) # raw video setup @@ -188,7 +186,7 @@ def _record_loop(self): try: while self.recording: - ret, frame = self.cap.read() + ret, frame = cap.read() if not ret: print('Frame grab failed. Stopping Recording') break @@ -201,10 +199,10 @@ def _record_loop(self): writer.write(frame_color) # Raw Video - if self.frame_index > self.raw_mm.shape[0]: + if frame_index > self.raw_mm.shape[0]: print('Reached Preallocated Size, Stopping Recording') break - self.raw_mm[self.frame_index] = frame + self.raw_mm[frame_index] = frame frame_index += 1 self.raw_mm.flush() From 448b7b53471385d007c2fb250113a1d83b545eee Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Mon, 10 Nov 2025 12:27:30 -1000 Subject: [PATCH 21/30] modified record.py to use boson_capture class --- examples/boson_record.py | 25 +++++++++++++++++++++++++ heatseek/boson_capture.py | 16 ++++------------ 2 files changed, 29 insertions(+), 12 deletions(-) create mode 100644 examples/boson_record.py diff --git a/examples/boson_record.py b/examples/boson_record.py new file mode 100644 index 0000000..acafa49 --- /dev/null +++ b/examples/boson_record.py @@ -0,0 +1,25 @@ +'''example file to record using flir boson''' + +from heatseek.boson_capture import BosonCapture +import argparse +from time import sleep + +if __name__=='__main__': + + parser = argparse.ArgumentParser() + parser.add_argument('--raw_output_path', + help='filepath to save raw radiometric output') + parser.add_argument('--video_output_path', + help='filepath to save normalized video output') + parser.add_argument('--bosonsdk_path', + help='filepath to boson sdk') + parser.add_argument('--recording_time', + help='time in s to record') + args = parser.parse_args() + + camera = BosonCapture() + camera.start_recording(raw=args.raw_output_path, + norm=args.video_output_path) + sleep(10) + camera.stop_recording() + camera.release_camera() diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index aef6336..3655440 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -14,7 +14,7 @@ from importlib import import_module import numpy as np import cv2 -from capture import Capture +from .capture import Capture class BosonCapture(Capture): @@ -132,7 +132,7 @@ def take_image(self): print('Taking Image- PLACEHOLDER') - def start_recording(self, raw='output.npy', norm='output.mp4'): + def start_recording(self, raw=None, norm=None): '''Begin thread for continuous recording Args: @@ -146,8 +146,8 @@ def start_recording(self, raw='output.npy', norm='output.mp4'): print('Recording in Progress!') return - self.raw_data_fpath = raw - self.viewable_video_fpath = norm + self.raw_data_fpath = raw or 'output.npy' + self.viewable_video_fpath = norm or 'output.mp4' self.recording = True self.recording_thread = threading.Thread( @@ -253,11 +253,3 @@ def release_camera(self): self.camera.Close() print('Release Camera') - - -# TEST CODE -boson_capture = BosonCapture() -boson_capture.start_recording() -sleep(10) -boson_capture.stop_recording() -boson_capture.release_camera() From 7cd69b6e894ccb727fe04dca3fca837f90b6e0eb Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Tue, 11 Nov 2025 12:17:21 -1000 Subject: [PATCH 22/30] Add doc strings --- examples/boson_record.py | 47 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/examples/boson_record.py b/examples/boson_record.py index acafa49..3126239 100644 --- a/examples/boson_record.py +++ b/examples/boson_record.py @@ -1,25 +1,64 @@ -'''example file to record using flir boson''' +'''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 \ + --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. +''' from heatseek.boson_capture import BosonCapture import argparse from time import sleep -if __name__=='__main__': + +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: + --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 + + Returns: + None + ''' parser = argparse.ArgumentParser() 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() camera.start_recording(raw=args.raw_output_path, norm=args.video_output_path) - sleep(10) + sleep(args.recording_time or 10) camera.stop_recording() camera.release_camera() + + +if __name__=='__main__': + main() From 40d2c592d6d5c2225755d6ac2873006072a876ba Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Tue, 11 Nov 2025 12:56:21 -1000 Subject: [PATCH 23/30] flake8/pylint boson_record.py --- examples/boson_record.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/examples/boson_record.py b/examples/boson_record.py index 3126239..500952f 100644 --- a/examples/boson_record.py +++ b/examples/boson_record.py @@ -1,35 +1,37 @@ '''Example script to record thermal imagery using FLIR Boson Radiometric camera. -This script demonstrates how to use the 'BosonCapture' class from +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: +Usage: python boson_record.py \ --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 +If no arguments are provided, default filenames are used and script records for 10s. -''' +''' -from heatseek.boson_capture import BosonCapture 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 + Parses arguments, initializes camera, starts recording for specified duration, then stops and releases camera. Command-line args: - --raw_output_path (str, opt): Filepath to save radiometric output (.npy) - --video_output_path (str, opt): Filepath to save normalized video (.mp4) + --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 @@ -51,7 +53,7 @@ def main(): type=int, help='time in s to record') args = parser.parse_args() - + camera = BosonCapture() camera.start_recording(raw=args.raw_output_path, norm=args.video_output_path) @@ -59,6 +61,6 @@ def main(): camera.stop_recording() camera.release_camera() - -if __name__=='__main__': + +if __name__ == '__main__': main() From 5fc1176bd5f437c5d5b19e203a378238a57386d6 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 12 Nov 2025 12:26:12 -1000 Subject: [PATCH 24/30] Fix relative import --- heatseek/boson_capture.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index 3655440..5c1ff10 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -14,7 +14,7 @@ from importlib import import_module import numpy as np import cv2 -from .capture import Capture +from heatseek.capture import Capture class BosonCapture(Capture): From 2444660c5c242ba276bb8ed665a1a2bc609cc700 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 12 Nov 2025 12:27:40 -1000 Subject: [PATCH 25/30] Remove boson scripts folder --- boson_scripts/README.md | 2 - boson_scripts/record.py | 158 ---------------------------------------- 2 files changed, 160 deletions(-) delete mode 100644 boson_scripts/README.md delete mode 100644 boson_scripts/record.py diff --git a/boson_scripts/README.md b/boson_scripts/README.md deleted file mode 100644 index 226ff60..0000000 --- a/boson_scripts/README.md +++ /dev/null @@ -1,2 +0,0 @@ -Download Boson 3.0 IDD & SDK: https://oem.flir.com/products/boson/?vertical=lwir&segment=oem -Modifications to make compatible with Raspberry Pi ARM architecture: https://flir.custhelp.com/app/answers/detail/a_id/6471/~/compiling-the-boson-sdk-for-arm64-linux-platforms diff --git a/boson_scripts/record.py b/boson_scripts/record.py deleted file mode 100644 index ef9a095..0000000 --- a/boson_scripts/record.py +++ /dev/null @@ -1,158 +0,0 @@ -''' -This script records radiometric data from the FLIR Boson thermal camera - -Usage: - $ python record.py \ - --raw_output_path path/to/save/raw/data.npy \ - --video_output_path path/to/save/video/data.mp4 \ - --bosonsdk_path path/to/boson/sdk/files -''' - - -from time import sleep -import argparse -import sys -import os -import numpy as np -import cv2 - - -# setup camera parameters -def setup(func, *args, delay=2, success_code=0, description=''): - '''run camera setup fucntion 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 - - Returns: - None - ''' - - print(f'\nSetting {description}...') - sleep(.5) - result = func(*args) - while result != success_code: - sleep(delay) - result = func(*args) - print('Success!') - - -def get_center_temp(frame_16bit): - '''get frame center temp in C and F - - Args: - frame_16bit: raw radiometric 16bit frame - - Returns: - center_c: center temperature in Celsius - center_f: center temperature in Fahrenheit - ''' - - center_raw = frame_16bit[int(frame_16bit.shape[0]/2), - int(frame_16bit.shape[1]/2)] - center_c = round((center_raw/100) - 273, 1) - center_f = round(center_c * 9/5 + 32, 1) - - return center_c, center_f - - -if __name__ == "__main__": - - # User set parameters - parser = argparse.ArgumentParser() - parser.add_argument('--raw_output_path', - default='output.npy', - help='filepath to save raw radiometric output') - parser.add_argument('--video_output_path', - default='output.mp4', - help='filepath to save normalized video output') - parser.add_argument('--bosonsdk_path', - default='~/BosonSDK/SDK_USER_PERMISSIONS', - help='filepath for Boson SDK') - args = parser.parse_args() - - # import Boson SDK - sdk_path = os.path.expanduser(args.bosonsdk_path) - sys.path.append(sdk_path) - from SDK_USER_PERMISSIONS import * - - # Connect to Camera - myCam = CamAPI.pyClient(manualport="/dev/ttyACM0") - - # Set Radiometric Parameters - setup(myCam.bosonSetGainMode, - FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN, - description='High Gain Mode') - setup(myCam.TLinearSetControl, - FLR_ENABLE_E.FLR_ENABLE, - description='TLinear Mode') - setup(myCam.sysctrlSetUsbVideoIR16Mode, - FLR_SYSCTRL_USBIR16_MODE_E.FLR_SYSCTRL_USBIR16_MODE_TLINEAR, - description='IR16 Mode') - setup(myCam.radiometrySetTransmissionWindow, - 100, - description='Window Transimission') - setup(myCam.TLinearRefreshLUT, - FLR_BOSON_GAINMODE_E.FLR_BOSON_HIGH_GAIN, - description='LUT Refresh') - setup(myCam.bosonRunFFC, - description='Flat Field Correction') - - # open camera - cap = cv2.VideoCapture(0, cv2.CAP_V4L2) - cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 256) - cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320) - cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) # don't auto-convert to RGB - cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y', '1', '6', ' ')) - - # output video settings - fourcc = cv2.VideoWriter_fourcc(*'mp4v') - writer = cv2.VideoWriter(args.video_output_path, fourcc, 30, (320, 256)) - headless = not os.environ.get("DISPLAY") - raw_frames = [] - print('\n\nRecording Started...') - - try: - while True: - ret, frame = cap.read() - if not ret: - print('frame grab failed, stopping') - break - - raw_frames.append(frame) - - # Viewable Window - min_val, max_val = np.min(frame), np.max(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) - writer.write(frame_color) - if not headless: - cv2.imshow('color', frame_color) - cv2.waitKey(1) - - # Display Center Temp - temp_c, temp_f = get_center_temp(frame) - sys.stdout.write(f'\rCenter temp: Temp C - {temp_c} | Temp F - {temp_f}') - sys.stdout.flush() - - except KeyboardInterrupt: - print('\n\nRecording Stopped...') - - finally: - raw_frames = np.stack(raw_frames, axis=0) - np.save(args.raw_output_path, raw_frames) - - print('\nRecordings Saved') - print(f'Radiometric data: {args.raw_output_path}') - print(f'Normalized video: {args.video_output_path}\n\n') - - cap.release() - writer.release() - if not headless: - cv2.destroyAllWindows() - - myCam.Close() From 85058a6915fa8b6ffde93b361bd7c90496e611e9 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Thu, 13 Nov 2025 08:39:05 -1000 Subject: [PATCH 26/30] Update port initialization --- examples/boson_record.py | 10 +++++++++- heatseek/boson_capture.py | 15 +++++++-------- heatseek/capture.py | 3 +-- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/examples/boson_record.py b/examples/boson_record.py index 500952f..6c8c7d9 100644 --- a/examples/boson_record.py +++ b/examples/boson_record.py @@ -40,6 +40,12 @@ def main(): ''' 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') @@ -54,7 +60,9 @@ def main(): help='time in s to record') args = parser.parse_args() - camera = BosonCapture() + 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) diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index 5c1ff10..4f4f97b 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -38,7 +38,7 @@ class BosonCapture(Capture): width (int): frame pixel width ''' - def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): + def __init__(self, serial_port=None, video_port=None, sdkpath=None): '''Initializes boson camera interface Loads Boson SDK, connects to camera and configures radiometric @@ -49,20 +49,20 @@ def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): skdpath (str, opt): path to FLIR Boson SDK folder. Defaults to ~/BosonSDK/SDK_USER_PERMISSIONS ''' - super().__init__(camera_id) + 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.cap = None self.raw_mm = None -# self.frame_index = None # import Boson SDK - path = os.path.expanduser(sdkpath) + 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( @@ -73,8 +73,7 @@ def __init__(self, camera_id=0, sdkpath='~/BosonSDK/SDK_USER_PERMISSIONS'): ) # create camera object & configure radiometric parameters - port = f'/dev/ttyACM{self.camera_id}' - self.camera = self.cam_api.pyClient(manualport=port) + self.camera = self.cam_api.pyClient(manualport=self.serial_port) self.setup() def setup(self): @@ -165,7 +164,7 @@ def _record_loop(self): ''' # video settings - cap = cv2.VideoCapture(self.camera_id, cv2.CAP_V4L2) + 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) diff --git a/heatseek/capture.py b/heatseek/capture.py index 19fa73b..2375128 100644 --- a/heatseek/capture.py +++ b/heatseek/capture.py @@ -17,7 +17,7 @@ class Capture(ABC): managing recordings across different camera types. ''' - def __init__(self, camera_id=0): + def __init__(self): '''Initialize capture interface This method should handle hardware initalization and @@ -27,7 +27,6 @@ def __init__(self, camera_id=0): camera_id (int, opt): Device identifier for camera. Defaults to 0. ''' - self.camera_id = camera_id @abstractmethod def setup(self): From 30cf6231128178f68acd4d70e39477a77728e829 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Thu, 13 Nov 2025 08:48:08 -1000 Subject: [PATCH 27/30] Update doc strings --- examples/boson_record.py | 9 ++++++++- heatseek/boson_capture.py | 8 ++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/examples/boson_record.py b/examples/boson_record.py index 6c8c7d9..9b54e90 100644 --- a/examples/boson_record.py +++ b/examples/boson_record.py @@ -7,6 +7,8 @@ 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 \ @@ -28,12 +30,17 @@ def main(): 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 + --recording_time (int, opt): Duration of recording in seconds. + Default is 10s. Returns: None diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index 4f4f97b..70944fd 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -24,7 +24,8 @@ class BosonCapture(Capture): camera. Attributes: - camera_id (int): used to select the appropriate USB port + 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 @@ -45,7 +46,10 @@ def __init__(self, serial_port=None, video_port=None, sdkpath=None): parameters Args: - camera_id (int, opt): USB port index for camera. Defaults to 0 + 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 ''' From 404e57827cbd825bd07dc94547127809f7707208 Mon Sep 17 00:00:00 2001 From: Sumega Mandadi Date: Wed, 19 Nov 2025 11:35:21 -1000 Subject: [PATCH 28/30] Fix quote type in doc strings --- examples/boson_record.py | 8 +++---- heatseek/boson_capture.py | 44 +++++++++++++++++++-------------------- heatseek/capture.py | 32 ++++++++++++++-------------- 3 files changed, 42 insertions(+), 42 deletions(-) diff --git a/examples/boson_record.py b/examples/boson_record.py index 9b54e90..9e1f4b5 100644 --- a/examples/boson_record.py +++ b/examples/boson_record.py @@ -1,4 +1,4 @@ -'''Example script to record thermal imagery using FLIR Boson +"""Example script to record thermal imagery using FLIR Boson Radiometric camera. This script demonstrates how to use the 'BosonCapture' class from @@ -16,7 +16,7 @@ If no arguments are provided, default filenames are used and script records for 10s. -''' +""" import argparse from time import sleep @@ -24,7 +24,7 @@ def main(): - '''Run exmaple thermal recording using BosonCapture + """Run exmaple thermal recording using BosonCapture Parses arguments, initializes camera, starts recording for specified duration, then stops and releases camera. @@ -44,7 +44,7 @@ def main(): Returns: None - ''' + """ parser = argparse.ArgumentParser() parser.add_argument('--serial_port', diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index 70944fd..b959fe5 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -1,11 +1,11 @@ -'''Implementation of Capture abstract class for FLIR Boson +"""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 @@ -18,7 +18,7 @@ class BosonCapture(Capture): - '''Camera interface for FLIR Boson Radiometric Thermal camera + """Camera interface for FLIR Boson Radiometric Thermal camera Implements 'Capture' abstract base class for Boson Radiometric thermal camera. @@ -37,10 +37,10 @@ class BosonCapture(Capture): 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 + """Initializes boson camera interface Loads Boson SDK, connects to camera and configures radiometric parameters @@ -52,7 +52,7 @@ def __init__(self, serial_port=None, video_port=None, sdkpath=None): 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' @@ -81,11 +81,11 @@ def __init__(self, serial_port=None, video_port=None, sdkpath=None): self.setup() def setup(self): - '''Configure radiometric parameters for capture + """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...') @@ -109,7 +109,7 @@ def setup(self): description='Flat Field Correction') def _configure(self, func, *args, delay=1, success_code=0, description=''): - '''run camera paramter configuration function until it succeeds + """run camera paramter configuration function until it succeeds Args: func: function to update camera parameter @@ -117,7 +117,7 @@ def _configure(self, func, *args, delay=1, success_code=0, description=''): 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) @@ -128,22 +128,22 @@ def _configure(self, func, *args, delay=1, success_code=0, description=''): print('Success!') def take_image(self): - '''Capture single image frame from camera. + """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 + """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!') @@ -160,12 +160,12 @@ def start_recording(self, raw=None, norm=None): print('Staring Recording...') def _record_loop(self): - '''Internal method: recording loop to continuously capture frames + """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) @@ -213,11 +213,11 @@ def _record_loop(self): self._finalize_recording() def stop_recording(self): - '''Stop ongoing recording session. + """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') @@ -231,11 +231,11 @@ def stop_recording(self): self.recording_thread = None def _finalize_recording(self): - '''Finalize recording. + """Finalize recording. Flushes and deletes memory-mapped raw frame array, closes video writer, and prints file locations. - ''' + """ self.recording = False @@ -249,10 +249,10 @@ def _finalize_recording(self): print(f'Viewable Video: {self.viewable_video_fpath}') def release_camera(self): - '''Release camera resources. + """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 index 2375128..de53e7d 100644 --- a/heatseek/capture.py +++ b/heatseek/capture.py @@ -1,24 +1,24 @@ -''' +""" 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 + """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 + """Initialize capture interface This method should handle hardware initalization and SDK loading @@ -26,45 +26,45 @@ def __init__(self): Args: camera_id (int, opt): Device identifier for camera. Defaults to 0. - ''' + """ @abstractmethod def setup(self): - '''Prepare camera for capture + """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 + """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 + """Begin video recording This method should capture frames continuously and save to disk - ''' + """ @abstractmethod def stop_recording(self): - '''Stop ongoing recording + """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 + """Release camera hardware and close This method should handle any releasing and/or shutting down of camera hardware - ''' + """ From 29a7d6a4191b4cbe6932f7c6e11d61474978a9d3 Mon Sep 17 00:00:00 2001 From: Akhil Nallacheruvu Date: Tue, 12 May 2026 10:36:58 -0700 Subject: [PATCH 29/30] Implementing Counting Methods --- .gitignore | 163 ++-- LICENSE | 42 +- README.md | 358 ++++----- environment.yml | 22 +- heatseek/cli.py | 252 +++--- heatseek/config/preproc_config.yaml | 18 +- heatseek/counting/README.md | 3 + heatseek/counting/__init__.py | 0 heatseek/counting/bg_sub/CountLine.py | 73 ++ heatseek/counting/bg_sub/README.md | 30 + heatseek/counting/bg_sub/__init__.py | 0 heatseek/counting/bg_sub/bat_functions.py | 384 +++++++++ heatseek/counting/bg_sub/crossing_tracks.py | 102 +++ .../counting/bg_sub/detections_to_tracks.py | 222 ++++++ heatseek/counting/bg_sub/koger_tracking.py | 744 ++++++++++++++++++ heatseek/counting/bg_sub/video_inference.py | 100 +++ heatseek/counting/dl/BatIterableDataset.py | 153 ++++ heatseek/counting/dl/CountLine.py | 72 ++ heatseek/counting/dl/README.md | 79 ++ heatseek/counting/dl/SegmentationDataset.py | 30 + heatseek/counting/dl/__init__.py | 0 heatseek/counting/dl/augmentations.py | 181 +++++ heatseek/counting/dl/bat_functions.py | 384 +++++++++ heatseek/counting/dl/bat_seg_models.py | 359 +++++++++ heatseek/counting/dl/crop_vids.py | 94 +++ heatseek/counting/dl/crossing_tracks.py | 143 ++++ heatseek/counting/dl/dataloaders.py | 16 + heatseek/counting/dl/detections_to_tracks.py | 222 ++++++ heatseek/counting/dl/extract_frames.py | 55 ++ heatseek/counting/dl/koger_tracking.py | 744 ++++++++++++++++++ heatseek/counting/dl/optimizer.py | 108 +++ heatseek/counting/dl/segment_frames.py | 88 +++ heatseek/counting/dl/trainer.py | 248 ++++++ heatseek/counting/dl/transforms.py | 117 +++ heatseek/counting/dl/video_inference.py | 180 +++++ heatseek/data_utils.py | 64 +- heatseek/density_annotator.py | 102 +-- heatseek/detect_track.py | 76 +- heatseek/preprocess.py | 170 ++-- heatseek/train.py | 80 +- ideas.txt | 18 +- requirements.txt | 22 +- setup.py | 44 +- 43 files changed, 5655 insertions(+), 707 deletions(-) create mode 100644 heatseek/counting/README.md create mode 100644 heatseek/counting/__init__.py create mode 100644 heatseek/counting/bg_sub/CountLine.py create mode 100644 heatseek/counting/bg_sub/README.md create mode 100644 heatseek/counting/bg_sub/__init__.py create mode 100644 heatseek/counting/bg_sub/bat_functions.py create mode 100644 heatseek/counting/bg_sub/crossing_tracks.py create mode 100644 heatseek/counting/bg_sub/detections_to_tracks.py create mode 100644 heatseek/counting/bg_sub/koger_tracking.py create mode 100644 heatseek/counting/bg_sub/video_inference.py create mode 100644 heatseek/counting/dl/BatIterableDataset.py create mode 100644 heatseek/counting/dl/CountLine.py create mode 100644 heatseek/counting/dl/README.md create mode 100644 heatseek/counting/dl/SegmentationDataset.py create mode 100644 heatseek/counting/dl/__init__.py create mode 100644 heatseek/counting/dl/augmentations.py create mode 100644 heatseek/counting/dl/bat_functions.py create mode 100644 heatseek/counting/dl/bat_seg_models.py create mode 100644 heatseek/counting/dl/crop_vids.py create mode 100644 heatseek/counting/dl/crossing_tracks.py create mode 100644 heatseek/counting/dl/dataloaders.py create mode 100644 heatseek/counting/dl/detections_to_tracks.py create mode 100644 heatseek/counting/dl/extract_frames.py create mode 100644 heatseek/counting/dl/koger_tracking.py create mode 100644 heatseek/counting/dl/optimizer.py create mode 100644 heatseek/counting/dl/segment_frames.py create mode 100644 heatseek/counting/dl/trainer.py create mode 100644 heatseek/counting/dl/transforms.py create mode 100644 heatseek/counting/dl/video_inference.py diff --git a/.gitignore b/.gitignore index e8f12a4..1488530 100644 --- a/.gitignore +++ b/.gitignore @@ -1,73 +1,90 @@ -# 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/bg_sub/detections_413_full_thermal_vid +heatseek/counting/bg_sub/detections_414_full_thermal_vid +heatseek/counting/dl/detections_413_full_train +heatseek/counting/dl/detections_414_full_train +heatseek/counting/dl/detections_413_single_train +heatseek/counting/dl/detections_414_single_train +heatseek/counting/dl/models_full_train +heatseek/counting/dl/models_single_train 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. - -

- Optical Flow - PB Example -

- -## 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 - -

- Density Annotator -

- - -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. + +

+ Optical Flow + PB Example +

+ +## 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 + +

+ Density Annotator +

+ + +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/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/README.md b/heatseek/counting/README.md new file mode 100644 index 0000000..a9af25c --- /dev/null +++ b/heatseek/counting/README.md @@ -0,0 +1,3 @@ +# Bat Track Counting + +For more information about counting bat tracks using deep learning, navigate to the dl directory. For more information about counting bat tracks using temperature-based background segmentation on a video taken on a FLIR Boson radiometric thermal camera with globally normalized pixel values, refer to the bg_sub directory. \ 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/bg_sub/CountLine.py b/heatseek/counting/bg_sub/CountLine.py new file mode 100644 index 0000000..fefc9db --- /dev/null +++ b/heatseek/counting/bg_sub/CountLine.py @@ -0,0 +1,73 @@ +import numpy as np + +class CountLine(): + ''' + Counts everytime a bat crosses a line defined in space (y = 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 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']) + # 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']) + 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/bg_sub/README.md b/heatseek/counting/bg_sub/README.md new file mode 100644 index 0000000..578e55b --- /dev/null +++ b/heatseek/counting/bg_sub/README.md @@ -0,0 +1,30 @@ +# Bat Track Counting Pipeline (Temperature-based Background Subtraction) +## Overview: +This file contains information related to generating crossing tracks using temperature-based background subtraction. This method skips the model training and relies on the pixel values from a radiometric camera image and their converted temperature value.The camera is assumed to be a FLIR Boson, and therefore, the raw temperature values are assumed to be encoded in centikelvins. Given this data, inference can directly be performed on the video to generate detection centroids, contours, and boundaries of bat blobs. From there, it can generate all the information related to the bat tracks. + +## 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 --vid_path path/to/inference/video --output_folder desired/path/of/output/folder + ``` + +2) Convert the detection data into raw tracks using `detections_to_tracks.py` + ```bash + cd path/to/working/dir + python -m detections_to_tracks --output_folder path/to/folder/with/detections + #should ideally be the same as output_folder in video_inference + ``` + +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 --raw_tracks_file path/to/raw/tracks/file --crossing_tracks_file desired/path/of/crossing_tracks/file + ``` + + An additional parameter must be added at the end of the second command. If the user wishes to generate a file containing tracks vertically crossing a horizontal midline and its corresponding visualization, they must add the additional parameter `--count_out`. If they instead wish to generate a file containing tracks horizontally crossing a vertical midline, they should use the parameter `--count_across` instead. + + + + + diff --git a/heatseek/counting/bg_sub/__init__.py b/heatseek/counting/bg_sub/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/heatseek/counting/bg_sub/bat_functions.py b/heatseek/counting/bg_sub/bat_functions.py new file mode 100644 index 0000000..38bb6b4 --- /dev/null +++ b/heatseek/counting/bg_sub/bat_functions.py @@ -0,0 +1,384 @@ +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 CountLine import CountLine +import 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): + """ 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 = [] + + max_distance_threshold = 30 + max_distance_threshold_noise = 30 + min_distance_threshold = 0 + max_unseen_time = 2 + min_new_track_distance = 3 + min_distance_big = 30 + +# #Create initial tracks based on the objects in the first frame +# raw_track_list = add_all_points_as_new_tracks( +# raw_track_list, positions[0], contours_list[0], sizes_list0, noise=0 +# ) + + #try to connect points to the next frame + if max_frame is None: + max_frame = len(positions) + + 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)) + +# for active_ind, track_ind in enumerate(active_list): +# if active_ind in row_ind: +# row_count = np.where(row_ind == active_ind)[0] +# raw_track_list[track_ind]['debug'].append( +# '{} dist {}, best {}'.format( +# frame_ind, +# distance[row_ind[row_count], +# col_ind[row_count]], +# np.min(distance[row_ind[row_count], +# :]) +# ) +# ) +# best_col = np.argmin(distance[row_ind[row_count], +# :]) +# row_count = np.where(col_ind == best_col)[0] +# raw_track_list[track_ind]['debug'].append( +# '{} row_ind {} col {} dist {} track {}'.format( +# frame_ind, row_ind[row_count], +# col_ind[row_count], +# distance[row_ind[row_count], +# col_ind[row_count]], +# active_list[row_ind[row_count][0]]) +# ) + + + # 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) + # result is 1 if forward crossing -1 is backward crossing + if count_out: + if out_frame_num: + crossing_track_list[-1]['crossed'] = out_frame_num * out_result + else: + crossing_track_list[-1]['crossed'] = 0 + if count_across: + if across_frame_num: + crossing_track_list[-1]['across_crossed'] = across_frame_num * across_result + else: + crossing_track_list[-1]['across_crossed'] = 0 + 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/bg_sub/crossing_tracks.py b/heatseek/counting/bg_sub/crossing_tracks.py new file mode 100644 index 0000000..b04ef90 --- /dev/null +++ b/heatseek/counting/bg_sub/crossing_tracks.py @@ -0,0 +1,102 @@ +from bat_functions import threshold_short_tracks, measure_crossing_bats +import numpy as np +import matplotlib.pyplot as plt +import argparse + +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 + + crossed_key = 'crossed' if count_out else 'across_crossed' + + 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()}') + print(f'Forward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] > 0)}') + print(f'Backward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] < 0)}') + + 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] + color = 'blue' if track[crossed_key] > 0 else 'red' + 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() + direction_labels = 'blue=downward, red=upward' + else: + ax.axvline(x=midline_x, color='green', linewidth=2, linestyle='--', label='midline') + direction_labels = 'blue=rightward, red=leftward' + + ax.set_title(f'Sample of {num_tracks} crossing tracks\n{direction_labels}') + ax.legend() + + ax = axes[1] + crossing_frames = [abs(t[crossed_key]) for t in crossing_tracks] + ax.hist(crossing_frames, bins=100) + ax.set_xlabel('Frame number') + ax.set_ylabel('Number of crossings') + ax.set_title('Bat crossings over time') + + print('Crossing frames:') + for t in crossing_tracks: + direction = 'forward' if t[crossed_key] > 0 else 'backward' + print(f' frame {abs(t[crossed_key])}: {direction}') + + plt.tight_layout() + plt.show() + +def main(): + parser = argparse.ArgumentParser(description='Get bat crossing tracks across the video') + parser.add_argument('--raw_tracks_file', type=str, help='Path to raw tracks file') + parser.add_argument('--crossing_tracks_file', type=str, help='Intended path for the crossing tracks 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() + + save_crossing_tracks_from_raw_tracks(args.raw_tracks_file, args.crossing_tracks_file, args.count_across, args.count_out) + visualize_crossing_tracks(args.crossing_tracks_file, count_out=args.count_out, count_across=args.count_across) + +if __name__ == '__main__': + main() diff --git a/heatseek/counting/bg_sub/detections_to_tracks.py b/heatseek/counting/bg_sub/detections_to_tracks.py new file mode 100644 index 0000000..cc80404 --- /dev/null +++ b/heatseek/counting/bg_sub/detections_to_tracks.py @@ -0,0 +1,222 @@ +import numpy as np +import os +import glob +import bat_functions as kbf +from multiprocessing import Pool +import argparse + +def track(camera_dict): + """ + 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). + + 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) + else: + print("Missing contour files.") + +def build_camera_dicts(output_folder, num_groups=10, fps=60, overlap_seconds=15): + """ + 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 15. + + 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) + + 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, processes=5): + """ + 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) + print(f'Tracking {len(camera_dicts)} groups') + with Pool(processes=processes) as pool: + pool.map(track, camera_dicts) + +def main(): + parser = argparse.ArgumentParser(description='Get centers and contours from detections from model inference') + parser.add_argument('--output_folder', type=str, help='Path of output folder containing the centers, contours, etc from video inference') + args = parser.parse_args() + run_tracking(args.output_folder) + combine_tracks(args.output_folder) + +if __name__ == '__main__': + main() diff --git a/heatseek/counting/bg_sub/koger_tracking.py b/heatseek/counting/bg_sub/koger_tracking.py new file mode 100644 index 0000000..cba6251 --- /dev/null +++ b/heatseek/counting/bg_sub/koger_tracking.py @@ -0,0 +1,744 @@ +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['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 + 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 fewwer new points than existing tracks. When it is the + later, 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 + +#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 = 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, + 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: upper distance threshold + max_distance_noise: upper distance threshold tracks with noise values + above 0 + min_distance: lowwer 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 + 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) + + 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)] = max_distance + # 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]) + 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]) + + 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/bg_sub/video_inference.py b/heatseek/counting/bg_sub/video_inference.py new file mode 100644 index 0000000..9b51da2 --- /dev/null +++ b/heatseek/counting/bg_sub/video_inference.py @@ -0,0 +1,100 @@ +import os +import cv2 +import logging +import numpy as np +import argparse +import bat_functions + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +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.') + + +def get_mask(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 + +def run_inference(video_file, min_val=28000, max_val=32000, + threshold=5.0, process_every_n_frames=1): + """Process every frame, storing only detection metadata — not images.""" + centers_list = [] + contours_list = [] + sizes_list = [] + rects_list = [] + + for record in iter_frames(video_file, process_every_n_frames): + mask = get_mask(record['image'], min_val, max_val, threshold) + 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) + # record['image'] goes out of scope here and gets GC'd + + if len(centers_list) % 1000 == 0: + logging.info(f'Processed {len(centers_list)} frames.') + + logging.info(f'Processed {len(centers_list)} frames.') + return centers_list, contours_list, sizes_list, rects_list + +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)) + print(f'Saved detections for {len(centers_list)} frames to {output_folder}') + +def main(): + parser = argparse.ArgumentParser(description='Process a thermal video and detect objects.') + parser.add_argument('--vid_path', help='Path to the input video file.') + parser.add_argument('--output_folder', help='Path to the output folder.') + args = parser.parse_args() + + centers_list, contours_list, sizes_list, rects_list = run_inference(args.vid_path) + save_detections(centers_list, contours_list, sizes_list, rects_list, args.output_folder) + +if __name__ == '__main__': + main() diff --git a/heatseek/counting/dl/BatIterableDataset.py b/heatseek/counting/dl/BatIterableDataset.py new file mode 100644 index 0000000..68e11f3 --- /dev/null +++ b/heatseek/counting/dl/BatIterableDataset.py @@ -0,0 +1,153 @@ +from torch.utils.data import IterableDataset +import cv2 + +class BatIterableDataset(IterableDataset): + """ An iterable PyTorch dataset that streams grayscale frames from one or + more video files for bat detection inference. + + Frames are read sequentially using OpenCV, converted to grayscale, and + optionally passed through an augmentation pipeline. Videos are processed + one at a time in the order provided; when a video is exhausted, the next + one starts automatically. + + Consecutive read failures are tolerated up to max_bad_reads before the + current video is closed and iteration ends. + + Attributes: + video_files (list[str]): Ordered list of video file paths. + augmentor (MaskCompose or None): Optional transform applied to each frame. + max_bad_reads (int): Max consecutive failed reads before closing a video. + total_frames_read (int): Running count of successfully read frames. + total_bad_reads (int): Running count of failed frame reads. + video_number (int): Index of the currently open video in video_files. + more_frames (bool): Whether the current video still has frames. + """ + + def __init__(self, video_files, augmentor=None, max_bad_reads=300): + """ + Initialize the dataset and open the first video file. + + Parameters: + video_files (list[str]): Ordered list of paths to video files. + The first file is opened immediately on construction. + augmentor (MaskCompose, optional): Transform pipeline applied to + each frame dict before yielding. Defaults to None. + max_bad_reads (int): Maximum consecutive failed cv2 reads before + the current video is considered exhausted. Defaults to 300. + + Raises: + AssertionError: If the first video file cannot be opened by OpenCV. + """ + self.vid_cap = cv2.VideoCapture(video_files[0]) + self.video_files = video_files + assert self.vid_cap.isOpened() + self.more_frames = True + self.max_bad_reads = max_bad_reads + self.total_frames_read = 0 + self.total_bad_reads = 0 + self.augmentor = augmentor + self.video_number = 0 + + def more_videos(self): + """ + Check whether there are additional unprocessed videos in the queue. + + Returns: + bool: True if video_number is within bounds of video_files. + """ + return self.video_number < len(self.video_files) + + def start_next_video(self): + """ + Release the current video capture and open the next video in the queue. + + Increments video_number, releases any open capture, and opens a new + cv2.VideoCapture for the next file. Does nothing if all videos have + already been processed. Prints a status message and frame read info + when a new video starts. + """ + if self.vid_cap.isOpened(): + self.vid_cap.release() + self.video_number += 1 + if self.video_number < len(self.video_files): + print('starting new video') + print(self.get_read_frame_info()) + self.vid_cap = cv2.VideoCapture(self.video_files[self.video_number]) + + def video_generator(self): + """ + Generator that yields preprocessed frames across all video files. + + For each video, reads frames sequentially. On a successful read the + frame is converted from BGR to grayscale, cropped by 2 pixels on each + edge, wrapped in a dict as {'image': frame}, and optionally transformed + by the augmentor before being yielded. + + Consecutive failed reads are counted; if max_bad_reads is reached the + video is released and iteration moves on. Transitions between videos + are handled automatically via start_next_video(). + + Yields: + dict or augmentor output: {'image': np.ndarray} if no augmentor, + otherwise the transformed output of augmentor({'image': frame}). + """ + + while(self.vid_cap.isOpened() or self.more_videos()): + if not self.vid_cap.isOpened(): + self.start_next_video() + good_read = False + num_bad_reads = 0 + while (not good_read and (num_bad_reads < self.max_bad_reads)): + grabbed, frame = self.vid_cap.read() + if grabbed: + good_read = True + self.total_frames_read += 1 + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + frame = {'image': frame[2:-2, 2:-2]} + if self.augmentor: + frame = self.augmentor(frame) + yield frame + else: + num_bad_reads += 1 + self.total_bad_reads += 1 + if not good_read: + self.vid_cap.release() + print("video capture closed") + + def __iter__(self): + """ + Return the frame generator, making this dataset compatible with + PyTorch DataLoader. + + Returns: + generator: The video_generator() iterator. + """ + return self.video_generator() + + def __del__(self): + """ + Release the OpenCV video capture on object destruction to free + file handles and decoder resources. + """ + if self.vid_cap.isOpened(): + self.vid_cap.release() + + def is_more_frames(self): + """ + Check whether the current video capture is still open and readable. + + Returns: + bool: True if vid_cap is currently open. + """ + return self.vid_cap.isOpened() + + def get_read_frame_info(self): + """ + Print a summary of total frames read and total bad reads so far. + + Example output: + 142 frames have been read with 3 bad reads + """ + print('{} frames have been read with {} bad reads'.format( + self.total_frames_read, self.total_bad_reads)) + \ No newline at end of file diff --git a/heatseek/counting/dl/CountLine.py b/heatseek/counting/dl/CountLine.py new file mode 100644 index 0000000..b66ee9d --- /dev/null +++ b/heatseek/counting/dl/CountLine.py @@ -0,0 +1,72 @@ +import numpy as np + +class CountLine(): + ''' + Counts everytime a bat crosses a line defined in space (y = 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 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']) + # 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']) + 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/dl/README.md b/heatseek/counting/dl/README.md new file mode 100644 index 0000000..70583a3 --- /dev/null +++ b/heatseek/counting/dl/README.md @@ -0,0 +1,79 @@ +# Bat Track Counting Pipeline (Deep Learning) +## Overview: +This file contains information related to generating crossing tracks using deep learning. This is the method that was originally implemented by Koger et al, but modified to fit our platform and data. Deep learning is useful when the data contains noisy objects (e.g. small bats in an RGB image). + +## Pipeline: +1) In order to train a model, we must create a training set consisting of images and segmentation masks as the model that is used in this pipeline is a UNet model. The steps for this stage are: + + 1) `crop_vids.py` - Cut the long videos into short clips of motion using the command: + ```bash + cd path/to/working/dir + python -m crop_vids --input path/to/input/file --output_dir path/to/desired/output/folder + ``` + There are additional parameters that may be set. Please refer to file for fine-tuning. + + The clips must go through human validation to ensure there are no false positives for motion. Once this is done, they must be arranged in a directory structure that looks like this: + + ``` + clips + |---- vid1_clips + |---------------- clip1.mp4 + |---------------- clip2.mp4 + . + . + . + |----- vid2_clips + |---------------- clip1.mp4 + |---------------- clip2.mp4 + . + . + . + . + . + . + + 2) `extract_frames.py` - Extract frames from the clips + ```bash + cd path/to/working/dir + python -m extract_frames --clips_dir path/to/directory/containing/clips + ``` + There are additional parameters that may be set. Please refer to file for fine-tuning. + + 3) `segment_frames.py` - Get masks for the frames using temperature-based background segmentation (assumes availability of globally normalized video with minimum and maximum pixel values corresponding to temperature from a radiometric thermal camera) + ```bash + cd path/to/working/dir + python -m segment_frames --clips_dir path/to/directory/containing/clips + ``` + There are additional parameters that may be set. Please refer to file for fine-tuning. + +2) Train the model on the training img, mask pairs using `trainer.py` + ```bash + cd path/to/working/dir + python -m trainer --clips_dir path/to/directory/containing/clips + ``` + There are additional parameters that may be set. Please refer to file for fine-tuning. + +3) 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 --vid_path path/to/inference/video --output_folder desired/path/of/output/folder --model_filepath path/to/model.tar --clips_dir path/to/directory/containing/clips + ``` + +4) Convert the detection data into raw tracks using `detections_to_tracks.py` + ```bash + cd path/to/working/dir + python -m detections_to_tracks --output_folder path/to/folder/with/detections + #should ideally be the same as output_folder in video_inference + ``` + +5) 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 --raw_tracks_file path/to/raw/tracks/file --crossing_tracks_file desired/path/of/crossing_tracks/file + ``` + + An additional parameter must be added at the end of the second command. If the user wishes to generate a file containing tracks vertically crossing a horizontal midline and its corresponding visualization, they must add the additional parameter `--count_out`. If they instead wish to generate a file containing tracks horizontally crossing a vertical midline, they should use the parameter `--count_across` instead. + + + + \ No newline at end of file diff --git a/heatseek/counting/dl/SegmentationDataset.py b/heatseek/counting/dl/SegmentationDataset.py new file mode 100644 index 0000000..63179e2 --- /dev/null +++ b/heatseek/counting/dl/SegmentationDataset.py @@ -0,0 +1,30 @@ +import torch.utils.data as data +import cv2 + +class SegmentationDataset(data.Dataset): + def __init__(self, pairs, transform=None, keep_orig=False): + """ + Args: + pairs (list of tuples): list of (frame_path, mask_path) pairs + transform (callable, optional): Optional transforms to be applied on a sample + keep_orig (bool): keep un-transformed image + """ + self.img_files = [p[0] for p in pairs] + self.mask_files = [p[1] for p in pairs] + self.keep_orig = keep_orig + self.transform = transform + + def __getitem__(self, index): + img_path = self.img_files[index] + mask_path = self.mask_files[index] + image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + sample = {'image': image[2:-2, 2:-2], 'mask': mask[2:-2, 2:-2]} + if self.keep_orig: + sample['orig'] = image + if self.transform: + sample = self.transform(sample) + return sample + + def __len__(self): + return len(self.img_files) diff --git a/heatseek/counting/dl/__init__.py b/heatseek/counting/dl/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/heatseek/counting/dl/augmentations.py b/heatseek/counting/dl/augmentations.py new file mode 100644 index 0000000..30eab98 --- /dev/null +++ b/heatseek/counting/dl/augmentations.py @@ -0,0 +1,181 @@ +import glob +import os +import cv2 +import numpy as np +import torch +import torchvision.transforms.functional as TF +import random + +def compute_mean_std(clips_dir, sample_size=500): + """ + Compute the mean and standard deviation of pixels in frames across the video clips + Parameters: + clips_dir (str) : path to the directory containing the clips + sample_size (int): number of frames to sample for computation + Returns: + mean (np.float64): the mean of the pixels + std (np.float64) : the standard deviation of values across pixels + """ + pixel_sum = 0.0 + pixel_sq_sum = 0.0 + pixel_count = 0 + + image_paths = glob.glob(os.path.join(clips_dir, '*', '*', 'frames', '*.jpg')) + + if len(image_paths) > sample_size: + image_paths = random.sample(image_paths, sample_size) + + for image_path in image_paths: + img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float64) + pixel_sum += img.sum() + pixel_sq_sum += (img ** 2).sum() + pixel_count += img.size + + mean = pixel_sum / pixel_count + std = np.sqrt(pixel_sq_sum / pixel_count - mean ** 2) + return mean, std + +class MaskToTensor: + def __call__(self, sample): + im = sample['image'] + im_tensor = torch.from_numpy(im).float() / 255 + im_tensor = torch.unsqueeze(im_tensor, 0) + sample['image'] = im_tensor + + if 'mask' in sample.keys(): + mask = sample['mask'] + mask_tensor = torch.from_numpy(np.asarray(mask, np.int64)) + mask_tensor = mask_tensor // 255 + sample['mask'] = mask_tensor + + return sample + +class MaskNormalize: + def __init__(self, mean, std): + """Normalize image, leave mask unchanged""" + self.mean = mean + self.std = std + + def __call__(self, sample): + im = sample['image'] + im_norm = TF.normalize(im, self.mean, self.std) + + sample['image'] = im_norm + return sample + +class MaskCompose: + def __init__(self, transform_list): + """Chain together transforms in transform. + Expect transform on both image and mask in dict + + Args: + transform_list (list): list of custom augmentations + """ + self.transform_list = transform_list + + def __call__(self, sample): + for transform in self.transform_list: + sample = transform(sample) + return sample + +class MaskImgAug: + def __call__(self, sample): + im = sample['image'].astype(np.uint8) + + # gaussian blur with 40% probability + if np.random.random() < 0.4: + sigma = np.random.uniform(0.0, 6.0) + ksize = int(6 * sigma + 1) + if ksize % 2 == 0: + ksize += 1 + im = cv2.GaussianBlur(im, (ksize, ksize), sigma) + + # poisson noise with 40% probability + if np.random.random() < 0.4: + lam = np.random.uniform(0, 30) + noise = np.random.poisson(lam, im.shape).astype(np.float64) + im = np.clip(im.astype(np.float64) + noise, 0, 255).astype(np.uint8) + + sample['image'] = im.astype(float) + return sample + +class MaskRandomCrop: + """Rotate by one of the given angles.""" + + def __init__(self, crop_size): + # size of square crop + self.crop_size = crop_size + + def __call__(self, sample): + im = sample['image'] + mask = sample['mask'] + + im_height = im.shape[0] + im_width = im.shape[1] + top = np.random.randint(im_height - self.crop_size) + left = np.random.randint(im_width - self.crop_size) + + im_crop = im[top:top+self.crop_size, left:left+self.crop_size] + mask_crop = mask[top:top+self.crop_size, left:left+self.crop_size] + + sample['image'] = im_crop + sample['mask'] = mask_crop + + return sample + +class Mask2dMultiplyAndAddToBrightness(): + def __init__(self, add, multiply): + """Add and multiply image values by given amount randomly within range. + + Expects image to be between 0 and 255. + + Args: + add: either number of tuple, if tuple randomly choose from range + multiply: either number ot tuple, if tuple randomly choose from range + """ + self.add = add + self.multiply = multiply + self.rng = np.random.default_rng() + + def __call__(self, sample): + im = sample['image'].astype(np.float64) + if isinstance(self.add, tuple): + add_val = self.rng.uniform(self.add[0], self.add[1]) + else: + add_val = self.add + if isinstance(self.multiply, tuple): + multiply_val = self.rng.uniform(self.multiply[0], self.multiply[1]) + else: + multiply_val = self.multiply + im += add_val + im *= multiply_val + im = np.maximum(im, 0) + im = np.minimum(im, 255) + + sample['image'] = im + + return sample + +class MaskContrast: + def __init__(self, contrast_factor): + """Change image contrast, leave mask unchanged + + Args: + contrast_factor: single value or tuple""" + self.contrast_factor = contrast_factor + self.rng = np.random.default_rng() + + def __call__(self, sample): + im = sample['image'] + if isinstance(self.contrast_factor, tuple): + contrast_factor = self.rng.uniform(self.contrast_factor[0], self.contrast_factor[1]) + else: + contrast_factor = self.contrast_factor + im_height, im_width = im.shape + aprox_mean = np.mean([im[0,0], im[-1, 0], im[0, -1], im[-1, -1], + im[im_height//2, im_width//2], im[-im_height//2, -im_width//2]]) + im_contrast = aprox_mean + contrast_factor * (im - aprox_mean) + im_contrast = np.maximum(im_contrast, 0) + im_contrast = np.minimum(im_contrast, 255) + sample['image'] = im_contrast + return sample diff --git a/heatseek/counting/dl/bat_functions.py b/heatseek/counting/dl/bat_functions.py new file mode 100644 index 0000000..38bb6b4 --- /dev/null +++ b/heatseek/counting/dl/bat_functions.py @@ -0,0 +1,384 @@ +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 CountLine import CountLine +import 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): + """ 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 = [] + + max_distance_threshold = 30 + max_distance_threshold_noise = 30 + min_distance_threshold = 0 + max_unseen_time = 2 + min_new_track_distance = 3 + min_distance_big = 30 + +# #Create initial tracks based on the objects in the first frame +# raw_track_list = add_all_points_as_new_tracks( +# raw_track_list, positions[0], contours_list[0], sizes_list0, noise=0 +# ) + + #try to connect points to the next frame + if max_frame is None: + max_frame = len(positions) + + 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)) + +# for active_ind, track_ind in enumerate(active_list): +# if active_ind in row_ind: +# row_count = np.where(row_ind == active_ind)[0] +# raw_track_list[track_ind]['debug'].append( +# '{} dist {}, best {}'.format( +# frame_ind, +# distance[row_ind[row_count], +# col_ind[row_count]], +# np.min(distance[row_ind[row_count], +# :]) +# ) +# ) +# best_col = np.argmin(distance[row_ind[row_count], +# :]) +# row_count = np.where(col_ind == best_col)[0] +# raw_track_list[track_ind]['debug'].append( +# '{} row_ind {} col {} dist {} track {}'.format( +# frame_ind, row_ind[row_count], +# col_ind[row_count], +# distance[row_ind[row_count], +# col_ind[row_count]], +# active_list[row_ind[row_count][0]]) +# ) + + + # 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) + # result is 1 if forward crossing -1 is backward crossing + if count_out: + if out_frame_num: + crossing_track_list[-1]['crossed'] = out_frame_num * out_result + else: + crossing_track_list[-1]['crossed'] = 0 + if count_across: + if across_frame_num: + crossing_track_list[-1]['across_crossed'] = across_frame_num * across_result + else: + crossing_track_list[-1]['across_crossed'] = 0 + 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/dl/bat_seg_models.py b/heatseek/counting/dl/bat_seg_models.py new file mode 100644 index 0000000..a3e0f07 --- /dev/null +++ b/heatseek/counting/dl/bat_seg_models.py @@ -0,0 +1,359 @@ +import torch +import torch.nn as nn + +class SuperSimpleSemSegNet(nn.Module): + def __init__(self, in_channel, out_channel): + super().__init__() + self.conv1 = torch.nn.Conv2d(in_channel, out_channel, + kernel_size=3, padding=1, stride=1) + self.ReLU = torch.nn.ReLU() + self.softmax = torch.nn.LogSoftmax(dim=1) + + def forward(self, x): + x = self.conv1(x) + x = self.ReLU(x) + x = self.softmax(x) + + return x + +class ThreeLayerSemSegNetWideView(nn.Module): + def __init__(self, in_channel, out_channel): + super().__init__() + self.conv1 = torch.nn.Conv2d(in_channel, 6, kernel_size=3, padding=1, stride=1) + self.conv1d100 = torch.nn.Conv2d(in_channel, 2, kernel_size=3, padding=101, + stride=1, dilation=101) + self.conv2d1 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=2, stride=1, dilation=2) + self.conv2d5 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=6, stride=1, dilation=6) + self.conv3 = torch.nn.Conv2d(8, out_channel, kernel_size=3, padding=1, stride=1) + self.ReLU1 = torch.nn.ReLU() + self.ReLU2 = torch.nn.ReLU() + self.softmax = torch.nn.LogSoftmax(dim=1) + self.batchnorm1 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + self.batchnorm2 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + + + def forward(self, x): + x1 = self.conv1(x) + x2 = self.conv1d100(x) + x = torch.cat((x1, x2), dim=1) + x = self.batchnorm1(x) + x = self.ReLU1(x) + x1 = self.conv2d1(x) + x2 = self.conv2d5(x) + x = torch.cat((x1, x2), dim=1) + x = self.batchnorm2(x) + x = self.ReLU2(x) + x = self.conv3(x) + x = self.softmax(x) + + return x + +class ThreeLayerSemSegNetWideViewHighDim(nn.Module): + """Each layer has more channels than the standard model""" + def __init__(self, in_channel, out_channel): + super().__init__() + self.conv1 = torch.nn.Conv2d(in_channel, 12, kernel_size=3, padding=1, stride=1) + self.conv1d100 = torch.nn.Conv2d(in_channel, 4, kernel_size=3, padding=101, + stride=1, dilation=101) + self.conv2d1 = torch.nn.Conv2d(16, 8, kernel_size=3, padding=2, stride=1, dilation=2) + self.conv2d5 = torch.nn.Conv2d(16, 8, kernel_size=3, padding=6, stride=1, dilation=6) + self.conv3 = torch.nn.Conv2d(16, out_channel, kernel_size=3, padding=1, stride=1) + self.ReLU1 = torch.nn.ReLU() + self.ReLU2 = torch.nn.ReLU() + self.softmax = torch.nn.LogSoftmax(dim=1) + self.batchnorm1 = torch.nn.BatchNorm2d(16, track_running_stats=False, momentum=1.0) + self.batchnorm2 = torch.nn.BatchNorm2d(16, track_running_stats=False, momentum=1.0) + + + def forward(self, x): + x1 = self.conv1(x) + x2 = self.conv1d100(x) + x = torch.cat((x1, x2), dim=1) + x = self.batchnorm1(x) + x = self.ReLU1(x) + x1 = self.conv2d1(x) + x2 = self.conv2d5(x) + x = torch.cat((x1, x2), dim=1) + x = self.batchnorm2(x) + x = self.ReLU2(x) + x = self.conv3(x) + x = self.softmax(x) + + return x + +class FourLayerSemSegNetWideView(nn.Module): + def __init__(self, in_channel, out_channel): + super().__init__() + self.conv1 = torch.nn.Conv2d(in_channel, 6, kernel_size=3, padding=1, stride=1) + self.conv1d100 = torch.nn.Conv2d(in_channel, 2, kernel_size=3, padding=101, + stride=1, dilation=101) + self.conv2d1 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=2, stride=1, dilation=2) + self.conv2d5 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=6, stride=1, dilation=6) + self.conv3d0 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=1, stride=1) + self.conv3d3 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=4, stride=1, dilation=4) + self.conv4 = torch.nn.Conv2d(8, out_channel, kernel_size=3, padding=1, stride=1) + self.ReLU1 = torch.nn.ReLU() + self.ReLU2 = torch.nn.ReLU() + self.ReLU3 = torch.nn.ReLU() + self.softmax = torch.nn.LogSoftmax(dim=1) + self.batchnorm1 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + self.batchnorm2 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + self.batchnorm3 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + + + def forward(self, x): + x1a = self.conv1(x) + x1b = self.conv1d100(x) + x1 = torch.cat((x1a, x1b), dim=1) + x1 = self.batchnorm1(x1) + x1 = self.ReLU1(x1) + x2a = self.conv2d1(x1) + x2b = self.conv2d5(x1) + x2 = torch.cat((x2a, x2b), dim=1) + x2 = self.batchnorm2(x2) + x2 = self.ReLU2(x2) + x3a = self.conv3d0(x2) + x3b = self.conv3d3(x2) + x3 = torch.cat((x3a, x3b), dim=1) + x3 = self.batchnorm3(x3) + x3 = self.ReLU3(x3) + x4 = self.conv4(x3) + xout = self.softmax(x4) + + return xout + +class ThreeLayerSemSegNet(nn.Module): + def __init__(self, in_channel, out_channel): + super().__init__() + self.conv1 = torch.nn.Conv2d(in_channel, 8, kernel_size=3, padding=1, stride=1) + self.conv2d1 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=2, stride=1, dilation=2) + self.conv2d5 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=6, stride=1, dilation=6) + self.conv3 = torch.nn.Conv2d(8, out_channel, kernel_size=3, padding=1, stride=1) + self.ReLU1 = torch.nn.ReLU() + self.ReLU2 = torch.nn.ReLU() + self.softmax = torch.nn.LogSoftmax(dim=1) + self.batchnorm1 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + self.batchnorm2 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) + + + def forward(self, x): + x = self.conv1(x) + x = self.batchnorm1(x) + x = self.ReLU1(x) + x1 = self.conv2d1(x) + x2 = self.conv2d5(x) + x = torch.cat((x1, x2), dim=1) + x = self.batchnorm2(x) + x = self.ReLU2(x) + x = self.conv3(x) + x = self.softmax(x) + + return x + + +class SimpleSemSegNet(nn.Module): + def __init__(self, in_channel, out_channel): + super().__init__() + + self.conv1 = self.contract_block(in_channel, 32, 7, 3) +# self.conv2 = self.contract_block(32, 64, 3, 1) +# self.conv3 = self.contract_block(64, 128, 3, 1) + +# self.upconv3 = self.expand_block(128, 64, 3, 1) +# self.upconv2 = self.expand_block(64*2, 32, 3, 1) + self.upconv1 = self.expand_block(32, out_channel, 3, 1) + self.log_softmax = torch.nn.LogSoftmax(dim=1) + + def forward(self, x): + conv1 = self.conv1(x) +# conv2 = self.conv2(conv1) +# conv3 = self.conv3(conv2) + +# upconv3 = self.upconv3(conv3) +# upconv2 = self.upconv2(torch.cat([upconv3, conv2], 1)) +# upconv1 = self.upconv1(torch.cat([upconv2, conv1], 1)) + upconv1 = self.upconv1(conv1) + + output = self.log_softmax(upconv1) + + return output + + def contract_block(self, in_channels, out_channels, kernel_size, padding): + contract = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, + stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, + stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + ) + return contract + + def expand_block(self, in_channels, out_channels, kernel_size, padding): + expand = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, + stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, + stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.ConvTranspose2d(out_channels, out_channels, kernel_size=3, + stride=2, padding=1, output_padding=1) + ) + return expand + +class UNET(nn.Module): + def __init__(self, in_channels, out_channels, should_pad=True): + super().__init__() + self.name = 'UNET' + if should_pad: + conv1_pad = 3 + gen_pad = 1 + else: + conv1_pad = 0 + gen_pad = 0 + self.conv1 = self.contract_block(in_channels, 32, 7, conv1_pad) + self.conv2 = self.contract_block(32, 64, 3, gen_pad) + self.conv3 = self.contract_block(64, 128, 3, gen_pad) + + self.upconv3 = self.expand_block(128, 64, 3, gen_pad) + self.upconv2 = self.expand_block(64*2, 32, 3, gen_pad) + self.upconv1 = self.expand_block(32*2, out_channels, 3, gen_pad) + + self.softmax = torch.nn.LogSoftmax(dim=1) + + def __call__(self, x): + + # downsampling part + conv1 = self.conv1(x) + conv2 = self.conv2(conv1) + conv3 = self.conv3(conv2) + upconv3 = self.upconv3(conv3) + + cat1_trim = 6 + cat2_trim = 18 + upconv2 = self.upconv2(torch.cat([upconv3, conv2[:, :, cat1_trim:-cat1_trim, cat1_trim:-cat1_trim]], 1)) + upconv1 = self.upconv1(torch.cat([upconv2, conv1[:, :, cat2_trim:-cat2_trim, cat2_trim:-cat2_trim]], 1)) + xout = self.softmax(upconv1) + + return xout + + def contract_block(self, in_channels, out_channels, kernel_size, padding): + + contract = nn.Sequential( + torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + ) + + return contract + + def expand_block(self, in_channels, out_channels, kernel_size, padding): + + expand = nn.Sequential(torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.Conv2d(out_channels, out_channels, kernel_size, stride=1, padding=padding), + torch.nn.BatchNorm2d(out_channels), + torch.nn.ReLU(), + torch.nn.ConvTranspose2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1) + ) + return expand + +class DoubleCNNBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, padding): + super().__init__() + self.name = 'DoubleCNNBlock' + + self.block = nn.Sequential( + torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, + stride=1, padding=padding), + torch.nn.ReLU(), + torch.nn.BatchNorm2d(out_channels), + torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, + stride=1, padding=padding), + torch.nn.ReLU(), + torch.nn.BatchNorm2d(out_channels) + ) + + def forward(self, x): + return self.block(x) + +class ContractBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, padding): + super().__init__() + self.name = "ContractBlock" + + self.maxpool_conv = nn.Sequential( + torch.nn.MaxPool2d(kernel_size=2, stride=2), + DoubleCNNBlock(in_channels, out_channels, kernel_size, padding) + ) + + def forward(self, x): + return self.maxpool_conv(x) + +class ExpandBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, padding): + super().__init__() + self.name = "ExpandBlock" + + self.expand = torch.nn.ConvTranspose2d(in_channels, in_channels//2, + kernel_size=3, stride=2, + padding=1, output_padding=1 + ) + self.conv = DoubleCNNBlock(in_channels, out_channels, kernel_size, padding) + + def forward(self, x1, x2): + x1_expand = self.expand(x1) + # Assumes dims have been corrected/padded to match elsewhere + x = torch.cat([x1_expand, x2], dim=1) + return self.conv(x) + + +class UNETTraditional(nn.Module): + def __init__(self, in_channels, out_channels, should_pad=True): + super().__init__() + self.name = 'UNETTraditional' + if should_pad: + conv1_pad = 3 + gen_pad = 1 + else: + conv1_pad = 0 + gen_pad = 0 + self.conv1 = DoubleCNNBlock(in_channels, 32, 7, conv1_pad) + self.conv2 = ContractBlock(32, 64, 3, gen_pad) + self.conv3 = ContractBlock(64, 128, 3, gen_pad) + + self.upconv1 = ExpandBlock(128, 64, 3, gen_pad) + self.upconv2 = ExpandBlock(64, 32, 3, gen_pad) + + self.outconv = nn.Conv2d(32, out_channels, kernel_size=1) + self.softmax = torch.nn.LogSoftmax(dim=1) + + def forward(self, x): + + # downsampling part + conv1 = self.conv1(x) + conv2 = self.conv2(conv1) + conv3 = self.conv3(conv2) + + conv2_crop = 4 + cropped_conv2 = conv2[:, :, conv2_crop:-conv2_crop, conv2_crop:-conv2_crop] + upconv1 = self.upconv1(conv3, cropped_conv2) + conv1_crop = 16 + cropped_conv1 = conv1[:, :, conv1_crop:-conv1_crop, conv1_crop:-conv1_crop] + upconv2 = self.upconv2(upconv1, cropped_conv1) + + outconv = self.outconv(upconv2) + xout = self.softmax(outconv) + + return xout \ No newline at end of file diff --git a/heatseek/counting/dl/crop_vids.py b/heatseek/counting/dl/crop_vids.py new file mode 100644 index 0000000..918f57d --- /dev/null +++ b/heatseek/counting/dl/crop_vids.py @@ -0,0 +1,94 @@ +import cv2 +import subprocess +import os +import logging +import argparse + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +def crop_vids(input_vid, output_dir, pixel_threshold, count_threshold, padding, min_gap): + """Scans the input video for motion and crops out clips of motion events. + Parameters: + - input_vid: path to the input video file + - output_dir: directory where the output clips will be saved + - threshold: sensitivity for motion detection (lower = more sensitive) + - padding: seconds of padding to add before and after each detected motion event + - min_gap: minimum gap in seconds to merge nearby motion events into a single clip + """ + os.makedirs(output_dir, exist_ok=True) + cap = cv2.VideoCapture(input_vid) + fps = cap.get(cv2.CAP_PROP_FPS) + prev_frame = None + motion_times = [] + logging.info("Scanning for motion") + + while True: #while there are frames to read + ret, frame = cap.read() #ret returns true if frame is read correctly, frame is the image array + + if not ret: #if there are no more frames left to read + break + + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) #convert the frame to grayscale for easier processing + + if prev_frame is not None: #if there is a previous frame to compare to (i.e. not the first frame) + diff = cv2.absdiff(prev_frame, gray) #calculate the absolute difference between the current frame and the previous frame + + # if diff.mean() > threshold: + if (diff > pixel_threshold).sum() > count_threshold: #if the number of pixels that changed significantly is greater than the count threshold, we consider it motion + t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000 #get the current time in milliseconds + motion_times.append(t) #add the time to the list of motion times + + prev_frame = gray #update the previous frame to the current frame for the next iteration + + cap.release() #release the video capture object + + if not motion_times: #if no motion was detected, print a message and exit + logging.info("No motion detected. Try lowering threshold") + return + + logging.info(f"Motion found at {len(motion_times)} frames, merging into clips") + clips = [] + start = motion_times[0] + end = motion_times[0] + + for t in motion_times[1:]: #iterate through the motion times starting from the second one + if t - end < min_gap: #if the time between the current motion time and the end of the last clip is less than the minimum gap, we consider it part of the same clip + end = t #update the end time of the current clip to the current motion time + else: + clips.append((start, end)) #if it's not part of the same clip, add the current clip to the list of clips + start = t #start a new clip with the current motion time as the start + end = t #and also set the end to the current motion time for now + + clips.append((start, end)) #add the last clip to the list of clips + logging.info(f"Found {len(clips)} clip(s):") + + for i, (s, e) in enumerate(clips): #iterate through the clips and print their start and end times + s_pad = max(0, s - padding) #add padding to the start time, ensuring it doesn't go below 0 + e_pad = e + padding #add padding to the end time + logging.info(f"Clip {i+1}: {s_pad:.1f}s → {e_pad:.1f}s") #print the clip number and its start and end times + + out_path = os.path.join(output_dir, f"clip_{i+1:03d}.mp4") #create the output path for the clip + subprocess.run([ #use ffmpeg to extract the clip from the original video + "ffmpeg", "-y", #-y to overwrite existing files without asking + "-ss", str(s_pad), #start time for the clip + "-to", str(e_pad), #end time for the clip + "-i", input_vid, #input video file + "-c", "copy", #copy the video and audio streams without re-encoding + out_path #output path for the clip + ]) + + logging.info(f"Done! Clips saved to {output_dir}") + +def main(): + parser = argparse.ArgumentParser(description="Crop motion events from a video into separate clips") + parser.add_argument("--input", required=True, help="Path to the input video file") + parser.add_argument("--output_dir", required=True, help="Directory to save the output clips") + parser.add_argument("--pixel_threshold", type=float, default=15, help="Sensitivity for motion detection (lower = more sensitive)") + parser.add_argument("--count_threshold", type=int, default=20, help="Minimum number of pixels that must change to consider it motion") + parser.add_argument("--padding", type=float, default=2.0, help="Seconds of padding to add before and after each detected motion event") + parser.add_argument("--min_gap", type=float, default=2.0, help="Minimum gap in seconds to merge nearby motion events into a single clip") + args = parser.parse_args() + crop_vids(args.input, args.output_dir, args.pixel_threshold, args.count_threshold, args.padding, args.min_gap) + +if __name__ == "__main__": + main() diff --git a/heatseek/counting/dl/crossing_tracks.py b/heatseek/counting/dl/crossing_tracks.py new file mode 100644 index 0000000..3c2676a --- /dev/null +++ b/heatseek/counting/dl/crossing_tracks.py @@ -0,0 +1,143 @@ +from bat_functions import threshold_short_tracks, measure_crossing_bats +import numpy as np +import matplotlib.pyplot as plt +import argparse + +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, midline_y=88, num_tracks=100): +# 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()}') +# print(f'Forward crossings: {sum(1 for t in crossing_tracks if t["crossed"] > 0)}') +# print(f'Backward crossings: {sum(1 for t in crossing_tracks if t["crossed"] < 0)}') + +# 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] +# color = 'blue' if track['crossed'] > 0 else 'red' +# ax.plot(track['track'][:, 0], track['track'][:, 1], color=color, alpha=0.3, linewidth=1.5) +# ax.axhline(y=midline_y, color='green', linewidth=2, linestyle='--', label='midline') +# ax.set_title(f'Sample of {num_tracks} crossing tracks\nblue=forward, red=backward') +# ax.invert_yaxis() +# ax.legend() + +# ax = axes[1] +# crossing_frames = [abs(t['crossed']) for t in crossing_tracks] +# ax.hist(crossing_frames, bins=100) +# ax.set_xlabel('Frame number') +# ax.set_ylabel('Number of crossings') +# ax.set_title('Bat crossings over time') + +# print('Crossing frames:') + +# for t in crossing_tracks: +# direction = 'forward' if t['crossed'] > 0 else 'backward' +# print(f' frame {abs(t["crossed"])}: {direction}') + +# plt.tight_layout() +# plt.show() + +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 + + crossed_key = 'crossed' if count_out else 'across_crossed' + + 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()}') + print(f'Forward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] > 0)}') + print(f'Backward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] < 0)}') + + 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] + color = 'blue' if track[crossed_key] > 0 else 'red' + 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() + direction_labels = 'blue=downward, red=upward' + else: + ax.axvline(x=midline_x, color='green', linewidth=2, linestyle='--', label='midline') + direction_labels = 'blue=rightward, red=leftward' + + ax.set_title(f'Sample of {num_tracks} crossing tracks\n{direction_labels}') + ax.legend() + + ax = axes[1] + crossing_frames = [abs(t[crossed_key]) for t in crossing_tracks] + ax.hist(crossing_frames, bins=100) + ax.set_xlabel('Frame number') + ax.set_ylabel('Number of crossings') + ax.set_title('Bat crossings over time') + + print('Crossing frames:') + for t in crossing_tracks: + direction = 'forward' if t[crossed_key] > 0 else 'backward' + print(f' frame {abs(t[crossed_key])}: {direction}') + + plt.tight_layout() + plt.show() + +def main(): + parser = argparse.ArgumentParser(description='Get bat crossing tracks across the video') + parser.add_argument('--raw_tracks_file', type=str, help='Path to raw tracks file') + parser.add_argument('--crossing_tracks_file', type=str, help='Intended path for the crossing tracks 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() + + save_crossing_tracks_from_raw_tracks(args.raw_tracks_file, args.crossing_tracks_file, args.count_across, args.count_out) + visualize_crossing_tracks(args.crossing_tracks_file, count_out=args.count_out, count_across=args.count_across) + +if __name__ == '__main__': + main() diff --git a/heatseek/counting/dl/dataloaders.py b/heatseek/counting/dl/dataloaders.py new file mode 100644 index 0000000..01de90b --- /dev/null +++ b/heatseek/counting/dl/dataloaders.py @@ -0,0 +1,16 @@ +import numpy as np +import datetime +from SegmentationDataset import SegmentationDataset +from transforms import split_train_val +import torch.utils.data as data + +def worker_init_fn(worker_id): + np.random.seed(datetime.datetime.now().microsecond + worker_id * 1000000) + +def load_dataloader(clips_dir, transform, split='train', batch_size=4, num_workers=7, pin_memory=True): + train_pairs, val_pairs = split_train_val(clips_dir) + pairs = train_pairs if split == 'train' else val_pairs + dataset = SegmentationDataset(pairs, transform) + return data.DataLoader(dataset, batch_size=batch_size, + shuffle=(split == 'train'), num_workers=num_workers, + pin_memory=pin_memory, worker_init_fn=worker_init_fn) diff --git a/heatseek/counting/dl/detections_to_tracks.py b/heatseek/counting/dl/detections_to_tracks.py new file mode 100644 index 0000000..cc80404 --- /dev/null +++ b/heatseek/counting/dl/detections_to_tracks.py @@ -0,0 +1,222 @@ +import numpy as np +import os +import glob +import bat_functions as kbf +from multiprocessing import Pool +import argparse + +def track(camera_dict): + """ + 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). + + 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) + else: + print("Missing contour files.") + +def build_camera_dicts(output_folder, num_groups=10, fps=60, overlap_seconds=15): + """ + 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 15. + + 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) + + 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, processes=5): + """ + 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) + print(f'Tracking {len(camera_dicts)} groups') + with Pool(processes=processes) as pool: + pool.map(track, camera_dicts) + +def main(): + parser = argparse.ArgumentParser(description='Get centers and contours from detections from model inference') + parser.add_argument('--output_folder', type=str, help='Path of output folder containing the centers, contours, etc from video inference') + args = parser.parse_args() + run_tracking(args.output_folder) + combine_tracks(args.output_folder) + +if __name__ == '__main__': + main() diff --git a/heatseek/counting/dl/extract_frames.py b/heatseek/counting/dl/extract_frames.py new file mode 100644 index 0000000..ba09b45 --- /dev/null +++ b/heatseek/counting/dl/extract_frames.py @@ -0,0 +1,55 @@ +import glob +import os +import cv2 +import logging +import argparse + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +def pull_out_frames(video_file, process_every_n_frames=1): + """Extract frames from a video file and optionally save them to disk. + Parameters: + - video_file: path to the input video file + - save_frames: whether to save the extracted frames to disk + - process_every_n_frames: only process every nth frame to speed up extraction + """ + video_dir = os.path.dirname(video_file) + video_name = os.path.splitext(os.path.basename(video_file))[0] + output_folder = os.path.join(video_dir, video_name, 'frames') + os.makedirs(output_folder, exist_ok=True) + logging.info(f'Saving frames to {output_folder}') + cap = cv2.VideoCapture(video_file) + frame_count = 0 + + while True: #read every frame until the end of the video + ret, frame = cap.read() + if not ret: + break + + if frame_count % process_every_n_frames == 0: + frame_name = f'{video_name}_{frame_count:05d}.jpg' + frame_file = os.path.join(output_folder, frame_name) + cv2.imwrite(frame_file, frame) + + frame_count += 1 + + cap.release() + logging.info(f'{video_name}: done. {frame_count} total frames.') + + + +def main(): + parser = argparse.ArgumentParser(description='Extract frames from videos') + parser.add_argument('--clips_dir', help='Parent folder containing video files') + parser.add_argument('--process_every_n_frames', type=int, default=1, help='Process every nth frame') + args = parser.parse_args() + + video_files = [] + video_files.extend(glob.glob(os.path.join(args.clips_dir, '**', '*.mp4'), recursive=True)) + video_files.sort() + + for video_file in video_files: + pull_out_frames(video_file, process_every_n_frames=args.process_every_n_frames) + +if __name__ == '__main__': + main() diff --git a/heatseek/counting/dl/koger_tracking.py b/heatseek/counting/dl/koger_tracking.py new file mode 100644 index 0000000..cba6251 --- /dev/null +++ b/heatseek/counting/dl/koger_tracking.py @@ -0,0 +1,744 @@ +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['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 + 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 fewwer new points than existing tracks. When it is the + later, 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 + +#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 = 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, + 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: upper distance threshold + max_distance_noise: upper distance threshold tracks with noise values + above 0 + min_distance: lowwer 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 + 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) + + 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)] = max_distance + # 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]) + 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]) + + 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/dl/optimizer.py b/heatseek/counting/dl/optimizer.py new file mode 100644 index 0000000..ca7aad0 --- /dev/null +++ b/heatseek/counting/dl/optimizer.py @@ -0,0 +1,108 @@ +import torch +from bat_seg_models import UNETTraditional +import os + +def lr_scheduler(epoch): + """ To pass to torch optim lr_scheduler. + + Args: + epoch: current epoch number + + return number to multiply base lr + """ + + if epoch < 15: + return 1.0 + elif epoch < 30: + return 0.2 + else: + return 0.05 + +def build_and_save_model(lr=0.01, accumulation_steps=4, epochs=91, momentum=0.9, batch_size=4): + """ + Instantiate a UNETTraditional model and construct its checkpoint file path. + + Builds a single-channel-input, two-class-output U-Net (no padding) and + generates a deterministic filename encoding the key hyperparameters so + that checkpoint files are self-describing and don't collide across runs. + + Args: + lr (float): Learning rate, embedded in the filename. Default: 0.01. + accumulation_steps (int): Gradient accumulation steps. Combined with + `batch_size` to record the effective batch size in the filename. + Default: 4. + epochs (int): Total training epochs, embedded in the filename. + Default: 91. + momentum (float): SGD momentum, embedded in the filename. Default: 0.9. + batch_size (int): Per-step batch size. Multiplied by + `accumulation_steps` to compute effective batch size. Default: 4. + + Returns: + model (UNETTraditional): Untrained model with in_channels=1, + out_channels=2, should_pad=False. + model_file (str): Path of the form + './models/model_UNETTraditional_epochs_{e}_batcheff_{b}_lr_{lr} + _momentum_{m}_aug_thermal-radiometric-heatseek.tar'. + The ./models/ directory is created if it does not exist. + + Note: + This function only constructs and names the model — it does not + train it or write any file to disk. Pass model_file to train() with + save_best_val=True to actually checkpoint the weights. + """ + + model = UNETTraditional(in_channels=1, out_channels=2, should_pad=False) + # refer to bat_seg_models.py to set any other models; Koger uses UNETTraditional for all experiments, so we hardcode that here for simplicity + model_name = "UNETTraditional" + aug_type = "thermal-radiometric-heatseek" + + folder = './models' + os.makedirs(folder, exist_ok=True) + + model_name = 'model_{}_epochs_{}_batcheff_{}_lr_{}_momentum_{}_aug_{}'.format( + model_name, epochs, accumulation_steps*batch_size, lr, momentum, aug_type) + + model_file = os.path.join(folder, model_name + '.tar') + return model, model_file + +def get_optimizer_and_scheduler(model, lr=0.01, momentum=0.9): + """ + Create three independent SGD optimizer/scheduler pairs for none, easy, + and hard augmentation training regimes. + + Each pair is fully independent — separate parameter groups, separate + state — so they can be used to train the same model sequentially under + different augmentation conditions without state leaking between regimes. + All three share the same `lr`, `momentum`, and `lr_scheduler` lambda. + + Args: + model (torch.nn.Module): Model whose parameters all three optimizers + will reference. + lr (float): Initial learning rate for all three SGD instances. + Default: 0.01. + momentum (float): Momentum for all three SGD instances. Default: 0.9. + + Returns: + optimizer_none (torch.optim.SGD): Optimizer for the no-augmentation regime. + scheduler_none (LambdaLR): Scheduler paired with optimizer_none. + optimizer_easy (torch.optim.SGD): Optimizer for the easy-augmentation regime. + scheduler_easy (LambdaLR): Scheduler paired with optimizer_easy. + optimizer_hard (torch.optim.SGD): Optimizer for the hard-augmentation regime. + scheduler_hard (LambdaLR): Scheduler paired with optimizer_hard. + + Note: + All three optimizers reference the same model.parameters() iterator + at construction time. If the model is replaced or re-initialized + between regimes, these optimizers will hold stale parameter references + and must be reconstructed. + """ + optimizer_none = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) + scheduler_none = torch.optim.lr_scheduler.LambdaLR(optimizer_none, lr_scheduler, last_epoch=-1) + + optimizer_easy = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) + scheduler_easy = torch.optim.lr_scheduler.LambdaLR(optimizer_easy, lr_scheduler, last_epoch=-1) + + optimizer_hard = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) + scheduler_hard = torch.optim.lr_scheduler.LambdaLR(optimizer_hard, lr_scheduler, last_epoch=-1) + + return optimizer_none, scheduler_none, optimizer_easy, scheduler_easy, optimizer_hard, scheduler_hard diff --git a/heatseek/counting/dl/segment_frames.py b/heatseek/counting/dl/segment_frames.py new file mode 100644 index 0000000..cc3249c --- /dev/null +++ b/heatseek/counting/dl/segment_frames.py @@ -0,0 +1,88 @@ +import cv2 +import numpy as np +import argparse +import os +import logging + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +def get_background_temp(image_path, min_val, max_val): + """Get background temperature of an image taken by a radiometric camera. + Parameters: + - image_path: path to the image file (e.g. frame.png) + - min_val: minimum temperature value corresponding to pixel value 0 + - max_val: maximum temperature value corresponding to pixel value 255 + """ + frame = cv2.imread(image_path) #read the frame + + if frame is None: + raise ValueError("Could not read image") + + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) #convert to grayscale dimensions + + k100 = (gray.astype(np.float32) / 255.0) * (max_val - min_val) + min_val #formula to convert pixel values to temperature in Kelvin (FLIR Boson) + temp_celsius = (k100 / 100.0) - 273.15 + background_temp = np.median(temp_celsius) + return background_temp #return the estimated background temperature in Celsius based on the median pixel value of the thermal map + +def get_mask(image_path, min_val, max_val, threshold): + """Get a segmentation mask of objects hotter than the background by a threshold. + Parameters: + - image_path: path to the image file + - min_val: minimum temperature value corresponding to pixel value 0 + - max_val: maximum temperature value corresponding to pixel value 255 + - background_temp: estimated background temperature in °C + - threshold: how many °C above background a pixel must be to be masked + """ + frame = cv2.imread(image_path) + + if frame is None: + raise ValueError("Could not read image") + + gray = cv2.cvtColor(frame, 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 = get_background_temp(image_path, min_val, max_val) + mask = (temp_celsius > (background_temp + threshold)).astype(np.uint8) * 255 #create a binary mask where pixels hotter than background + threshold are set to 255 (white) and others are 0 (black) + return mask + + + +def main(): + parser = argparse.ArgumentParser(description='Generate thermal segmentation mask for a single image') + parser.add_argument('--clips_dir', type=str, help='Path to the clips folder') + parser.add_argument('--min_val', type=float, default=28000) + parser.add_argument('--max_val', type=float, default=32000) + parser.add_argument('--threshold', type=float, default=5.0) + args = parser.parse_args() + + for session_folder in os.listdir(args.clips_dir): + session_path = os.path.join(args.clips_dir, session_folder) + + if not os.path.isdir(session_path): + continue + + for clip_folder in os.listdir(session_path): + clip_path = os.path.join(session_path, clip_folder) + + if not os.path.isdir(clip_path): + continue + + frames_path = os.path.join(clip_path, 'frames') + masks_path = os.path.join(clip_path, 'masks') + + if not os.path.isdir(frames_path): + continue + + os.makedirs(masks_path, exist_ok=True) + + for filename in os.listdir(frames_path): + image_path = os.path.join(frames_path, filename) + mask = get_mask(image_path, args.min_val, args.max_val, args.threshold) + output_file = os.path.join(masks_path, filename.replace('.jpg', '.png')) + cv2.imwrite(output_file, mask) + + logging.info(f"Finished processing {clip_folder}") + +if __name__ == "__main__": + main() diff --git a/heatseek/counting/dl/trainer.py b/heatseek/counting/dl/trainer.py new file mode 100644 index 0000000..9979b19 --- /dev/null +++ b/heatseek/counting/dl/trainer.py @@ -0,0 +1,248 @@ +import time +import numpy as np +import torch +import argparse +from optimizer import build_and_save_model +from dataloaders import load_dataloader +from transforms import get_transforms +from optimizer import get_optimizer_and_scheduler +import logging + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +def get_conf_matrix(conf_matrix, num_classes, pred_batch, mask_batch): + """ Calculate confusion matrix and add to passed confusion matrix. + + Args: + conf_matrix (np array): confusion matrix size + (num_classes + 1, num_classes + 1) + num_classes (int): number of classes being segmented + pred_batch (np array): NxM or BxNxM + mask_batch (np array): NxM or BxNxM + """ + + N = num_classes + 1 + + if len(pred_batch.shape) == 2: + pred_batch = np.expand_dims(pred_batch, 0) + mask_batch = np.expand_dims(mask_batch, 0) + for pred, mask in zip(pred_batch, mask_batch): + conf_matrix += np.bincount( + N * pred.reshape(-1) + mask.reshape(-1), minlength=N ** 2 + ).reshape(N, N) + return conf_matrix + +def evaluate(num_classes, conf_matrix): + """ + Compute per-class accuracy and IoU from a confusion matrix. + + The confusion matrix is expected to have an extra row/column for an + ignore/background class (shape: (num_classes + 1, num_classes + 1)). + Only the first `num_classes` rows and columns are evaluated; the final + row/column is treated as a void/ignore label and excluded. + + Parameters: + num_classes (int): Number of foreground classes to evaluate. + conf_matrix (np.ndarray): Square confusion matrix of shape + (num_classes + 1, num_classes + 1), where entry [i, j] is the + number of pixels of true class i predicted as class j. + + Returns: + iou (np.ndarray): Per-class Intersection over Union of shape + (num_classes,). Classes with no ground-truth pixels are NaN. + acc (np.ndarray): Per-class accuracy (recall) of shape + (num_classes,). Classes with no ground-truth pixels are NaN. + """ + + acc = np.full(num_classes, np.nan, dtype=np.float64) + iou = np.full(num_classes, np.nan, dtype=np.float64) + tp = conf_matrix.diagonal()[:-1].astype(np.float64) + pos_gt = np.sum(conf_matrix[:-1, :-1], axis=0).astype(np.float64) + pos_pred = np.sum(conf_matrix[:-1, :-1], axis=1).astype(np.float64) + acc_valid = pos_gt > 0 + acc[acc_valid] = tp[acc_valid] / pos_gt[acc_valid] + union = pos_gt + pos_pred - tp + iou[acc_valid] = tp[acc_valid] / union[acc_valid] + return iou, acc + + +def acc_metric(predb, yb): + """ + Compute batch accuracy for a multi-class classifier. + + Parameters: + predb (torch.Tensor): Raw logits or probabilities of shape + (N, C), where N is batch size and C is number of classes. + yb (torch.Tensor): Ground-truth class indices of shape (N,). + + Returns: + torch.Tensor: Scalar mean accuracy over the batch, in [0.0, 1.0]. + """ + return (predb.argmax(dim=1) == yb).float().mean() + +def train(model, num_classes, train_dl, val_dl, loss_fn, optimizer, + epochs=91, lr_scheduler=None, save_best_val=False, + model_file='model.tar', accumulation_steps=4, val_epoch=2): + """ + Train a segmentation model with gradient accumulation and periodic validation. + + Parameters: + model (torch.nn.Module): Segmentation model to train. + num_classes (int): Number of foreground classes. + train_dl (DataLoader): Training dataloader. Batches must be dicts + with keys 'image' and 'mask'. + val_dl (DataLoader): Validation dataloader. Same format as train_dl. + loss_fn (callable): Loss function with signature + loss_fn(outputs, masks) -> scalar tensor. + optimizer (torch.optim.Optimizer): Optimizer instance. + epochs (int): Total number of training epochs. Default: 91. + lr_scheduler (optional): Scheduler with a .step() method called + after each optimizer step. Default: None. + save_best_val (bool): If True, saves model state dict to + model_file whenever validation loss improves. Default: False. + model_file (str): Path for the checkpoint file. Default: 'model.tar'. + accumulation_steps (int): Number of batches to accumulate gradients + over before calling optimizer.step(). Default: 4. + val_epoch (int): Validation runs on epochs where + epoch % val_epoch == 0. Default: 2. + + Returns: + train_loss (list[float]): Mean training loss recorded each epoch. + valid_loss (list[float]): Mean validation loss recorded on each + validated epoch. + """ + start = time.time() + train_loss, valid_loss = [], [] + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + model.to(device) + + max_batchnum = (len(train_dl) // accumulation_steps) * accumulation_steps + logging.info("Using {} batches ({} images)".format( + max_batchnum, max_batchnum * train_dl.batch_size)) + + top_val_loss = float('inf') + + for epoch in range(epochs): + logging.info('Epoch {}/{}'.format(epoch, epochs - 1)) + optimizer.zero_grad() + + for phase in ['train', 'val']: + if phase == 'train': + model.train(True) + dataloader = train_dl + else: + if epoch % val_epoch != 0: + continue + model.train(False) + dataloader = val_dl + conf_matrix = np.zeros((num_classes + 1, num_classes + 1), dtype=np.int64) + + running_loss = 0.0 + + for batch_ind, batch in enumerate(dataloader): + if batch_ind >= max_batchnum: + break + + im_batch = batch['image'].to(device) + masks = batch['mask'][:, 24:-24, 24:-24].to(device) + + if phase == 'train': + outputs = model(im_batch) + loss = loss_fn(outputs, masks) + loss.backward() + if (batch_ind + 1) % accumulation_steps == 0: + optimizer.step() + if lr_scheduler: + lr_scheduler.step() + model.zero_grad() + else: + with torch.no_grad(): + outputs = model(im_batch) + loss = loss_fn(outputs, masks) + np_preds = np.argmax(outputs.cpu().numpy(), axis=1) + np_masks = masks.cpu().numpy() + conf_matrix = get_conf_matrix(conf_matrix, num_classes, + np_preds, np_masks) + + running_loss += loss.item() * dataloader.batch_size + + epoch_loss = running_loss / len(dataloader.dataset) + logging.info('{} Loss: {:.4f}'.format(phase, epoch_loss)) + + if phase == 'val': + valid_loss.append(epoch_loss) + if epoch_loss < top_val_loss: + top_val_loss = epoch_loss + if save_best_val: + torch.save(model.state_dict(), model_file) + logging.info('Saved new best model. Val loss: {:.4f}'.format(epoch_loss)) + iou, acc = evaluate(num_classes, conf_matrix) + logging.info('IOU: {}, Acc: {}'.format(iou, acc)) + else: + train_loss.append(epoch_loss) + + time_elapsed = time.time() - start + logging.info('Training complete in {:.0f}m {:.0f}s'.format( + time_elapsed // 60, time_elapsed % 60)) + + return train_loss, valid_loss + +def main(): + parser = argparse.ArgumentParser(description='Train UNETTraditional on thermal segmentation dataset') + parser.add_argument('--clips_dir', type=str, help='path to clips directory') + parser.add_argument('--save_model', default=True, help='whether to save the trained model') + parser.add_argument('--lr', type=float, default=0.01, help='learning rate for optimizer') + parser.add_argument('--accumulation_steps', type=int, default=4, help='number of batches to accumulate gradients over') + parser.add_argument('--epochs', type=int, default=91, help='number of epochs to train for') + parser.add_argument('--momentum', type=float, default=0.9, help='momentum for SGD optimizer (not applicable for UNETTraditional)') + parser.add_argument('--batch_size', type=int, default=4, help='batch size for training (effective batch size is batch_size * accumulation_steps)') + parser.add_argument('--num_workers', type=int, default=7, help='number of workers for dataloader') + parser.add_argument('--pin_memory', type=bool, default=True, help='whether to pin memory in dataloader') + args = parser.parse_args() + + total_train_loss = [] + total_val_loss = [] + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + loss_fn = torch.nn.NLLLoss(weight=torch.tensor([.01, .99]).to(device)) + model, model_file = build_and_save_model(args.lr, args.accumulation_steps, args.epochs) + optimizer_none, scheduler_none, optimizer_easy, scheduler_easy, optimizer_hard, scheduler_hard = get_optimizer_and_scheduler(model, args.lr, args.momentum) + + no_aug_epochs = 1 + easy_epochs = 30 + hard_epochs = 60 + + none_transforms, easy_transforms, hard_transforms = get_transforms(args.clips_dir) + + train_dataloader_no_aug = load_dataloader(args.clips_dir, none_transforms, split='train', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) + train_dataloader_easy_aug = load_dataloader(args.clips_dir, easy_transforms, split='train', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) + train_dataloader_hard_aug = load_dataloader(args.clips_dir, hard_transforms, split='train', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) + val_dataloader = load_dataloader(args.clips_dir, hard_transforms, split='val', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) + + train_loss, val_loss = train(model, 2, train_dataloader_no_aug, val_dataloader, + loss_fn, optimizer_none, epochs=no_aug_epochs, + lr_scheduler=scheduler_none, save_best_val=True, + model_file=model_file, val_epoch=1) + total_train_loss.extend(train_loss) + total_val_loss.extend(val_loss) + + train_loss, val_loss = train(model, 2, train_dataloader_easy_aug, val_dataloader, + loss_fn, optimizer_easy, epochs=easy_epochs, + lr_scheduler=scheduler_easy, save_best_val=True, + model_file=model_file) + total_train_loss.extend(train_loss) + total_val_loss.extend(val_loss) + + train_loss, val_loss = train(model, 2, train_dataloader_hard_aug, val_dataloader, + loss_fn, optimizer_hard, epochs=hard_epochs, + lr_scheduler=scheduler_hard, save_best_val=True, + model_file=model_file) + total_train_loss.extend(train_loss) + total_val_loss.extend(val_loss) + + if args.save_model: + torch.save(model.state_dict(), model_file) + +if __name__ == "__main__": + main() diff --git a/heatseek/counting/dl/transforms.py b/heatseek/counting/dl/transforms.py new file mode 100644 index 0000000..4e37fd6 --- /dev/null +++ b/heatseek/counting/dl/transforms.py @@ -0,0 +1,117 @@ +from augmentations import compute_mean_std, MaskCompose, MaskRandomCrop, MaskToTensor, MaskNormalize, Mask2dMultiplyAndAddToBrightness, MaskContrast, MaskImgAug +import os + +def get_transforms(clips_dir): + """ + Build the three augmentation pipelines used in curriculum training. + + Computes per-dataset mean and std from the clips in `clips_dir` and + constructs none, easy, and hard transform pipelines of increasing + augmentation intensity. All three pipelines share the same crop size + and normalization strategy: mean-centering by the dataset mean (scaled + to [0,1]) with no std scaling (std=1.0), preserving the radiometric + intensity range. + + Augmentation progression: + none — crop and normalize only; used for initial warmup training. + easy — adds mild brightness/contrast jitter and imgaug transforms; + brightness multiply in (0.85, 1.0), additive shift in (-15, 0). + hard — same structure as easy but with wider jitter ranges; + brightness multiply in (0.7, 1.0), additive shift in (-25, 0), + contrast factor in (0.6, 1.2). + + Parameters: + clips_dir (str): Directory containing training clips, passed + to compute_mean_std() to derive normalization statistics. + + Returns: + none_data_transforms (MaskCompose): Minimal pipeline for warmup phase. + easy_data_transforms (MaskCompose): Moderate augmentation pipeline. + hard_data_transforms (MaskCompose): Heavy augmentation pipeline for + final training phase. + + Note: + std=1.0 means normalization only subtracts the mean without rescaling + by standard deviation. This is intentional for radiometric data where + the absolute intensity scale is meaningful and should be preserved + relative to the mean. + """ + + mean, std = compute_mean_std(clips_dir) + + none_data_transforms = MaskCompose([ + MaskRandomCrop(224), + MaskToTensor(), + MaskNormalize(mean=mean/255, std=1.0) + ]) + + easy_data_transforms = MaskCompose([ + MaskRandomCrop(224), + Mask2dMultiplyAndAddToBrightness(multiply=(0.85, 1.0), add=(-15, 0)), + MaskContrast(contrast_factor=(0.8, 1.2)), + MaskImgAug(), + MaskToTensor(), + MaskNormalize(mean=mean/255, std=1.0) + ]) + + hard_data_transforms = MaskCompose([ + MaskRandomCrop(224), + Mask2dMultiplyAndAddToBrightness(multiply=(0.7, 1.0), add=(-25, 0)), + MaskContrast(contrast_factor=(0.6, 1.2)), + MaskImgAug(), + MaskToTensor(), + MaskNormalize(mean=mean/255, std=1.0) + ]) + + return none_data_transforms, easy_data_transforms, hard_data_transforms + +def split_train_val(clips_dir, train_split=0.8, total_samples=5000): + """ + Split frame/mask pairs into training and validation sets, sampled evenly + across all clips. + Parameters: + clips_dir (str): Parent directory containing video subdirectories, + each containing clip folders with frames/ and masks/ + train_split (float): Fraction of samples for training. Default: 0.8 + total_samples (int): Total number of frame/mask pairs to use. Default: 5000 + Returns: + train_pairs (list of tuples): (frame_path, mask_path) pairs for training + val_pairs (list of tuples): (frame_path, mask_path) pairs for validation + """ + # Collect all valid frame/mask pairs across the full directory structure + all_pairs = [] + + for video_dir in sorted(os.listdir(clips_dir)): + video_path = os.path.join(clips_dir, video_dir) + + if not os.path.isdir(video_path): + continue + + for clip_dir in sorted(os.listdir(video_path)): + clip_path = os.path.join(video_path, clip_dir) + + if not os.path.isdir(clip_path): + continue + + frames_path = os.path.join(clip_path, 'frames') + masks_path = os.path.join(clip_path, 'masks') + + if not os.path.isdir(frames_path) or not os.path.isdir(masks_path): + continue + + for filename in sorted(os.listdir(frames_path)): + frame_path = os.path.join(frames_path, filename) + mask_path = os.path.join(masks_path, filename.replace('.jpg', '.png')) + + if os.path.exists(mask_path): + all_pairs.append((frame_path, mask_path)) + + if len(all_pairs) > total_samples: + step = len(all_pairs) / total_samples + all_pairs = [all_pairs[int(i * step)] for i in range(total_samples)] + + train_end = int(train_split * len(all_pairs)) + train_pairs = all_pairs[:train_end] + val_pairs = all_pairs[train_end:] + + return train_pairs, val_pairs diff --git a/heatseek/counting/dl/video_inference.py b/heatseek/counting/dl/video_inference.py new file mode 100644 index 0000000..cb5c345 --- /dev/null +++ b/heatseek/counting/dl/video_inference.py @@ -0,0 +1,180 @@ +import os +import time +import cv2 +import numpy as np +import torch +import torch.utils.data as data +from bat_seg_models import UNETTraditional +from augmentations import compute_mean_std, MaskCompose, MaskToTensor, MaskNormalize +import bat_functions +from BatIterableDataset import BatIterableDataset +import argparse + +def denorm_image(im, mean): + """ + Reverses mean normalization on a float image array and converts it to uint8. + + Adds back the per-channel mean (scaled to [0, 1]) to a normalized image, + rescales pixel values to [0, 255], clips to valid range, and casts to uint8. + + Parameters: + im (np.ndarray): Normalized image array of shape (H, W, C) with float + values expected in roughly [-1, 1] or [0, 1] range. + mean (np.ndarray or list): Per-channel mean values in [0, 255] range, + shape (C,), used during the original normalization step. + + Returns: + np.ndarray: Denormalized image as uint8 with pixel values in [0, 255], + same spatial shape as input. + """ + im += mean / 255 + im *= 255 + im = np.maximum(im, 0) + im = np.minimum(im, 255) + return im.astype(np.uint8) + +def load_model(model_file): + """ + Load a trained UNETTraditional model from a state dict file and prepare + it for inference. + + Automatically selects CUDA if available, otherwise falls back to CPU. + The model is set to evaluation mode (gradients and dropout disabled) + before being returned. + + Parameters: + model_file (str or path-like): Path to a saved state dict file + (.pth or .pt) produced by torch.save(). + + Returns: + model (UNETTraditional): Trained model in eval mode, moved to the + appropriate device (CUDA or CPU). + """ + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = UNETTraditional(1, 2, should_pad=False) + model.load_state_dict(torch.load(model_file, map_location=device)) + model.to(device) + model.train(False) + return model + +def get_augmentor(mean): + """ + Build an image augmentation pipeline that converts an image/mask pair + to tensors and applies mean normalization. + + Parameters: + mean (np.ndarray or float): Per-channel mean values in [0, 255] range, + divided by 255 internally to normalize to [0, 1] scale. + + Returns: + MaskCompose: A composed transform that applies, in order: + 1. MaskToTensor — converts image and mask to torch.Tensor. + 2. MaskNormalize — subtracts mean/255 from the image (std=1.0). + """ + return MaskCompose([ + MaskToTensor(), + MaskNormalize(mean=mean/255, std=1.0) + ]) + +def run_inference(video_path, output_folder, model, mean, + bat_prob_thresh=0.6, batch_size=2, + save_every_n_frames=18000, early_stop=None): + """Run model inference on a video file and save detections to disk.""" + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + os.makedirs(output_folder, exist_ok=True) + os.makedirs(os.path.join(output_folder, 'example-frames'), exist_ok=True) + + augmentor = get_augmentor(mean) + bat_dataset = BatIterableDataset([video_path], augmentor=augmentor) + dataloader = data.DataLoader(bat_dataset, batch_size=batch_size, + shuffle=False, num_workers=0, pin_memory=True) + + centers_list = [] + contours_list = [] + sizes_list = [] + rects_list = [] + num_frames = 0 + t0 = None + + for batch_ind, batch in enumerate(dataloader): + if batch_ind == 0: + print('started...') + t0 = time.time() + if early_stop and batch_ind >= early_stop: + break + + im_batch = batch['image'].to(device) + + with torch.no_grad(): + outputs = model(im_batch) + masks = (outputs[:, 1].cpu().numpy() > np.log(bat_prob_thresh)).astype(np.uint8) + + for ind, mask in enumerate(masks): + 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 save_every_n_frames and num_frames % save_every_n_frames == 0: + im_name = f'frame_{num_frames:07d}.jpg' + im_file = os.path.join(output_folder, 'example-frames', im_name) + im = denorm_image(np.squeeze(batch['image'][ind].numpy()).copy(), mean) + cv2.imwrite(im_file, im) + num_frames += 1 + + if batch_ind % 1000 == 0 and t0 is not None: + print(f'batch {batch_ind}, frame {num_frames}, time {time.time()-t0:.1f}s') + + total_time = time.time() - t0 + print(f'{total_time:.1f}s total, {total_time/max(batch_ind,1)/batch_size:.3f}s per frame') + bat_dataset.get_read_frame_info() + + return centers_list, contours_list, sizes_list, rects_list + +def save_detections(centers_list, contours_list, sizes_list, rects_list, + output_folder, num_contour_files=15): + """Save per-frame detections to disk.""" + 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)) + print(f'Saved detections for {len(centers_list)} frames to {output_folder}') + +def run_video_inference(video_path, output_folder, model_file, clips_dir, + bat_prob_thresh=0.6, batch_size=2, + save_every_n_frames=18000, early_stop=None): + """Top level function — load model, run inference, save detections.""" + model = load_model(model_file) + mean, _ = compute_mean_std(clips_dir) + centers_list, contours_list, sizes_list, rects_list = run_inference( + video_path, output_folder, model, mean, + bat_prob_thresh, batch_size, save_every_n_frames, early_stop) + save_detections(centers_list, contours_list, sizes_list, rects_list, output_folder) + +def main(): + parser = argparse.ArgumentParser(description='Get centers and contours from detections from model inference') + parser.add_argument('--vid_path', type=str, help='Path of video') + parser.add_argument('--output_folder', type=str, help='Path to the folder where you want your outputs saved') + parser.add_argument('--model_filepath', type=str, help='Filepath of trained model') + parser.add_argument('--clips_dir', type=str, help='Directory containing the clips (faster way to calculate mean and std)') + args = parser.parse_args() + run_video_inference(video_path=args.vid_path, output_folder=args.output_folder, model_file=args.model_filepath, clips_dir=args.clips_dir) + +if __name__ == '__main__': + main() + \ No newline at end of file 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", + ], + }, +) From d669e91fb5c2715b7f5ee4469ea86cb9c1671aa9 Mon Sep 17 00:00:00 2001 From: Akhil Nallacheruvu Date: Mon, 20 Jul 2026 08:19:29 -0700 Subject: [PATCH 30/30] Restructured dirs and keeping thresholding with configs --- .gitignore | 18 +- heatseek/counting/{bg_sub => }/CountLine.py | 18 +- heatseek/counting/README.md | 129 ++- heatseek/counting/{dl => }/bat_functions.py | 65 +- heatseek/counting/bg_sub/README.md | 30 - heatseek/counting/bg_sub/__init__.py | 0 heatseek/counting/bg_sub/bat_functions.py | 384 --------- heatseek/counting/bg_sub/crossing_tracks.py | 102 --- heatseek/counting/bg_sub/video_inference.py | 100 --- .../config/crossing_tracks_config.yaml | 2 + .../config/detections_to_tracks_config.yaml | 8 + .../config/video_inference_config.yaml | 10 + heatseek/counting/crossing_tracks.py | 274 +++++++ .../{bg_sub => }/detections_to_tracks.py | 171 +++- heatseek/counting/dl/BatIterableDataset.py | 153 ---- heatseek/counting/dl/CountLine.py | 72 -- heatseek/counting/dl/README.md | 79 -- heatseek/counting/dl/SegmentationDataset.py | 30 - heatseek/counting/dl/__init__.py | 0 heatseek/counting/dl/augmentations.py | 181 ----- heatseek/counting/dl/bat_seg_models.py | 359 --------- heatseek/counting/dl/crop_vids.py | 94 --- heatseek/counting/dl/crossing_tracks.py | 143 ---- heatseek/counting/dl/dataloaders.py | 16 - heatseek/counting/dl/detections_to_tracks.py | 222 ------ heatseek/counting/dl/extract_frames.py | 55 -- heatseek/counting/dl/koger_tracking.py | 744 ------------------ heatseek/counting/dl/optimizer.py | 108 --- heatseek/counting/dl/segment_frames.py | 88 --- heatseek/counting/dl/trainer.py | 248 ------ heatseek/counting/dl/transforms.py | 117 --- heatseek/counting/dl/video_inference.py | 180 ----- .../counting/{bg_sub => }/koger_tracking.py | 108 ++- heatseek/counting/video_inference.py | 225 ++++++ 34 files changed, 917 insertions(+), 3616 deletions(-) rename heatseek/counting/{bg_sub => }/CountLine.py (76%) rename heatseek/counting/{dl => }/bat_functions.py (84%) delete mode 100644 heatseek/counting/bg_sub/README.md delete mode 100644 heatseek/counting/bg_sub/__init__.py delete mode 100644 heatseek/counting/bg_sub/bat_functions.py delete mode 100644 heatseek/counting/bg_sub/crossing_tracks.py delete mode 100644 heatseek/counting/bg_sub/video_inference.py create mode 100644 heatseek/counting/config/crossing_tracks_config.yaml create mode 100644 heatseek/counting/config/detections_to_tracks_config.yaml create mode 100644 heatseek/counting/config/video_inference_config.yaml create mode 100644 heatseek/counting/crossing_tracks.py rename heatseek/counting/{bg_sub => }/detections_to_tracks.py (56%) delete mode 100644 heatseek/counting/dl/BatIterableDataset.py delete mode 100644 heatseek/counting/dl/CountLine.py delete mode 100644 heatseek/counting/dl/README.md delete mode 100644 heatseek/counting/dl/SegmentationDataset.py delete mode 100644 heatseek/counting/dl/__init__.py delete mode 100644 heatseek/counting/dl/augmentations.py delete mode 100644 heatseek/counting/dl/bat_seg_models.py delete mode 100644 heatseek/counting/dl/crop_vids.py delete mode 100644 heatseek/counting/dl/crossing_tracks.py delete mode 100644 heatseek/counting/dl/dataloaders.py delete mode 100644 heatseek/counting/dl/detections_to_tracks.py delete mode 100644 heatseek/counting/dl/extract_frames.py delete mode 100644 heatseek/counting/dl/koger_tracking.py delete mode 100644 heatseek/counting/dl/optimizer.py delete mode 100644 heatseek/counting/dl/segment_frames.py delete mode 100644 heatseek/counting/dl/trainer.py delete mode 100644 heatseek/counting/dl/transforms.py delete mode 100644 heatseek/counting/dl/video_inference.py rename heatseek/counting/{bg_sub => }/koger_tracking.py (89%) create mode 100644 heatseek/counting/video_inference.py diff --git a/.gitignore b/.gitignore index 1488530..df7e3cf 100644 --- a/.gitignore +++ b/.gitignore @@ -80,11 +80,13 @@ Thumbs.db # Data Folders heatseek/counting/clips -heatseek/counting/bg_sub/detections_413_full_thermal_vid -heatseek/counting/bg_sub/detections_414_full_thermal_vid -heatseek/counting/dl/detections_413_full_train -heatseek/counting/dl/detections_414_full_train -heatseek/counting/dl/detections_413_single_train -heatseek/counting/dl/detections_414_single_train -heatseek/counting/dl/models_full_train -heatseek/counting/dl/models_single_train +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/heatseek/counting/bg_sub/CountLine.py b/heatseek/counting/CountLine.py similarity index 76% rename from heatseek/counting/bg_sub/CountLine.py rename to heatseek/counting/CountLine.py index fefc9db..f94d927 100644 --- a/heatseek/counting/bg_sub/CountLine.py +++ b/heatseek/counting/CountLine.py @@ -2,7 +2,7 @@ class CountLine(): ''' - Counts everytime a bat crosses a line defined in space (y = constant) + 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): @@ -25,6 +25,12 @@ def __init__(self, line_value, line_dim=1, total_frames=None): 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): ''' @@ -42,8 +48,9 @@ def is_crossing(self, track, track_ind): 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 = (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: @@ -61,8 +68,9 @@ def is_crossing(self, track, track_ind): 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 = (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 diff --git a/heatseek/counting/README.md b/heatseek/counting/README.md index a9af25c..dd319b7 100644 --- a/heatseek/counting/README.md +++ b/heatseek/counting/README.md @@ -1,3 +1,128 @@ -# Bat Track Counting +# Counting Pipeline +This pipeline consists of 3 stages: -For more information about counting bat tracks using deep learning, navigate to the dl directory. For more information about counting bat tracks using temperature-based background segmentation on a video taken on a FLIR Boson radiometric thermal camera with globally normalized pixel values, refer to the bg_sub directory. \ No newline at end of file +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/dl/bat_functions.py b/heatseek/counting/bat_functions.py similarity index 84% rename from heatseek/counting/dl/bat_functions.py rename to heatseek/counting/bat_functions.py index 38bb6b4..d8efa36 100644 --- a/heatseek/counting/dl/bat_functions.py +++ b/heatseek/counting/bat_functions.py @@ -5,8 +5,8 @@ import pandas as pd from scipy import signal from scipy.optimize import linear_sum_assignment -from CountLine import CountLine -import koger_tracking as ktf +from heatseek.counting.CountLine import CountLine +import heatseek.counting.koger_tracking as ktf def get_blob_info(binary_image, background=None, size_threshold=0): @@ -100,7 +100,8 @@ def add_all_points_as_new_tracks(raw_track_list, positions, contours, def find_tracks(first_frame_ind, positions, contours_files=None, contours_list=None, sizes_list=None, max_frame=None, verbose=True, - tracks_file=None): + 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: @@ -115,22 +116,12 @@ def find_tracks(first_frame_ind, positions, """ raw_track_list = [] - - max_distance_threshold = 30 - max_distance_threshold_noise = 30 - min_distance_threshold = 0 - max_unseen_time = 2 - min_new_track_distance = 3 + min_new_track_distance = 1 min_distance_big = 30 -# #Create initial tracks based on the objects in the first frame -# raw_track_list = add_all_points_as_new_tracks( -# raw_track_list, positions[0], contours_list[0], sizes_list0, noise=0 -# ) - #try to connect points to the next frame if max_frame is None: - max_frame = len(positions) + max_frame = len(positions) #positions = centers contours_file_ind = 0 previous_contours_seen = 0 @@ -227,31 +218,6 @@ def find_tracks(first_frame_ind, positions, row_ind, col_ind = linear_sum_assignment(np.log(distance + 1)) -# for active_ind, track_ind in enumerate(active_list): -# if active_ind in row_ind: -# row_count = np.where(row_ind == active_ind)[0] -# raw_track_list[track_ind]['debug'].append( -# '{} dist {}, best {}'.format( -# frame_ind, -# distance[row_ind[row_count], -# col_ind[row_count]], -# np.min(distance[row_ind[row_count], -# :]) -# ) -# ) -# best_col = np.argmin(distance[row_ind[row_count], -# :]) -# row_count = np.where(col_ind == best_col)[0] -# raw_track_list[track_ind]['debug'].append( -# '{} row_ind {} col {} dist {} track {}'.format( -# frame_ind, row_ind[row_count], -# col_ind[row_count], -# distance[row_ind[row_count], -# col_ind[row_count]], -# active_list[row_ind[row_count][0]]) -# ) - - # 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 @@ -355,24 +321,25 @@ def measure_crossing_bats(track_list, frame_height=None, frame_width=None, 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) - # result is 1 if forward crossing -1 is backward crossing + if count_out: - if out_frame_num: - crossing_track_list[-1]['crossed'] = out_frame_num * out_result - else: - crossing_track_list[-1]['crossed'] = 0 + crossing_track_list[-1]['crossed'] = out_result + crossing_track_list[-1]['crossed_frame'] = out_frame_num if count_across: - if across_frame_num: - crossing_track_list[-1]['across_crossed'] = across_frame_num * across_result - else: - crossing_track_list[-1]['across_crossed'] = 0 + 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) diff --git a/heatseek/counting/bg_sub/README.md b/heatseek/counting/bg_sub/README.md deleted file mode 100644 index 578e55b..0000000 --- a/heatseek/counting/bg_sub/README.md +++ /dev/null @@ -1,30 +0,0 @@ -# Bat Track Counting Pipeline (Temperature-based Background Subtraction) -## Overview: -This file contains information related to generating crossing tracks using temperature-based background subtraction. This method skips the model training and relies on the pixel values from a radiometric camera image and their converted temperature value.The camera is assumed to be a FLIR Boson, and therefore, the raw temperature values are assumed to be encoded in centikelvins. Given this data, inference can directly be performed on the video to generate detection centroids, contours, and boundaries of bat blobs. From there, it can generate all the information related to the bat tracks. - -## 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 --vid_path path/to/inference/video --output_folder desired/path/of/output/folder - ``` - -2) Convert the detection data into raw tracks using `detections_to_tracks.py` - ```bash - cd path/to/working/dir - python -m detections_to_tracks --output_folder path/to/folder/with/detections - #should ideally be the same as output_folder in video_inference - ``` - -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 --raw_tracks_file path/to/raw/tracks/file --crossing_tracks_file desired/path/of/crossing_tracks/file - ``` - - An additional parameter must be added at the end of the second command. If the user wishes to generate a file containing tracks vertically crossing a horizontal midline and its corresponding visualization, they must add the additional parameter `--count_out`. If they instead wish to generate a file containing tracks horizontally crossing a vertical midline, they should use the parameter `--count_across` instead. - - - - - diff --git a/heatseek/counting/bg_sub/__init__.py b/heatseek/counting/bg_sub/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/heatseek/counting/bg_sub/bat_functions.py b/heatseek/counting/bg_sub/bat_functions.py deleted file mode 100644 index 38bb6b4..0000000 --- a/heatseek/counting/bg_sub/bat_functions.py +++ /dev/null @@ -1,384 +0,0 @@ -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 CountLine import CountLine -import 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): - """ 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 = [] - - max_distance_threshold = 30 - max_distance_threshold_noise = 30 - min_distance_threshold = 0 - max_unseen_time = 2 - min_new_track_distance = 3 - min_distance_big = 30 - -# #Create initial tracks based on the objects in the first frame -# raw_track_list = add_all_points_as_new_tracks( -# raw_track_list, positions[0], contours_list[0], sizes_list0, noise=0 -# ) - - #try to connect points to the next frame - if max_frame is None: - max_frame = len(positions) - - 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)) - -# for active_ind, track_ind in enumerate(active_list): -# if active_ind in row_ind: -# row_count = np.where(row_ind == active_ind)[0] -# raw_track_list[track_ind]['debug'].append( -# '{} dist {}, best {}'.format( -# frame_ind, -# distance[row_ind[row_count], -# col_ind[row_count]], -# np.min(distance[row_ind[row_count], -# :]) -# ) -# ) -# best_col = np.argmin(distance[row_ind[row_count], -# :]) -# row_count = np.where(col_ind == best_col)[0] -# raw_track_list[track_ind]['debug'].append( -# '{} row_ind {} col {} dist {} track {}'.format( -# frame_ind, row_ind[row_count], -# col_ind[row_count], -# distance[row_ind[row_count], -# col_ind[row_count]], -# active_list[row_ind[row_count][0]]) -# ) - - - # 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) - # result is 1 if forward crossing -1 is backward crossing - if count_out: - if out_frame_num: - crossing_track_list[-1]['crossed'] = out_frame_num * out_result - else: - crossing_track_list[-1]['crossed'] = 0 - if count_across: - if across_frame_num: - crossing_track_list[-1]['across_crossed'] = across_frame_num * across_result - else: - crossing_track_list[-1]['across_crossed'] = 0 - 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/bg_sub/crossing_tracks.py b/heatseek/counting/bg_sub/crossing_tracks.py deleted file mode 100644 index b04ef90..0000000 --- a/heatseek/counting/bg_sub/crossing_tracks.py +++ /dev/null @@ -1,102 +0,0 @@ -from bat_functions import threshold_short_tracks, measure_crossing_bats -import numpy as np -import matplotlib.pyplot as plt -import argparse - -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 - - crossed_key = 'crossed' if count_out else 'across_crossed' - - 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()}') - print(f'Forward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] > 0)}') - print(f'Backward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] < 0)}') - - 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] - color = 'blue' if track[crossed_key] > 0 else 'red' - 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() - direction_labels = 'blue=downward, red=upward' - else: - ax.axvline(x=midline_x, color='green', linewidth=2, linestyle='--', label='midline') - direction_labels = 'blue=rightward, red=leftward' - - ax.set_title(f'Sample of {num_tracks} crossing tracks\n{direction_labels}') - ax.legend() - - ax = axes[1] - crossing_frames = [abs(t[crossed_key]) for t in crossing_tracks] - ax.hist(crossing_frames, bins=100) - ax.set_xlabel('Frame number') - ax.set_ylabel('Number of crossings') - ax.set_title('Bat crossings over time') - - print('Crossing frames:') - for t in crossing_tracks: - direction = 'forward' if t[crossed_key] > 0 else 'backward' - print(f' frame {abs(t[crossed_key])}: {direction}') - - plt.tight_layout() - plt.show() - -def main(): - parser = argparse.ArgumentParser(description='Get bat crossing tracks across the video') - parser.add_argument('--raw_tracks_file', type=str, help='Path to raw tracks file') - parser.add_argument('--crossing_tracks_file', type=str, help='Intended path for the crossing tracks 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() - - save_crossing_tracks_from_raw_tracks(args.raw_tracks_file, args.crossing_tracks_file, args.count_across, args.count_out) - visualize_crossing_tracks(args.crossing_tracks_file, count_out=args.count_out, count_across=args.count_across) - -if __name__ == '__main__': - main() diff --git a/heatseek/counting/bg_sub/video_inference.py b/heatseek/counting/bg_sub/video_inference.py deleted file mode 100644 index 9b51da2..0000000 --- a/heatseek/counting/bg_sub/video_inference.py +++ /dev/null @@ -1,100 +0,0 @@ -import os -import cv2 -import logging -import numpy as np -import argparse -import bat_functions - -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') - -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.') - - -def get_mask(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 - -def run_inference(video_file, min_val=28000, max_val=32000, - threshold=5.0, process_every_n_frames=1): - """Process every frame, storing only detection metadata — not images.""" - centers_list = [] - contours_list = [] - sizes_list = [] - rects_list = [] - - for record in iter_frames(video_file, process_every_n_frames): - mask = get_mask(record['image'], min_val, max_val, threshold) - 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) - # record['image'] goes out of scope here and gets GC'd - - if len(centers_list) % 1000 == 0: - logging.info(f'Processed {len(centers_list)} frames.') - - logging.info(f'Processed {len(centers_list)} frames.') - return centers_list, contours_list, sizes_list, rects_list - -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)) - print(f'Saved detections for {len(centers_list)} frames to {output_folder}') - -def main(): - parser = argparse.ArgumentParser(description='Process a thermal video and detect objects.') - parser.add_argument('--vid_path', help='Path to the input video file.') - parser.add_argument('--output_folder', help='Path to the output folder.') - args = parser.parse_args() - - centers_list, contours_list, sizes_list, rects_list = run_inference(args.vid_path) - save_detections(centers_list, contours_list, sizes_list, rects_list, args.output_folder) - -if __name__ == '__main__': - main() 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/bg_sub/detections_to_tracks.py b/heatseek/counting/detections_to_tracks.py similarity index 56% rename from heatseek/counting/bg_sub/detections_to_tracks.py rename to heatseek/counting/detections_to_tracks.py index cc80404..4fed421 100644 --- a/heatseek/counting/bg_sub/detections_to_tracks.py +++ b/heatseek/counting/detections_to_tracks.py @@ -1,11 +1,15 @@ import numpy as np import os import glob -import bat_functions as kbf +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): +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. @@ -23,6 +27,12 @@ def track(camera_dict): 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: @@ -39,18 +49,21 @@ def track(camera_dict): 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, + raw_tracks = kbf.find_tracks(first_frame, centers, contours_files=contours_files, sizes_list=sizes, tracks_file=tracks_file, - max_frame=max_frame) + 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=60, overlap_seconds=15): +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. @@ -72,7 +85,7 @@ def build_camera_dicts(output_folder, num_groups=10, fps=60, overlap_seconds=15) 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 15. + segments. Defaults to 10. Returns: list[dict]: One dict per unprocessed segment, each containing: @@ -84,6 +97,12 @@ def build_camera_dicts(output_folder, num_groups=10, fps=60, overlap_seconds=15) """ 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 = [] @@ -166,6 +185,7 @@ def combine_overlapping_tracks(output_folder, first_group=0, last_group=None, sa 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}') @@ -189,7 +209,8 @@ def combine_tracks(output_folder): 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, processes=5): +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. @@ -206,17 +227,143 @@ def run_tracking(output_folder, processes=5): Returns: None """ - camera_dicts = build_camera_dicts(output_folder) + 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.map(track, camera_dicts) + 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('--output_folder', type=str, help='Path of output folder containing the centers, contours, etc from video inference') + parser.add_argument('--config', help='Path to Raw Tracking Config File') args = parser.parse_args() - run_tracking(args.output_folder) - combine_tracks(args.output_folder) + 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/dl/BatIterableDataset.py b/heatseek/counting/dl/BatIterableDataset.py deleted file mode 100644 index 68e11f3..0000000 --- a/heatseek/counting/dl/BatIterableDataset.py +++ /dev/null @@ -1,153 +0,0 @@ -from torch.utils.data import IterableDataset -import cv2 - -class BatIterableDataset(IterableDataset): - """ An iterable PyTorch dataset that streams grayscale frames from one or - more video files for bat detection inference. - - Frames are read sequentially using OpenCV, converted to grayscale, and - optionally passed through an augmentation pipeline. Videos are processed - one at a time in the order provided; when a video is exhausted, the next - one starts automatically. - - Consecutive read failures are tolerated up to max_bad_reads before the - current video is closed and iteration ends. - - Attributes: - video_files (list[str]): Ordered list of video file paths. - augmentor (MaskCompose or None): Optional transform applied to each frame. - max_bad_reads (int): Max consecutive failed reads before closing a video. - total_frames_read (int): Running count of successfully read frames. - total_bad_reads (int): Running count of failed frame reads. - video_number (int): Index of the currently open video in video_files. - more_frames (bool): Whether the current video still has frames. - """ - - def __init__(self, video_files, augmentor=None, max_bad_reads=300): - """ - Initialize the dataset and open the first video file. - - Parameters: - video_files (list[str]): Ordered list of paths to video files. - The first file is opened immediately on construction. - augmentor (MaskCompose, optional): Transform pipeline applied to - each frame dict before yielding. Defaults to None. - max_bad_reads (int): Maximum consecutive failed cv2 reads before - the current video is considered exhausted. Defaults to 300. - - Raises: - AssertionError: If the first video file cannot be opened by OpenCV. - """ - self.vid_cap = cv2.VideoCapture(video_files[0]) - self.video_files = video_files - assert self.vid_cap.isOpened() - self.more_frames = True - self.max_bad_reads = max_bad_reads - self.total_frames_read = 0 - self.total_bad_reads = 0 - self.augmentor = augmentor - self.video_number = 0 - - def more_videos(self): - """ - Check whether there are additional unprocessed videos in the queue. - - Returns: - bool: True if video_number is within bounds of video_files. - """ - return self.video_number < len(self.video_files) - - def start_next_video(self): - """ - Release the current video capture and open the next video in the queue. - - Increments video_number, releases any open capture, and opens a new - cv2.VideoCapture for the next file. Does nothing if all videos have - already been processed. Prints a status message and frame read info - when a new video starts. - """ - if self.vid_cap.isOpened(): - self.vid_cap.release() - self.video_number += 1 - if self.video_number < len(self.video_files): - print('starting new video') - print(self.get_read_frame_info()) - self.vid_cap = cv2.VideoCapture(self.video_files[self.video_number]) - - def video_generator(self): - """ - Generator that yields preprocessed frames across all video files. - - For each video, reads frames sequentially. On a successful read the - frame is converted from BGR to grayscale, cropped by 2 pixels on each - edge, wrapped in a dict as {'image': frame}, and optionally transformed - by the augmentor before being yielded. - - Consecutive failed reads are counted; if max_bad_reads is reached the - video is released and iteration moves on. Transitions between videos - are handled automatically via start_next_video(). - - Yields: - dict or augmentor output: {'image': np.ndarray} if no augmentor, - otherwise the transformed output of augmentor({'image': frame}). - """ - - while(self.vid_cap.isOpened() or self.more_videos()): - if not self.vid_cap.isOpened(): - self.start_next_video() - good_read = False - num_bad_reads = 0 - while (not good_read and (num_bad_reads < self.max_bad_reads)): - grabbed, frame = self.vid_cap.read() - if grabbed: - good_read = True - self.total_frames_read += 1 - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - frame = {'image': frame[2:-2, 2:-2]} - if self.augmentor: - frame = self.augmentor(frame) - yield frame - else: - num_bad_reads += 1 - self.total_bad_reads += 1 - if not good_read: - self.vid_cap.release() - print("video capture closed") - - def __iter__(self): - """ - Return the frame generator, making this dataset compatible with - PyTorch DataLoader. - - Returns: - generator: The video_generator() iterator. - """ - return self.video_generator() - - def __del__(self): - """ - Release the OpenCV video capture on object destruction to free - file handles and decoder resources. - """ - if self.vid_cap.isOpened(): - self.vid_cap.release() - - def is_more_frames(self): - """ - Check whether the current video capture is still open and readable. - - Returns: - bool: True if vid_cap is currently open. - """ - return self.vid_cap.isOpened() - - def get_read_frame_info(self): - """ - Print a summary of total frames read and total bad reads so far. - - Example output: - 142 frames have been read with 3 bad reads - """ - print('{} frames have been read with {} bad reads'.format( - self.total_frames_read, self.total_bad_reads)) - \ No newline at end of file diff --git a/heatseek/counting/dl/CountLine.py b/heatseek/counting/dl/CountLine.py deleted file mode 100644 index b66ee9d..0000000 --- a/heatseek/counting/dl/CountLine.py +++ /dev/null @@ -1,72 +0,0 @@ -import numpy as np - -class CountLine(): - ''' - Counts everytime a bat crosses a line defined in space (y = 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 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']) - # 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']) - 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/dl/README.md b/heatseek/counting/dl/README.md deleted file mode 100644 index 70583a3..0000000 --- a/heatseek/counting/dl/README.md +++ /dev/null @@ -1,79 +0,0 @@ -# Bat Track Counting Pipeline (Deep Learning) -## Overview: -This file contains information related to generating crossing tracks using deep learning. This is the method that was originally implemented by Koger et al, but modified to fit our platform and data. Deep learning is useful when the data contains noisy objects (e.g. small bats in an RGB image). - -## Pipeline: -1) In order to train a model, we must create a training set consisting of images and segmentation masks as the model that is used in this pipeline is a UNet model. The steps for this stage are: - - 1) `crop_vids.py` - Cut the long videos into short clips of motion using the command: - ```bash - cd path/to/working/dir - python -m crop_vids --input path/to/input/file --output_dir path/to/desired/output/folder - ``` - There are additional parameters that may be set. Please refer to file for fine-tuning. - - The clips must go through human validation to ensure there are no false positives for motion. Once this is done, they must be arranged in a directory structure that looks like this: - - ``` - clips - |---- vid1_clips - |---------------- clip1.mp4 - |---------------- clip2.mp4 - . - . - . - |----- vid2_clips - |---------------- clip1.mp4 - |---------------- clip2.mp4 - . - . - . - . - . - . - - 2) `extract_frames.py` - Extract frames from the clips - ```bash - cd path/to/working/dir - python -m extract_frames --clips_dir path/to/directory/containing/clips - ``` - There are additional parameters that may be set. Please refer to file for fine-tuning. - - 3) `segment_frames.py` - Get masks for the frames using temperature-based background segmentation (assumes availability of globally normalized video with minimum and maximum pixel values corresponding to temperature from a radiometric thermal camera) - ```bash - cd path/to/working/dir - python -m segment_frames --clips_dir path/to/directory/containing/clips - ``` - There are additional parameters that may be set. Please refer to file for fine-tuning. - -2) Train the model on the training img, mask pairs using `trainer.py` - ```bash - cd path/to/working/dir - python -m trainer --clips_dir path/to/directory/containing/clips - ``` - There are additional parameters that may be set. Please refer to file for fine-tuning. - -3) 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 --vid_path path/to/inference/video --output_folder desired/path/of/output/folder --model_filepath path/to/model.tar --clips_dir path/to/directory/containing/clips - ``` - -4) Convert the detection data into raw tracks using `detections_to_tracks.py` - ```bash - cd path/to/working/dir - python -m detections_to_tracks --output_folder path/to/folder/with/detections - #should ideally be the same as output_folder in video_inference - ``` - -5) 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 --raw_tracks_file path/to/raw/tracks/file --crossing_tracks_file desired/path/of/crossing_tracks/file - ``` - - An additional parameter must be added at the end of the second command. If the user wishes to generate a file containing tracks vertically crossing a horizontal midline and its corresponding visualization, they must add the additional parameter `--count_out`. If they instead wish to generate a file containing tracks horizontally crossing a vertical midline, they should use the parameter `--count_across` instead. - - - - \ No newline at end of file diff --git a/heatseek/counting/dl/SegmentationDataset.py b/heatseek/counting/dl/SegmentationDataset.py deleted file mode 100644 index 63179e2..0000000 --- a/heatseek/counting/dl/SegmentationDataset.py +++ /dev/null @@ -1,30 +0,0 @@ -import torch.utils.data as data -import cv2 - -class SegmentationDataset(data.Dataset): - def __init__(self, pairs, transform=None, keep_orig=False): - """ - Args: - pairs (list of tuples): list of (frame_path, mask_path) pairs - transform (callable, optional): Optional transforms to be applied on a sample - keep_orig (bool): keep un-transformed image - """ - self.img_files = [p[0] for p in pairs] - self.mask_files = [p[1] for p in pairs] - self.keep_orig = keep_orig - self.transform = transform - - def __getitem__(self, index): - img_path = self.img_files[index] - mask_path = self.mask_files[index] - image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) - mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) - sample = {'image': image[2:-2, 2:-2], 'mask': mask[2:-2, 2:-2]} - if self.keep_orig: - sample['orig'] = image - if self.transform: - sample = self.transform(sample) - return sample - - def __len__(self): - return len(self.img_files) diff --git a/heatseek/counting/dl/__init__.py b/heatseek/counting/dl/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/heatseek/counting/dl/augmentations.py b/heatseek/counting/dl/augmentations.py deleted file mode 100644 index 30eab98..0000000 --- a/heatseek/counting/dl/augmentations.py +++ /dev/null @@ -1,181 +0,0 @@ -import glob -import os -import cv2 -import numpy as np -import torch -import torchvision.transforms.functional as TF -import random - -def compute_mean_std(clips_dir, sample_size=500): - """ - Compute the mean and standard deviation of pixels in frames across the video clips - Parameters: - clips_dir (str) : path to the directory containing the clips - sample_size (int): number of frames to sample for computation - Returns: - mean (np.float64): the mean of the pixels - std (np.float64) : the standard deviation of values across pixels - """ - pixel_sum = 0.0 - pixel_sq_sum = 0.0 - pixel_count = 0 - - image_paths = glob.glob(os.path.join(clips_dir, '*', '*', 'frames', '*.jpg')) - - if len(image_paths) > sample_size: - image_paths = random.sample(image_paths, sample_size) - - for image_path in image_paths: - img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float64) - pixel_sum += img.sum() - pixel_sq_sum += (img ** 2).sum() - pixel_count += img.size - - mean = pixel_sum / pixel_count - std = np.sqrt(pixel_sq_sum / pixel_count - mean ** 2) - return mean, std - -class MaskToTensor: - def __call__(self, sample): - im = sample['image'] - im_tensor = torch.from_numpy(im).float() / 255 - im_tensor = torch.unsqueeze(im_tensor, 0) - sample['image'] = im_tensor - - if 'mask' in sample.keys(): - mask = sample['mask'] - mask_tensor = torch.from_numpy(np.asarray(mask, np.int64)) - mask_tensor = mask_tensor // 255 - sample['mask'] = mask_tensor - - return sample - -class MaskNormalize: - def __init__(self, mean, std): - """Normalize image, leave mask unchanged""" - self.mean = mean - self.std = std - - def __call__(self, sample): - im = sample['image'] - im_norm = TF.normalize(im, self.mean, self.std) - - sample['image'] = im_norm - return sample - -class MaskCompose: - def __init__(self, transform_list): - """Chain together transforms in transform. - Expect transform on both image and mask in dict - - Args: - transform_list (list): list of custom augmentations - """ - self.transform_list = transform_list - - def __call__(self, sample): - for transform in self.transform_list: - sample = transform(sample) - return sample - -class MaskImgAug: - def __call__(self, sample): - im = sample['image'].astype(np.uint8) - - # gaussian blur with 40% probability - if np.random.random() < 0.4: - sigma = np.random.uniform(0.0, 6.0) - ksize = int(6 * sigma + 1) - if ksize % 2 == 0: - ksize += 1 - im = cv2.GaussianBlur(im, (ksize, ksize), sigma) - - # poisson noise with 40% probability - if np.random.random() < 0.4: - lam = np.random.uniform(0, 30) - noise = np.random.poisson(lam, im.shape).astype(np.float64) - im = np.clip(im.astype(np.float64) + noise, 0, 255).astype(np.uint8) - - sample['image'] = im.astype(float) - return sample - -class MaskRandomCrop: - """Rotate by one of the given angles.""" - - def __init__(self, crop_size): - # size of square crop - self.crop_size = crop_size - - def __call__(self, sample): - im = sample['image'] - mask = sample['mask'] - - im_height = im.shape[0] - im_width = im.shape[1] - top = np.random.randint(im_height - self.crop_size) - left = np.random.randint(im_width - self.crop_size) - - im_crop = im[top:top+self.crop_size, left:left+self.crop_size] - mask_crop = mask[top:top+self.crop_size, left:left+self.crop_size] - - sample['image'] = im_crop - sample['mask'] = mask_crop - - return sample - -class Mask2dMultiplyAndAddToBrightness(): - def __init__(self, add, multiply): - """Add and multiply image values by given amount randomly within range. - - Expects image to be between 0 and 255. - - Args: - add: either number of tuple, if tuple randomly choose from range - multiply: either number ot tuple, if tuple randomly choose from range - """ - self.add = add - self.multiply = multiply - self.rng = np.random.default_rng() - - def __call__(self, sample): - im = sample['image'].astype(np.float64) - if isinstance(self.add, tuple): - add_val = self.rng.uniform(self.add[0], self.add[1]) - else: - add_val = self.add - if isinstance(self.multiply, tuple): - multiply_val = self.rng.uniform(self.multiply[0], self.multiply[1]) - else: - multiply_val = self.multiply - im += add_val - im *= multiply_val - im = np.maximum(im, 0) - im = np.minimum(im, 255) - - sample['image'] = im - - return sample - -class MaskContrast: - def __init__(self, contrast_factor): - """Change image contrast, leave mask unchanged - - Args: - contrast_factor: single value or tuple""" - self.contrast_factor = contrast_factor - self.rng = np.random.default_rng() - - def __call__(self, sample): - im = sample['image'] - if isinstance(self.contrast_factor, tuple): - contrast_factor = self.rng.uniform(self.contrast_factor[0], self.contrast_factor[1]) - else: - contrast_factor = self.contrast_factor - im_height, im_width = im.shape - aprox_mean = np.mean([im[0,0], im[-1, 0], im[0, -1], im[-1, -1], - im[im_height//2, im_width//2], im[-im_height//2, -im_width//2]]) - im_contrast = aprox_mean + contrast_factor * (im - aprox_mean) - im_contrast = np.maximum(im_contrast, 0) - im_contrast = np.minimum(im_contrast, 255) - sample['image'] = im_contrast - return sample diff --git a/heatseek/counting/dl/bat_seg_models.py b/heatseek/counting/dl/bat_seg_models.py deleted file mode 100644 index a3e0f07..0000000 --- a/heatseek/counting/dl/bat_seg_models.py +++ /dev/null @@ -1,359 +0,0 @@ -import torch -import torch.nn as nn - -class SuperSimpleSemSegNet(nn.Module): - def __init__(self, in_channel, out_channel): - super().__init__() - self.conv1 = torch.nn.Conv2d(in_channel, out_channel, - kernel_size=3, padding=1, stride=1) - self.ReLU = torch.nn.ReLU() - self.softmax = torch.nn.LogSoftmax(dim=1) - - def forward(self, x): - x = self.conv1(x) - x = self.ReLU(x) - x = self.softmax(x) - - return x - -class ThreeLayerSemSegNetWideView(nn.Module): - def __init__(self, in_channel, out_channel): - super().__init__() - self.conv1 = torch.nn.Conv2d(in_channel, 6, kernel_size=3, padding=1, stride=1) - self.conv1d100 = torch.nn.Conv2d(in_channel, 2, kernel_size=3, padding=101, - stride=1, dilation=101) - self.conv2d1 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=2, stride=1, dilation=2) - self.conv2d5 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=6, stride=1, dilation=6) - self.conv3 = torch.nn.Conv2d(8, out_channel, kernel_size=3, padding=1, stride=1) - self.ReLU1 = torch.nn.ReLU() - self.ReLU2 = torch.nn.ReLU() - self.softmax = torch.nn.LogSoftmax(dim=1) - self.batchnorm1 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - self.batchnorm2 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - - - def forward(self, x): - x1 = self.conv1(x) - x2 = self.conv1d100(x) - x = torch.cat((x1, x2), dim=1) - x = self.batchnorm1(x) - x = self.ReLU1(x) - x1 = self.conv2d1(x) - x2 = self.conv2d5(x) - x = torch.cat((x1, x2), dim=1) - x = self.batchnorm2(x) - x = self.ReLU2(x) - x = self.conv3(x) - x = self.softmax(x) - - return x - -class ThreeLayerSemSegNetWideViewHighDim(nn.Module): - """Each layer has more channels than the standard model""" - def __init__(self, in_channel, out_channel): - super().__init__() - self.conv1 = torch.nn.Conv2d(in_channel, 12, kernel_size=3, padding=1, stride=1) - self.conv1d100 = torch.nn.Conv2d(in_channel, 4, kernel_size=3, padding=101, - stride=1, dilation=101) - self.conv2d1 = torch.nn.Conv2d(16, 8, kernel_size=3, padding=2, stride=1, dilation=2) - self.conv2d5 = torch.nn.Conv2d(16, 8, kernel_size=3, padding=6, stride=1, dilation=6) - self.conv3 = torch.nn.Conv2d(16, out_channel, kernel_size=3, padding=1, stride=1) - self.ReLU1 = torch.nn.ReLU() - self.ReLU2 = torch.nn.ReLU() - self.softmax = torch.nn.LogSoftmax(dim=1) - self.batchnorm1 = torch.nn.BatchNorm2d(16, track_running_stats=False, momentum=1.0) - self.batchnorm2 = torch.nn.BatchNorm2d(16, track_running_stats=False, momentum=1.0) - - - def forward(self, x): - x1 = self.conv1(x) - x2 = self.conv1d100(x) - x = torch.cat((x1, x2), dim=1) - x = self.batchnorm1(x) - x = self.ReLU1(x) - x1 = self.conv2d1(x) - x2 = self.conv2d5(x) - x = torch.cat((x1, x2), dim=1) - x = self.batchnorm2(x) - x = self.ReLU2(x) - x = self.conv3(x) - x = self.softmax(x) - - return x - -class FourLayerSemSegNetWideView(nn.Module): - def __init__(self, in_channel, out_channel): - super().__init__() - self.conv1 = torch.nn.Conv2d(in_channel, 6, kernel_size=3, padding=1, stride=1) - self.conv1d100 = torch.nn.Conv2d(in_channel, 2, kernel_size=3, padding=101, - stride=1, dilation=101) - self.conv2d1 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=2, stride=1, dilation=2) - self.conv2d5 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=6, stride=1, dilation=6) - self.conv3d0 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=1, stride=1) - self.conv3d3 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=4, stride=1, dilation=4) - self.conv4 = torch.nn.Conv2d(8, out_channel, kernel_size=3, padding=1, stride=1) - self.ReLU1 = torch.nn.ReLU() - self.ReLU2 = torch.nn.ReLU() - self.ReLU3 = torch.nn.ReLU() - self.softmax = torch.nn.LogSoftmax(dim=1) - self.batchnorm1 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - self.batchnorm2 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - self.batchnorm3 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - - - def forward(self, x): - x1a = self.conv1(x) - x1b = self.conv1d100(x) - x1 = torch.cat((x1a, x1b), dim=1) - x1 = self.batchnorm1(x1) - x1 = self.ReLU1(x1) - x2a = self.conv2d1(x1) - x2b = self.conv2d5(x1) - x2 = torch.cat((x2a, x2b), dim=1) - x2 = self.batchnorm2(x2) - x2 = self.ReLU2(x2) - x3a = self.conv3d0(x2) - x3b = self.conv3d3(x2) - x3 = torch.cat((x3a, x3b), dim=1) - x3 = self.batchnorm3(x3) - x3 = self.ReLU3(x3) - x4 = self.conv4(x3) - xout = self.softmax(x4) - - return xout - -class ThreeLayerSemSegNet(nn.Module): - def __init__(self, in_channel, out_channel): - super().__init__() - self.conv1 = torch.nn.Conv2d(in_channel, 8, kernel_size=3, padding=1, stride=1) - self.conv2d1 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=2, stride=1, dilation=2) - self.conv2d5 = torch.nn.Conv2d(8, 4, kernel_size=3, padding=6, stride=1, dilation=6) - self.conv3 = torch.nn.Conv2d(8, out_channel, kernel_size=3, padding=1, stride=1) - self.ReLU1 = torch.nn.ReLU() - self.ReLU2 = torch.nn.ReLU() - self.softmax = torch.nn.LogSoftmax(dim=1) - self.batchnorm1 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - self.batchnorm2 = torch.nn.BatchNorm2d(8, track_running_stats=False, momentum=1.0) - - - def forward(self, x): - x = self.conv1(x) - x = self.batchnorm1(x) - x = self.ReLU1(x) - x1 = self.conv2d1(x) - x2 = self.conv2d5(x) - x = torch.cat((x1, x2), dim=1) - x = self.batchnorm2(x) - x = self.ReLU2(x) - x = self.conv3(x) - x = self.softmax(x) - - return x - - -class SimpleSemSegNet(nn.Module): - def __init__(self, in_channel, out_channel): - super().__init__() - - self.conv1 = self.contract_block(in_channel, 32, 7, 3) -# self.conv2 = self.contract_block(32, 64, 3, 1) -# self.conv3 = self.contract_block(64, 128, 3, 1) - -# self.upconv3 = self.expand_block(128, 64, 3, 1) -# self.upconv2 = self.expand_block(64*2, 32, 3, 1) - self.upconv1 = self.expand_block(32, out_channel, 3, 1) - self.log_softmax = torch.nn.LogSoftmax(dim=1) - - def forward(self, x): - conv1 = self.conv1(x) -# conv2 = self.conv2(conv1) -# conv3 = self.conv3(conv2) - -# upconv3 = self.upconv3(conv3) -# upconv2 = self.upconv2(torch.cat([upconv3, conv2], 1)) -# upconv1 = self.upconv1(torch.cat([upconv2, conv1], 1)) - upconv1 = self.upconv1(conv1) - - output = self.log_softmax(upconv1) - - return output - - def contract_block(self, in_channels, out_channels, kernel_size, padding): - contract = torch.nn.Sequential( - torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, - stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, - stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - ) - return contract - - def expand_block(self, in_channels, out_channels, kernel_size, padding): - expand = torch.nn.Sequential( - torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, - stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, - stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.ConvTranspose2d(out_channels, out_channels, kernel_size=3, - stride=2, padding=1, output_padding=1) - ) - return expand - -class UNET(nn.Module): - def __init__(self, in_channels, out_channels, should_pad=True): - super().__init__() - self.name = 'UNET' - if should_pad: - conv1_pad = 3 - gen_pad = 1 - else: - conv1_pad = 0 - gen_pad = 0 - self.conv1 = self.contract_block(in_channels, 32, 7, conv1_pad) - self.conv2 = self.contract_block(32, 64, 3, gen_pad) - self.conv3 = self.contract_block(64, 128, 3, gen_pad) - - self.upconv3 = self.expand_block(128, 64, 3, gen_pad) - self.upconv2 = self.expand_block(64*2, 32, 3, gen_pad) - self.upconv1 = self.expand_block(32*2, out_channels, 3, gen_pad) - - self.softmax = torch.nn.LogSoftmax(dim=1) - - def __call__(self, x): - - # downsampling part - conv1 = self.conv1(x) - conv2 = self.conv2(conv1) - conv3 = self.conv3(conv2) - upconv3 = self.upconv3(conv3) - - cat1_trim = 6 - cat2_trim = 18 - upconv2 = self.upconv2(torch.cat([upconv3, conv2[:, :, cat1_trim:-cat1_trim, cat1_trim:-cat1_trim]], 1)) - upconv1 = self.upconv1(torch.cat([upconv2, conv1[:, :, cat2_trim:-cat2_trim, cat2_trim:-cat2_trim]], 1)) - xout = self.softmax(upconv1) - - return xout - - def contract_block(self, in_channels, out_channels, kernel_size, padding): - - contract = nn.Sequential( - torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - ) - - return contract - - def expand_block(self, in_channels, out_channels, kernel_size, padding): - - expand = nn.Sequential(torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.Conv2d(out_channels, out_channels, kernel_size, stride=1, padding=padding), - torch.nn.BatchNorm2d(out_channels), - torch.nn.ReLU(), - torch.nn.ConvTranspose2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1) - ) - return expand - -class DoubleCNNBlock(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, padding): - super().__init__() - self.name = 'DoubleCNNBlock' - - self.block = nn.Sequential( - torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, - stride=1, padding=padding), - torch.nn.ReLU(), - torch.nn.BatchNorm2d(out_channels), - torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, - stride=1, padding=padding), - torch.nn.ReLU(), - torch.nn.BatchNorm2d(out_channels) - ) - - def forward(self, x): - return self.block(x) - -class ContractBlock(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, padding): - super().__init__() - self.name = "ContractBlock" - - self.maxpool_conv = nn.Sequential( - torch.nn.MaxPool2d(kernel_size=2, stride=2), - DoubleCNNBlock(in_channels, out_channels, kernel_size, padding) - ) - - def forward(self, x): - return self.maxpool_conv(x) - -class ExpandBlock(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, padding): - super().__init__() - self.name = "ExpandBlock" - - self.expand = torch.nn.ConvTranspose2d(in_channels, in_channels//2, - kernel_size=3, stride=2, - padding=1, output_padding=1 - ) - self.conv = DoubleCNNBlock(in_channels, out_channels, kernel_size, padding) - - def forward(self, x1, x2): - x1_expand = self.expand(x1) - # Assumes dims have been corrected/padded to match elsewhere - x = torch.cat([x1_expand, x2], dim=1) - return self.conv(x) - - -class UNETTraditional(nn.Module): - def __init__(self, in_channels, out_channels, should_pad=True): - super().__init__() - self.name = 'UNETTraditional' - if should_pad: - conv1_pad = 3 - gen_pad = 1 - else: - conv1_pad = 0 - gen_pad = 0 - self.conv1 = DoubleCNNBlock(in_channels, 32, 7, conv1_pad) - self.conv2 = ContractBlock(32, 64, 3, gen_pad) - self.conv3 = ContractBlock(64, 128, 3, gen_pad) - - self.upconv1 = ExpandBlock(128, 64, 3, gen_pad) - self.upconv2 = ExpandBlock(64, 32, 3, gen_pad) - - self.outconv = nn.Conv2d(32, out_channels, kernel_size=1) - self.softmax = torch.nn.LogSoftmax(dim=1) - - def forward(self, x): - - # downsampling part - conv1 = self.conv1(x) - conv2 = self.conv2(conv1) - conv3 = self.conv3(conv2) - - conv2_crop = 4 - cropped_conv2 = conv2[:, :, conv2_crop:-conv2_crop, conv2_crop:-conv2_crop] - upconv1 = self.upconv1(conv3, cropped_conv2) - conv1_crop = 16 - cropped_conv1 = conv1[:, :, conv1_crop:-conv1_crop, conv1_crop:-conv1_crop] - upconv2 = self.upconv2(upconv1, cropped_conv1) - - outconv = self.outconv(upconv2) - xout = self.softmax(outconv) - - return xout \ No newline at end of file diff --git a/heatseek/counting/dl/crop_vids.py b/heatseek/counting/dl/crop_vids.py deleted file mode 100644 index 918f57d..0000000 --- a/heatseek/counting/dl/crop_vids.py +++ /dev/null @@ -1,94 +0,0 @@ -import cv2 -import subprocess -import os -import logging -import argparse - -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') - -def crop_vids(input_vid, output_dir, pixel_threshold, count_threshold, padding, min_gap): - """Scans the input video for motion and crops out clips of motion events. - Parameters: - - input_vid: path to the input video file - - output_dir: directory where the output clips will be saved - - threshold: sensitivity for motion detection (lower = more sensitive) - - padding: seconds of padding to add before and after each detected motion event - - min_gap: minimum gap in seconds to merge nearby motion events into a single clip - """ - os.makedirs(output_dir, exist_ok=True) - cap = cv2.VideoCapture(input_vid) - fps = cap.get(cv2.CAP_PROP_FPS) - prev_frame = None - motion_times = [] - logging.info("Scanning for motion") - - while True: #while there are frames to read - ret, frame = cap.read() #ret returns true if frame is read correctly, frame is the image array - - if not ret: #if there are no more frames left to read - break - - gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) #convert the frame to grayscale for easier processing - - if prev_frame is not None: #if there is a previous frame to compare to (i.e. not the first frame) - diff = cv2.absdiff(prev_frame, gray) #calculate the absolute difference between the current frame and the previous frame - - # if diff.mean() > threshold: - if (diff > pixel_threshold).sum() > count_threshold: #if the number of pixels that changed significantly is greater than the count threshold, we consider it motion - t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000 #get the current time in milliseconds - motion_times.append(t) #add the time to the list of motion times - - prev_frame = gray #update the previous frame to the current frame for the next iteration - - cap.release() #release the video capture object - - if not motion_times: #if no motion was detected, print a message and exit - logging.info("No motion detected. Try lowering threshold") - return - - logging.info(f"Motion found at {len(motion_times)} frames, merging into clips") - clips = [] - start = motion_times[0] - end = motion_times[0] - - for t in motion_times[1:]: #iterate through the motion times starting from the second one - if t - end < min_gap: #if the time between the current motion time and the end of the last clip is less than the minimum gap, we consider it part of the same clip - end = t #update the end time of the current clip to the current motion time - else: - clips.append((start, end)) #if it's not part of the same clip, add the current clip to the list of clips - start = t #start a new clip with the current motion time as the start - end = t #and also set the end to the current motion time for now - - clips.append((start, end)) #add the last clip to the list of clips - logging.info(f"Found {len(clips)} clip(s):") - - for i, (s, e) in enumerate(clips): #iterate through the clips and print their start and end times - s_pad = max(0, s - padding) #add padding to the start time, ensuring it doesn't go below 0 - e_pad = e + padding #add padding to the end time - logging.info(f"Clip {i+1}: {s_pad:.1f}s → {e_pad:.1f}s") #print the clip number and its start and end times - - out_path = os.path.join(output_dir, f"clip_{i+1:03d}.mp4") #create the output path for the clip - subprocess.run([ #use ffmpeg to extract the clip from the original video - "ffmpeg", "-y", #-y to overwrite existing files without asking - "-ss", str(s_pad), #start time for the clip - "-to", str(e_pad), #end time for the clip - "-i", input_vid, #input video file - "-c", "copy", #copy the video and audio streams without re-encoding - out_path #output path for the clip - ]) - - logging.info(f"Done! Clips saved to {output_dir}") - -def main(): - parser = argparse.ArgumentParser(description="Crop motion events from a video into separate clips") - parser.add_argument("--input", required=True, help="Path to the input video file") - parser.add_argument("--output_dir", required=True, help="Directory to save the output clips") - parser.add_argument("--pixel_threshold", type=float, default=15, help="Sensitivity for motion detection (lower = more sensitive)") - parser.add_argument("--count_threshold", type=int, default=20, help="Minimum number of pixels that must change to consider it motion") - parser.add_argument("--padding", type=float, default=2.0, help="Seconds of padding to add before and after each detected motion event") - parser.add_argument("--min_gap", type=float, default=2.0, help="Minimum gap in seconds to merge nearby motion events into a single clip") - args = parser.parse_args() - crop_vids(args.input, args.output_dir, args.pixel_threshold, args.count_threshold, args.padding, args.min_gap) - -if __name__ == "__main__": - main() diff --git a/heatseek/counting/dl/crossing_tracks.py b/heatseek/counting/dl/crossing_tracks.py deleted file mode 100644 index 3c2676a..0000000 --- a/heatseek/counting/dl/crossing_tracks.py +++ /dev/null @@ -1,143 +0,0 @@ -from bat_functions import threshold_short_tracks, measure_crossing_bats -import numpy as np -import matplotlib.pyplot as plt -import argparse - -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, midline_y=88, num_tracks=100): -# 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()}') -# print(f'Forward crossings: {sum(1 for t in crossing_tracks if t["crossed"] > 0)}') -# print(f'Backward crossings: {sum(1 for t in crossing_tracks if t["crossed"] < 0)}') - -# 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] -# color = 'blue' if track['crossed'] > 0 else 'red' -# ax.plot(track['track'][:, 0], track['track'][:, 1], color=color, alpha=0.3, linewidth=1.5) -# ax.axhline(y=midline_y, color='green', linewidth=2, linestyle='--', label='midline') -# ax.set_title(f'Sample of {num_tracks} crossing tracks\nblue=forward, red=backward') -# ax.invert_yaxis() -# ax.legend() - -# ax = axes[1] -# crossing_frames = [abs(t['crossed']) for t in crossing_tracks] -# ax.hist(crossing_frames, bins=100) -# ax.set_xlabel('Frame number') -# ax.set_ylabel('Number of crossings') -# ax.set_title('Bat crossings over time') - -# print('Crossing frames:') - -# for t in crossing_tracks: -# direction = 'forward' if t['crossed'] > 0 else 'backward' -# print(f' frame {abs(t["crossed"])}: {direction}') - -# plt.tight_layout() -# plt.show() - -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 - - crossed_key = 'crossed' if count_out else 'across_crossed' - - 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()}') - print(f'Forward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] > 0)}') - print(f'Backward crossings: {sum(1 for t in crossing_tracks if t[crossed_key] < 0)}') - - 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] - color = 'blue' if track[crossed_key] > 0 else 'red' - 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() - direction_labels = 'blue=downward, red=upward' - else: - ax.axvline(x=midline_x, color='green', linewidth=2, linestyle='--', label='midline') - direction_labels = 'blue=rightward, red=leftward' - - ax.set_title(f'Sample of {num_tracks} crossing tracks\n{direction_labels}') - ax.legend() - - ax = axes[1] - crossing_frames = [abs(t[crossed_key]) for t in crossing_tracks] - ax.hist(crossing_frames, bins=100) - ax.set_xlabel('Frame number') - ax.set_ylabel('Number of crossings') - ax.set_title('Bat crossings over time') - - print('Crossing frames:') - for t in crossing_tracks: - direction = 'forward' if t[crossed_key] > 0 else 'backward' - print(f' frame {abs(t[crossed_key])}: {direction}') - - plt.tight_layout() - plt.show() - -def main(): - parser = argparse.ArgumentParser(description='Get bat crossing tracks across the video') - parser.add_argument('--raw_tracks_file', type=str, help='Path to raw tracks file') - parser.add_argument('--crossing_tracks_file', type=str, help='Intended path for the crossing tracks 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() - - save_crossing_tracks_from_raw_tracks(args.raw_tracks_file, args.crossing_tracks_file, args.count_across, args.count_out) - visualize_crossing_tracks(args.crossing_tracks_file, count_out=args.count_out, count_across=args.count_across) - -if __name__ == '__main__': - main() diff --git a/heatseek/counting/dl/dataloaders.py b/heatseek/counting/dl/dataloaders.py deleted file mode 100644 index 01de90b..0000000 --- a/heatseek/counting/dl/dataloaders.py +++ /dev/null @@ -1,16 +0,0 @@ -import numpy as np -import datetime -from SegmentationDataset import SegmentationDataset -from transforms import split_train_val -import torch.utils.data as data - -def worker_init_fn(worker_id): - np.random.seed(datetime.datetime.now().microsecond + worker_id * 1000000) - -def load_dataloader(clips_dir, transform, split='train', batch_size=4, num_workers=7, pin_memory=True): - train_pairs, val_pairs = split_train_val(clips_dir) - pairs = train_pairs if split == 'train' else val_pairs - dataset = SegmentationDataset(pairs, transform) - return data.DataLoader(dataset, batch_size=batch_size, - shuffle=(split == 'train'), num_workers=num_workers, - pin_memory=pin_memory, worker_init_fn=worker_init_fn) diff --git a/heatseek/counting/dl/detections_to_tracks.py b/heatseek/counting/dl/detections_to_tracks.py deleted file mode 100644 index cc80404..0000000 --- a/heatseek/counting/dl/detections_to_tracks.py +++ /dev/null @@ -1,222 +0,0 @@ -import numpy as np -import os -import glob -import bat_functions as kbf -from multiprocessing import Pool -import argparse - -def track(camera_dict): - """ - 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). - - 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) - else: - print("Missing contour files.") - -def build_camera_dicts(output_folder, num_groups=10, fps=60, overlap_seconds=15): - """ - 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 15. - - 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) - - 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, processes=5): - """ - 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) - print(f'Tracking {len(camera_dicts)} groups') - with Pool(processes=processes) as pool: - pool.map(track, camera_dicts) - -def main(): - parser = argparse.ArgumentParser(description='Get centers and contours from detections from model inference') - parser.add_argument('--output_folder', type=str, help='Path of output folder containing the centers, contours, etc from video inference') - args = parser.parse_args() - run_tracking(args.output_folder) - combine_tracks(args.output_folder) - -if __name__ == '__main__': - main() diff --git a/heatseek/counting/dl/extract_frames.py b/heatseek/counting/dl/extract_frames.py deleted file mode 100644 index ba09b45..0000000 --- a/heatseek/counting/dl/extract_frames.py +++ /dev/null @@ -1,55 +0,0 @@ -import glob -import os -import cv2 -import logging -import argparse - -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') - -def pull_out_frames(video_file, process_every_n_frames=1): - """Extract frames from a video file and optionally save them to disk. - Parameters: - - video_file: path to the input video file - - save_frames: whether to save the extracted frames to disk - - process_every_n_frames: only process every nth frame to speed up extraction - """ - video_dir = os.path.dirname(video_file) - video_name = os.path.splitext(os.path.basename(video_file))[0] - output_folder = os.path.join(video_dir, video_name, 'frames') - os.makedirs(output_folder, exist_ok=True) - logging.info(f'Saving frames to {output_folder}') - cap = cv2.VideoCapture(video_file) - frame_count = 0 - - while True: #read every frame until the end of the video - ret, frame = cap.read() - if not ret: - break - - if frame_count % process_every_n_frames == 0: - frame_name = f'{video_name}_{frame_count:05d}.jpg' - frame_file = os.path.join(output_folder, frame_name) - cv2.imwrite(frame_file, frame) - - frame_count += 1 - - cap.release() - logging.info(f'{video_name}: done. {frame_count} total frames.') - - - -def main(): - parser = argparse.ArgumentParser(description='Extract frames from videos') - parser.add_argument('--clips_dir', help='Parent folder containing video files') - parser.add_argument('--process_every_n_frames', type=int, default=1, help='Process every nth frame') - args = parser.parse_args() - - video_files = [] - video_files.extend(glob.glob(os.path.join(args.clips_dir, '**', '*.mp4'), recursive=True)) - video_files.sort() - - for video_file in video_files: - pull_out_frames(video_file, process_every_n_frames=args.process_every_n_frames) - -if __name__ == '__main__': - main() diff --git a/heatseek/counting/dl/koger_tracking.py b/heatseek/counting/dl/koger_tracking.py deleted file mode 100644 index cba6251..0000000 --- a/heatseek/counting/dl/koger_tracking.py +++ /dev/null @@ -1,744 +0,0 @@ -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['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 - 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 fewwer new points than existing tracks. When it is the - later, 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 - -#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 = 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, - 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: upper distance threshold - max_distance_noise: upper distance threshold tracks with noise values - above 0 - min_distance: lowwer 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 - 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) - - 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)] = max_distance - # 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]) - 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]) - - 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/dl/optimizer.py b/heatseek/counting/dl/optimizer.py deleted file mode 100644 index ca7aad0..0000000 --- a/heatseek/counting/dl/optimizer.py +++ /dev/null @@ -1,108 +0,0 @@ -import torch -from bat_seg_models import UNETTraditional -import os - -def lr_scheduler(epoch): - """ To pass to torch optim lr_scheduler. - - Args: - epoch: current epoch number - - return number to multiply base lr - """ - - if epoch < 15: - return 1.0 - elif epoch < 30: - return 0.2 - else: - return 0.05 - -def build_and_save_model(lr=0.01, accumulation_steps=4, epochs=91, momentum=0.9, batch_size=4): - """ - Instantiate a UNETTraditional model and construct its checkpoint file path. - - Builds a single-channel-input, two-class-output U-Net (no padding) and - generates a deterministic filename encoding the key hyperparameters so - that checkpoint files are self-describing and don't collide across runs. - - Args: - lr (float): Learning rate, embedded in the filename. Default: 0.01. - accumulation_steps (int): Gradient accumulation steps. Combined with - `batch_size` to record the effective batch size in the filename. - Default: 4. - epochs (int): Total training epochs, embedded in the filename. - Default: 91. - momentum (float): SGD momentum, embedded in the filename. Default: 0.9. - batch_size (int): Per-step batch size. Multiplied by - `accumulation_steps` to compute effective batch size. Default: 4. - - Returns: - model (UNETTraditional): Untrained model with in_channels=1, - out_channels=2, should_pad=False. - model_file (str): Path of the form - './models/model_UNETTraditional_epochs_{e}_batcheff_{b}_lr_{lr} - _momentum_{m}_aug_thermal-radiometric-heatseek.tar'. - The ./models/ directory is created if it does not exist. - - Note: - This function only constructs and names the model — it does not - train it or write any file to disk. Pass model_file to train() with - save_best_val=True to actually checkpoint the weights. - """ - - model = UNETTraditional(in_channels=1, out_channels=2, should_pad=False) - # refer to bat_seg_models.py to set any other models; Koger uses UNETTraditional for all experiments, so we hardcode that here for simplicity - model_name = "UNETTraditional" - aug_type = "thermal-radiometric-heatseek" - - folder = './models' - os.makedirs(folder, exist_ok=True) - - model_name = 'model_{}_epochs_{}_batcheff_{}_lr_{}_momentum_{}_aug_{}'.format( - model_name, epochs, accumulation_steps*batch_size, lr, momentum, aug_type) - - model_file = os.path.join(folder, model_name + '.tar') - return model, model_file - -def get_optimizer_and_scheduler(model, lr=0.01, momentum=0.9): - """ - Create three independent SGD optimizer/scheduler pairs for none, easy, - and hard augmentation training regimes. - - Each pair is fully independent — separate parameter groups, separate - state — so they can be used to train the same model sequentially under - different augmentation conditions without state leaking between regimes. - All three share the same `lr`, `momentum`, and `lr_scheduler` lambda. - - Args: - model (torch.nn.Module): Model whose parameters all three optimizers - will reference. - lr (float): Initial learning rate for all three SGD instances. - Default: 0.01. - momentum (float): Momentum for all three SGD instances. Default: 0.9. - - Returns: - optimizer_none (torch.optim.SGD): Optimizer for the no-augmentation regime. - scheduler_none (LambdaLR): Scheduler paired with optimizer_none. - optimizer_easy (torch.optim.SGD): Optimizer for the easy-augmentation regime. - scheduler_easy (LambdaLR): Scheduler paired with optimizer_easy. - optimizer_hard (torch.optim.SGD): Optimizer for the hard-augmentation regime. - scheduler_hard (LambdaLR): Scheduler paired with optimizer_hard. - - Note: - All three optimizers reference the same model.parameters() iterator - at construction time. If the model is replaced or re-initialized - between regimes, these optimizers will hold stale parameter references - and must be reconstructed. - """ - optimizer_none = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) - scheduler_none = torch.optim.lr_scheduler.LambdaLR(optimizer_none, lr_scheduler, last_epoch=-1) - - optimizer_easy = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) - scheduler_easy = torch.optim.lr_scheduler.LambdaLR(optimizer_easy, lr_scheduler, last_epoch=-1) - - optimizer_hard = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) - scheduler_hard = torch.optim.lr_scheduler.LambdaLR(optimizer_hard, lr_scheduler, last_epoch=-1) - - return optimizer_none, scheduler_none, optimizer_easy, scheduler_easy, optimizer_hard, scheduler_hard diff --git a/heatseek/counting/dl/segment_frames.py b/heatseek/counting/dl/segment_frames.py deleted file mode 100644 index cc3249c..0000000 --- a/heatseek/counting/dl/segment_frames.py +++ /dev/null @@ -1,88 +0,0 @@ -import cv2 -import numpy as np -import argparse -import os -import logging - -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') - -def get_background_temp(image_path, min_val, max_val): - """Get background temperature of an image taken by a radiometric camera. - Parameters: - - image_path: path to the image file (e.g. frame.png) - - min_val: minimum temperature value corresponding to pixel value 0 - - max_val: maximum temperature value corresponding to pixel value 255 - """ - frame = cv2.imread(image_path) #read the frame - - if frame is None: - raise ValueError("Could not read image") - - gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) #convert to grayscale dimensions - - k100 = (gray.astype(np.float32) / 255.0) * (max_val - min_val) + min_val #formula to convert pixel values to temperature in Kelvin (FLIR Boson) - temp_celsius = (k100 / 100.0) - 273.15 - background_temp = np.median(temp_celsius) - return background_temp #return the estimated background temperature in Celsius based on the median pixel value of the thermal map - -def get_mask(image_path, min_val, max_val, threshold): - """Get a segmentation mask of objects hotter than the background by a threshold. - Parameters: - - image_path: path to the image file - - min_val: minimum temperature value corresponding to pixel value 0 - - max_val: maximum temperature value corresponding to pixel value 255 - - background_temp: estimated background temperature in °C - - threshold: how many °C above background a pixel must be to be masked - """ - frame = cv2.imread(image_path) - - if frame is None: - raise ValueError("Could not read image") - - gray = cv2.cvtColor(frame, 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 = get_background_temp(image_path, min_val, max_val) - mask = (temp_celsius > (background_temp + threshold)).astype(np.uint8) * 255 #create a binary mask where pixels hotter than background + threshold are set to 255 (white) and others are 0 (black) - return mask - - - -def main(): - parser = argparse.ArgumentParser(description='Generate thermal segmentation mask for a single image') - parser.add_argument('--clips_dir', type=str, help='Path to the clips folder') - parser.add_argument('--min_val', type=float, default=28000) - parser.add_argument('--max_val', type=float, default=32000) - parser.add_argument('--threshold', type=float, default=5.0) - args = parser.parse_args() - - for session_folder in os.listdir(args.clips_dir): - session_path = os.path.join(args.clips_dir, session_folder) - - if not os.path.isdir(session_path): - continue - - for clip_folder in os.listdir(session_path): - clip_path = os.path.join(session_path, clip_folder) - - if not os.path.isdir(clip_path): - continue - - frames_path = os.path.join(clip_path, 'frames') - masks_path = os.path.join(clip_path, 'masks') - - if not os.path.isdir(frames_path): - continue - - os.makedirs(masks_path, exist_ok=True) - - for filename in os.listdir(frames_path): - image_path = os.path.join(frames_path, filename) - mask = get_mask(image_path, args.min_val, args.max_val, args.threshold) - output_file = os.path.join(masks_path, filename.replace('.jpg', '.png')) - cv2.imwrite(output_file, mask) - - logging.info(f"Finished processing {clip_folder}") - -if __name__ == "__main__": - main() diff --git a/heatseek/counting/dl/trainer.py b/heatseek/counting/dl/trainer.py deleted file mode 100644 index 9979b19..0000000 --- a/heatseek/counting/dl/trainer.py +++ /dev/null @@ -1,248 +0,0 @@ -import time -import numpy as np -import torch -import argparse -from optimizer import build_and_save_model -from dataloaders import load_dataloader -from transforms import get_transforms -from optimizer import get_optimizer_and_scheduler -import logging - -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') - -def get_conf_matrix(conf_matrix, num_classes, pred_batch, mask_batch): - """ Calculate confusion matrix and add to passed confusion matrix. - - Args: - conf_matrix (np array): confusion matrix size - (num_classes + 1, num_classes + 1) - num_classes (int): number of classes being segmented - pred_batch (np array): NxM or BxNxM - mask_batch (np array): NxM or BxNxM - """ - - N = num_classes + 1 - - if len(pred_batch.shape) == 2: - pred_batch = np.expand_dims(pred_batch, 0) - mask_batch = np.expand_dims(mask_batch, 0) - for pred, mask in zip(pred_batch, mask_batch): - conf_matrix += np.bincount( - N * pred.reshape(-1) + mask.reshape(-1), minlength=N ** 2 - ).reshape(N, N) - return conf_matrix - -def evaluate(num_classes, conf_matrix): - """ - Compute per-class accuracy and IoU from a confusion matrix. - - The confusion matrix is expected to have an extra row/column for an - ignore/background class (shape: (num_classes + 1, num_classes + 1)). - Only the first `num_classes` rows and columns are evaluated; the final - row/column is treated as a void/ignore label and excluded. - - Parameters: - num_classes (int): Number of foreground classes to evaluate. - conf_matrix (np.ndarray): Square confusion matrix of shape - (num_classes + 1, num_classes + 1), where entry [i, j] is the - number of pixels of true class i predicted as class j. - - Returns: - iou (np.ndarray): Per-class Intersection over Union of shape - (num_classes,). Classes with no ground-truth pixels are NaN. - acc (np.ndarray): Per-class accuracy (recall) of shape - (num_classes,). Classes with no ground-truth pixels are NaN. - """ - - acc = np.full(num_classes, np.nan, dtype=np.float64) - iou = np.full(num_classes, np.nan, dtype=np.float64) - tp = conf_matrix.diagonal()[:-1].astype(np.float64) - pos_gt = np.sum(conf_matrix[:-1, :-1], axis=0).astype(np.float64) - pos_pred = np.sum(conf_matrix[:-1, :-1], axis=1).astype(np.float64) - acc_valid = pos_gt > 0 - acc[acc_valid] = tp[acc_valid] / pos_gt[acc_valid] - union = pos_gt + pos_pred - tp - iou[acc_valid] = tp[acc_valid] / union[acc_valid] - return iou, acc - - -def acc_metric(predb, yb): - """ - Compute batch accuracy for a multi-class classifier. - - Parameters: - predb (torch.Tensor): Raw logits or probabilities of shape - (N, C), where N is batch size and C is number of classes. - yb (torch.Tensor): Ground-truth class indices of shape (N,). - - Returns: - torch.Tensor: Scalar mean accuracy over the batch, in [0.0, 1.0]. - """ - return (predb.argmax(dim=1) == yb).float().mean() - -def train(model, num_classes, train_dl, val_dl, loss_fn, optimizer, - epochs=91, lr_scheduler=None, save_best_val=False, - model_file='model.tar', accumulation_steps=4, val_epoch=2): - """ - Train a segmentation model with gradient accumulation and periodic validation. - - Parameters: - model (torch.nn.Module): Segmentation model to train. - num_classes (int): Number of foreground classes. - train_dl (DataLoader): Training dataloader. Batches must be dicts - with keys 'image' and 'mask'. - val_dl (DataLoader): Validation dataloader. Same format as train_dl. - loss_fn (callable): Loss function with signature - loss_fn(outputs, masks) -> scalar tensor. - optimizer (torch.optim.Optimizer): Optimizer instance. - epochs (int): Total number of training epochs. Default: 91. - lr_scheduler (optional): Scheduler with a .step() method called - after each optimizer step. Default: None. - save_best_val (bool): If True, saves model state dict to - model_file whenever validation loss improves. Default: False. - model_file (str): Path for the checkpoint file. Default: 'model.tar'. - accumulation_steps (int): Number of batches to accumulate gradients - over before calling optimizer.step(). Default: 4. - val_epoch (int): Validation runs on epochs where - epoch % val_epoch == 0. Default: 2. - - Returns: - train_loss (list[float]): Mean training loss recorded each epoch. - valid_loss (list[float]): Mean validation loss recorded on each - validated epoch. - """ - start = time.time() - train_loss, valid_loss = [], [] - - device = 'cuda' if torch.cuda.is_available() else 'cpu' - model.to(device) - - max_batchnum = (len(train_dl) // accumulation_steps) * accumulation_steps - logging.info("Using {} batches ({} images)".format( - max_batchnum, max_batchnum * train_dl.batch_size)) - - top_val_loss = float('inf') - - for epoch in range(epochs): - logging.info('Epoch {}/{}'.format(epoch, epochs - 1)) - optimizer.zero_grad() - - for phase in ['train', 'val']: - if phase == 'train': - model.train(True) - dataloader = train_dl - else: - if epoch % val_epoch != 0: - continue - model.train(False) - dataloader = val_dl - conf_matrix = np.zeros((num_classes + 1, num_classes + 1), dtype=np.int64) - - running_loss = 0.0 - - for batch_ind, batch in enumerate(dataloader): - if batch_ind >= max_batchnum: - break - - im_batch = batch['image'].to(device) - masks = batch['mask'][:, 24:-24, 24:-24].to(device) - - if phase == 'train': - outputs = model(im_batch) - loss = loss_fn(outputs, masks) - loss.backward() - if (batch_ind + 1) % accumulation_steps == 0: - optimizer.step() - if lr_scheduler: - lr_scheduler.step() - model.zero_grad() - else: - with torch.no_grad(): - outputs = model(im_batch) - loss = loss_fn(outputs, masks) - np_preds = np.argmax(outputs.cpu().numpy(), axis=1) - np_masks = masks.cpu().numpy() - conf_matrix = get_conf_matrix(conf_matrix, num_classes, - np_preds, np_masks) - - running_loss += loss.item() * dataloader.batch_size - - epoch_loss = running_loss / len(dataloader.dataset) - logging.info('{} Loss: {:.4f}'.format(phase, epoch_loss)) - - if phase == 'val': - valid_loss.append(epoch_loss) - if epoch_loss < top_val_loss: - top_val_loss = epoch_loss - if save_best_val: - torch.save(model.state_dict(), model_file) - logging.info('Saved new best model. Val loss: {:.4f}'.format(epoch_loss)) - iou, acc = evaluate(num_classes, conf_matrix) - logging.info('IOU: {}, Acc: {}'.format(iou, acc)) - else: - train_loss.append(epoch_loss) - - time_elapsed = time.time() - start - logging.info('Training complete in {:.0f}m {:.0f}s'.format( - time_elapsed // 60, time_elapsed % 60)) - - return train_loss, valid_loss - -def main(): - parser = argparse.ArgumentParser(description='Train UNETTraditional on thermal segmentation dataset') - parser.add_argument('--clips_dir', type=str, help='path to clips directory') - parser.add_argument('--save_model', default=True, help='whether to save the trained model') - parser.add_argument('--lr', type=float, default=0.01, help='learning rate for optimizer') - parser.add_argument('--accumulation_steps', type=int, default=4, help='number of batches to accumulate gradients over') - parser.add_argument('--epochs', type=int, default=91, help='number of epochs to train for') - parser.add_argument('--momentum', type=float, default=0.9, help='momentum for SGD optimizer (not applicable for UNETTraditional)') - parser.add_argument('--batch_size', type=int, default=4, help='batch size for training (effective batch size is batch_size * accumulation_steps)') - parser.add_argument('--num_workers', type=int, default=7, help='number of workers for dataloader') - parser.add_argument('--pin_memory', type=bool, default=True, help='whether to pin memory in dataloader') - args = parser.parse_args() - - total_train_loss = [] - total_val_loss = [] - - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - loss_fn = torch.nn.NLLLoss(weight=torch.tensor([.01, .99]).to(device)) - model, model_file = build_and_save_model(args.lr, args.accumulation_steps, args.epochs) - optimizer_none, scheduler_none, optimizer_easy, scheduler_easy, optimizer_hard, scheduler_hard = get_optimizer_and_scheduler(model, args.lr, args.momentum) - - no_aug_epochs = 1 - easy_epochs = 30 - hard_epochs = 60 - - none_transforms, easy_transforms, hard_transforms = get_transforms(args.clips_dir) - - train_dataloader_no_aug = load_dataloader(args.clips_dir, none_transforms, split='train', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) - train_dataloader_easy_aug = load_dataloader(args.clips_dir, easy_transforms, split='train', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) - train_dataloader_hard_aug = load_dataloader(args.clips_dir, hard_transforms, split='train', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) - val_dataloader = load_dataloader(args.clips_dir, hard_transforms, split='val', batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_memory) - - train_loss, val_loss = train(model, 2, train_dataloader_no_aug, val_dataloader, - loss_fn, optimizer_none, epochs=no_aug_epochs, - lr_scheduler=scheduler_none, save_best_val=True, - model_file=model_file, val_epoch=1) - total_train_loss.extend(train_loss) - total_val_loss.extend(val_loss) - - train_loss, val_loss = train(model, 2, train_dataloader_easy_aug, val_dataloader, - loss_fn, optimizer_easy, epochs=easy_epochs, - lr_scheduler=scheduler_easy, save_best_val=True, - model_file=model_file) - total_train_loss.extend(train_loss) - total_val_loss.extend(val_loss) - - train_loss, val_loss = train(model, 2, train_dataloader_hard_aug, val_dataloader, - loss_fn, optimizer_hard, epochs=hard_epochs, - lr_scheduler=scheduler_hard, save_best_val=True, - model_file=model_file) - total_train_loss.extend(train_loss) - total_val_loss.extend(val_loss) - - if args.save_model: - torch.save(model.state_dict(), model_file) - -if __name__ == "__main__": - main() diff --git a/heatseek/counting/dl/transforms.py b/heatseek/counting/dl/transforms.py deleted file mode 100644 index 4e37fd6..0000000 --- a/heatseek/counting/dl/transforms.py +++ /dev/null @@ -1,117 +0,0 @@ -from augmentations import compute_mean_std, MaskCompose, MaskRandomCrop, MaskToTensor, MaskNormalize, Mask2dMultiplyAndAddToBrightness, MaskContrast, MaskImgAug -import os - -def get_transforms(clips_dir): - """ - Build the three augmentation pipelines used in curriculum training. - - Computes per-dataset mean and std from the clips in `clips_dir` and - constructs none, easy, and hard transform pipelines of increasing - augmentation intensity. All three pipelines share the same crop size - and normalization strategy: mean-centering by the dataset mean (scaled - to [0,1]) with no std scaling (std=1.0), preserving the radiometric - intensity range. - - Augmentation progression: - none — crop and normalize only; used for initial warmup training. - easy — adds mild brightness/contrast jitter and imgaug transforms; - brightness multiply in (0.85, 1.0), additive shift in (-15, 0). - hard — same structure as easy but with wider jitter ranges; - brightness multiply in (0.7, 1.0), additive shift in (-25, 0), - contrast factor in (0.6, 1.2). - - Parameters: - clips_dir (str): Directory containing training clips, passed - to compute_mean_std() to derive normalization statistics. - - Returns: - none_data_transforms (MaskCompose): Minimal pipeline for warmup phase. - easy_data_transforms (MaskCompose): Moderate augmentation pipeline. - hard_data_transforms (MaskCompose): Heavy augmentation pipeline for - final training phase. - - Note: - std=1.0 means normalization only subtracts the mean without rescaling - by standard deviation. This is intentional for radiometric data where - the absolute intensity scale is meaningful and should be preserved - relative to the mean. - """ - - mean, std = compute_mean_std(clips_dir) - - none_data_transforms = MaskCompose([ - MaskRandomCrop(224), - MaskToTensor(), - MaskNormalize(mean=mean/255, std=1.0) - ]) - - easy_data_transforms = MaskCompose([ - MaskRandomCrop(224), - Mask2dMultiplyAndAddToBrightness(multiply=(0.85, 1.0), add=(-15, 0)), - MaskContrast(contrast_factor=(0.8, 1.2)), - MaskImgAug(), - MaskToTensor(), - MaskNormalize(mean=mean/255, std=1.0) - ]) - - hard_data_transforms = MaskCompose([ - MaskRandomCrop(224), - Mask2dMultiplyAndAddToBrightness(multiply=(0.7, 1.0), add=(-25, 0)), - MaskContrast(contrast_factor=(0.6, 1.2)), - MaskImgAug(), - MaskToTensor(), - MaskNormalize(mean=mean/255, std=1.0) - ]) - - return none_data_transforms, easy_data_transforms, hard_data_transforms - -def split_train_val(clips_dir, train_split=0.8, total_samples=5000): - """ - Split frame/mask pairs into training and validation sets, sampled evenly - across all clips. - Parameters: - clips_dir (str): Parent directory containing video subdirectories, - each containing clip folders with frames/ and masks/ - train_split (float): Fraction of samples for training. Default: 0.8 - total_samples (int): Total number of frame/mask pairs to use. Default: 5000 - Returns: - train_pairs (list of tuples): (frame_path, mask_path) pairs for training - val_pairs (list of tuples): (frame_path, mask_path) pairs for validation - """ - # Collect all valid frame/mask pairs across the full directory structure - all_pairs = [] - - for video_dir in sorted(os.listdir(clips_dir)): - video_path = os.path.join(clips_dir, video_dir) - - if not os.path.isdir(video_path): - continue - - for clip_dir in sorted(os.listdir(video_path)): - clip_path = os.path.join(video_path, clip_dir) - - if not os.path.isdir(clip_path): - continue - - frames_path = os.path.join(clip_path, 'frames') - masks_path = os.path.join(clip_path, 'masks') - - if not os.path.isdir(frames_path) or not os.path.isdir(masks_path): - continue - - for filename in sorted(os.listdir(frames_path)): - frame_path = os.path.join(frames_path, filename) - mask_path = os.path.join(masks_path, filename.replace('.jpg', '.png')) - - if os.path.exists(mask_path): - all_pairs.append((frame_path, mask_path)) - - if len(all_pairs) > total_samples: - step = len(all_pairs) / total_samples - all_pairs = [all_pairs[int(i * step)] for i in range(total_samples)] - - train_end = int(train_split * len(all_pairs)) - train_pairs = all_pairs[:train_end] - val_pairs = all_pairs[train_end:] - - return train_pairs, val_pairs diff --git a/heatseek/counting/dl/video_inference.py b/heatseek/counting/dl/video_inference.py deleted file mode 100644 index cb5c345..0000000 --- a/heatseek/counting/dl/video_inference.py +++ /dev/null @@ -1,180 +0,0 @@ -import os -import time -import cv2 -import numpy as np -import torch -import torch.utils.data as data -from bat_seg_models import UNETTraditional -from augmentations import compute_mean_std, MaskCompose, MaskToTensor, MaskNormalize -import bat_functions -from BatIterableDataset import BatIterableDataset -import argparse - -def denorm_image(im, mean): - """ - Reverses mean normalization on a float image array and converts it to uint8. - - Adds back the per-channel mean (scaled to [0, 1]) to a normalized image, - rescales pixel values to [0, 255], clips to valid range, and casts to uint8. - - Parameters: - im (np.ndarray): Normalized image array of shape (H, W, C) with float - values expected in roughly [-1, 1] or [0, 1] range. - mean (np.ndarray or list): Per-channel mean values in [0, 255] range, - shape (C,), used during the original normalization step. - - Returns: - np.ndarray: Denormalized image as uint8 with pixel values in [0, 255], - same spatial shape as input. - """ - im += mean / 255 - im *= 255 - im = np.maximum(im, 0) - im = np.minimum(im, 255) - return im.astype(np.uint8) - -def load_model(model_file): - """ - Load a trained UNETTraditional model from a state dict file and prepare - it for inference. - - Automatically selects CUDA if available, otherwise falls back to CPU. - The model is set to evaluation mode (gradients and dropout disabled) - before being returned. - - Parameters: - model_file (str or path-like): Path to a saved state dict file - (.pth or .pt) produced by torch.save(). - - Returns: - model (UNETTraditional): Trained model in eval mode, moved to the - appropriate device (CUDA or CPU). - """ - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - model = UNETTraditional(1, 2, should_pad=False) - model.load_state_dict(torch.load(model_file, map_location=device)) - model.to(device) - model.train(False) - return model - -def get_augmentor(mean): - """ - Build an image augmentation pipeline that converts an image/mask pair - to tensors and applies mean normalization. - - Parameters: - mean (np.ndarray or float): Per-channel mean values in [0, 255] range, - divided by 255 internally to normalize to [0, 1] scale. - - Returns: - MaskCompose: A composed transform that applies, in order: - 1. MaskToTensor — converts image and mask to torch.Tensor. - 2. MaskNormalize — subtracts mean/255 from the image (std=1.0). - """ - return MaskCompose([ - MaskToTensor(), - MaskNormalize(mean=mean/255, std=1.0) - ]) - -def run_inference(video_path, output_folder, model, mean, - bat_prob_thresh=0.6, batch_size=2, - save_every_n_frames=18000, early_stop=None): - """Run model inference on a video file and save detections to disk.""" - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - os.makedirs(output_folder, exist_ok=True) - os.makedirs(os.path.join(output_folder, 'example-frames'), exist_ok=True) - - augmentor = get_augmentor(mean) - bat_dataset = BatIterableDataset([video_path], augmentor=augmentor) - dataloader = data.DataLoader(bat_dataset, batch_size=batch_size, - shuffle=False, num_workers=0, pin_memory=True) - - centers_list = [] - contours_list = [] - sizes_list = [] - rects_list = [] - num_frames = 0 - t0 = None - - for batch_ind, batch in enumerate(dataloader): - if batch_ind == 0: - print('started...') - t0 = time.time() - if early_stop and batch_ind >= early_stop: - break - - im_batch = batch['image'].to(device) - - with torch.no_grad(): - outputs = model(im_batch) - masks = (outputs[:, 1].cpu().numpy() > np.log(bat_prob_thresh)).astype(np.uint8) - - for ind, mask in enumerate(masks): - 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 save_every_n_frames and num_frames % save_every_n_frames == 0: - im_name = f'frame_{num_frames:07d}.jpg' - im_file = os.path.join(output_folder, 'example-frames', im_name) - im = denorm_image(np.squeeze(batch['image'][ind].numpy()).copy(), mean) - cv2.imwrite(im_file, im) - num_frames += 1 - - if batch_ind % 1000 == 0 and t0 is not None: - print(f'batch {batch_ind}, frame {num_frames}, time {time.time()-t0:.1f}s') - - total_time = time.time() - t0 - print(f'{total_time:.1f}s total, {total_time/max(batch_ind,1)/batch_size:.3f}s per frame') - bat_dataset.get_read_frame_info() - - return centers_list, contours_list, sizes_list, rects_list - -def save_detections(centers_list, contours_list, sizes_list, rects_list, - output_folder, num_contour_files=15): - """Save per-frame detections to disk.""" - 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)) - print(f'Saved detections for {len(centers_list)} frames to {output_folder}') - -def run_video_inference(video_path, output_folder, model_file, clips_dir, - bat_prob_thresh=0.6, batch_size=2, - save_every_n_frames=18000, early_stop=None): - """Top level function — load model, run inference, save detections.""" - model = load_model(model_file) - mean, _ = compute_mean_std(clips_dir) - centers_list, contours_list, sizes_list, rects_list = run_inference( - video_path, output_folder, model, mean, - bat_prob_thresh, batch_size, save_every_n_frames, early_stop) - save_detections(centers_list, contours_list, sizes_list, rects_list, output_folder) - -def main(): - parser = argparse.ArgumentParser(description='Get centers and contours from detections from model inference') - parser.add_argument('--vid_path', type=str, help='Path of video') - parser.add_argument('--output_folder', type=str, help='Path to the folder where you want your outputs saved') - parser.add_argument('--model_filepath', type=str, help='Filepath of trained model') - parser.add_argument('--clips_dir', type=str, help='Directory containing the clips (faster way to calculate mean and std)') - args = parser.parse_args() - run_video_inference(video_path=args.vid_path, output_folder=args.output_folder, model_file=args.model_filepath, clips_dir=args.clips_dir) - -if __name__ == '__main__': - main() - \ No newline at end of file diff --git a/heatseek/counting/bg_sub/koger_tracking.py b/heatseek/counting/koger_tracking.py similarity index 89% rename from heatseek/counting/bg_sub/koger_tracking.py rename to heatseek/counting/koger_tracking.py index cba6251..c766391 100644 --- a/heatseek/counting/bg_sub/koger_tracking.py +++ b/heatseek/counting/koger_tracking.py @@ -46,6 +46,7 @@ def update_unmatched_track(track, debug=False): ''' 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) @@ -90,7 +91,8 @@ def update_matched_track(track, frame_index, new_position, if track['noise'] > 0: # Maybe this should go to 0 after found - track['noise'] -= 1 + # track['noise'] -= 1 + track['noise'] = 0 if 'size' in track: track['size'].append(size) if 'rects' in track: @@ -147,8 +149,8 @@ def filter_tracks_without_new_points(track_list, distance, row_ind, """ 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 fewwer new points than existing tracks. When it is the - later, the longer track takes presedence. + 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 @@ -464,6 +466,23 @@ def finalize_tracks(track_list): 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 @@ -481,6 +500,7 @@ def calculate_distances(new_positions, track_list, active_list): """ #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) @@ -502,7 +522,7 @@ def calculate_active_list(track_list, max_unseen_time, frame_num, debug=False): 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, +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): @@ -518,10 +538,10 @@ def calculate_max_distance(track_list, active_list, max_distance, Args: track_list: list of all tracks active_list: list of tracks that could be added to - max_distance: upper distance threshold + max_distance_threshold: upper distance threshold max_distance_noise: upper distance threshold tracks with noise values above 0 - min_distance: lowwer distance threshold + 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 @@ -530,25 +550,41 @@ def calculate_max_distance(track_list, active_list, max_distance, # only check distances to established tracks defined by a noise value of # 0 or below - 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) + #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") + # 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)] = max_distance - # 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 + # # 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: @@ -559,20 +595,22 @@ def calculate_max_distance(track_list, active_list, max_distance, # 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]) - 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: + # 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]) - 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]) + # #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 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()