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
 15
 16import attrs
 17from numpy.random import Generator as RandomGenerator
 18
 19from vkit.element import Shapable, Box, 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 Page(Shapable):
 82    image: Image
 83    page_image_collection: PageImageCollection
 84    page_bottom_layer_image: Image
 85    page_text_line_collection: PageTextLineCollection
 86    page_seal_impression_text_line_collection: PageSealImpressionTextLineCollection
 87    page_char_polygon_collection: PageCharPolygonCollection
 88    page_text_line_polygon_collection: PageTextLinePolygonCollection
 89    page_disconnected_text_region_collection: PageDisconnectedTextRegionCollection
 90    page_non_text_region_collection: PageNonTextRegionCollection
 91
 92    @property
 93    def height(self):
 94        return self.image.height
 95
 96    @property
 97    def width(self):
 98        return self.image.width
 99
100
101@attrs.define
102class PageAssemblerStepOutput:
103    page: Page
104
105
106class PageAssemblerStep(
107    PipelineStep[
108        PageAssemblerStepConfig,
109        PageAssemblerStepInput,
110        PageAssemblerStepOutput,
111    ]
112):  # yapf: disable
113
114    def run(self, input: PageAssemblerStepInput, rng: RandomGenerator):
115        page_layout_step_output = input.page_layout_step_output
116        page_layout = page_layout_step_output.page_layout
117
118        page_background_step_output = input.page_background_step_output
119        background_image = page_background_step_output.background_image
120
121        page_image_step_output = input.page_image_step_output
122        page_image_collection = page_image_step_output.page_image_collection
123        page_bottom_layer_image = page_image_step_output.page_bottom_layer_image
124
125        page_barcode_step_output = input.page_barcode_step_output
126
127        page_text_line_step_output = input.page_text_line_step_output
128        page_text_line_collection = page_text_line_step_output.page_text_line_collection
129        page_seal_impression_text_line_collection = \
130            page_text_line_step_output.page_seal_impression_text_line_collection
131
132        page_non_text_symbol_step_output = input.page_non_text_symbol_step_output
133
134        page_text_line_bounding_box_step_output = input.page_text_line_bounding_box_step_output
135        text_line_bounding_box_score_maps = page_text_line_bounding_box_step_output.score_maps
136        text_line_bounding_box_colors = page_text_line_bounding_box_step_output.colors
137
138        page_text_line_label_step_output = input.page_text_line_label_step_output
139        page_char_polygon_collection = \
140            page_text_line_label_step_output.page_char_polygon_collection
141        page_text_line_polygon_collection = \
142            page_text_line_label_step_output.page_text_line_polygon_collection
143
144        # Page background.
145        assert background_image.mat.shape == (page_layout.height, page_layout.width, 3)
146        assembled_image = background_image.copy()
147
148        # Page images.
149        for page_image in page_image_collection.page_images:
150            page_image.box.fill_image(
151                assembled_image,
152                page_image.image,
153                alpha=page_image.alpha,
154            )
155
156        # Page barcodes.
157        for barcode_qr_score_map in page_barcode_step_output.barcode_qr_score_maps:
158            assembled_image[barcode_qr_score_map] = (0, 0, 0)
159        for barcode_code39_score_map in page_barcode_step_output.barcode_code39_score_maps:
160            assembled_image[barcode_code39_score_map] = (0, 0, 0)
161
162        # Page text line bounding boxes.
163        for text_line_bounding_box_score_map, text_line_bounding_box_color in zip(
164            text_line_bounding_box_score_maps, text_line_bounding_box_colors
165        ):
166            assembled_image[text_line_bounding_box_score_map] = text_line_bounding_box_color
167
168        # Page text lines.
169        for text_line in page_text_line_collection.text_lines:
170            if text_line.score_map:
171                text_line.score_map.fill_image(assembled_image, text_line.glyph_color)
172            else:
173                text_line.mask.fill_image(assembled_image, text_line.image)
174
175        # Page non-text symbols.
176        for image, box, alpha in zip(
177            page_non_text_symbol_step_output.images,
178            page_non_text_symbol_step_output.boxes,
179            page_non_text_symbol_step_output.alphas,
180        ):
181            box.fill_image(assembled_image, value=image, alpha=alpha)
182
183        # Page seal impressions.
184        for seal_impression, seal_impression_resource in zip(
185            page_seal_impression_text_line_collection.seal_impressions,
186            page_seal_impression_text_line_collection.seal_impression_resources,
187        ):
188            alpha = seal_impression.alpha
189            color = seal_impression.color
190
191            # Prepare foreground (text) and background.
192            background_mask = seal_impression.background_mask
193            text_line_filled_score_map = fill_text_line_to_seal_impression(
194                seal_impression,
195                seal_impression_resource.text_line_slot_indices,
196                seal_impression_resource.text_lines,
197                seal_impression_resource.internal_text_line,
198            )
199
200            # Rotate, shift, and trim.
201            rotated_result = rotate.distort(
202                {'angle': seal_impression_resource.angle},
203                mask=background_mask,
204                score_map=text_line_filled_score_map,
205            )
206            assert rotated_result.mask
207            background_mask = rotated_result.mask
208            assert rotated_result.score_map
209            text_line_filled_score_map = rotated_result.score_map
210            assert background_mask.shape == text_line_filled_score_map.shape
211
212            box_center_point = seal_impression_resource.box.get_center_point()
213            up = box_center_point.y - background_mask.height // 2
214            down = up + background_mask.height - 1
215            left = box_center_point.x - background_mask.width // 2
216            right = left + background_mask.width - 1
217
218            if up < 0 or down >= assembled_image.height \
219                    or left < 0 or right >= assembled_image.width:
220                extract_up = 0
221                if up < 0:
222                    extract_up = abs(up)
223                    up = 0
224
225                extract_down = background_mask.height - 1
226                if down >= assembled_image.height:
227                    extract_down -= down + 1 - assembled_image.height
228                    down = assembled_image.height - 1
229
230                extract_left = 0
231                if left < 0:
232                    extract_left = abs(left)
233                    left = 0
234
235                extract_right = background_mask.width - 1
236                if right >= assembled_image.width:
237                    extract_right -= right + 1 - assembled_image.width
238                    right = assembled_image.width - 1
239
240                extract_box = Box(
241                    up=extract_up,
242                    down=extract_down,
243                    left=extract_left,
244                    right=extract_right,
245                )
246                background_mask = extract_box.extract_mask(background_mask)
247                text_line_filled_score_map = extract_box.extract_score_map(
248                    text_line_filled_score_map
249                )
250
251            # Rendering.
252            box = Box(up=up, down=down, left=left, right=right)
253            box.fill_image(assembled_image, value=color, image_mask=background_mask, alpha=alpha)
254            box.fill_image(assembled_image, value=color, alpha=text_line_filled_score_map)
255
256        # For char-level polygon regression.
257        page_disconnected_text_region_collection = PageDisconnectedTextRegionCollection(
258            page_layout.disconnected_text_regions
259        )
260
261        # For sampling negative text region area.
262        page_non_text_region_collection = PageNonTextRegionCollection(page_layout.non_text_regions)
263
264        page = Page(
265            image=assembled_image,
266            page_image_collection=page_image_collection,
267            page_bottom_layer_image=page_bottom_layer_image,
268            page_text_line_collection=page_text_line_collection,
269            page_seal_impression_text_line_collection=page_seal_impression_text_line_collection,
270            page_char_polygon_collection=page_char_polygon_collection,
271            page_text_line_polygon_collection=page_text_line_polygon_collection,
272            page_disconnected_text_region_collection=page_disconnected_text_region_collection,
273            page_non_text_region_collection=page_non_text_region_collection,
274        )
275        return PageAssemblerStepOutput(page=page)
276
277
278page_assembler_step_factory = PipelineStepFactory(PageAssemblerStep)