diff --git a/.gitignore b/.gitignore index d334382..fc1eda8 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,5 @@ orisys.egg-info/ /old_orisys/ /dist_nuitka/ /build_nuitka -/data/ \ No newline at end of file +/data/ +/models/ \ No newline at end of file diff --git a/RAFT_test.py b/RAFT_test.py index 6896f4a..7b454fb 100644 --- a/RAFT_test.py +++ b/RAFT_test.py @@ -50,14 +50,27 @@ def track_raft(prev_frame, frame, model, device): curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0) prev_tensor = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device) curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device) + t0 = time.perf_counter() with torch.no_grad(): flow = model(prev_tensor, curr_tensor) predicted_flows = flow[-1] - + t1 = time.perf_counter() + print(f"RAFT time: {t1 - t0}") flow_np = predicted_flows[0].cpu().numpy() # or .squeeze(0) # Transpose to (H, W, 2) format expected by draw_arrows flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2) + + + # torch.onnx.export( + # model, # model to export + # (prev_tensor, curr_tensor), # inputs of the model, + # "my_model.onnx", # filename of the ONNX model + # input_names=["prev_tensor", "curr_tensor"], # Rename inputs for the ONNX model + # dynamo=True # True or False to select the exporter to use + # ) + # print("ONNX model exported successfully") + return flow_np def main_simple(): @@ -66,7 +79,7 @@ def main_simple(): model = setup_simple_raft() device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") - cap = cv2.VideoCapture(0) + cap = cv2.VideoCapture("s1.mp4") cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080) diff --git a/config/compound_eye.png b/config/compound_eye.png index 3cfe523..e3ee26c 100644 Binary files a/config/compound_eye.png and b/config/compound_eye.png differ diff --git a/config/prototype_of3.json b/config/prototype_of3.json index c8a0345..9b2dd1d 100644 --- a/config/prototype_of3.json +++ b/config/prototype_of3.json @@ -2,11 +2,11 @@ "name": "prototype", "version": "1.0", "camera": { - "resolution_raw": [1920, 1080], + "resolution_raw": [960, 540], "fps": 30, "units": 15, "units_used": 9, - "resolution_stitch": [400, 400] + "resolution_stitch": [200, 200] }, "stitcher": { "min_area_frac": 0.01, diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index 5028119..f984a97 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -2,7 +2,8 @@ import cv2 import numpy as np import torch import torchvision.transforms as T -from torchvision.models.optical_flow import raft_large +from torchvision.models.optical_flow import raft_large, raft_small +import time class Tracker: def __init__(self, @@ -54,6 +55,7 @@ class Tracker: self.__upper_threshold = config_tracker["upper_threshold"] # error greater than threshold will add new reference frame ### RAFT self.__raft_device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"RAFT device: {self.__raft_device}") self.__raft_model = self.setup_simple_raft() self.__raft_ref_frame = None self.__raft_ref_tensor = None @@ -250,7 +252,7 @@ class Tracker: model = model.eval() # Optimize matmul precision for faster inference (TF32 on Ampere+ GPUs) - torch.set_float32_matmul_precision('medium') + # torch.set_float32_matmul_precision('medium') # Enable CUDA-specific optimizations if self.__raft_device == 'cuda': @@ -263,7 +265,7 @@ class Tracker: print(f"RAFT model loaded on {self.__raft_device}") return model - + def _preprocess_frame_to_tensor(self, frame, resize=None): """Efficiently convert numpy frame to RAFT-ready tensor""" # Ensure contiguous memory for faster tensor conversion @@ -296,7 +298,6 @@ class Tracker: # Add batch dimension tensor = tensor.unsqueeze(0) - # Keep in float32 - autocast will handle mixed precision during inference return tensor def track_raft(self, frame): @@ -319,16 +320,22 @@ class Tracker: prev_tensor = self.__raft_ref_tensor curr_tensor = self._preprocess_frame_to_tensor(frame) - # OPTIMIZATION: Remove CUDA stream since immediate sync eliminates async benefit - # The bottleneck is inference, not data transfer (already optimized with pin_memory) - with torch.no_grad(): + # Synchronize before timing to ensure preprocessing is complete + if self.__raft_device == 'cuda': + torch.cuda.synchronize() + + with torch.inference_mode(): if self.__raft_device == 'cuda': with torch.amp.autocast('cuda'): flow = self.__raft_model(prev_tensor, curr_tensor) else: flow = self.__raft_model(prev_tensor, curr_tensor) predicted_flows = flow[-1] - + + # Synchronize after inference to get accurate timing + if self.__raft_device == 'cuda': + torch.cuda.synchronize() + # Optimize tensor->numpy conversion: contiguous() ensures optimal memory layout flow_np = predicted_flows[0].contiguous().cpu().numpy()