working version

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

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

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."

node-hub/dora-cotracker/dora_cotracker/__main__.py

@@ -0,0 +1,5 @@
+from .main import main
+
+
+if __name__ == "__main__":
+    main()

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()

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"

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()