vkit.pipeline.text_detection.page_assembler

  1# Copyright 2022 vkit-x Administrator. All Rights Reserved.
  2#
  3# This project (vkit-x/vkit) is dual-licensed under commercial and SSPL licenses.
  4#
  5# The commercial license gives you the full rights to create and distribute software
  6# on your own terms without any SSPL license obligations. For more information,
  7# please see the "LICENSE_COMMERCIAL.txt" file.
  8#
  9# This project is also available under Server Side Public License (SSPL).
 10# The SSPL licensing is ideal for use cases such as open source projects with
 11# SSPL distribution, student/academic purposes, hobby projects, internal research
 12# projects without external distribution, or other projects where all SSPL
 13# obligations can be met. For more information, please see the "LICENSE_SSPL.txt" file.
 14from typing import Sequence, List
 15
 16import attrs
 17from numpy.random import Generator as RandomGenerator
 18
 19from vkit.element import Shapable, Box, Polygon, Image
 20from vkit.engine.seal_impression import fill_text_line_to_seal_impression
 21from vkit.mechanism.distortion import rotate
 22from ..interface import PipelineStep, PipelineStepFactory
 23from .page_layout import (
 24    PageLayoutStepOutput,
 25    DisconnectedTextRegion,
 26    NonTextRegion,
 27)
 28from .page_background import PageBackgroundStepOutput
 29from .page_image import PageImageStepOutput, PageImageCollection
 30from .page_barcode import PageBarcodeStepOutput
 31from .page_text_line import (
 32    PageTextLineStepOutput,
 33    PageTextLineCollection,
 34    PageSealImpressionTextLineCollection,
 35)
 36from .page_non_text_symbol import PageNonTextSymbolStepOutput
 37from .page_text_line_label import (
 38    PageTextLineLabelStepOutput,
 39    PageCharPolygonCollection,
 40    PageTextLinePolygonCollection,
 41)
 42from .page_text_line_bounding_box import PageTextLineBoundingBoxStepOutput
 43
 44
 45@attrs.define
 46class PageAssemblerStepConfig:
 47    pass
 48
 49
 50@attrs.define
 51class PageAssemblerStepInput:
 52    page_layout_step_output: PageLayoutStepOutput
 53    page_background_step_output: PageBackgroundStepOutput
 54    page_image_step_output: PageImageStepOutput
 55    page_barcode_step_output: PageBarcodeStepOutput
 56    page_text_line_step_output: PageTextLineStepOutput
 57    page_non_text_symbol_step_output: PageNonTextSymbolStepOutput
 58    page_text_line_bounding_box_step_output: PageTextLineBoundingBoxStepOutput
 59    page_text_line_label_step_output: PageTextLineLabelStepOutput
 60
 61
 62@attrs.define
 63class PageDisconnectedTextRegionCollection:
 64    disconnected_text_regions: Sequence[DisconnectedTextRegion]
 65
 66    def to_polygons(self):
 67        for disconnected_text_region in self.disconnected_text_regions:
 68            yield disconnected_text_region.polygon
 69
 70
 71@attrs.define
 72class PageNonTextRegionCollection:
 73    non_text_regions: Sequence[NonTextRegion]
 74
 75    def to_polygons(self):
 76        for non_text_region in self.non_text_regions:
 77            yield non_text_region.polygon
 78
 79
 80@attrs.define
 81class PageSealImpressionCharPolygonCollection:
 82    char_polygons: Sequence[Polygon]
 83
 84
 85@attrs.define
 86class Page(Shapable):
 87    image: Image
 88    page_image_collection: PageImageCollection
 89    page_bottom_layer_image: Image
 90    page_text_line_collection: PageTextLineCollection
 91    page_seal_impression_text_line_collection: PageSealImpressionTextLineCollection
 92    page_char_polygon_collection: PageCharPolygonCollection
 93    page_text_line_polygon_collection: PageTextLinePolygonCollection
 94    page_disconnected_text_region_collection: PageDisconnectedTextRegionCollection
 95    page_non_text_region_collection: PageNonTextRegionCollection
 96    page_seal_impression_char_polygon_collection: PageSealImpressionCharPolygonCollection
 97
 98    @property
 99    def height(self):
