update torch conditional import
This commit is contained in:
@@ -17,7 +17,7 @@
|
||||
"nominal_ov_pos": 100,
|
||||
"n": [3,3],
|
||||
"row_tol": 0.6,
|
||||
"box_size": [250, 250]
|
||||
"box_size": [220, 220]
|
||||
},
|
||||
"tracker": {
|
||||
"adaptive": true,
|
||||
|
||||
+6
-2
@@ -127,7 +127,7 @@ def draw_magnitude_map(vfield, threshold=3.0, colormap=cv2.COLORMAP_JET) -> np.n
|
||||
|
||||
|
||||
|
||||
npoints = 100 # 图表数据长度
|
||||
npoints = 400 # 图表数据长度
|
||||
plot_threshold = 0 # only plot force values above this (below → 0)
|
||||
pg.setConfigOption('background', bg_color_hex)
|
||||
pg.setConfigOption('foreground', bg_color_hex)
|
||||
@@ -154,7 +154,7 @@ def main():
|
||||
pyqt_writer = cv2.VideoWriter(args.video.replace(".mp4", "_pyqt.mp4"), cv2.VideoWriter_fourcc(*'mp4v'), fps_video, (1600, 1200))
|
||||
depth_writer = cv2.VideoWriter(args.video.replace(".mp4", "_depth.mp4"), cv2.VideoWriter_fourcc(*'mp4v'),fps_video, (384, 384))
|
||||
|
||||
sensor1 = orisys.Sensor(args.video, config_name="./config/prototype_of3.json", cal_path= args.video.replace(".mp4", ".npy"), isstitch=True, cuda=True, verbose=True)
|
||||
sensor1 = orisys.Sensor(args.video, config_name="./config/prototype_raft.json", cal_path= args.video.replace(".mp4", ".npy"), isstitch=True, cuda=True, verbose=True)
|
||||
|
||||
start_time = time.time()
|
||||
is_save = False
|
||||
@@ -203,6 +203,10 @@ def main():
|
||||
data[:-1, 1, i] = data[1:, 1, i] # shift force data left
|
||||
data[-1, 0, :] = frame_count/fps_video # set current time for all plots
|
||||
data[-1, 1, 0] = fn
|
||||
if abs(fx) - data[-2, 1, 1] > 100000:
|
||||
fx = data[-2, 1, 1] + data[-2, 1, 1]*0.5
|
||||
if abs(fy) - data[-2, 1, 2] > 100000:
|
||||
fy = data[-2, 1, 2] + data[-2, 1, 2]*0.5
|
||||
data[-1, 1, 1] = abs(fx)
|
||||
data[-1, 1, 2] = abs(fy)
|
||||
# Moving average on force data (column 1): simple loop, no singularities at ends
|
||||
|
||||
@@ -12,8 +12,8 @@ import cv2
|
||||
import time
|
||||
from enum import Enum, auto
|
||||
from typing import Tuple
|
||||
from threading import Thread, Event
|
||||
from queue import Queue, Empty
|
||||
# from threading import Thread, Event
|
||||
# from queue import Queue, Empty
|
||||
|
||||
from .Stitcher import Stitcher
|
||||
from .tracker import Tracker
|
||||
|
||||
@@ -25,11 +25,15 @@ class Tracker:
|
||||
dis_preset = config_tracker["dis_preset"]
|
||||
loop_closure_interval = config_tracker["loop_closure_interval"]
|
||||
self.__algorithm = config_tracker["algorithm"]
|
||||
|
||||
# Conditionally import torch only if RAFT algorithm is used
|
||||
self.torch = None
|
||||
if self.__algorithm == "raft":
|
||||
import torch
|
||||
self.torch = torch
|
||||
|
||||
|
||||
if self.__algorithm == "raft" and (not self.cuda or not torch.cuda.is_available()):
|
||||
if self.__algorithm == "raft" and (not self.cuda or not self.torch.cuda.is_available()):
|
||||
self.__algorithm = "dis"
|
||||
if self.verbose:
|
||||
print("No CUDA available, falling back to OPENCV DIS optical flow")
|
||||
@@ -217,7 +221,7 @@ class Tracker:
|
||||
frame = np.ascontiguousarray(frame)
|
||||
|
||||
# Upload to device as uint8 first to save bandwidth (4x less data than float32)
|
||||
tensor = torch.from_numpy(frame).permute(2, 0, 1)
|
||||
tensor = self.torch.from_numpy(frame).permute(2, 0, 1)
|
||||
|
||||
# Pin memory for faster CPU-GPU transfer (similar to DataLoader pin_memory=True)
|
||||
if self.__raft_device == 'cuda':
|
||||
@@ -258,9 +262,9 @@ class Tracker:
|
||||
if self.verbose:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
with torch.inference_mode():
|
||||
with self.torch.inference_mode():
|
||||
if self.__raft_device == 'cuda':
|
||||
with torch.amp.autocast('cuda'):
|
||||
with self.torch.amp.autocast('cuda'):
|
||||
flow = self.__raft_model(prev_tensor, curr_tensor)
|
||||
else:
|
||||
flow = self.__raft_model(prev_tensor, curr_tensor)
|
||||
|
||||
Reference in New Issue
Block a user