Postprocesses layout predictions by cleaning up clusters and mapping cells.
| 156 | |
| 157 | |
| 158 | class LayoutPostprocessor: |
| 159 | """Postprocesses layout predictions by cleaning up clusters and mapping cells.""" |
| 160 | |
| 161 | # Cluster type-specific parameters for overlap resolution |
| 162 | OVERLAP_PARAMS = { |
| 163 | "regular": {"area_threshold": 1.3, "conf_threshold": 0.05}, |
| 164 | "picture": {"area_threshold": 2.0, "conf_threshold": 0.3}, |
| 165 | "wrapper": {"area_threshold": 2.0, "conf_threshold": 0.2}, |
| 166 | } |
| 167 | |
| 168 | WRAPPER_TYPES = { |
| 169 | DocItemLabel.FORM, |
| 170 | DocItemLabel.KEY_VALUE_REGION, |
| 171 | DocItemLabel.TABLE, |
| 172 | DocItemLabel.DOCUMENT_INDEX, |
| 173 | } |
| 174 | SPECIAL_TYPES = WRAPPER_TYPES.union({DocItemLabel.PICTURE}) |
| 175 | |
| 176 | CONFIDENCE_THRESHOLDS = { |
| 177 | DocItemLabel.CAPTION: 0.5, |
| 178 | DocItemLabel.FOOTNOTE: 0.5, |
| 179 | DocItemLabel.FORMULA: 0.5, |
| 180 | DocItemLabel.LIST_ITEM: 0.5, |
| 181 | DocItemLabel.PAGE_FOOTER: 0.5, |
| 182 | DocItemLabel.PAGE_HEADER: 0.5, |
| 183 | DocItemLabel.PICTURE: 0.5, |
| 184 | DocItemLabel.SECTION_HEADER: 0.45, |
| 185 | DocItemLabel.TABLE: 0.5, |
| 186 | DocItemLabel.TEXT: 0.5, # 0.45, |
| 187 | DocItemLabel.TITLE: 0.45, |
| 188 | DocItemLabel.CODE: 0.45, |
| 189 | DocItemLabel.CHECKBOX_SELECTED: 0.45, |
| 190 | DocItemLabel.CHECKBOX_UNSELECTED: 0.45, |
| 191 | DocItemLabel.FORM: 0.45, |
| 192 | DocItemLabel.KEY_VALUE_REGION: 0.45, |
| 193 | DocItemLabel.DOCUMENT_INDEX: 0.45, |
| 194 | } |
| 195 | |
| 196 | LABEL_REMAPPING = { |
| 197 | # DocItemLabel.DOCUMENT_INDEX: DocItemLabel.TABLE, |
| 198 | DocItemLabel.TITLE: DocItemLabel.SECTION_HEADER, |
| 199 | } |
| 200 | |
| 201 | def __init__(self, cells: List[Cell], clusters: List[Cluster], page_size: Size): |
| 202 | """Initialize processor with cells and clusters.""" |
| 203 | """Initialize processor with cells and spatial indices.""" |
| 204 | self.cells = cells |
| 205 | self.page_size = page_size |
| 206 | self.regular_clusters = [ |
| 207 | c for c in clusters if c.label not in self.SPECIAL_TYPES |
| 208 | ] |
| 209 | self.special_clusters = [c for c in clusters if c.label in self.SPECIAL_TYPES] |
| 210 | |
| 211 | # Build spatial indices once |
| 212 | self.regular_index = SpatialClusterIndex(self.regular_clusters) |
| 213 | self.picture_index = SpatialClusterIndex( |
| 214 | [c for c in self.special_clusters if c.label == DocItemLabel.PICTURE] |
| 215 | ) |