3D Canvas added with yolo-depth

Author: MathematicLove

README.md

@@ -32,9 +32,42 @@ those keypoints, so recognition is independent of hand size, distance, or handed
 Keys: `S` save, `C` clear, `U` undo, `1..6` color, `[` / `]` thickness,
 `SPACE` 2D/3D, `R` reset view, `K` skeleton, `D` debug, `H` help, `Q` quit.
 
+## 3D Drawing
+
+The **3D Drawing** mode uses the YOLO depth model (`models/yolo26s-depth.pt`), which
+returns a per pixel depth map in meters. Every drawn point takes its Z from the metric
+depth under the fingertip, so moving the hand closer or further really draws into the
+scene: the stroke is a 3D curve in space, not a flat picture placed on a board. The
+first drawn point sets the zero plane, everything after that is relative to it.
+Depth runs in a background thread, so tracking and drawing stay at camera speed.
+
+Raw monocular depth flickers, so the pen does not follow it directly: the global scale
+drift of the map is cancelled against the static background, the fingertip depth goes
+through a one euro filter with a noise deadband and a speed limit, and the newest points
+of a stroke are low passed along the curve. A circle drawn without moving the hand comes
+out flat, while a deliberate push of 25 cm is kept in full.
+
+By default strokes are colored by a depth gradient - warm near the camera, cool further
+away. `1..6` switch to a solid color, `G` goes back to the gradient.
+
+Gestures and keys are the same as in 2D, with two additions:
+
+- **Two ring fingers (hold)** - toggle depth colors for the whole frame, drawing keeps working
+- **Open palm (hold)** - reset the view
+
+In 3D the fist rotates the drawing itself and two fists scale it, `SPACE` toggles
+depth colors, `R` resets the view.
+
 ## Run
 
 ```bash
 pip install -r requirements.txt
 python src/main.py
 ```
+
+Pick the mode in the menu, or skip it:
+
+```bash
+python src/main.py --no-menu --mode 3d
+python src/main.py --no-menu --mode 3d --depth-size 448 --depth-every 1
+```

requirements.txt

@@ -1,4 +1,6 @@
 mediapipe>=0.10.30
+ultralytics>=8.4.19
+torch>=2.4
 opencv-python>=4.10
 numpy>=1.26
 pillow>=10.0

src/app3d.py

