diff --git a/RAFT_test.py b/RAFT_test.py index b06d9c8..020ae58 100644 --- a/RAFT_test.py +++ b/RAFT_test.py @@ -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() diff --git a/src/orisys/Stitcher.py b/src/orisys/Stitcher.py index e1a9403..259067d 100644 --- a/src/orisys/Stitcher.py +++ b/src/orisys/Stitcher.py @@ -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 diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index f0a0fb1..8e70cf7 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -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