Skip to content

Commit

Permalink
Add Ball Candidate Panel to Twix (#1394)
Browse files Browse the repository at this point in the history
* Move grayscale sample extraction

* Add sample boxes to ball detection overlay

* Implement ball detection panel

* Bump Twix version
  • Loading branch information
h3ndrk authored Jul 20, 2024
1 parent 595ecfd commit 08ca138
Show file tree
Hide file tree
Showing 8 changed files with 277 additions and 46 deletions.
2 changes: 1 addition & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

23 changes: 22 additions & 1 deletion crates/types/src/ycbcr422_image.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,18 +6,21 @@ use std::{
};

use color_eyre::eyre::{self, WrapErr};
use geometry::circle::Circle;
use image::{io::Reader, RgbImage};
use serde::{Deserialize, Serialize};

use coordinate_systems::Pixel;
use linear_algebra::Point2;
use linear_algebra::{vector, Point2};
use path_serde::{PathDeserialize, PathIntrospect, PathSerialize};

use crate::{
color::{Rgb, YCbCr422, YCbCr444},
jpeg::JpegImage,
};

pub const SAMPLE_SIZE: usize = 32;

#[derive(
Clone, Debug, Default, Deserialize, Serialize, PathSerialize, PathIntrospect, PathDeserialize,
)]
Expand Down Expand Up @@ -242,4 +245,22 @@ impl YCbCr422Image {
ycbcr444
})
}

pub fn sample_grayscale(&self, candidate: Circle<Pixel>) -> Sample {
let top_left = candidate.center - vector![candidate.radius, candidate.radius];
let image_pixels_per_sample_pixel = candidate.radius * 2.0 / SAMPLE_SIZE as f32;

let mut sample = Sample::default();
for (y, column) in sample.iter_mut().enumerate() {
for (x, pixel) in column.iter_mut().enumerate() {
let x = (top_left.x() + x as f32 * image_pixels_per_sample_pixel) as u32;
let y = (top_left.y() + y as f32 * image_pixels_per_sample_pixel) as u32;
*pixel = self.try_at(x, y).map_or(128.0, |pixel| pixel.y as f32);
}
}

sample
}
}

pub type Sample = [[f32; SAMPLE_SIZE]; SAMPLE_SIZE];
50 changes: 14 additions & 36 deletions crates/vision/src/ball_detection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,9 @@ use types::{
multivariate_normal_distribution::MultivariateNormalDistribution,
parameters::BallDetectionParameters,
perspective_grid_candidates::PerspectiveGridCandidates,
ycbcr422_image::YCbCr422Image,
ycbcr422_image::{Sample, YCbCr422Image, SAMPLE_SIZE},
};

pub const SAMPLE_SIZE: usize = 32;
pub type Sample = [[f32; SAMPLE_SIZE]; SAMPLE_SIZE];

struct NeuralNetworks {
preclassifier: CompiledNN,
classifier: CompiledNN,
Expand Down Expand Up @@ -184,22 +181,6 @@ fn position_sample(network: &mut CompiledNN, sample: &Sample) -> Circle<Pixel> {
}
}

fn sample_grayscale(image: &YCbCr422Image, candidate: Circle<Pixel>) -> Sample {
let top_left = candidate.center - vector![candidate.radius, candidate.radius];
let image_pixels_per_sample_pixel = candidate.radius * 2.0 / SAMPLE_SIZE as f32;

let mut sample = Sample::default();
for (y, column) in sample.iter_mut().enumerate() {
for (x, pixel) in column.iter_mut().enumerate() {
let x = (top_left.x() + x as f32 * image_pixels_per_sample_pixel) as u32;
let y = (top_left.y() + y as f32 * image_pixels_per_sample_pixel) as u32;
*pixel = image.try_at(x, y).map_or(128.0, |pixel| pixel.y as f32);
}
}

sample
}