100        return self.image.height
101
102    @property
103    def width(self):
104        return self.image.width
105
106
107@attrs.define
108class PageAssemblerStepOutput:
109    page: Page
110
111
112class PageAssemblerStep(
113    PipelineStep[
114        PageAssemblerStepConfig,
115        PageAssemblerStepInput,
116        PageAssemblerStepOutput,
117    ]
118):  # yapf: disable
119
120    def run(self, input: PageAssemblerStepInput, rng: RandomGenerator):
121        page_layout_step_output = input.page_layout_step_output
122        page_layout = page_layout_step_output.page_layout
123
124        page_background_step_output = input.page_background_step_output
125        background_image = page_background_step_output.background_image
126
127        page_image_step_output = input.page_image_step_output
128        page_image_collection = page_image_step_output.page_image_collection
129        page_bottom_layer_image = page_image_step_output.page_bottom_layer_image
130
131        page_barcode_step_output = input.page_barcode_step_output
132
133        page_text_line_step_output = input.page_text_line_step_output
134        page_text_line_collection = page_text_line_step_output.page_text_line_collection
135        page_seal_impression_text_line_collection = \
136            page_text_line_step_output.page_seal_impression_text_line_collection
137
138        page_non_text_symbol_step_output = input.page_non_text_symbol_step_output
139
140        page_text_line_bounding_box_step_output = input.page_text_line_bounding_box_step_output
141        text_line_bounding_box_score_maps = page_text_line_bounding_box_step_output.score_maps
142        text_line_bounding_box_colors = page_text_line_bounding_box_step_output.colors
143
144        page_text_line_label_step_output = input.page_text_line_label_step_output
145        page_char_polygon_collection = \
146            page_text_line_label_step_output.page_char_polygon_collection
147        page_text_line_polygon_collection = \
148            page_text_line_label_step_output.page_text_line_polygon_collection
149
150        # Page background.
151        assert background_image.mat.shape == (page_layout.height, page_layout.width, 3)
152        assembled_image = background_image.copy()
153
154        # Page images.
155        for page_image in page_image_collection.page_images:
156            page_image.box.fill_image(
157                assembled_image,
158                page_image.image,
159                alpha=page_image.alpha,
160            )
161
162        # Page barcodes.
163        for barcode_qr_score_map in page_barcode_step_output.barcode_qr_score_maps:
164            assembled_image[barcode_qr_score_map] = (0, 0, 0)
165        for barcode_code39_score_map in page_barcode_step_output.barcode_code39_score_maps:
166            assembled_image[barcode_code39_score_map] = (0, 0, 0)
167
168        # Page text line bounding boxes.
169        for text_line_bounding_box_score_map, text_line_bounding_box_color in zip(
170            text_line_bounding_box_score_maps, text_line_bounding_box_colors
171        ):
172            assembled_image[text_line_bounding_box_score_map] = text_line_bounding_box_color
173
174        # Page text lines.
175        for text_line in page_text_line_collection.text_lines:
176            if text_line.score_map:
177                text_line.score_map.fill_image(assembled_image, text_line.glyph_color)
178            else:
179                text_line.mask.fill_image(assembled_image, text_line.image)
180
181        # Page non-text symbols.
182        for image, box, alpha in zip(
183            page_non_text_symbol_step_output.images,
184            page_non_text_symbol_step_output.boxes,
185            page_non_text_symbol_step_output.alphas,
186        ):
187            box.fill_image(assembled_image, value=image, alpha=alpha)
188
189        # Page seal impressions.
190        page_seal_impression_char_polygons: List[Polygon] = []
191        for seal_impression, seal_impression_resource in zip(
192            page_seal_impression_text_line_collection.seal_impressions,
193            page_seal_impression_text_line_collection.seal_impression_resources,
194        ):
195            alpha = seal_impression.alpha
196            color = seal_impression.color
197
198            # Prepare foreground (text) and background.
199            background_mask = seal_impression.background_mask
200            text_line_filled_score_map, char_polygons = fill_text_line_to_seal_impression(
201                seal_impression,
202                seal_impression_resource.text_line_slot_indices,
203                seal_impression_resource.text_lines,
204                seal_impression_resource.internal_text_line,
205            )
206
207            # Rotate, shift, and trim.
208            rotated_result = rotate.distort(
209                {'angle': seal_impression_resource.angle},
210                mask=background_mask,
211                score_map=text_line_filled_score_map,
212                polygons=char_polygons,
213            )
214            assert rotated_result.mask
215            background_mask = rotated_result.mask
216            assert rotated_result.score_map
217            text_line_filled_score_map = rotated_result.score_map
218            assert background_mask.shape == text_line_filled_score_map.shape
219            assert rotated_result.polygons
220
221            # Placement box.
222            box_center_point = seal_impression_resource.box.get_center_point()
223            up = box_center_point.y - background_mask.height // 2
224            down = up + background_mask.height - 1
225            left = box_center_point.x - background_mask.width // 2
226            right = left + background_mask.width - 1
227
228            if up < 0 or down >= assembled_image.height \
229                    or left < 0 or right >= assembled_image.width:
230                # Simply ignore this seal impression.
231                continue
232
233            # Rendering.
234            box = Box(up=up, down=down, left=left, right=right)
235            box.fill_image(assembled_image, value=color, image_mask=background_mask, alpha=alpha)
236            box.fill_image(assembled_image, value=color, alpha=text_line_filled_score_map)
237
238            # Shift and keep char polyghons.
239            char_polygons = [
240                char_polygon.to_shifted_polygon(
241                    offset_y=up,
242                    offset_x=left,
243                ) for char_polygon in rotated_result.polygons
244            ]
245            page_seal_impression_char_polygons.extend(char_polygons)
246
247        # For char-level polygon regression.
248        page_disconnected_text_region_collection = PageDisconnectedTextRegionCollection(
249            page_layout.disconnected_text_regions
250        )
251
252        # For sampling negative text region area.
253        page_non_text_region_collection = PageNonTextRegionCollection(page_layout.non_text_regions)
254
255        # For generating char-level seal impression labeling.
256        page_seal_impression_char_polygon_collection = PageSealImpressionCharPolygonCollection(
257            char_polygons=page_seal_impression_char_polygons,
258        )
259
260        page = Page(
261            image=assembled_image,
262            page_image_collection=page_image_collection,
263            page_bottom_layer_image=page_bottom_layer_image,
264            page_text_line_collection=page_text_line_collection,
265            page_seal_impression_text_line_collection=page_seal_impression_text_line_collection,
266            page_char_polygon_collection=page_char_polygon_collection,
267            page_text_line_polygon_collection=page_text_line_polygon_collection,
268            page_disconnected_text_region_collection=page_disconnected_text_region_collection,
269            page_non_text_region_collection=page_non_text_region_collection,
270            page_seal_impression_char_polygon_collection=(
271                page_seal_impression_char_polygon_collection
272            ),
273        )
274        return PageAssemblerStepOutput(page=page)
275
276
277page_assembler_step_factory = PipelineStepFactory(PageAssemblerStep)