add multi-sensor support
This commit is contained in:
bluepulsar
+45
-19
@@ -16,7 +16,7 @@ class Sensor:
|
||||
Sensor helper that can decompose optical flow via Fourier methods.
|
||||
Stores the decomposed components and global norms as attributes after compute().
|
||||
"""
|
||||
def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True, config_name="prototype", override=False, box=None):
|
||||
def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True, config_name="prototype", cal_path=None,override=False, box=None):
|
||||
"""
|
||||
Initialize Sensor with video source and configuration.
|
||||
|
||||
@@ -38,7 +38,7 @@ class Sensor:
|
||||
if vid_src is None:
|
||||
raise Exception("User must specify a video source.")
|
||||
|
||||
self.vid_src, self.isstitch, self.cuda, self.verbose, self.config_name = vid_src, isstitch, cuda, verbose, config_name
|
||||
self.vid_src, self.isstitch, self.cuda, self.verbose, self.config_name, self.cal_path = vid_src, isstitch, cuda, verbose, config_name, cal_path
|
||||
|
||||
# Initialize state variables
|
||||
self.stitch_valid = self.imgupdate = self.vupdate = False
|
||||
@@ -58,7 +58,7 @@ class Sensor:
|
||||
self.__open_video_source(override, box)
|
||||
|
||||
# Logging state
|
||||
self.last_log_time = time.time()
|
||||
self.last_log_time = time.perf_counter()
|
||||
self.fps = 0
|
||||
|
||||
if self.verbose:
|
||||
@@ -102,9 +102,7 @@ class Sensor:
|
||||
|
||||
if override:
|
||||
if box is not None:
|
||||
|
||||
self.isstitch = False
|
||||
self.__stitcher.isstitch = False
|
||||
self.__stitcher.isstitch = self.isstitch = False
|
||||
self.__stitcher.boxes = box
|
||||
else:
|
||||
print("Override is True but no box provided.")
|
||||
@@ -137,7 +135,7 @@ class Sensor:
|
||||
print(" Resolution after stitch: ", str(self.config_resolution_stitch))
|
||||
print(" Index for stitching: ", str(self.config_index))
|
||||
|
||||
def get_img(self):
|
||||
def get_img(self, log_timing=False):
|
||||
"""
|
||||
Capture frame from video source, convert to grayscale, and perform stitching.
|
||||
|
||||
@@ -147,14 +145,32 @@ class Sensor:
|
||||
ret, frame = self.cap.read()
|
||||
if not ret:
|
||||
print("Failed to get_img from video source.")
|
||||
self.cap.release()
|
||||
exit(1)
|
||||
return None
|
||||
|
||||
if log_timing:
|
||||
timings = {}
|
||||
t_start = time.perf_counter()
|
||||
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
if log_timing:
|
||||
timings['cvt_color'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
self.__stitcher.get_stitch_unit(gray_frame)
|
||||
self.__get_tiles(gray_frame)
|
||||
if log_timing:
|
||||
timings['get_tiles'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Step 1: Optical flow tracking
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
img = np.ascontiguousarray(self.__stitcher.stitch_rt_optimized())
|
||||
|
||||
if log_timing:
|
||||
timings['stitch'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
img_resized = cv2.resize(img, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
@@ -171,6 +187,18 @@ class Sensor:
|
||||
# Convert back to uint8
|
||||
self.img = self.img_filtered.astype(np.uint8)
|
||||
self.imgupdate = True
|
||||
|
||||
if log_timing:
|
||||
timings['get_img'] = (time.perf_counter() - t_start) * 1000
|
||||
print("\n[get_img timing]")
|
||||
print(f" cvt_color: {timings['cvt_color']:7.2f} ms")
|
||||
print(f" get_tiles: {timings['get_tiles']:7.2f} ms")
|
||||
print(f" stitch: {timings['stitch']:7.2f} ms")
|
||||
print(f" ─────────────────────────")
|
||||
print(f" TOTAL: {timings['get_img']:7.2f} ms ({1000/timings['get_img']:.1f} FPS)")
|
||||
|
||||
# Store timings for external access
|
||||
self._last_timings = timings
|
||||
return self.img
|
||||
|
||||
def compute_deformation(self, log_timing=False):
|
||||
@@ -199,8 +227,6 @@ class Sensor:
|
||||
if log_timing:
|
||||
timings['optical_flow'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
|
||||
|
||||
if hasattr(flow, 'download'):
|
||||
flow = flow.download()
|
||||
|
||||
@@ -239,8 +265,8 @@ class Sensor:
|
||||
timings['get_density'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Calculate FPS
|
||||
self.fps = 1/(time.time() - self.last_log_time)
|
||||
self.last_log_time = time.time()
|
||||
self.fps = 1/((time.perf_counter() - self.last_log_time))
|
||||
self.last_log_time = time.perf_counter()
|
||||
|
||||
# Print timing summary
|
||||
if log_timing:
|
||||
@@ -351,7 +377,7 @@ class Sensor:
|
||||
n_boxes, v_mean0, timeout = 0, 0, 5
|
||||
start_time = cv2.getTickCount()
|
||||
|
||||
# print("Waiting for camera stabilization...")
|
||||
# Wait for camera stabilization
|
||||
while n_boxes < self.config_units:
|
||||
ret, frame = self.cap.read()
|
||||
v_mean = np.sum(cv2.mean(frame))
|
||||
@@ -361,8 +387,7 @@ class Sensor:
|
||||
|
||||
if (cv2.getTickCount() - start_time) / cv2.getTickFrequency() > timeout:
|
||||
print("Timeout for finding valid vision units.")
|
||||
self.isstitch = False
|
||||
self.__stitcher.isstitch = False
|
||||
self.__stitcher.isstitch = self.isstitch = False
|
||||
break
|
||||
v_mean0 = v_mean
|
||||
|
||||
@@ -390,13 +415,16 @@ class Sensor:
|
||||
- Creates new calibration and saves to ./config/cal.npy if not found
|
||||
- Calibration contains transformation parameters for tile alignment
|
||||
"""
|
||||
cal_path = "./config/cal.npy"
|
||||
if self.cal_path is None:
|
||||
self.cal_path = "./config/cal.npy"
|
||||
cal_path = self.cal_path
|
||||
|
||||
if os.path.exists(cal_path):
|
||||
print("Loading existing calibration from", cal_path)
|
||||
self.stitch_calibration = np.load(cal_path, allow_pickle=True)
|
||||
self.__stitcher.stitch_calibration = self.stitch_calibration
|
||||
img_stitched_height = self.__stitcher.stitch_rt_optimized()
|
||||
|
||||
else:
|
||||
print("Calibration file not found, creating new calibration...")
|
||||
img_stitched_height, self.stitch_calibration = self.__stitcher.stitch_cal()
|
||||
@@ -404,6 +432,4 @@ class Sensor:
|
||||
np.save(cal_path, self.stitch_calibration)
|
||||
print("Calibration saved to", cal_path)
|
||||
|
||||
# print("Image stitched with shape:", img_stitched_height.shape)
|
||||
return img_stitched_height
|
||||
|
||||
|
||||
Reference in New Issue
Block a user