fn evaluate_candidates(
candidates: &[Circle<Pixel>],
image: &YCbCr422Image,
Expand All @@ -221,7 +202,7 @@ fn evaluate_candidates(
center: candidate.center,
radius: candidate.radius * ball_radius_enlargement_factor,
};
let sample = sample_grayscale(image, enlarged_candidate);
let sample = image.sample_grayscale(enlarged_candidate);
let preclassifier_confidence = preclassify_sample(preclassifier, &sample);

let mut classifier_confidence = None;
Expand Down Expand Up @@ -401,13 +382,12 @@ mod tests {
fn preclassify_ball() {
let mut network = CompiledNN::default();
network.compile(CLASSIFIER_PATH);
let sample = sample_grayscale(
&YCbCr422Image::load_from_444_png(Path::new(BALL_SAMPLE_PATH)).unwrap(),
Circle {
let sample = YCbCr422Image::load_from_444_png(Path::new(BALL_SAMPLE_PATH))
.unwrap()
.sample_grayscale(Circle {
center: point![16.0, 16.0],
radius: 16.0,
},
);
});
let confidence = preclassify_sample(&mut network, &sample);

println!("{confidence:?}");
Expand All @@ -418,13 +398,12 @@ mod tests {
fn classify_ball() {
let mut network = CompiledNN::default();
network.compile(PRECLASSIFIER_PATH);
let sample = sample_grayscale(
&YCbCr422Image::load_from_444_png(Path::new(BALL_SAMPLE_PATH)).unwrap(),
Circle {
let sample = YCbCr422Image::load_from_444_png(Path::new(BALL_SAMPLE_PATH))
.unwrap()
.sample_grayscale(Circle {
center: point![16.0, 16.0],
radius: 16.0,
},
);
});
let confidence = classify_sample(&mut network, &sample);

println!("{confidence:?}");
Expand All @@ -435,13 +414,12 @@ mod tests {
fn position_ball() {
let mut network = CompiledNN::default();
network.compile(POSITIONER_PATH);
let sample = sample_grayscale(
&YCbCr422Image::load_from_444_png(Path::new(BALL_SAMPLE_PATH)).unwrap(),
Circle {
let sample = YCbCr422Image::load_from_444_png(Path::new(BALL_SAMPLE_PATH))
.unwrap()
.sample_grayscale(Circle {
center: point![16.0, 16.0],
radius: 16.0,
},
);
});
let circle = position_sample(&mut network, &sample);

assert_relative_eq!(
Expand Down
2 changes: 1 addition & 1 deletion tools/twix/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "twix"
version = "0.7.0"
version = "0.7.1"
edition.workspace = true
license.workspace = true
homepage.workspace = true
Expand Down
7 changes: 4 additions & 3 deletions tools/twix/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@ use log::error;
use nao::Nao;
use panel::Panel;
use panels::{
BehaviorSimulatorPanel, EnumPlotPanel, ImageColorSelectPanel, ImagePanel, ImageSegmentsPanel,
LookAtPanel, ManualCalibrationPanel, MapPanel, ParameterPanel, PlotPanel, RemotePanel,
TextPanel, VisionTunerPanel,
BallCandidatePanel, BehaviorSimulatorPanel, EnumPlotPanel, ImageColorSelectPanel, ImagePanel,
ImageSegmentsPanel, LookAtPanel, ManualCalibrationPanel, MapPanel, ParameterPanel, PlotPanel,
RemotePanel, TextPanel, VisionTunerPanel,
};

use repository::{get_repository_root, Repository};
Expand Down Expand Up @@ -155,6 +155,7 @@ impl ReachableNaos {
}

impl_selectable_panel!(
BallCandidatePanel,
BehaviorSimulatorPanel,
ImagePanel,
ImageSegmentsPanel,
Expand Down
208 changes: 208 additions & 0 deletions tools/twix/src/panels/ball_candidates.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,208 @@
use std::sync::Arc;

use coordinate_systems::Pixel;
use eframe::{
egui::{Color32, Pos2, Rect, Response, Stroke, Ui, Vec2, Widget},
emath::RectTransform,
};
use geometry::circle::Circle;
use linear_algebra::{point, vector};
use serde_json::{json, Value};
use types::{
ball_detection::CandidateEvaluation,
ycbcr422_image::{YCbCr422Image, SAMPLE_SIZE},
};

use crate::{
nao::Nao,
panel::Panel,
twix_painter::{Orientation, TwixPainter},
value_buffer::BufferHandle,
};

use super::image::cycler_selector::{VisionCycler, VisionCyclerSelector};

pub struct BallCandidatePanel {
nao: Arc<Nao>,
cycler: VisionCycler,
ball_radius_enlargement_factor: BufferHandle<f32>,
ball_candidates: BufferHandle<Option<Vec<CandidateEvaluation>>>,
image: BufferHandle<YCbCr422Image>,
}

impl Panel for BallCandidatePanel {
const NAME: &'static str = "Ball Candidates";

fn new(nao: Arc<Nao>, value: Option<&Value>) -> Self {
let cycler = value
.and_then(|value| {
let string = value.get("cycler")?.as_str()?;
VisionCycler::try_from(string).ok()
})
.unwrap_or(VisionCycler::Top);

let cycler_path = cycler.as_snake_case_path();
let ball_radius_enlargement_factor = nao.subscribe_value(format!(
"parameters.ball_detection.{cycler_path}.ball_radius_enlargement_factor",
));
let cycler_path = cycler.as_path();
let ball_candidates =
nao.subscribe_value(format!("{cycler_path}.additional_outputs.ball_candidates"));
let image = nao.subscribe_value(format!("{cycler_path}.main_outputs.image"));
Self {
nao,
cycler,
ball_radius_enlargement_factor,
ball_candidates,
image,
}
}

fn save(&self) -> Value {
json!({
"cycler": self.cycler.as_path(),
})
}
}

impl Widget for &mut BallCandidatePanel {
fn ui(self, ui: &mut Ui) -> Response {
ui.vertical(|ui| {
ui.horizontal(|ui| {
let mut cycler_selector = VisionCyclerSelector::new(&mut self.cycler);
if cycler_selector.ui(ui).changed() {
self.resubscribe();
}
});
ui.separator();
if let Some((ball_radius_enlargement_factor, ball_candidates, image)) = self
.ball_radius_enlargement_factor
.get_last_value()
.ok()
.flatten()
.and_then(|ball_radius_enlargement_factor| {
self.ball_candidates
.get_last_value()
.ok()
.flatten()
.flatten()
.map(|ball_candidates| (ball_radius_enlargement_factor, ball_candidates))
})
.and_then(|(ball_radius_enlargement_factor, ball_candidates)| {
self.image
.get_last_value()
.ok()
.flatten()
.map(|image| (ball_radius_enlargement_factor, ball_candidates, image))
})
{
ui.horizontal_wrapped(|ui| {
for candidate in ball_candidates {
ui.add(CandidateSample {
ball_radius_enlargement_factor,
candidate,
image: image.clone(),
});
}
});
} else {
ui.label("Some outputs are missing");
}
})
.response
}
}

impl BallCandidatePanel {
fn resubscribe(&mut self) {
let cycler_path = self.cycler.as_snake_case_path();
self.ball_radius_enlargement_factor = self.nao.subscribe_value(format!(
"parameters.ball_detection.{cycler_path}.ball_radius_enlargement_factor",
));
let cycler_path = self.cycler.as_path();
self.ball_candidates = self
.nao
.subscribe_value(format!("{cycler_path}.additional_outputs.ball_candidates"));
self.image = self
.nao
.subscribe_value(format!("{cycler_path}.main_outputs.image"));
}
}

struct CandidateSample {
ball_radius_enlargement_factor: f32,
candidate: CandidateEvaluation,
image: YCbCr422Image,
}

impl Widget for CandidateSample {
fn ui(self, ui: &mut Ui) -> Response {
let enlarged_candidate = Circle {
center: self.candidate.candidate_circle.center,
radius: self.candidate.candidate_circle.radius * self.ball_radius_enlargement_factor,
};

let sample = self.image.sample_grayscale(enlarged_candidate);

const SAMPLE_SIZE_F32: f32 = SAMPLE_SIZE as f32;
const SCALING: f32 = 3.0;
ui.allocate_ui(
Vec2::new(SAMPLE_SIZE_F32 * SCALING, SAMPLE_SIZE_F32 * SCALING),
|ui| {
let (response, painter) = TwixPainter::<Pixel>::allocate(
ui,
vector![SAMPLE_SIZE_F32, SAMPLE_SIZE_F32],
point![0.0, 0.0],
Orientation::LeftHanded,
);

for (y, sample_row) in sample.iter().enumerate() {
let y = y as f32;
for (x, sample_value) in sample_row.iter().enumerate() {
let x = x as f32;
painter.rect_filled(
point![x, y],
point![x + 1.0, y + 1.0],
Color32::from_gray(*sample_value as u8),
);
}
}

if let Some(corrected_circle) = self.candidate.corrected_circle {
let candidate_circle = self.candidate.candidate_circle;
let transform = RectTransform::from_to(
Rect::from_center_size(
Pos2::new(candidate_circle.center.x(), candidate_circle.center.y()),
Vec2::new(
candidate_circle.radius * 2.0 * self.ball_radius_enlargement_factor,
candidate_circle.radius * 2.0 * self.ball_radius_enlargement_factor,
),
),
Rect::from_min_size(
Pos2::ZERO,
Vec2::new(SAMPLE_SIZE_F32, SAMPLE_SIZE_F32),
),
);
let corrected_center_in_sample = transform.transform_pos(Pos2::new(
corrected_circle.center.x(),
corrected_circle.center.y(),
));
let corrected_center_in_sample =
point![corrected_center_in_sample.x, corrected_center_in_sample.y];
let corrected_radius_in_sample = corrected_circle.radius
/ (self.candidate.candidate_circle.radius
* self.ball_radius_enlargement_factor)
* (SAMPLE_SIZE_F32 / 2.0);
painter.circle_stroke(
corrected_center_in_sample,
corrected_radius_in_sample,
Stroke::new(0.5, Color32::GREEN),
);
}

response
},
)
.response
}
}
Loading

0 comments on commit 08ca138

Please sign in to comment.