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 time
import orisys
from torchvision.utils import flow_to_image
import matplotlib.pyplot as plt
# from torchvision.utils import flow_to_image
# import matplotlib.pyplot as plt
def plot(imgs, **imshow_kwargs):
if not isinstance(imgs[0], list):
# Make a 2d grid even if there's just 1 row
imgs = [imgs]
# def plot(imgs, **imshow_kwargs):
# if not isinstance(imgs[0], list):
# # Make a 2d grid even if there's just 1 row
# imgs = [imgs]
num_rows = len(imgs)
num_cols = len(imgs[0])
_, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
for row_idx, row in enumerate(imgs):
for col_idx, img in enumerate(row):
ax = axs[row_idx, col_idx]
img = F.to_pil_image(img.to("cpu"))
ax.imshow(np.asarray(img), **imshow_kwargs)
ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
# num_rows = len(imgs)
# num_cols = len(imgs[0])
# _, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
# for row_idx, row in enumerate(imgs):
# for col_idx, img in enumerate(row):
# ax = axs[row_idx, col_idx]
# img = F.to_pil_image(img.to("cpu"))
# ax.imshow(np.asarray(img), **imshow_kwargs)
# ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
plt.tight_layout()
# plt.tight_layout()
def preprocess(batch, resize=None):
"""
@@ -66,7 +66,7 @@ def setup_simple_raft():
# return raft
# except:
# pass
model = raft_small(pretrained=True, progress=False).to('cuda')
model = raft_large(pretrained=True, progress=False).to('cuda')
model = model.eval()
# Alternative: Use LiteFlowNet (lighter alternative)
print("RAFT loaded successfully")
@@ -96,13 +96,13 @@ def compute_dense_optical_flow_simple(prev_gray, curr_gray):
def main_simple():
# 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()
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
# cap = cv2.VideoCapture(0)
# cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
# cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
prev_gray = None
fps_buffer = []
@@ -115,11 +115,21 @@ def main_simple():
# # Convert to grayscale
# gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
sensor.get_img()
frame = sensor.img.copy()
ret, frame = cap.read()
if not ret:
break
dx = 300
dy = 0
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:
prev_gray = frame.copy
prev_gray = frame.copy()
gray = frame
# Compute optical flow
if prev_gray is not None:
@@ -174,13 +184,13 @@ def main_simple():
# Transpose to (H, W, 2) format expected by draw_arrows
flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2)
if flow_track is None:
flow_track = np.zeros_like(flow_np)
else:
if abs(np.mean(flow_np)) >0.1:
flow_track = flow_track + flow_np
# if flow_track is None:
# flow_track = np.zeros_like(flow_np)
# else:
# if abs(np.mean(flow_np)) >0.1:
# flow_track = flow_track + flow_np
flow_np = flow_track
# flow_np = flow_track
# Flow is already in pixel coordinates, no denormalization needed
# Keep as float32 for draw_arrows (it can handle float values)
@@ -191,8 +201,8 @@ def main_simple():
arrows = orisys.util.draw_arrows(
gray,
flow_np,
threshold=1,
grid_spacing=10,
threshold=2,
grid_spacing=40,
arrow_scale=1.0
)
@@ -202,6 +212,7 @@ def main_simple():
# cv2.imshow('Original', frame)
# cv2.imshow('Grayscale', gray)
arrows = cv2.resize(arrows, (1200, 400))
cv2.imshow('Optical Flow', arrows)
# prev_gray = gray.copy()
src/orisys/Stitcher.py
+4 -3
View File
@@ -170,9 +170,10 @@ class Stitcher:
def extract_tiles(self):
if not self.isstitch:
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 = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1])
tile_idx = 7
# if tile_idx is None:
# tile_idx = 7
# Ensure index is always iterable (list or array)
self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx
self.ncol = 1
src/orisys/sensor.py
+4 -3
View File
@@ -166,8 +166,8 @@ class Sensor:
return None
self.frame = frame
# Convert to grayscale
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
gray = frame
# Extract and stitch tiles
self.stitcher.get_stitch_unit(gray)
self._extract_tiles(gray)
@@ -215,7 +215,8 @@ class Sensor:
if self.img is None:
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'):
flow = flow.download()
self.flow = flow