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