use image::DynamicImage;
use crate::component::Component;
use crate::element::Element;
use crate::config::Config;
pub fn merge(
img: &DynamicImage,
comps: &[Component],
texts: &[Element],
config: &Config,
is_paragraph: bool,
is_remove_bar: bool,
) -> Vec<Element> {
let img_shape = (img.height(), img.width());
let img_height = img_shape.0 as f64;
let comp_eles: Vec<Element> = comps.iter()
.map(|c| {
Element::new(0, &c.bbox, &c.category, None)
})
.collect();
let mut text_eles = texts.to_vec();
text_eles = refine_texts(&text_eles, img_shape, config.text_max_height);
let mut elements = refine_elements(&comp_eles, &text_eles, (2, 2), 0.8);
if is_remove_bar {
elements = remove_top_bar(&elements, img_height, config.top_bottom_bar);
elements = remove_bottom_bar(&elements, img_height, config.top_bottom_bar);
}
if is_paragraph {
elements = merge_text_lines(&elements, config.text_max_word_gap as i32);
}
reassign_ids(&mut elements);
check_containment(&mut elements);
elements
}
pub fn refine_texts(texts: &[Element], img_shape: (u32, u32), text_max_height_ratio: f64) -> Vec<Element> {
texts.iter()
.filter(|t| {
let content = t.text_content.as_deref().unwrap_or("");
let content_len = content.len();
if content_len == 0 {
return false; }
if content_len == 1 {
let c = content.chars().next().unwrap();
if c.is_ascii_digit() {
return true;
}
if c.is_ascii_alphabetic() {
return true;
}
if ('\u{4E00}'..='\u{9FFF}').contains(&c) {
return true;
}
return false;
}
if content_len > 5 && t.width > t.height * 2 {
return true;
}
let height_ratio = (t.height as f64) / (img_shape.0 as f64);
height_ratio < text_max_height_ratio
})
.cloned()
.collect()
}
pub fn refine_elements(
comps: &[Element],
texts: &[Element],
bias: (i32, i32),
containment_ratio: f64,
) -> Vec<Element> {
let mut elements: Vec<Element> = Vec::new();
for comp in comps {
let mut is_covered_by_text = false;
for text in texts.iter() {
let (inter, _, ioa, _) = comp.calc_intersection(text, bias);
if inter > 0 {
if ioa >= containment_ratio
&& comp.class != "Block"
&& comp.class != "Image"
{
is_covered_by_text = true;
break;
}
}
}
if !is_covered_by_text && comp.area() > 0 {
elements.push(comp.clone());
}
}
for text in texts.iter() {
elements.push(text.clone());
}
elements
}
pub fn remove_top_bar(elements: &[Element], img_height: f64, top_bottom_bar: (f64, f64)) -> Vec<Element> {
let max_height = img_height * 0.04;
let top_threshold = (img_height * top_bottom_bar.0) as i32;
elements.iter()
.filter(|e| !(e.position.row_min < top_threshold && (e.height as f64) < max_height))
.cloned()
.collect()
}
pub fn remove_bottom_bar(elements: &[Element], img_height: f64, top_bottom_bar: (f64, f64)) -> Vec<Element> {
let bottom_start = (img_height * top_bottom_bar.1) as i32;
elements.iter()
.filter(|e| {
!(e.position.row_min > bottom_start
&& (e.height as f64) < img_height * 0.03
&& (e.width as f64) < img_height * 0.03)
})
.cloned()
.collect()
}
pub fn merge_text_lines(elements: &[Element], max_line_gap: i32) -> Vec<Element> {
let mut texts: Vec<Element> = Vec::new();
let mut non_texts: Vec<Element> = Vec::new();
for ele in elements {
if ele.class == "Text" {
texts.push(ele.clone());
} else {
non_texts.push(ele.clone());
}
}
let mut changed = true;
while changed {
changed = false;
let mut temp: Vec<Element> = Vec::new();
'outer: for ta in &texts {
for tb in &mut temp {
let (inter, _, _, _) = ta.calc_intersection(tb, (0, max_line_gap));
if inter > 0 {
tb.element_merge(ta);
changed = true;
continue 'outer;
}
}
temp.push(ta.clone());
}
texts = temp;
}
non_texts.extend(texts);
non_texts
}
pub fn reassign_ids(elements: &mut [Element]) {
for (i, ele) in elements.iter_mut().enumerate() {
ele.id = i;
}
}
pub fn check_containment(elements: &mut [Element]) {
let n = elements.len();
for i in 0..n {
for j in (i + 1)..n {
let rel = elements[i].element_relation(&elements[j], (2, 2));
if rel == -1 {
let i_id = elements[i].id;
let current_parent = elements[i].parent;
let should_assign = match current_parent {
None => true,
Some(pid) => {
elements[j].area() < elements[pid].area()
}
};
if should_assign {
if let Some(pid) = current_parent {
if let Some(children) = elements[pid].children.as_mut() {
children.retain(|&c| c != i_id);
}
}
elements[j].children.get_or_insert(vec![]).push(elements[i].id);
elements[i].parent = Some(elements[j].id);
}
} else if rel == 1 {
let j_id = elements[j].id;
let current_parent = elements[j].parent;
let should_assign = match current_parent {
None => true,
Some(pid) => {
elements[i].area() < elements[pid].area()
}
};
if should_assign {
if let Some(pid) = current_parent {
if let Some(children) = elements[pid].children.as_mut() {
children.retain(|&c| c != j_id);
}
}
elements[i].children.get_or_insert(vec![]).push(elements[j].id);
elements[j].parent = Some(elements[i].id);
}
}
}
}
}
pub fn synthesize_orphan_text_regions(
elements: &mut Vec<Element>,
raw_texts: &[Element],
coverage_threshold: f64,
line_gap: i32,
) {
if raw_texts.is_empty() {
return;
}
let non_texts: Vec<&Element> = elements.iter()
.filter(|e| e.class != "Text")
.collect();
let orphan_indices: Vec<usize> = raw_texts.iter().enumerate()
.filter(|(_, text)| {
let max_iou = non_texts.iter()
.map(|nt| {
let (_, iou, _, _) = text.calc_intersection(nt, (2, 2));
iou
})
.fold(0.0f64, f64::max);
max_iou < coverage_threshold
})
.map(|(idx, _)| idx)
.collect();
if orphan_indices.is_empty() {
return;
}
let mut sorted_orphan: Vec<(usize, i32)> = orphan_indices.iter()
.map(|&idx| (idx, raw_texts[idx].position.row_min))
.collect();
sorted_orphan.sort_by_key(|&(_, row)| row);
let mut groups: Vec<Vec<usize>> = Vec::new();
let mut used = vec![false; sorted_orphan.len()];
for i in 0..sorted_orphan.len() {
if used[i] { continue; }
let mut group = vec![sorted_orphan[i].0];
used[i] = true;
for j in (i + 1)..sorted_orphan.len() {
if used[j] { continue; }
let last_idx = *group.last().unwrap();
let last = &raw_texts[last_idx];
let candidate = &raw_texts[sorted_orphan[j].0];
let vert_gap = candidate.position.row_min - last.position.row_max;
let horiz_gap = if candidate.position.column_min > last.position.column_max {
candidate.position.column_min - last.position.column_max
} else if last.position.column_min > candidate.position.column_max {
last.position.column_min - candidate.position.column_max
} else {
0 };
if vert_gap <= line_gap && horiz_gap <= 60 {
group.push(sorted_orphan[j].0);
used[j] = true;
}
}
groups.push(group);
}
let mut new_blocks: Vec<Element> = Vec::new();
for (_g_idx, group) in groups.iter().enumerate() {
if group.is_empty() { continue; }
let col_min = group.iter().map(|&idx| raw_texts[idx].position.column_min).min().unwrap_or(0);
let row_min = group.iter().map(|&idx| raw_texts[idx].position.row_min).min().unwrap_or(0);
let col_max = group.iter().map(|&idx| raw_texts[idx].position.column_max).max().unwrap_or(0);
let row_max = group.iter().map(|&idx| raw_texts[idx].position.row_max).max().unwrap_or(0);
let pad = 6i32;
let block = Element::from_parts(
0,
(col_min - pad).max(0),
(row_min - pad).max(0),
col_max + pad,
row_max + pad,
"Block",
);
new_blocks.push(block);
}
if !new_blocks.is_empty() {
for e in elements.iter_mut() {
e.parent = None;
e.children = None;
}
elements.extend(new_blocks);
reassign_ids(elements);
check_containment(elements);
}
}