Add some more camera controls

crates/brush-viewer/src/orbit_controls.rs

@@ -1,35 +1,71 @@
 use core::f32;
+use std::ops::Range;
 
-use glam::{Affine3A, Mat3A, Vec2, Vec3A};
+use glam::{Affine3A, EulerRot, Quat, Vec2, Vec3A};
 
 pub struct OrbitControls {
-    pub transform: Affine3A,
+    pub position: Vec3A,
+    pub yaw: f32,
+    pub pitch: f32,
 
     pub focus: Vec3A,
     pub dirty: bool,
 
     pan_momentum: Vec2,
     rotate_momentum: Vec2,
 
-    min_radius: f32,
-    max_radius: f32,
+    radius_range: Range<f32>,
+    yaw_range: Range<f32>,
+    pitch_range: Range<f32>,
 }
 
 impl OrbitControls {
-    pub fn new(transform: Affine3A, min_radius: f32, max_radius: f32) -> Self {
+    pub fn new(
+        radius: f32,
+        radius_range: Range<f32>,
+        yaw_range: Range<f32>,
+        pitch_range: Range<f32>,
+    ) -> Self {
         Self {
-            transform,
+            position: -Vec3A::Z * radius,
             focus: Vec3A::ZERO,
             pan_momentum: Vec2::ZERO,
             rotate_momentum: Vec2::ZERO,
             dirty: false,
-            min_radius,
-            max_radius
+            radius_range,
+            yaw_range,
+            pitch_range,
+            yaw: 0.0,
+            pitch: 0.0,
         }
     }
 
     pub fn radius(&self) -> f32 {
-        (self.transform.translation - self.focus).length()
+        (self.position - self.focus).length()
+    }
+
+    fn clamp_smooth(val: f32, range: Range<f32>) -> f32 {
+        let mut val = val;
+        if val < range.start {
+            val = val * 0.5 + range.start * 0.5;
+        }
+
+        if val > range.end {
+            val = val * 0.5 + range.end * 0.5;
+        }
+        val
+    }
+
+    fn clamp_rotation(quat: Quat, pitch_range: Range<f32>, yaw_range: Range<f32>) -> Quat {
+        // Get current angles
+        let (pitch, yaw, _) = quat.to_euler(EulerRot::YXZ);
+
+        // Clamp them
+        let clamped_pitch = Self::clamp_smooth(pitch, pitch_range);
+        let clamped_yaw = Self::clamp_smooth(yaw, yaw_range);
+
+        // Make new quaternion with clamped angles
+        Quat::from_euler(EulerRot::YXZ, clamped_yaw, clamped_pitch, 0.0)
     }
 
     pub fn pan_orbit_camera(
@@ -40,8 +76,11 @@ impl OrbitControls {
         window: Vec2,
         delta_time: f32,
     ) -> bool {
-        let mut rotation = self.transform.matrix3;
-        let mut radius = (self.transform.translation - self.focus).length();
+        let mut yaw = self.yaw;
+        let mut pitch = self.pitch;
+
+        let mut radius = self.radius();
+
         // Adjust momentum with the new input
         self.pan_momentum += pan;
         self.rotate_momentum += rotate;
@@ -57,9 +96,18 @@ impl OrbitControls {
 
         let delta_x = rotate_velocity.x * std::f32::consts::PI * 2.0 / window.x;
         let delta_y = rotate_velocity.y * std::f32::consts::PI / window.y;
-        let yaw = Mat3A::from_rotation_y(delta_x);
-        let pitch = Mat3A::from_rotation_x(-delta_y);
-        rotation = yaw * rotation * pitch;
+
+        yaw = Self::clamp_smooth(yaw + delta_x, self.yaw_range.clone());
+        pitch = Self::clamp_smooth(pitch - delta_y, self.pitch_range.clone());
+        self.yaw = yaw;
+        self.pitch = pitch;
+
+        // let (mut pitch, mut yaw, _) = Quat::from_mat3a(&new_rotation).to_euler(glam::EulerRot::YXZ);
+        // yaw = Self::clamp_smooth(yaw, self.yaw_range.clone());
+        // pitch = Self::clamp_smooth(pitch, self.pitch_range.clone());
+        // rotation = Mat3A::from_quat(Quat::from_euler(glam::EulerRot::XYZ, pitch, yaw, 0.0));
+
+        let rotation = Quat::from_rotation_y(self.yaw) * Quat::from_rotation_x(self.pitch);
 
         let scaled_pan = pan_velocity * Vec2::new(1.0 / window.x, 1.0 / window.y);
 
@@ -69,18 +117,8 @@ impl OrbitControls {
         let translation = (right + up) * radius;
         self.focus += translation;
         radius -= scroll * radius * 0.2;
-
-        // smooth clamp to min/max radius.
-        if radius < self.min_radius {
-            radius = radius * 0.5 + self.min_radius * 0.5;
-        }
-
-        if radius > self.max_radius {
-            radius = radius * 0.5 + self.max_radius * 0.5;
-        }
-
-        self.transform.translation = self.focus + rotation * Vec3A::new(0.0, 0.0, -radius);
-        self.transform.matrix3 = rotation;
+        radius = Self::clamp_smooth(radius, self.radius_range.clone());
+        self.position = self.focus + rotation * Vec3A::new(0.0, 0.0, -radius);
 
         scroll.abs() > 0.0
             || pan.length_squared() > 0.0
@@ -89,4 +127,11 @@ impl OrbitControls {
             || self.rotate_momentum.length_squared() > 0.001
             || self.dirty
     }
+
+    pub(crate) fn transform(&self) -> Affine3A {
+        Affine3A::from_rotation_translation(
+            Quat::from_rotation_y(self.yaw) * Quat::from_rotation_x(self.pitch),
+            self.position.into(),
+        )
+    }
 }

crates/brush-viewer/src/panels/scene.rs

@@ -134,7 +134,7 @@ impl ScenePanel {
             delta_time.as_secs_f32(),
         );
 
-        let total_transform = context.model_transform * context.controls.transform;
+        let total_transform = context.model_transform * context.controls.transform();
         context.camera.position = total_transform.translation.into();
         context.camera.rotation = Quat::from_mat3a(&total_transform.matrix3);
 

crates/brush-viewer/src/viewer.rs

@@ -1,3 +1,5 @@
+use core::f32;
+use std::ops::Range;
 use std::{pin::Pin, sync::Arc};
 
 use async_fn_stream::try_fn_stream;
@@ -235,8 +237,11 @@ impl DataSource {
 
 struct CameraSettings {
     focal: f64,
-    min_radius: f32,
-    max_radius: f32,
+    radius: f32,
+
+    yaw_range: Range<f32>,
+    pitch_range: Range<f32>,
+    radius_range: Range<f32>,
 }
 
 impl ViewerContext {
@@ -248,11 +253,11 @@ impl ViewerContext {
     ) -> Self {
         let model_transform = Affine3A::IDENTITY;
 
-        let avg_radius = (cam_settings.min_radius + cam_settings.max_radius) / 2.0;
         let controls = OrbitControls::new(
-            Affine3A::from_translation(-Vec3::Z * avg_radius),
-            cam_settings.min_radius,
-            cam_settings.max_radius,
+            cam_settings.radius,
+            cam_settings.radius_range,
+            cam_settings.yaw_range,
+            cam_settings.pitch_range,
         );
 
         // Camera position will be controller by controls.
@@ -264,6 +269,26 @@ impl ViewerContext {
             glam::vec2(0.5, 0.5),
         );
 
+        // TODO: Generalize this "inner control" logic.
+        let (inner_send, inner_control) = ::tokio::sync::mpsc::unbounded_channel();
+        let ctx_spawn = ctx.clone();
+        let mut controller = controller;
+        task::spawn(async move {
+            // Loop until there are no more messages, processing is done.
+            while let Some(m) = controller.recv().await {
+                ctx_spawn.request_repaint();
+
+                // Give back to the runtime for a second.
+                // This only really matters in the browser.
+                tokio::task::yield_now().await;
+
+                // If channel is closed, bail.
+                if inner_send.send(m).is_err() {
+                    break;
+                }
+            }
+        });
+
         Self {
             camera,
             controls,
@@ -273,7 +298,7 @@ impl ViewerContext {
             dataset: Dataset::empty(),
             rec_process_msg: None,
             send_train_msg: None,
-            rec_ui_control_msg: controller,
+            rec_ui_control_msg: inner_control,
         }
     }
 
@@ -286,10 +311,17 @@ impl ViewerContext {
     pub fn focus_view(&mut self, cam: &Camera) {
         // set the controls transform.
         let cam_transform = Affine3A::from_rotation_translation(cam.rotation, cam.position);
-        self.controls.transform = self.model_transform.inverse() * cam_transform;
+        let transform = self.model_transform.inverse() * cam_transform;
+        let (_, rotation, position) = transform.to_scale_rotation_translation();
+
+        let (pitch, yaw, _) = rotation.to_euler(glam::EulerRot::ZYX);
+        self.controls.position = position.into();
+        self.controls.yaw = yaw;
+        self.controls.pitch = pitch;
+
         // Copy the camera, mostly to copy over the intrinsics and such.
-        self.controls.focus = self.controls.transform.translation
-            + self.controls.transform.matrix3
+        self.controls.focus = transform.translation
+            + transform.matrix3
                 * Vec3A::Z
                 * self.dataset.train.bounds(0.0, 0.0).extent.length()
                 * 0.5;
@@ -319,8 +351,11 @@ impl ViewerContext {
 
         self.dataset = Dataset::empty();
         let ctx = self.ctx.clone();
+        ctx.request_repaint();
 
         let fut = async move {
+            ctx.request_repaint();
+
             // Map errors to a viewer message containing thee error.
             let mut stream = process_loop(
                 source,
@@ -405,6 +440,10 @@ impl Viewer {
             .get("focal")
             .and_then(|f| f.parse().ok())
             .unwrap_or(0.5);
+        let radius = search_params
+            .get("radius")
+            .and_then(|f| f.parse().ok())
+            .unwrap_or(4.0);
         let min_radius = search_params
             .get("min_radius")
             .and_then(|f| f.parse().ok())
@@ -414,10 +453,34 @@ impl Viewer {
             .and_then(|f| f.parse().ok())
             .unwrap_or(10.0);
 
+        let min_yaw = search_params
+            .get("min_yaw")
+            .and_then(|f| f.parse::<f32>().ok())
+            .map(|d| d.to_radians())
+            .unwrap_or(f32::MIN);
+        let max_yaw = search_params
+            .get("max_yaw")
+            .and_then(|f| f.parse::<f32>().ok())
+            .map(|d| d.to_radians())
+            .unwrap_or(f32::MAX);
+
+        let min_pitch = search_params
+            .get("min_pitch")
+            .and_then(|f| f.parse::<f32>().ok())
+            .map(|d| d.to_radians())
+            .unwrap_or(f32::MIN);
+        let max_pitch = search_params
+            .get("max_pitch")
+            .and_then(|f| f.parse::<f32>().ok())
+            .map(|d| d.to_radians())
+            .unwrap_or(f32::MAX);
+
         let settings = CameraSettings {
             focal,
-            min_radius,
-            max_radius,
+            radius,
+            radius_range: min_radius..max_radius,
+            yaw_range: min_yaw..max_yaw,
+            pitch_range: min_pitch..max_pitch,
         };
 
         let context = ViewerContext::new(device.clone(), cc.egui_ctx.clone(), settings, controller);
@@ -526,8 +589,6 @@ impl eframe::App for Viewer {
                         Tile::Container(_) => {}
                     }
                 }
-
-                ctx.request_repaint();
             }
         }