diff --git a/node-hub/dora-cotracker/README.md b/node-hub/dora-cotracker/README.md new file mode 100644 index 00000000..b3a3dc7f --- /dev/null +++ b/node-hub/dora-cotracker/README.md @@ -0,0 +1,221 @@ +# dora-cotracker + +A Dora node that implements real-time object tracking using Facebook's CoTracker model. The node supports both interactive point selection via clicking and programmatic point input through Dora's messaging system. + +## Features + +- Real-time object tracking using CoTracker +- Support for multiple tracking points +- Interactive point selection via mouse clicks +- Programmatic point input through Dora messages +- Visualization of tracked points with unique identifiers + +## Getting Started + +### Installation + +Install using uv: + +```bash +uv venv -p 3.11 --seed +uv pip install -e . +``` + +## Demo Video + +Watch a demonstration of the dora-cotracker node in action: + +[![Dora CoTracker Demo](https://img.youtube.com/vi/1VmC1BNq6J0/0.jpg)](https://youtu.be/1VmC1BNq6J0) + +The video shows: +- Setting up the node +- Interactive point selection +- Real-time tracking performance + +### Basic Usage + +1. Create a YAML configuration file (e.g., `demo.yml`): + +```yaml +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 dora-cotracker + path: dora-cotracker + inputs: + image: camera/image + points_to_track: input/points_to_track + outputs: + - tracked_image + - tracked_points + + - id: display + build: pip install dora-rerun + path: dora-rerun + inputs: + image: camera/image + tracked_image: tracker/tracked_image +``` + +*Note* - this only has the cv2 as an input source. see below to add your nodes workflow and pass points directly. + +2. Run the demo: + +```bash +dora run demo.yml --uv +``` + +## Usage Examples + +### 1. Interactive Point Selection +Click points directly in the "Raw Feed" window to start tracking them: +- Left-click to add tracking points +- Points will be tracked automatically across frames +- Each point is assigned a unique identifier (C0, C1, etc. for clicked points and I0, I1, etc for input points) + +### 2. Dynamic Point Integration +The node can receive tracking points from other models or nodes in your pipeline. Common use cases include: + +- Tracking YOLO detection centroids +- Following pose estimation keypoints +- Monitoring segmentation mask centers +- Custom object detection points + +example showing how to send tracking points through Dora messages using a custom input node: + +```python +import numpy as np +import pyarrow as pa +from dora import Node + +class PointInputNode: + def __init__(self): + self.node = Node("point-input") + + def send_points(self, points): + """ + Send points to tracker + Args: + points: Nx2 array of (x,y) coordinates + """ + points = np.array(points, dtype=np.float32) + self.node.send_output( + "points_to_track", + pa.array(points.ravel()), + { + "num_points": len(points), + "dtype": "float32", + "shape": (len(points), 2) + } + ) + + def run(self): + # Example: Track 3 points + points = np.array([ + [320, 240], # Center + [160, 120], # Top-left + [480, 360] # Bottom-right + ]) + self.send_points(points) +``` + + + +To connect your existing node that outputs tracking points with the CoTracker node, add the following to your YAML configuration: + +```yaml +nodes: + # Your existing point source node (e.g., YOLO detector, pose estimator, etc.) + - id: point_source + build: pip install your-node # Replace with your node's name + path: your-point-source-node # Replace with your node's path + inputs: + image: camera/image # If your node needs image input + outputs: + - points_to_track # Must output points in required format + + # CoTracker node configuration + - id: tracker + build: pip install dora-cotracker + path: dora-cotracker + inputs: + image: camera/image + points_to_track: point_source/points_to_track # Connect to your point source + outputs: + - tracked_image + - tracked_points + + # Optional visualization + - id: display + build: pip install dora-rerun + path: dora-rerun + inputs: + image: camera/image + tracked_image: tracker/tracked_image +``` + +Your point source node must output points in the following format: +- Topic name: `points_to_track` +- Data: Flattened numpy array of x,y coordinates +- Metadata: + ```python + { + "num_points": len(points), # Number of points + "dtype": "float32", # Data type + "shape": (N, 2) # N points, 2 coordinates each + } + ``` + +Example point source implementations: +- YOLO detection centroids +- Pose estimation keypoints +- Face landmark detectors +- Custom object detectors + +For dynamic updates, send new points whenever your source node processes a new frame. The tracker will maintain temporal consistency between updates. +** +## API Reference + +### Input Topics +- `image`: Input video stream (RGB format) +- `points_to_track`: Points to track + - Format: Flattened array of x,y coordinates + - Metadata: + - `num_points`: Number of points + - `dtype`: "float32" + - `shape`: (N, 2) where N is number of points + +### Output Topics +- `tracked_image`: Visualization with tracked points +- `tracked_points`: Current positions of tracked points + - Same format as input points + +## Development + +Format code with ruff: +```bash +uv pip install ruff +uv run ruff check . --fix +``` + +Run tests: +```bash +uv pip install pytest +uv run pytest +``` + +## License + +dora-cotracker's code are released under the MIT License \ No newline at end of file diff --git a/node-hub/dora-cotracker/demo.yml b/node-hub/dora-cotracker/demo.yml new file mode 100644 index 00000000..6bb36707 --- /dev/null +++ b/node-hub/dora-cotracker/demo.yml @@ -0,0 +1,41 @@ +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 dora-cotracker + path: dora-cotracker + inputs: + image: camera/image + # points_to_track: input/points_to_track # uncomment this if using input node + outputs: + - tracked_image + - tracked_points + + - id: plot + build: pip install dora-rerun + path: dora-rerun + inputs: + image: camera/image + tracked_image: tracker/tracked_image + + + # replace with your own node that outputs tracking points # uncomment if input via node + # (e.g., YOLO detector, pose estimator, etc.) + # - id: point_source + # build: pip install your-node # Replace with your node's name + # path: your-point-source-node # Replace with your node's path + # inputs: + # image: camera/image # If your node needs image input + # outputs: + # - points_to_track # Must output points in required format \ 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..5e4c08a4 --- /dev/null +++ b/node-hub/dora-cotracker/dora_cotracker/main.py @@ -0,0 +1,149 @@ +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" + self.model = torch.hub.load("facebookresearch/co-tracker", "cotracker3_online") + self.model = self.model.to(self.device) + self.model.step = 8 + self.buffer_size = self.model.step * 2 + self.window_frames = deque(maxlen=self.buffer_size) + self.is_first_step = True + self.clicked_points = [] + self.input_points = [] + + def mouse_callback(self, event, x, y, flags, param): + if event == cv2.EVENT_LBUTTONDOWN: + self.clicked_points.append([x, y]) + self.is_first_step = True + # print(f"Clicked point added at: ({x}, {y})") + + def process_tracking(self, frame): + """Process frame for tracking""" + if len(self.window_frames) == self.buffer_size: + all_points = self.input_points + self.clicked_points + + if not all_points: + print("No points to track") + return None, None + + video_chunk = torch.tensor( + np.stack(list(self.window_frames)), + device=self.device + ).float() + video_chunk = video_chunk / 255.0 + # Reshape to [B,T,C,H,W] + video_chunk = video_chunk.permute(0, 3, 1, 2)[None] + query_points = torch.tensor(all_points, device=self.device).float() + time_dim = torch.zeros(len(all_points), 1, device=self.device) + queries = torch.cat([time_dim, query_points], dim=1).unsqueeze(0) + # Track points + pred_tracks, pred_visibility = self.model( + video_chunk, + is_first_step=self.is_first_step, + grid_size=0, + queries=queries, + add_support_grid=False + ) + self.is_first_step = False + + if pred_tracks is not None and pred_visibility is not None: + tracks = pred_tracks[0, -1].cpu().numpy() + visibility = pred_visibility[0, -1].cpu().numpy() + visible_tracks = [] + for pt, vis in zip(tracks, visibility): + if vis > 0.5: + visible_tracks.append([int(pt[0]), int(pt[1])]) + visible_tracks = np.array(visible_tracks, dtype=np.float32) + + frame_viz = frame.copy() + num_input_stream = len(self.input_points) + # Draw input points in red + for i, (pt, vis) in enumerate(zip(tracks[:num_input_stream], visibility[:num_input_stream])): + if vis > 0.5: + x, y = int(pt[0]), int(pt[1]) + cv2.circle(frame_viz, (x, y), radius=3, + color=(0, 255, 0), thickness=-1) + cv2.putText(frame_viz, f"I{i}", (x + 5, y - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + + # Draw clicked points in red + for i, (pt, vis) in enumerate(zip(tracks[num_input_stream:], visibility[num_input_stream:])): + if vis > 0.5: + x, y = int(pt[0]), int(pt[1]) + cv2.circle(frame_viz, (x, y), radius=3, + color=(0, 0, 255), thickness=-1) + cv2.putText(frame_viz, f"C{i}", (x + 5, y - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) + + # 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) + } + ) + + return frame, frame_viz + + return None, None + + def run(self): + """Main run loop""" + cv2.namedWindow("Raw Feed", cv2.WINDOW_NORMAL) + cv2.setMouseCallback("Raw Feed", self.mouse_callback) + + for event in self.node: + if event["type"] == "INPUT": + if 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) + original_frame, tracked_frame = self.process_tracking(frame) + 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 + ) + + display_frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + cv2.imshow("Raw Feed", display_frame) + cv2.waitKey(1) + + if event["id"] == "points_to_track": + # Handle points from input_stream node + metadata = event["metadata"] + points_array = event["value"].to_numpy() + num_points = metadata["num_points"] + self.input_points = points_array.reshape((num_points, 2)).tolist() + self.is_first_step = True + print(f"Received {num_points} points from input_stream") + + + +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..3359a888 --- /dev/null +++ b/node-hub/dora-cotracker/pyproject.toml @@ -0,0 +1,34 @@ +[project] +name = "dora-cotracker" +version = "0.1.0" +authors = [ + { name = "Shashwat Patil", email = "shashwatpatil974@gmail.com" } +] +description = "A Dora node implementing real-time object tracking using Facebook's CoTracker model" +license = { text = "MIT" } +readme = "README.md" +requires-python = ">=3.10" + +dependencies = [ + "dora-rs>=0.3.9", + "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", +] + +[dependency-groups] +dev = ["pytest >=8.1.1", "ruff >=0.9.1"] + +[project.scripts] +dora-cotracker = "dora_cotracker.main:main" + +[tool.ruff.lint] +extend-select = [ + "PERF", # Performance + "RET", # Return statements + "RSE", # Runtime errors + "NPY", # NumPy + "N", # Naming +] 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()