diff --git a/node-hub/dora-cotracker/README.md b/node-hub/dora-cotracker/README.md new file mode 100644 index 00000000..2d5f1217 --- /dev/null +++ b/node-hub/dora-cotracker/README.md @@ -0,0 +1,40 @@ +# dora-cotracker + +## Getting started + +- Install it with uv: + +```bash +uv venv -p 3.11 --seed +uv pip install -e . +``` + +## Contribution Guide + +- Format with [ruff](https://docs.astral.sh/ruff/): + +```bash +uv pip install ruff +uv run ruff check . --fix +``` + +- Lint with ruff: + +```bash +uv run ruff check . +``` + +- Test with [pytest](https://github.com/pytest-dev/pytest) + +```bash +uv pip install pytest +uv run pytest . # Test +``` + +## YAML Specification + +## Examples + +## License + +dora-cotracker's code are released under the MIT License diff --git a/node-hub/dora-cotracker/demo.yml b/node-hub/dora-cotracker/demo.yml new file mode 100644 index 00000000..102d1e76 --- /dev/null +++ b/node-hub/dora-cotracker/demo.yml @@ -0,0 +1,39 @@ +nodes: + - id: camera + build: pip install opencv-video-capture + path: opencv-video-capture + inputs: + tick: dora/timer/millis/100 + outputs: + - image + env: + CAPTURE_PATH: "0" + ENCODING: "rgb8" + IMAGE_WIDTH: "640" + IMAGE_HEIGHT: "480" + + - id: tracker + build: pip install -e . + path: dora-cotracker + inputs: + image: camera/image + # points_to_track: debug/points_to_track + outputs: + - tracked_image + - tracked_points + + - id: plot + build: pip install dora-rerun + path: dora-rerun + inputs: + image: camera/image + tracked_image: tracker/tracked_image + # points: tracker/tracked_points + + - id: debug + build: pip install -e . + path: dora-sgp_debug_node + inputs: + points: tracker/tracked_points + # outputs: + # - points_to_track \ No newline at end of file diff --git a/node-hub/dora-cotracker/dora_cotracker/__init__.py b/node-hub/dora-cotracker/dora_cotracker/__init__.py new file mode 100644 index 00000000..ac3cbef9 --- /dev/null +++ b/node-hub/dora-cotracker/dora_cotracker/__init__.py @@ -0,0 +1,11 @@ +import os + +# Define the path to the README file relative to the package directory +readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md") + +# Read the content of the README file +try: + with open(readme_path, "r", encoding="utf-8") as f: + __doc__ = f.read() +except FileNotFoundError: + __doc__ = "README file not found." diff --git a/node-hub/dora-cotracker/dora_cotracker/__main__.py b/node-hub/dora-cotracker/dora_cotracker/__main__.py new file mode 100644 index 00000000..bcbfde6d --- /dev/null +++ b/node-hub/dora-cotracker/dora_cotracker/__main__.py @@ -0,0 +1,5 @@ +from .main import main + + +if __name__ == "__main__": + main() diff --git a/node-hub/dora-cotracker/dora_cotracker/main.py b/node-hub/dora-cotracker/dora_cotracker/main.py new file mode 100644 index 00000000..b12d9613 --- /dev/null +++ b/node-hub/dora-cotracker/dora_cotracker/main.py @@ -0,0 +1,128 @@ +import numpy as np +import pyarrow as pa +from dora import Node +import cv2 +import torch +from collections import deque + +class VideoTrackingNode: + def __init__(self): + self.node = Node("video-tracking-node") + + # Initialize CoTracker + self.device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Using device: {self.device}") + self.model = torch.hub.load("facebookresearch/co-tracker", "cotracker3_online") + self.model = self.model.to(self.device) + + # Initialize tracking variables + self.buffer_size = self.model.step * 2 + self.window_frames = deque(maxlen=self.buffer_size) + self.is_first_step = True + self.grid_size = 10 # Smaller grid for better visualization + self.grid_query_frame = 0 + self.frame_count = 0 + + def process_tracking(self, frame): + """Process frame for tracking""" + if len(self.window_frames) == self.buffer_size: + try: + # Stack frames and convert to tensor + video_chunk = torch.tensor( + np.stack(list(self.window_frames)), + device=self.device + ).float() + + # Normalize pixel values to [0, 1] + video_chunk = video_chunk / 255.0 + + # Reshape to [B,T,C,H,W] + video_chunk = video_chunk.permute(0, 3, 1, 2)[None] + + # Run tracking with grid parameters + pred_tracks, pred_visibility = self.model( + video_chunk, + is_first_step=self.is_first_step, + grid_size=self.grid_size, + grid_query_frame=self.grid_query_frame + ) + self.is_first_step = False + + if pred_tracks is not None and pred_visibility is not None: + # Get the latest tracks and visibility + tracks = pred_tracks[0, -1].cpu().numpy() + visibility = pred_visibility[0, -1].cpu().numpy() + + # Filter high-confidence points + visible_mask = visibility > 0.5 + visible_tracks = tracks[visible_mask] + + # Send tracked points + if len(visible_tracks) > 0: + self.node.send_output( + "tracked_points", + pa.array(visible_tracks.ravel()), + { + "num_points": len(visible_tracks), + "dtype": "float32", + "shape": (len(visible_tracks), 2) + } + ) + + # Visualize tracked points + frame_viz = frame.copy() + for pt, vis in zip(tracks, visibility): + if vis > 0.5: # Only draw high-confidence points + x, y = int(pt[0]), int(pt[1]) + cv2.circle(frame_viz, (x, y), radius=3, + color=(0, 255, 0), thickness=-1) + + return frame, frame_viz + else: + print("Debug - Model returned None values") + + except Exception as e: + print(f"Error in processing: {str(e)}") + import traceback + traceback.print_exc() + + return None, None + + def run(self): + """Main run loop""" + try: + for event in self.node: + if event["type"] == "INPUT" and event["id"] == "image": + metadata = event["metadata"] + frame = event["value"].to_numpy().reshape(( + metadata["height"], + metadata["width"], + 3 + )) + + # Add frame to tracking window + self.window_frames.append(frame) + + # Process tracking + original_frame, tracked_frame = self.process_tracking(frame) + + # Only publish when we have processed frames + if original_frame is not None and tracked_frame is not None: + self.node.send_output("image", + pa.array(original_frame.ravel()), + metadata + ) + self.node.send_output("tracked_image", + pa.array(tracked_frame.ravel()), + metadata + ) + + finally: + pass + +def main(): + tracker = VideoTrackingNode() + tracker.run() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/node-hub/dora-cotracker/pyproject.toml b/node-hub/dora-cotracker/pyproject.toml new file mode 100644 index 00000000..fde7fe1b --- /dev/null +++ b/node-hub/dora-cotracker/pyproject.toml @@ -0,0 +1,26 @@ +[project] +name = "dora-cotracker" +version = "0.0.1" +authors = [{ name = "Your Name", email = "email@email.com" }] +description = "dora-cotracker" +license = { text = "MIT" } +readme = "README.md" +requires-python = ">=3.10" + +dependencies = [ + "dora-rs>=0.3.9", + "gradio>=4.0.0", + "torch>=2.0.0", + "numpy>=1.24.0", + "opencv-python>=4.8.0", + "pyarrow>=14.0.1", + "cotracker @ git+https://github.com/facebookresearch/co-tracker.git", + "imageio>=2.31.0", + "imageio-ffmpeg>=0.4.9", +] + +[dependency-groups] +dev = ["pytest >=8.1.1", "ruff >=0.9.1"] + +[project.scripts] +dora-cotracker = "dora_cotracker.main:main" diff --git a/node-hub/dora-cotracker/tests/test_dora_cotracker.py b/node-hub/dora-cotracker/tests/test_dora_cotracker.py new file mode 100644 index 00000000..50d8eb83 --- /dev/null +++ b/node-hub/dora-cotracker/tests/test_dora_cotracker.py @@ -0,0 +1,9 @@ +import pytest + + +def test_import_main(): + from dora_cotracker.main import main + + # Check that everything is working, and catch dora Runtime Exception as we're not running in a dora dataflow. + with pytest.raises(RuntimeError): + main()