Fix focus distance of orbit not being reset after focusing a view.

crates/brush-app/src/app.rs

@@ -180,6 +180,12 @@ impl AppContext {
         self.match_controls_to(&view.camera);
         self.controls.stop_movement();
         self.view_aspect = Some(view.image.width() as f32 / view.image.height() as f32);
+
+        self.controls.focus_distance = self
+            .dataset
+            .train
+            .estimate_extent()
+            .map_or(4.0, |x| x / 3.0);
     }
 
     pub fn connect_to(&mut self, process: RunningProcess) {

crates/brush-app/src/orbit_controls.rs

@@ -6,8 +6,8 @@ use glam::{Quat, Vec2, Vec3};
 pub struct CameraController {
     pub position: Vec3,
     pub rotation: Quat,
+    pub focus_distance: f32,
     roll: Quat,
-    focus_distance: f32,
     fly_velocity: Vec3,
     orbit_velocity: Vec2,
 }

crates/brush-dataset/src/scene_loader.rs

@@ -11,28 +11,14 @@ pub struct SceneLoader<B: Backend> {
     receiver: Receiver<SceneBatch<B>>,
 }
 
-use glam::Vec3;
-
-fn find_two_smallest(v: Vec3) -> (f32, f32) {
-    let mut arr = v.to_array();
-    arr.sort_by(|a, b| a.partial_cmp(b).expect("NaN"));
-    (arr[0], arr[1])
-}
-
-fn estimate_scene_extent(scene: &Scene) -> f32 {
-    let bounds = scene.bounds();
-    let smallest = find_two_smallest(bounds.extent * 2.0);
-    smallest.0.hypot(smallest.1)
-}
-
 impl<B: Backend> SceneLoader<B> {
     pub fn new(scene: &Scene, seed: u64, device: &B::Device) -> Self {
         let scene = scene.clone();
         // The bounded size == number of batches to prefetch.
         let (tx, rx) = mpsc::channel(5);
         let device = device.clone();
 
-        let scene_extent = estimate_scene_extent(&scene);
+        let scene_extent = scene.estimate_extent().unwrap_or(1.0);
 
         let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
 

crates/brush-train/src/scene.rs

@@ -43,6 +43,12 @@ fn camera_distance_penalty(cam_local_to_world: Affine3A, reference: Affine3A) ->
     penalty
 }
 
+fn find_two_smallest(v: Vec3) -> (f32, f32) {
+    let mut arr = v.to_array();
+    arr.sort_by(|a, b| a.partial_cmp(b).expect("NaN"));
+    (arr[0], arr[1])
+}
+
 impl Scene {
     pub fn new(views: Vec<SceneView>) -> Self {
         Self {
@@ -83,4 +89,15 @@ impl Scene {
             })
             .map(|(index, _)| index) // We return the index instead of the camera
     }
+
+    pub fn estimate_extent(&self) -> Option<f32> {
+        if self.views.len() < 5 {
+            None
+        } else {
+            // TODO: This is really sensitive to outliers.
+            let bounds = self.bounds();
+            let smallest = find_two_smallest(bounds.extent * 2.0);
+            Some(smallest.0.hypot(smallest.1))
+        }
+    }
 }