diff --git a/examples/boson_record.py b/examples/boson_record.py index 9e1f4b5..2bbd26e 100644 --- a/examples/boson_record.py +++ b/examples/boson_record.py @@ -9,8 +9,9 @@ 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 \ + --autonorm_path path/to/output.mp4 \ + --globalnorm_path path/to/output.mp4 \ + --metadata_path path/to/metadata.yaml \ --bosonsdk_path path/to/BosonSDK/SDK_USER_PERMISSIONS \ --recording_time 60 @@ -19,7 +20,7 @@ """ import argparse -from time import sleep +from time import time, sleep from heatseek.boson_capture import BosonCapture @@ -34,13 +35,15 @@ def main(): 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) + --autonorm_path (str, opt): Filepath to save + normalized viewable video (.mp4) + --globalnorm_path (str, opt): Filepath to save + normalized radiometric (.mp4) + --metadata_path (str, opt): Filepath to save + radiometric metadata --bosonsdk_path (str, opt) Filepath to boson sdk folder --recording_time (int, opt): Duration of recording in seconds. - Default is 10s. + Default is 3hrs Returns: None @@ -53,12 +56,15 @@ def main(): parser.add_argument('--video_port', type=str, help='path to video port. Default is /dev/video0') - parser.add_argument('--raw_output_path', + parser.add_argument('--autonorm_path', type=str, - help='filepath to save raw radiometric output') - parser.add_argument('--video_output_path', + help='filepath to save noramlized viewable video') + parser.add_argument('--globalnorm_path', type=str, - help='filepath to save normalized video output') + help='filepath to save normalized radiometric output') + parser.add_argument('--metadata_path', + type=str, + help='filepath to save radiometric metadata') parser.add_argument('--bosonsdk_path', type=str, help='filepath to boson sdk') @@ -67,14 +73,32 @@ def main(): help='time in s to record') args = parser.parse_args() + # Camera Setup camera = BosonCapture(serial_port=args.serial_port, video_port=args.video_port, sdkpath=args.bosonsdk_path) - camera.start_recording(raw=args.raw_output_path, - norm=args.video_output_path) - sleep(args.recording_time or 10) - camera.stop_recording() - camera.release_camera() + camera.start_recording(autonorm=args.autonorm_path, + globalnorm=args.globalnorm_path, + meta=args.metadata_path) + max_recording_time = args.recording_time or 10800 # 3hrs default + start_time = time() + + # Recording Loop + try: + print('Recording ... Press CTRL+C to stop manually') + while True: + elapsed = time() - start_time + if elapsed >= max_recording_time: + print('Reached max recording time. Stopping recording') + break + sleep(.01) + + except KeyboardInterrupt: + print('Stopping Recording') + + finally: + camera.stop_recording() + camera.release_camera() if __name__ == '__main__': diff --git a/heatseek/boson_capture.py b/heatseek/boson_capture.py index b959fe5..ff143e7 100644 --- a/heatseek/boson_capture.py +++ b/heatseek/boson_capture.py @@ -11,7 +11,10 @@ import sys import os from time import sleep +import subprocess from importlib import import_module +from datetime import datetime +import yaml import numpy as np import cv2 from heatseek.capture import Capture @@ -46,7 +49,7 @@ def __init__(self, serial_port=None, video_port=None, sdkpath=None): parameters Args: - serial_port (str, opt): path to serial port. + serial_port (str, opt): path to serial port. Defaults to \dev\ttyACM0. video_port (str, opt): path to video port. Defaults to \dev\video0. @@ -60,8 +63,12 @@ def __init__(self, serial_port=None, video_port=None, sdkpath=None): self.recording = False self.height = 256 self.width = 320 - self.raw_data_fpath = None - self.viewable_video_fpath = None + self.GLOBAL_MIN = 28000 + self.GLOBAL_MAX = 32000 + self.n_frames = 0 + self.autonorm_fpath = None + self.globalnorm_fpath = None + self.metadata_fpath = None self.recording_thread = None self.raw_mm = None @@ -135,22 +142,25 @@ def take_image(self): print('Taking Image- PLACEHOLDER') - def start_recording(self, raw=None, norm=None): + def start_recording(self, autonorm=None, globalnorm=None, meta=None): """Begin thread for continuous recording Args: - raw (str, opt): filepath to save raw radiometric data. - Defaults to output.npy - norm (str, opt): filepath to save normalized mp4 video. - Defaults to output.mp4 + autonorm (str, opt): filepath to save frame by frame + noramlized video. Defaults to autonorm_{timestamp}.mp4 + globalnorm (str, opt): filepath to save globally + normalized mp4 video. Defaults to globalnorm_{timestamp}.mp4 """ if self.recording: print('Recording in Progress!') return - self.raw_data_fpath = raw or 'output.npy' - self.viewable_video_fpath = norm or 'output.mp4' + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + + self.autonorm_fpath = autonorm or f'autonorm_{timestamp}.mp4' + self.globalnorm_fpath = globalnorm or f'globalnorm_{timestamp}.mkv' + self.metadata_fpath = meta or f'metadata_{timestamp}.yaml' self.recording = True self.recording_thread = threading.Thread( @@ -159,6 +169,13 @@ def start_recording(self, raw=None, norm=None): self.recording_thread.start() print('Staring Recording...') + def _get_center_temp(self, frame): + """Internal method: returns pixel value of frame center + """ + + center = frame[int(self.height/2), int(self.width/2)] + return center + def _record_loop(self): """Internal method: recording loop to continuously capture frames @@ -167,25 +184,38 @@ def _record_loop(self): memory is full """ - # video settings + # MP4 Setup cap = cv2.VideoCapture(self.video_port, cv2.CAP_V4L2) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height) cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width) cap.set(cv2.CAP_PROP_CONVERT_RGB, 0) cap.set(cv2.CAP_PROP_FOURCC, - cv2.VideoWriter_fourcc('Y', '1', '6', ' ')) - fourcc = cv2.VideoWriter_fourcc(*'mp4v') - writer = cv2.VideoWriter(self.viewable_video_fpath, - fourcc, - 60, - (self.width, self.height)) + cv2.VideoWriter_fourcc('Y', '1', '6', ' ')) + mpv4_fourcc = cv2.VideoWriter_fourcc(*'mp4v') + autonorm_writer = cv2.VideoWriter(self.autonorm_fpath, + mpv4_fourcc, + 60, + (self.width, self.height)) + + # MKV Setup + proc = subprocess.Popen([ + "ffmpeg", "-y", + "-f", "rawvideo", + "-pixel_format", "gray", + "-video_size", f"{self.width}x{self.height}", + "-framerate", "60", + "-i", "-", + "-c:v", "libx264", + self.globalnorm_fpath + ], stdin=subprocess.PIPE) # 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 + # self.raw_mm = np.memmap('raw_test.raw', + # dtype=np.uint16, + # mode='w+', + # shape=(600, self.height, self.width)) + + self.n_frames = 0 try: while self.recording: @@ -194,24 +224,39 @@ def _record_loop(self): 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) + # Autonorm Frame (MP4) + frame_8bit_autonorm = cv2.normalize(frame, None, 0, 255, + norm_type=cv2.NORM_MINMAX).astype(np.uint8) + frame_color = cv2.applyColorMap(frame_8bit_autonorm, + cv2.COLORMAP_INFERNO).astype(np.uint8) + autonorm_writer.write(frame_color) + + # Globalnorm Frame (MKV) + frame_32bit = frame.astype(np.float32) + frame_8bit_globalnorm = (255.0 * (frame_32bit-self.GLOBAL_MIN)/(self.GLOBAL_MAX-self.GLOBAL_MIN)).astype(np.uint8) + proc.stdin.write(frame_8bit_globalnorm.tobytes()) # Raw Video - if frame_index > self.raw_mm.shape[0]: - print('Reached Preallocated Size, Stopping Recording') - break - self.raw_mm[frame_index] = frame - frame_index += 1 - self.raw_mm.flush() + # if frame_index < self.raw_mm.shape[0]: + # if self.n_frames > self.raw_mm.shape[0]: + # print('Reached Preallocated Size, Stopping Recording') + # break + # self.raw_mm[self.n_frames] = frame + # frame_index += 1 + # self.n_frames += 1 + # self.raw_mm.flush() + + self.n_frames += 1 finally: + cap.release() + autonorm_writer.release() + self._finalize_recording() + proc.stdin.close() + proc.wait() + def stop_recording(self): """Stop ongoing recording session. @@ -239,14 +284,33 @@ def _finalize_recording(self): self.recording = False - if self.raw_mm is not None: - self.raw_mm.flush() - del self.raw_mm - self.raw_mm = None + # if self.raw_mm is not None: + # self.raw_mm.flush() + # del self.raw_mm + # self.raw_mm = None + + self._write_radiometric_metadata() print('Recording Successfully Completed') - print(f'Radiometric data saved: {self.raw_data_fpath}') - print(f'Viewable Video: {self.viewable_video_fpath}') + print(f'Viewable Video: {self.autonorm_fpath}') + print(f'Globally normalized Video: {self.globalnorm_fpath}') + print(f'Radiometric Metadata: {self.metadata_fpath}') + + def _write_radiometric_metadata(self): + """Saves out all radiometric video metadata + """ + + meta = { + 'width': self.width, + 'height': self.height, + 'n_frames': self.n_frames, + 'min': self.GLOBAL_MIN, + 'max': self.GLOBAL_MAX, + 'dtype': 'uint16', + } + + with open(self.metadata_fpath, "w", encoding="utf-8") as f: + yaml.safe_dump(meta, f, sort_keys=False) def release_camera(self): """Release camera resources. diff --git a/heatseek/cli.py b/heatseek/cli.py index f31a2a3..05d395f 100644 --- a/heatseek/cli.py +++ b/heatseek/cli.py @@ -5,7 +5,7 @@ import requests from .data_utils import download_dataset, write_data_yaml from .train import train -from .preprocess import reduce_background +from .preprocess import reduce_background, reduce_background_radiometric from .detect_track import detect_and_track from .density_annotator import annotate_folder def main(): @@ -58,6 +58,17 @@ def main(): help="YAML config for optical‐flow thresholds", ) + # preprocess + pre_r = subs.add_parser("preprocess_radiometric", help="Reduce background in a radiometric video") + pre_r.add_argument("--input", required=True, help="Input video path") + pre_r.add_argument("--output", required=True, help="Output video path") + pre_r.add_argument("--metadata", required=True, help="Radiometic video metadata path") + pre_r.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") @@ -112,6 +123,9 @@ def main(): elif args.cmd == "preprocess": reduce_background(args.input, args.output, args.config) + elif args.cmd == "preprocess_radiometric": + reduce_background_radiometric(args.input, args.output, args.metadata, args.config) + elif args.cmd == "track": detect_and_track(args.input, args.output, args.weights) diff --git a/heatseek/config/preproc_config.yaml b/heatseek/config/preproc_config.yaml index ea0127f..d420de8 100644 --- a/heatseek/config/preproc_config.yaml +++ b/heatseek/config/preproc_config.yaml @@ -1,5 +1,5 @@ # preproc_config.yaml -motion_thresh: 2.75 +motion_thresh: 0.5 flow_params: pyr_scale: 0.5 levels: 3 diff --git a/heatseek/preprocess.py b/heatseek/preprocess.py index 1a7cdbb..03434d0 100644 --- a/heatseek/preprocess.py +++ b/heatseek/preprocess.py @@ -1,3 +1,4 @@ +import subprocess import cv2 import numpy as np from tqdm import tqdm @@ -83,4 +84,105 @@ def reduce_background(in_path: str, out_path: str, yaml_path: str = "heatseek/co cap.release() out.release() - print(f"[video_preproc] saved → {out_path}") \ No newline at end of file + print(f"[video_preproc] saved → {out_path}") + + +def reduce_background_radiometric(in_path: str, + out_path: str, + meta_path: str, + yaml_path: str = "heatseek/config/preproc_config.yaml"): + """ + Optical flow background reduction using parameters from a YAML config on radiometric data. + + Args: + in_path (str): Path to input video file. + out_path (str): Path to save processed video. + meta_path (str): Path to radiometric metadata. + yaml_path (str): Path to YAML configuration file containing motion_thresh and flow_params. + """ + + # Load Config + 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) + + # Load Metadata + # with open(meta_path, 'r') as f: + # meta = yaml.safe_load(f) + # + # fps = 60 + # width = meta.get('width', 320) + # height = meta.get('height', 256) + # n_frames = meta.get('n_frames', 600) + # fourcc = cv2.VideoWriter_fourcc(*"mp4v") + # out = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) + + cap = cv2.VideoCapture(in_path) + if not cap.isOpened(): + raise RuntimeError(f'Failed to 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)) + n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + # Start FFmpeg process + proc = subprocess.Popen([ + "ffmpeg", "-y", + "-f", "rawvideo", + "-pixel_format", "bgr24", + "-video_size", f"{width}x{height}", + "-framerate", f"{int(fps)}", + "-i", "-", # stdin + "-c:v", "libx264", + "-pix_fmt", "yuv420p", # standard for h264 + out_path + ], stdin=subprocess.PIPE) + + # Read first frame + ret, prev_frame = cap.read() + if not ret: + raise RuntimeError("Failed to read the first frame from the video.") + + prev_gray = prev_frame[:,:,0] + + for _ in tqdm(range(n_frames - 1), desc="bg-reduce"): + ret, frame = cap.read() + if not ret: + break + + gray = frame[:,:,0] + + 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 + dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5) + mask = (dist > 1.5).astype(np.uint8) * 255 + mask_color = cv2.merge([mask, mask, mask]) + + fg = cv2.bitwise_and(frame, mask_color) + + # Write frame to FFmpeg stdin + proc.stdin.write(fg.tobytes()) + + prev_gray = gray + + cap.release() + proc.stdin.close() + proc.wait() + + print(f'video saved - {out_path}')