From 233ae841038e177b2ded3e1f90c4db18da10b69a Mon Sep 17 00:00:00 2001 From: bluepulsar Date: Thu, 12 Feb 2026 10:15:54 +0800 Subject: [PATCH] update torch conditional import --- config/prototype_raft.json | 2 +- promote_video.py | 8 ++++++-- src/orisys/sensor.py | 4 ++-- src/orisys/tracker.py | 12 ++++++++---- 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/config/prototype_raft.json b/config/prototype_raft.json index 097e2bb..bbc0619 100644 --- a/config/prototype_raft.json +++ b/config/prototype_raft.json @@ -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, diff --git a/promote_video.py b/promote_video.py index bbbf127..6f7b29a 100644 --- a/promote_video.py +++ b/promote_video.py @@ -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 diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index ef08047..cfdbcda 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -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 diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index 9e31628..bc484da 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -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)