orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

raft working

Browse Source
This commit is contained in:
bluepulsar
2026-01-31 12:24:15 +08:00
parent def7dd6916
commit e2100204aa
3 changed files with 51 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
RAFT_test.py
+43 -32
View File
@@ -8,25 +8,25 @@ from torchvision.models.optical_flow import raft_large, raft_small
import numpy as np import numpy as np
import time import time
import orisys import orisys
from torchvision.utils import flow_to_image # from torchvision.utils import flow_to_image
import matplotlib.pyplot as plt # import matplotlib.pyplot as plt
def plot(imgs, **imshow_kwargs): # def plot(imgs, **imshow_kwargs):
if not isinstance(imgs[0], list): # if not isinstance(imgs[0], list):
# Make a 2d grid even if there's just 1 row # # Make a 2d grid even if there's just 1 row
imgs = [imgs] # imgs = [imgs]
num_rows = len(imgs) # num_rows = len(imgs)
num_cols = len(imgs[0]) # num_cols = len(imgs[0])
_, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False) # _, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
for row_idx, row in enumerate(imgs): # for row_idx, row in enumerate(imgs):
for col_idx, img in enumerate(row): # for col_idx, img in enumerate(row):
ax = axs[row_idx, col_idx] # ax = axs[row_idx, col_idx]
img = F.to_pil_image(img.to("cpu")) # img = F.to_pil_image(img.to("cpu"))
ax.imshow(np.asarray(img), **imshow_kwargs) # ax.imshow(np.asarray(img), **imshow_kwargs)
ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[]) # ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
plt.tight_layout() # plt.tight_layout()
def preprocess(batch, resize=None): def preprocess(batch, resize=None):
""" """
@@ -66,7 +66,7 @@ def setup_simple_raft():
# return raft # return raft
# except: # except:
# pass # pass
model = raft_small(pretrained=True, progress=False).to('cuda') model = raft_large(pretrained=True, progress=False).to('cuda')
model = model.eval() model = model.eval()
# Alternative: Use LiteFlowNet (lighter alternative) # Alternative: Use LiteFlowNet (lighter alternative)
print("RAFT loaded successfully") print("RAFT loaded successfully")
@@ -96,13 +96,13 @@ def compute_dense_optical_flow_simple(prev_gray, curr_gray):
def main_simple(): def main_simple():
# Initialize camera # Initialize camera
sensor = orisys.Sensor('s1.mp4', config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) # sensor = orisys.Sensor('s1.mp4', config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
model = setup_simple_raft() model = setup_simple_raft()
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}") print(f"Using device: {device}")
# cap = cv2.VideoCapture(0) cap = cv2.VideoCapture(0)
# cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
# cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
prev_gray = None prev_gray = None
fps_buffer = [] fps_buffer = []
@@ -115,11 +115,21 @@ def main_simple():
# # Convert to grayscale # # Convert to grayscale
# gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
sensor.get_img() ret, frame = cap.read()
frame = sensor.img.copy() if not ret:
break
dx = 300
dy = 0
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame1 = frame[420:660, 840:1080, :]
frame2 = frame[420+dy:660+dy, 840+dx:1080+dx, :]
frame3 = frame[420-dy:660-dy, 840-dx:1080-dx, :]
frame = np.concatenate([frame3, frame1, frame2], axis=1)
# cv2.imshow("frame", frame3)
# cv2.waitKey(0)
if prev_gray is None: if prev_gray is None:
prev_gray = frame.copy prev_gray = frame.copy()
gray = frame gray = frame
# Compute optical flow # Compute optical flow
if prev_gray is not None: if prev_gray is not None:
@@ -174,13 +184,13 @@ def main_simple():
# Transpose to (H, W, 2) format expected by draw_arrows # Transpose to (H, W, 2) format expected by draw_arrows
flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2) flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2)
if flow_track is None: # if flow_track is None:
flow_track = np.zeros_like(flow_np) # flow_track = np.zeros_like(flow_np)
else: # else:
if abs(np.mean(flow_np)) >0.1: # if abs(np.mean(flow_np)) >0.1:
flow_track = flow_track + flow_np # flow_track = flow_track + flow_np
flow_np = flow_track # flow_np = flow_track
# Flow is already in pixel coordinates, no denormalization needed # Flow is already in pixel coordinates, no denormalization needed
# Keep as float32 for draw_arrows (it can handle float values) # Keep as float32 for draw_arrows (it can handle float values)
@@ -191,8 +201,8 @@ def main_simple():
arrows = orisys.util.draw_arrows( arrows = orisys.util.draw_arrows(
gray, gray,
flow_np, flow_np,
threshold=1, threshold=2,
grid_spacing=10, grid_spacing=40,
arrow_scale=1.0 arrow_scale=1.0
) )
@@ -202,6 +212,7 @@ def main_simple():
# cv2.imshow('Original', frame) # cv2.imshow('Original', frame)
# cv2.imshow('Grayscale', gray) # cv2.imshow('Grayscale', gray)
arrows = cv2.resize(arrows, (1200, 400))
cv2.imshow('Optical Flow', arrows) cv2.imshow('Optical Flow', arrows)
# prev_gray = gray.copy() # prev_gray = gray.copy()
src/orisys/Stitcher.py
+3 -2
View File
@@ -170,9 +170,10 @@ class Stitcher:
def extract_tiles(self): def extract_tiles(self):
if not self.isstitch: if not self.isstitch:
tile_idx = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1]) # tile_idx = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1])
if tile_idx is None:
tile_idx = 7 tile_idx = 7
# if tile_idx is None:
# tile_idx = 7
# Ensure index is always iterable (list or array) # Ensure index is always iterable (list or array)
self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx
self.ncol = 1 self.ncol = 1
src/orisys/sensor.py
+4 -3
View File
@@ -166,8 +166,8 @@ class Sensor:
return None return None
self.frame = frame self.frame = frame
# Convert to grayscale # Convert to grayscale
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
gray = frame
# Extract and stitch tiles # Extract and stitch tiles
self.stitcher.get_stitch_unit(gray) self.stitcher.get_stitch_unit(gray)
self._extract_tiles(gray) self._extract_tiles(gray)
@@ -215,7 +215,8 @@ class Sensor:
if self.img is None: if self.img is None:
raise Exception("Warning: Get image first") raise Exception("Warning: Get image first")
flow = self.__tracker.track(self.img) # flow = self.__tracker.track(self.img)
flow = self.__tracker.track_raft(self.img)
if hasattr(flow, 'download'): if hasattr(flow, 'download'):
flow = flow.download() flow = flow.download()
self.flow = flow self.flow = flow