vkit.engine.char_heatmap.default

  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 Optional
 15
 16import attrs
 17from numpy.random import Generator as RandomGenerator
 18import numpy as np
 19import cv2 as cv
 20
 21from vkit.element import Mask, ScoreMap, ElementSetOperationMode
 22from ..interface import (
 23    Engine,
 24    EngineExecutorFactory,
 25    NoneTypeEngineInitResource,
 26)
 27from .type import CharHeatmapEngineRunConfig, CharHeatmap
 28
 29
 30def build_np_distance(radius: int):
 31    # Make it symmetric.
 32    side_length = radius * 2 + 1
 33
 34    # Build distances to the center point.
 35    np_offset = np.abs(np.arange(side_length, dtype=np.float32) - radius)
 36    np_vert_offset = np.repeat(np_offset[:, None], side_length, axis=1)
 37    np_hori_offset = np.repeat(np_offset[None, :], side_length, axis=0)
 38    np_distance = np.sqrt(np.square(np_vert_offset) + np.square(np_hori_offset))
 39
 40    return np_distance
 41
 42
 43@attrs.define
 44class CharHeatmapDefaultEngineInitConfig:
 45    # Adjust the std. The std gets smaller as the distance factor gets larger.
 46    # The activated area shrinks to the center as the std gets smaller.
 47    # https://colab.research.google.com/drive/1TQ1-BTisMYZHIRVVNpVwDFPviXYMhT7A
 48    gaussian_map_distance_factor: float = 2.25
 49    gaussian_map_char_radius: int = 25
 50    gaussian_map_preserving_score_min: float = 0.9
 51    weight_neutralized_score_map: float = 0.4
 52
 53
 54@attrs.define
 55class CharHeatmapDefaultDebug:
 56    score_map_max: ScoreMap
 57    score_map_min: ScoreMap
 58    char_overlapped_mask: Mask
 59    char_neutralized_score_map: ScoreMap
 60    neutralized_mask: Mask
 61    neutralized_score_map: ScoreMap
 62
 63
 64class CharHeatmapDefaultEngine(
 65    Engine[
 66        CharHeatmapDefaultEngineInitConfig,
 67        NoneTypeEngineInitResource,
 68        CharHeatmapEngineRunConfig,
 69        CharHeatmap,
 70    ]
 71):  # yapf: disable
 72
 73    @classmethod
 74    def get_type_name(cls) -> str:
 75        return 'default'
 76
 77    def generate_np_gaussian_map(self):
 78        char_radius = self.init_config.gaussian_map_char_radius
 79        np_distance = build_np_distance(char_radius)
 80        side_length = np_distance.shape[0]
 81
 82        np_norm_distance = np_distance / char_radius
 83        np_gaussian_map = np.exp(
 84            -0.5 * np.square(self.init_config.gaussian_map_distance_factor * np_norm_distance)
 85        )
 86
 87        # For perspective transformation.
 88        char_begin = 0
 89        char_end = side_length - 1
 90        np_char_src_points = np.asarray(
 91            [
 92                (char_begin, char_begin),
 93                (char_end, char_begin),
 94                (char_end, char_end),
 95                (char_begin, char_end),
 96            ],
 97            dtype=np.float32,
 98        )
 99
100        return np_gaussian_map, np_char_src_points
101
102    def __init__(
103        self,
104        init_config: CharHeatmapDefaultEngineInitConfig,
105        init_resource: Optional[NoneTypeEngineInitResource] = None,
106    ):
107        super().__init__(init_config, init_resource)
108
109        self.np_gaussian_map, self.np_char_points = self.generate_np_gaussian_map()
110
111    def run(
112        self,
113        run_config: CharHeatmapEngineRunConfig,
114        rng: Optional[RandomGenerator] = None,
115    ) -> CharHeatmap:
116        height = run_config.height
117        width = run_config.width
118        char_polygons = run_config.char_polygons
119
120        shape = (height, width)
121
122        # Intermediate score maps.
123        score_map_max = ScoreMap.from_shape(shape)
124        score_map_min = ScoreMap.from_shape(shape, value=1.0)
125
126        for char_polygon in char_polygons:
127            # Get transformation matrix.
128            np_trans_mat = cv.getPerspectiveTransform(
129                self.np_char_points,
130                char_polygon.internals.np_self_relative_points,
131                cv.DECOMP_SVD,
132            )
133
134            # Transform gaussian map.
135            char_bounding_box = char_polygon.bounding_box
136
137            np_gaussian_map = cv.warpPerspective(
138                self.np_gaussian_map,
139                np_trans_mat,
140                (char_bounding_box.width, char_bounding_box.height),
141            )
142            score_map = ScoreMap(mat=np_gaussian_map, box=char_bounding_box)
143
144            # Fill score maps.
145            char_polygon.fill_score_map(score_map_max, score_map, keep_max_value=True)
146            char_polygon.fill_score_map(score_map_min, score_map, keep_min_value=True)
147
148        # Set the char overlapped area while preserving score gte threshold.
149        char_overlapped_mask = Mask.from_polygons(
150            shape,
151            char_polygons,
152            ElementSetOperationMode.INTERSECT,
153        )
154
155        preserving_score_min = self.init_config.gaussian_map_preserving_score_min
156        preserving_mask = Mask(mat=(score_map_max.mat >= preserving_score_min).astype(np.uint8))
157
158        neutralized_mask = Mask.from_masks(
159            shape,
160            [
161                char_overlapped_mask,
162                preserving_mask.to_inverted_mask(),
163            ],
164            ElementSetOperationMode.INTERSECT,
165        )
166
167        # Estimate the neutralized score.
168        np_delta: np.ndarray = score_map_max.mat - score_map_min.mat  # type: ignore
169        np_delta = np.clip(np_delta, 0.0, 1.0)
170        char_neutralized_score_map = ScoreMap(mat=np_delta)
171
172        neutralized_score_map = score_map_max.copy()
173        neutralized_mask.fill_score_map(neutralized_score_map, char_neutralized_score_map)
174
175        weight = self.init_config.weight_neutralized_score_map
176        score_map = ScoreMap(
177            mat=((1 - weight) * score_map_max.mat + weight * neutralized_score_map.mat)
178        )
179
180        debug = None
181        if run_config.enable_debug:
182            debug = CharHeatmapDefaultDebug(
183                score_map_max=score_map_max,
184                score_map_min=score_map_min,
185                char_overlapped_mask=char_overlapped_mask,
186                char_neutralized_score_map=char_neutralized_score_map,
187                neutralized_mask=neutralized_mask,
188                neutralized_score_map=neutralized_score_map,
189            )
190
191        return CharHeatmap(score_map=score_map, debug=debug)
192
193
194char_heatmap_default_engine_executor_factory = EngineExecutorFactory(CharHeatmapDefaultEngine)