@@ -0,0 +1,275 @@
+"""3D drawing mode built on monocular depth."""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import cv2
+import numpy as np
+
+import ui
+from canvas3d import Canvas3D, View
+from depth import DepthEstimator, DepthWorker
+from gestures import Gesture
+from main import App
+
+RING_HOLD_SECONDS = 0.35
+PALM_RESET_SECONDS = 0.6
+DEPTH_MIX = 0.78
+SCENE_ALPHA = 0.02
+
+
+class App3D(App):
+    """Draw real 3D strokes using estimated depth."""
+
+    def __init__(self, args: argparse.Namespace) -> None:
+        super().__init__(args)
+        self.canvas = Canvas3D(self.w, self.h, args.output)
+        self.depth = DepthWorker(args.depth_model, imgsz=args.depth_size)
+        print(f"[i] depth model: {args.depth_model}", flush=True)
+        print(f"[i] depth device: {self.depth.device}  imgsz: {args.depth_size}", flush=True)
+
+        self.depth_map: np.ndarray | None = None
+        self.depth_norm: np.ndarray | None = None
+        self.depth_view = False
+        self._view = View()
+        self._depth_seq = -1
+        self._scene_ref: float | None = None
+        self._scene_gain = 1.0
+        self._depth_announced = False
+        self._frame_index = 0
+        self._ring_since: float | None = None
+        self._ring_latched = False
+        self._palm_since: float | None = None
+        self._palm_latched = False
+
+        self.yaw, self.pitch, self.roll, self.scale = 0.0, 0.0, 0.0, 1.0
+        self.yaw_s.set(0.0)
+        self.pitch_s.set(0.0)
+        self.scale_s.set(1.0)
+
+    def view(self) -> View:
+        """Current smoothed camera."""
+        return self._view
+
+    def step_view(self) -> View:
+        """Advance the smoothed camera."""
+        self._view = View(self.yaw_s(self.yaw), self.pitch_s(self.pitch),
+                          self.roll, self.scale_s(self.scale))
+        self.canvas.set_view(self._view)
+        return self._view
+
+    def on_frame(self, frame: np.ndarray) -> None:
+        """Feed the depth worker."""
+        self.step_view()
+        self._frame_index += 1
+        if self._frame_index % max(1, self.args.depth_every) == 0:
+            self.depth.submit(frame)
+        latest = self.depth.latest
+        if latest is None or self.depth.frames == self._depth_seq:
+            return
+        self._depth_seq = self.depth.frames
+        self.depth_map = latest
+        self.depth_norm = self.depth.latest_norm
+        self.track_scene(latest)
+        if not self._depth_announced:
+            self._depth_announced = True
+            self.toast("Depth ready. Move your hand closer or further to draw in depth")
+
+    def track_scene(self, depth_map: np.ndarray) -> None:
+        """Cancel global depth drift."""
+        sample = depth_map[::8, ::8]
+        sample = sample[np.isfinite(sample)]
+        if sample.size == 0:
+            return
+        median = float(np.median(sample))
+        if median <= 0.0:
+            return
+        if self._scene_ref is None:
+            self._scene_ref = median
+        else:
+            self._scene_ref += SCENE_ALPHA * (median - self._scene_ref)
+        self._scene_gain = float(np.clip(self._scene_ref / median, 0.5, 2.0))
+
+    def depth_at(self, point) -> float:
+        """Metric depth under a screen point."""
+        if self.depth_map is None:
+            return float("nan")
+        d = DepthEstimator.sample(self.depth_map, float(point[0]), float(point[1]),
+                                  default=float("nan"))
+        return d * self._scene_gain
+
+    def depth_readout(self) -> str:
+        """Depth status line."""
+        if self.depth_map is None:
+            return "depth: warming up"
+        if self.canvas.depth_ref is None:
+            return f"depth: {self.canvas.last_depth:.2f} m   draw to set the zero plane"
+        offset = (self.canvas.last_depth - self.canvas.depth_ref) * 100.0
+        turn = "yaw %+.0f pitch %+.0f" % (np.degrees(self._view.yaw), np.degrees(self._view.pitch))
+        return f"depth: {self.canvas.last_depth:.2f} m   z {offset:+.0f} cm   {turn}"
+
+    def handle_ring(self, states: list) -> None:
+        """Toggle depth colors on two ring fingers."""
+        rings = [s for s in states if s.gesture == Gesture.RING and s.stable]
+        if len(rings) < 2:
+            self._ring_since = None
+            self._ring_latched = False
+            return
+        now = time.time()
+        self._ring_since = self._ring_since or now
+        if not self._ring_latched and now - self._ring_since >= RING_HOLD_SECONDS:
+            self._ring_latched = True
+            self.depth_view = not self.depth_view
+            self.toast("Depth colors " + ("on" if self.depth_view else "off"))
+
+    def handle_palm(self, states: list) -> None:
+        """Reset the view on open palm."""
+        palms = [s for s in states if s.gesture == Gesture.OPEN_PALM and s.stable]
+        if not palms:
+            self._palm_since = None
+            self._palm_latched = False
+            return
+        now = time.time()
+        self._palm_since = self._palm_since or now
+        if not self._palm_latched and now - self._palm_since >= PALM_RESET_SECONDS:
+            self._palm_latched = True
+            self.reset_view()
+            self.toast("View reset")
+
+    def reset_view(self) -> None:
+        """Back to the front view."""
+        self.yaw, self.pitch, self.roll, self.scale = 0.0, 0.0, 0.0, 1.0
+        self.yaw_s.set(0.0)
+        self.pitch_s.set(0.0)
+        self.scale_s.set(1.0)
+        self._grab = None
+        self._grab_span = None
+        self.step_view()
+
+    def rotate_view(self, grabbing: list) -> None:
+        """Rotate and scale the drawing."""
+        p = grabbing[0].pinch_point
+        if self._grab is None:
+            self._grab = (p.copy(), self.yaw, self.pitch)
+        anchor, yaw0, pitch0 = self._grab
+        dx = (p[0] - anchor[0]) / self.w
+        dy = (p[1] - anchor[1]) / self.h
+        self.yaw = yaw0 + dx * 2.0 * np.pi * 1.1
+        self.pitch = float(np.clip(pitch0 + dy * np.pi * 1.1, -1.3, 1.3))
+
+        if len(grabbing) >= 2:
+            span = float(np.linalg.norm(grabbing[0].pinch_point - grabbing[1].pinch_point))
+            if self._grab_span is None:
+                self._grab_span = (span, self.scale)
+            span0, scale0 = self._grab_span
+            if span0 > 1e-3:
+                self.scale = float(np.clip(scale0 * (span / span0), 0.25, 4.0))
+        else:
+            self._grab_span = None
+
+    def handle_draw_mode(self, states: list, dt: float) -> None:
+        """Draw in 3D, grab to rotate."""
+        self.handle_ring(states)
+        self.handle_palm(states)
+
+        drawing = [s for s in states if s.gesture == Gesture.DRAW]
+        grabbing = [s for s in states if s.gesture == Gesture.GRAB]
+
+        if grabbing:
+            self.stop_stroke(force=True)
+            self.rotate_view(grabbing)
+            self.grab_since = None
+            return
+
+        self._grab = None
+        self._grab_span = None
+
+        if drawing:
+            point = self.tip_filter(drawing[0].cursor, dt)
+            self.canvas.begin(self._view)
+            self.canvas.add_point(point, self.depth_at(point), dt, self._view)
+            self._draw_lost_since = None
+            self.grab_since = None
+            return
+
+        self.stop_stroke()
+        self.grab_since = None
+
+    def clear_canvas(self) -> None:
+        """Wipe the drawing."""
+        self.stop_stroke(force=True)
+        self.canvas.clear()
+        self._grab = None
+        self._grab_span = None
+        self.grab_since = None
+        self.toast("Canvas cleared")
+
+    def background(self, frame: np.ndarray) -> np.ndarray:
+        """Camera or depth colored frame."""
+        if not self.depth_view or self.depth_norm is None:
+            return frame
+        colored = DepthEstimator.colorize(self.depth_norm)
+        return cv2.addWeighted(colored, DEPTH_MIX, frame, 1.0 - DEPTH_MIX, 0.0)
+
+    def render(self, frame: np.ndarray, hands, states) -> np.ndarray:
+        """Compose the output frame."""
+        view = self._view
+        out = self.canvas.render_over(self.background(frame), view)
+
+        if self.show_skeleton:
+            for hand, st in zip(hands, states):
+                ui.draw_hand(out, hand, st)
+        pen = self.canvas.pen_color()
+        active = self.active_state(states)
+        if active is not None:
+            ui.draw_cursor(out, active.cursor, active.gesture, pen, self.canvas.thickness)
+
+        shown = Gesture.NONE if active is None else active.gesture
+        if len([s for s in states if s.gesture == Gesture.GRAB]) >= 2:
+            shown = Gesture.ZOOM
+        ui.draw_hud(
+            out,
+            gesture=shown,
+            hands_info=[(s.handedness, s.gesture) for s in states],
+            mode_3d=False,
+            fps=self.fps,
+            device=self.device,
+            color=pen,
+            thickness=self.canvas.thickness,
+            strokes=len(self.canvas.strokes),
+            clear_progress=0.0,
+            toast=self.toast_text if time.time() < self.toast_until else "",
+            mode_label="3D DRAW",
+            hint="Finger draw    Fist rotate    Two rings depth",
+        )
+        ui.put_text(out, self.depth_readout(), (16, 84), 16, (170, 220, 255))
+        if self.show_debug and states:
+            ui.draw_debug(out, hands, states)
+        if self.show_help:
+            ui.draw_help(out, ui.HELP_LINES_3D)
+        return out
+
+    def handle_key(self, key: int) -> bool:
+        """Handle one keypress."""
+        if key == ord("g"):
+            self.canvas.use_gradient()
+            self.toast("Depth gradient color")
+            return True
+        if key == ord(" "):
+            self.depth_view = not self.depth_view
+            self.toast("Depth colors " + ("on" if self.depth_view else "off"))
+            return True
+        if key == ord("r"):
+            self.reset_view()
+            self.toast("View reset")
+            return True
+        return super().handle_key(key)
+
+    def run(self) -> None:
+        """Run the main loop."""
+        try:
+            super().run()
+        finally:
+            self.depth.close()