lbf/samplers/
ls_sampler.rs1use rand::Rng;
2use rand_distr::Distribution;
3use rand_distr::Normal;
4use std::f32::consts::PI;
5
6use crate::samplers::rotation_distr::NormalRotDistr;
7use jagua_rs::entities::Item;
8use jagua_rs::geometry::DTransformation;
9use jagua_rs::geometry::primitives::Rect;
10
11pub const SD_TRANSL: (f32, f32) = (0.01, 0.0005);
13
14pub const SD_ROT: (f32, f32) = (2.0 * PI / 180.0, 0.5 * PI / 180.0);
16
17pub struct LSSampler {
21 normal_x: Normal<f32>,
22 normal_y: Normal<f32>,
23 normal_r: NormalRotDistr,
24 reflected: bool,
25 sd_transl: f32,
26 sd_rot: f32,
27 sd_transl_range: (f32, f32),
28 sd_rot_range: (f32, f32),
29 pub(crate) n_samples: usize,
30}
31
32impl LSSampler {
33 pub fn new(
34 item: &Item,
35 ref_transform: DTransformation,
36 sd_transl_range: (f32, f32),
37 sd_rot_range: (f32, f32),
38 ) -> Self {
39 let sd_transl = sd_transl_range.0;
40 let sd_rot = sd_rot_range.0;
41
42 let normal_x = Normal::new(ref_transform.translation().0, sd_transl).unwrap();
43 let normal_y = Normal::new(ref_transform.translation().1, sd_transl).unwrap();
44 let normal_r = NormalRotDistr::from_item(item, ref_transform.rotation(), sd_rot);
45
46 Self {
47 normal_x,
48 normal_y,
49 normal_r,
50 reflected: ref_transform.reflected,
51 sd_transl,
52 sd_rot,
53 sd_transl_range,
54 sd_rot_range,
55 n_samples: 0,
56 }
57 }
58
59 pub fn from_defaults(item: &Item, ref_transform: DTransformation, bbox: Rect) -> Self {
61 let max_dim = f32::max(bbox.width(), bbox.height());
62 let sd_transl_range = (SD_TRANSL.0 * max_dim, SD_TRANSL.1 * max_dim);
63 Self::new(item, ref_transform, sd_transl_range, SD_ROT)
64 }
65
66 pub fn shift_mean(&mut self, ref_transform: DTransformation) {
68 self.normal_x = Normal::new(ref_transform.translation().0, self.sd_transl).unwrap();
69 self.normal_y = Normal::new(ref_transform.translation().1, self.sd_transl).unwrap();
70 self.normal_r.set_mean(ref_transform.rotation());
71 self.reflected = ref_transform.reflected;
72 }
73
74 pub fn set_stddev(&mut self, stddev_transl: f32, stddev_rot: f32) {
76 assert!(stddev_transl >= 0.0 && stddev_rot >= 0.0);
77
78 self.sd_transl = stddev_transl;
79 self.sd_rot = stddev_rot;
80 self.normal_x = Normal::new(self.normal_x.mean(), self.sd_transl).unwrap();
81 self.normal_y = Normal::new(self.normal_y.mean(), self.sd_transl).unwrap();
82 self.normal_r.set_stddev(self.sd_rot);
83 }
84
85 pub fn decay_stddev(&mut self, progress_pct: f32) {
93 let calc_stddev = |(init, end): (f32, f32), pct: f32| init * (end / init).powf(pct);
94 self.set_stddev(
95 calc_stddev(self.sd_transl_range, progress_pct),
96 calc_stddev(self.sd_rot_range, progress_pct),
97 );
98 }
99
100 pub fn sample(&mut self, rng: &mut impl Rng) -> DTransformation {
102 self.n_samples += 1;
103
104 DTransformation::new(
105 self.normal_r.sample(rng),
106 (self.normal_x.sample(rng), self.normal_y.sample(rng)),
107 )
108 .with_reflection(self.reflected)
109 }
110}