Skip to content

Commit

Permalink
clippy fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
chris-ha458 committed Jan 7, 2024
1 parent 688c21f commit 6c102ce
Show file tree
Hide file tree
Showing 6 changed files with 19 additions and 26 deletions.
10 changes: 6 additions & 4 deletions ocrs-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -219,16 +219,18 @@ fn main() -> Result<(), Box<dyn Error>> {
let detection_model_src = args
.detection_model
.as_ref()
.map(|path| ModelSource::Path(path))
.unwrap_or(ModelSource::Url(DETECTION_MODEL));
.map_or(ModelSource::Url(DETECTION_MODEL), |path| {
ModelSource::Path(path)
});
let detection_model = load_model(detection_model_src)
.file_error_context("Failed to load text detection model", detection_model_src)?;

let recognition_model_src = args
.recognition_model
.as_ref()
.map(|path| ModelSource::Path(path))
.unwrap_or(ModelSource::Url(RECOGNITION_MODEL));
.map_or(ModelSource::Url(RECOGNITION_MODEL), |path| {
ModelSource::Path(path)
});
let recognition_model = load_model(recognition_model_src).file_error_context(
"Failed to load text recognition model",
recognition_model_src,
Expand Down
2 changes: 1 addition & 1 deletion ocrs-cli/src/output.rs
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ pub fn generate_annotated_png(args: GeneratePngArgs) -> NdTensor<f32, 3> {
let floor_point = |p: PointF| Point::from_yx(p.y as i32, p.x as i32);

// Draw line bounding rects from layout analysis step.
for line in line_rects.iter() {
for line in line_rects {
let line_points: Vec<_> = line
.iter()
.flat_map(|word_rect| word_rect.corners().into_iter())
Expand Down
8 changes: 3 additions & 5 deletions ocrs/src/detection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,9 @@ impl TextDetector {
// Add batch dim
let image = image.reshaped([1, img_chans, img_height, img_width]);

let (in_height, in_width) = match self.input_shape[..] {
[_, _, Dimension::Fixed(h), Dimension::Fixed(w)] => (h, w),
_ => {
return Err("failed to get model dims".into());
}
let [_, _, Dimension::Fixed(in_height), Dimension::Fixed(in_width)] = self.input_shape[..]
else {
return Err("failed to get model dims".into());
};

// Pad small images to the input size of the text detection model. This is
Expand Down
9 changes: 4 additions & 5 deletions ocrs/src/layout_analysis.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,27 +92,26 @@ pub fn find_block_separators(words: &[RotatedRect]) -> Vec<Rect> {
lines.sort_by_key(|l| l.first().unwrap().bounding_rect().top().round() as i32);

let mut all_word_spacings = Vec::new();
for line in lines.iter() {
for line in lines {
if line.len() > 1 {
let mut spacings: Vec<_> = zip(line.iter(), line.iter().skip(1))
.map(|(cur, next)| {
(next.bounding_rect().left() - cur.bounding_rect().right()).round() as i32
})
.collect();
spacings.sort();
spacings.sort_unstable();
all_word_spacings.extend_from_slice(&spacings);
}
}
all_word_spacings.sort();
all_word_spacings.sort_unstable();

let median_word_spacing = all_word_spacings
.get(all_word_spacings.len() / 2)
.copied()
.unwrap_or(10);
let median_height = words
.get(words.len() / 2)
.map(|r| r.height())
.unwrap_or(10.)
.map_or(10.0, |r| r.height())
.round() as i32;

// Scoring function for empty rectangles. Taken from Section 3.D in [1].
Expand Down
2 changes: 1 addition & 1 deletion ocrs/src/layout_analysis/empty_rects.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ where
queue.push(Partition {
score: score(boundary),
boundary,
obstacles: obstacles.to_vec(),
obstacles: obstacles.clone(),
});
}

Expand Down
14 changes: 4 additions & 10 deletions ocrs/src/recognition.rs
Original file line number Diff line number Diff line change
Expand Up @@ -157,14 +157,12 @@ fn polygon_slice_bounding_rect(
let trunc_edge_start = e
.to_f32()
.y_for_x(min_x as f32)
.map(|y| Point::from_yx(y.round() as i32, min_x))
.unwrap_or(e.start);
.map_or(e.start, |y| Point::from_yx(y.round() as i32, min_x));

let trunc_edge_end = e
.to_f32()
.y_for_x(max_x as f32)
.map(|y| Point::from_yx(y.round() as i32, max_x))
.unwrap_or(e.end);
.map_or(e.end, |y| Point::from_yx(y.round() as i32, max_x));

Some(Line::from_endpoints(trunc_edge_start, trunc_edge_end))
})
Expand Down Expand Up @@ -370,7 +368,7 @@ impl TextRecognizer {
let min_width = 10.;
let max_width = 800.;
let aspect_ratio = orig_width as f32 / orig_height as f32;
(height as f32 * aspect_ratio).max(min_width).min(max_width) as u32
(height as f32 * aspect_ratio).clamp(min_width, max_width) as u32
}

// Group lines into batches which will have similar widths after resizing
Expand Down Expand Up @@ -489,11 +487,7 @@ mod tests {
// Vary the orientation of words. The output of `line_polygon`
// should be invariant to different orientations of a RotatedRect
// that cover the same pixels.
let up = if i % 2 == 0 {
Vec2::from_yx(-1., 0.)
} else {
Vec2::from_yx(1., 0.)
};
let up = Vec2::from_yx(if i % 2 == 0 { -1. } else { 1. }, 0.);
RotatedRect::new(center, up, width, height)
})
.collect();
Expand Down

0 comments on commit 6c102ce

Please sign in to comment.