vkit.engine.char_mask.external_ellipse
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, List 15import math 16import itertools 17 18import attrs 19from numpy.random import Generator as RandomGenerator 20import numpy as np 21import cv2 as cv 22 23from vkit.element import Box, Polygon, Mask 24from vkit.mechanism.distortion.geometric.affine import affine_np_points 25from ..char_heatmap.default import build_np_distance 26from ..interface import ( 27 Engine, 28 EngineExecutorFactory, 29 NoneTypeEngineInitResource, 30) 31from .type import CharMaskEngineRunConfig, CharMask 32 33 34@attrs.define 35class CharMaskExternalEllipseEngineInitConfig: 36 internal_side_length: int = 40 37 38 39class CharMaskExternalEllipseEngine( 40 Engine[ 41 CharMaskExternalEllipseEngineInitConfig, 42 NoneTypeEngineInitResource, 43 CharMaskEngineRunConfig, 44 CharMask, 45 ] 46): # yapf: disable 47 48 @classmethod 49 def get_type_name(cls) -> str: 50 return 'external_ellipse' 51 52 def __init__( 53 self, 54 init_config: CharMaskExternalEllipseEngineInitConfig, 55 init_resource: Optional[NoneTypeEngineInitResource] = None, 56 ): 57 super().__init__(init_config, init_resource) 58 59 internal_side_length = init_config.internal_side_length 60 external_radius = math.ceil(internal_side_length / math.sqrt(2)) 61 62 # Build distances to the center point. 63 np_distance = build_np_distance(external_radius) 64 65 # Build mask. 66 self.np_external_mask = (np_distance <= external_radius).astype(np.uint8) 67 external_side_length = self.np_external_mask.shape[0] 68 69 # For perspective transformation. 70 char_pad = (external_side_length - internal_side_length) // 2 71 char_begin = char_pad 72 char_end = char_pad + internal_side_length - 1 73 self.np_char_points = np.asarray( 74 [ 75 (char_begin, char_begin), 76 (char_end, char_begin), 77 (char_end, char_end), 78 (char_begin, char_end), 79 ], 80 dtype=np.float32, 81 ) 82 83 external_begin = 0 84 external_end = external_side_length - 1 85 self.np_external_points = np.asarray( 86 [ 87 (external_begin, external_begin), 88 (external_end, external_begin), 89 (external_end, external_end), 90 (external_begin, external_end), 91 ], 92 dtype=np.float32, 93 ) 94 95 def run( 96 self, 97 run_config: CharMaskEngineRunConfig, 98 rng: Optional[RandomGenerator] = None, 99 ) -> CharMask: 100 char_polygons = run_config.char_polygons 101 char_bounding_boxes = run_config.char_bounding_boxes 102 char_bounding_polygons = run_config.char_bounding_polygons 103 104 if char_bounding_boxes or char_bounding_polygons: 105 assert not (char_bounding_boxes and char_bounding_polygons) 106 107 if char_bounding_boxes: 108 assert len(char_bounding_boxes) == len(char_polygons) 109 char_bounding_elements = char_bounding_boxes 110 111 elif char_bounding_polygons: 112 assert len(char_bounding_polygons) == len(char_polygons) 113 char_bounding_elements = char_bounding_polygons 114 115 else: 116 bounding_box = Box( 117 up=0, 118 down=run_config.height - 1, 119 left=0, 120 right=run_config.width - 1, 121 ) 122 char_bounding_elements = itertools.repeat(bounding_box) 123 124 combined_chars_mask = Mask.from_shape((run_config.height, run_config.width)) 125 char_masks: List[Mask] = [] 126 127 for char_polygon, char_bounding_element in zip(char_polygons, char_bounding_elements): 128 # 1. Find the transformed external points. 129 assert char_polygon.num_points == 4 130 np_trans_mat = cv.getPerspectiveTransform( 131 self.np_char_points, 132 char_polygon.internals.np_self_relative_points, 133 cv.DECOMP_SVD, 134 ) 135 np_transformed_external_points = affine_np_points( 136 np_trans_mat, 137 self.np_external_points, 138 ) 139 140 # Make self-relative and keep the offset. 141 y_offset = np_transformed_external_points[:, 1].min() 142 x_offset = np_transformed_external_points[:, 0].min() 143 np_transformed_external_points[:, 1] -= y_offset 144 np_transformed_external_points[:, 0] -= x_offset 145 146 # 2. Transform the external mask. 147 np_transformed_external_points = np_transformed_external_points.astype(np.float32) 148 np_trans_mat = cv.getPerspectiveTransform( 149 self.np_external_points, 150 np_transformed_external_points, 151 cv.DECOMP_SVD, 152 ) 153 y_max = np_transformed_external_points[:, 1].max() 154 transformed_height = math.ceil(y_max) 155 x_max = np_transformed_external_points[:, 0].max() 156 transformed_width = math.ceil(x_max) 157 np_transformed_external_mask = cv.warpPerspective( 158 self.np_external_mask, 159 np_trans_mat, 160 (transformed_width, transformed_height), 161 ) 162 163 # 3. Place char mask. 164 smooth_y_min = min(point.smooth_y for point in char_polygon.points) 165 smooth_x_min = min(point.smooth_x for point in char_polygon.points) 166 167 target_up = round(smooth_y_min + y_offset) 168 target_down = target_up + transformed_height - 1 169 target_left = round(smooth_x_min + x_offset) 170 target_right = target_left + transformed_width - 1 171 172 if isinstance(char_bounding_element, Box): 173 char_bounding_box = char_bounding_element 174 else: 175 char_bounding_box = char_bounding_element.bounding_box 176 177 trimmed_up = 0 178 if target_up < char_bounding_box.up: 179 trimmed_up = char_bounding_box.up - target_up 180 target_up = char_bounding_box.up 181 182 trimmed_down = transformed_height - 1 183 if target_down > char_bounding_box.down: 184 trimmed_down -= (target_down - char_bounding_box.down) 185 target_down = char_bounding_box.down 186 187 trimmed_left = 0 188 if target_left < char_bounding_box.left: 189 trimmed_left = char_bounding_box.left - target_left 190 target_left = char_bounding_box.left 191 192 trimmed_right = transformed_width - 1 193 if target_right > char_bounding_box.right: 194 trimmed_right -= (target_right - char_bounding_box.right) 195 target_right = char_bounding_box.right 196 197 target_box = Box( 198 up=target_up, 199 down=target_down, 200 left=target_left, 201 right=target_right, 202 ) 203 np_transformed_external_mask = \ 204 np_transformed_external_mask[ 205 trimmed_up:trimmed_down + 1, 206 trimmed_left:trimmed_right + 1 207 ] 208 char_mask = Mask(mat=np_transformed_external_mask, box=target_box) 209 210 if isinstance(char_bounding_element, Polygon): 211 char_inverted_mask = char_bounding_element.mask.to_inverted_mask() 212 char_inverted_mask = target_box.extract_mask(char_inverted_mask) 213 char_inverted_mask.fill_mask(char_mask, 0) 214 215 char_masks.append(char_mask) 216 217 # Fill. 218 char_mask.fill_mask(combined_chars_mask, 1, keep_max_value=True) 219 220 return CharMask( 221 combined_chars_mask=combined_chars_mask, 222 char_masks=char_masks, 223 ) 224 225 226char_mask_external_ellipse_engine_executor_factory = EngineExecutorFactory( 227 CharMaskExternalEllipseEngine 228)
class
CharMaskExternalEllipseEngineInitConfig:
class
CharMaskExternalEllipseEngine(vkit.engine.interface.Engine[vkit.engine.char_mask.external_ellipse.CharMaskExternalEllipseEngineInitConfig, vkit.engine.interface.NoneTypeEngineInitResource, vkit.engine.char_mask.type.CharMaskEngineRunConfig, vkit.engine.char_mask.type.CharMask]):
40class CharMaskExternalEllipseEngine( 41 Engine[ 42 CharMaskExternalEllipseEngineInitConfig, 43 NoneTypeEngineInitResource, 44 CharMaskEngineRunConfig, 45 CharMask, 46 ] 47): # yapf: disable 48 49 @classmethod 50 def get_type_name(cls) -> str: 51 return 'external_ellipse' 52 53 def __init__( 54 self, 55 init_config: CharMaskExternalEllipseEngineInitConfig, 56 init_resource: Optional[NoneTypeEngineInitResource] = None, 57 ): 58 super().__init__(init_config, init_resource) 59 60 internal_side_length = init_config.internal_side_length 61 external_radius = math.ceil(internal_side_length / math.sqrt(2)) 62 63 # Build distances to the center point. 64 np_distance = build_np_distance(external_radius) 65 66 # Build mask. 67 self.np_external_mask = (np_distance <= external_radius).astype(np.uint8) 68 external_side_length = self.np_external_mask.shape[0] 69 70 # For perspective transformation. 71 char_pad = (external_side_length - internal_side_length) // 2 72 char_begin = char_pad 73 char_end = char_pad + internal_side_length - 1 74 self.np_char_points = np.asarray( 75 [ 76 (char_begin, char_begin), 77 (char_end, char_begin), 78 (char_end, char_end), 79 (char_begin, char_end), 80 ], 81 dtype=np.float32, 82 ) 83 84 external_begin = 0 85 external_end = external_side_length - 1 86 self.np_external_points = np.asarray( 87 [ 88 (external_begin, external_begin), 89 (external_end, external_begin), 90 (external_end, external_end), 91 (external_begin, external_end), 92 ], 93 dtype=np.float32, 94 ) 95 96 def run( 97 self, 98 run_config: CharMaskEngineRunConfig, 99 rng: Optional[RandomGenerator] = None, 100 ) -> CharMask: 101 char_polygons = run_config.char_polygons 102 char_bounding_boxes = run_config.char_bounding_boxes 103 char_bounding_polygons = run_config.char_bounding_polygons 104 105 if char_bounding_boxes or char_bounding_polygons: 106 assert not (char_bounding_boxes and char_bounding_polygons) 107 108 if char_bounding_boxes: 109 assert len(char_bounding_boxes) == len(char_polygons) 110 char_bounding_elements = char_bounding_boxes 111 112 elif char_bounding_polygons: 113 assert len(char_bounding_polygons) == len(char_polygons) 114 char_bounding_elements = char_bounding_polygons 115 116 else: 117 bounding_box = Box( 118 up=0, 119 down=run_config.height - 1, 120 left=0, 121 right=run_config.width - 1, 122 ) 123 char_bounding_elements = itertools.repeat(bounding_box) 124 125 combined_chars_mask = Mask.from_shape((run_config.height, run_config.width)) 126 char_masks: List[Mask] = [] 127 128 for char_polygon, char_bounding_element in zip(char_polygons, char_bounding_elements): 129 # 1. Find the transformed external points. 130 assert char_polygon.num_points == 4 131 np_trans_mat = cv.getPerspectiveTransform( 132 self.np_char_points, 133 char_polygon.internals.np_self_relative_points, 134 cv.DECOMP_SVD, 135 ) 136 np_transformed_external_points = affine_np_points( 137 np_trans_mat, 138 self.np_external_points, 139 ) 140 141 # Make self-relative and keep the offset. 142 y_offset = np_transformed_external_points[:, 1].min() 143 x_offset = np_transformed_external_points[:, 0].min() 144 np_transformed_external_points[:, 1] -= y_offset 145 np_transformed_external_points[:, 0] -= x_offset 146 147 # 2. Transform the external mask. 148 np_transformed_external_points = np_transformed_external_points.astype(np.float32) 149 np_trans_mat = cv.getPerspectiveTransform( 150 self.np_external_points, 151 np_transformed_external_points, 152 cv.DECOMP_SVD, 153 ) 154 y_max = np_transformed_external_points[:, 1].max() 155 transformed_height = math.ceil(y_max) 156 x_max = np_transformed_external_points[:, 0].max() 157 transformed_width = math.ceil(x_max) 158 np_transformed_external_mask = cv.warpPerspective( 159 self.np_external_mask, 160 np_trans_mat, 161 (transformed_width, transformed_height), 162 ) 163 164 # 3. Place char mask. 165 smooth_y_min = min(point.smooth_y for point in char_polygon.points) 166 smooth_x_min = min(point.smooth_x for point in char_polygon.points) 167 168 target_up = round(smooth_y_min + y_offset) 169 target_down = target_up + transformed_height - 1 170 target_left = round(smooth_x_min + x_offset) 171 target_right = target_left + transformed_width - 1 172 173 if isinstance(char_bounding_element, Box): 174 char_bounding_box = char_bounding_element 175 else: 176 char_bounding_box = char_bounding_element.bounding_box 177 178 trimmed_up = 0 179 if target_up < char_bounding_box.up: 180 trimmed_up = char_bounding_box.up - target_up 181 target_up = char_bounding_box.up 182 183 trimmed_down = transformed_height - 1 184 if target_down > char_bounding_box.down: 185 trimmed_down -= (target_down - char_bounding_box.down) 186 target_down = char_bounding_box.down 187 188 trimmed_left = 0 189 if target_left < char_bounding_box.left: 190 trimmed_left = char_bounding_box.left - target_left 191 target_left = char_bounding_box.left 192 193 trimmed_right = transformed_width - 1 194 if target_right > char_bounding_box.right: 195 trimmed_right -= (target_right - char_bounding_box.right) 196 target_right = char_bounding_box.right 197 198 target_box = Box( 199 up=target_up, 200 down=target_down, 201 left=target_left, 202 right=target_right, 203 ) 204 np_transformed_external_mask = \ 205 np_transformed_external_mask[ 206 trimmed_up:trimmed_down + 1, 207 trimmed_left:trimmed_right + 1 208 ] 209 char_mask = Mask(mat=np_transformed_external_mask, box=target_box) 210 211 if isinstance(char_bounding_element, Polygon): 212 char_inverted_mask = char_bounding_element.mask.to_inverted_mask() 213 char_inverted_mask = target_box.extract_mask(char_inverted_mask) 214 char_inverted_mask.fill_mask(char_mask, 0) 215 216 char_masks.append(char_mask) 217 218 # Fill. 219 char_mask.fill_mask(combined_chars_mask, 1, keep_max_value=True) 220 221 return CharMask( 222 combined_chars_mask=combined_chars_mask, 223 char_masks=char_masks, 224 )
Abstract base class for generic types.
A generic type is typically declared by inheriting from this class parameterized with one or more type variables. For example, a generic mapping type might be defined as::
class Mapping(Generic[KT, VT]): def __getitem__(self, key: KT) -> VT: ... # Etc.
This class can then be used as follows::
def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT: try: return mapping[key] except KeyError: return default
CharMaskExternalEllipseEngine( init_config: vkit.engine.char_mask.external_ellipse.CharMaskExternalEllipseEngineInitConfig, init_resource: Union[vkit.engine.interface.NoneTypeEngineInitResource, NoneType] = None)
53 def __init__( 54 self, 55 init_config: CharMaskExternalEllipseEngineInitConfig, 56 init_resource: Optional[NoneTypeEngineInitResource] = None, 57 ): 58 super().__init__(init_config, init_resource) 59 60 internal_side_length = init_config.internal_side_length 61 external_radius = math.ceil(internal_side_length / math.sqrt(2)) 62 63 # Build distances to the center point. 64 np_distance = build_np_distance(external_radius) 65 66 # Build mask. 67 self.np_external_mask = (np_distance <= external_radius).astype(np.uint8) 68 external_side_length = self.np_external_mask.shape[0] 69 70 # For perspective transformation. 71 char_pad = (external_side_length - internal_side_length) // 2 72 char_begin = char_pad 73 char_end = char_pad + internal_side_length - 1 74 self.np_char_points = np.asarray( 75 [ 76 (char_begin, char_begin), 77 (char_end, char_begin), 78 (char_end, char_end), 79 (char_begin, char_end), 80 ], 81 dtype=np.float32, 82 ) 83 84 external_begin = 0 85 external_end = external_side_length - 1 86 self.np_external_points = np.asarray( 87 [ 88 (external_begin, external_begin), 89 (external_end, external_begin), 90 (external_end, external_end), 91 (external_begin, external_end), 92 ], 93 dtype=np.float32, 94 )
def
run( self, run_config: vkit.engine.char_mask.type.CharMaskEngineRunConfig, rng: Union[numpy.random._generator.Generator, NoneType] = None) -> vkit.engine.char_mask.type.CharMask:
96 def run( 97 self, 98 run_config: CharMaskEngineRunConfig, 99 rng: Optional[RandomGenerator] = None, 100 ) -> CharMask: 101 char_polygons = run_config.char_polygons 102 char_bounding_boxes = run_config.char_bounding_boxes 103 char_bounding_polygons = run_config.char_bounding_polygons 104 105 if char_bounding_boxes or char_bounding_polygons: 106 assert not (char_bounding_boxes and char_bounding_polygons) 107 108 if char_bounding_boxes: 109 assert len(char_bounding_boxes) == len(char_polygons) 110 char_bounding_elements = char_bounding_boxes 111 112 elif char_bounding_polygons: 113 assert len(char_bounding_polygons) == len(char_polygons) 114 char_bounding_elements = char_bounding_polygons 115 116 else: 117 bounding_box = Box( 118 up=0, 119 down=run_config.height - 1, 120 left=0, 121 right=run_config.width - 1, 122 ) 123 char_bounding_elements = itertools.repeat(bounding_box) 124 125 combined_chars_mask = Mask.from_shape((run_config.height, run_config.width)) 126 char_masks: List[Mask] = [] 127 128 for char_polygon, char_bounding_element in zip(char_polygons, char_bounding_elements): 129 # 1. Find the transformed external points. 130 assert char_polygon.num_points == 4 131 np_trans_mat = cv.getPerspectiveTransform( 132 self.np_char_points, 133 char_polygon.internals.np_self_relative_points, 134 cv.DECOMP_SVD, 135 ) 136 np_transformed_external_points = affine_np_points( 137 np_trans_mat, 138 self.np_external_points, 139 ) 140 141 # Make self-relative and keep the offset. 142 y_offset = np_transformed_external_points[:, 1].min() 143 x_offset = np_transformed_external_points[:, 0].min() 144 np_transformed_external_points[:, 1] -= y_offset 145 np_transformed_external_points[:, 0] -= x_offset 146 147 # 2. Transform the external mask. 148 np_transformed_external_points = np_transformed_external_points.astype(np.float32) 149 np_trans_mat = cv.getPerspectiveTransform( 150 self.np_external_points, 151 np_transformed_external_points, 152 cv.DECOMP_SVD, 153 ) 154 y_max = np_transformed_external_points[:, 1].max() 155 transformed_height = math.ceil(y_max) 156 x_max = np_transformed_external_points[:, 0].max() 157 transformed_width = math.ceil(x_max) 158 np_transformed_external_mask = cv.warpPerspective( 159 self.np_external_mask, 160 np_trans_mat, 161 (transformed_width, transformed_height), 162 ) 163 164 # 3. Place char mask. 165 smooth_y_min = min(point.smooth_y for point in char_polygon.points) 166 smooth_x_min = min(point.smooth_x for point in char_polygon.points) 167 168 target_up = round(smooth_y_min + y_offset) 169 target_down = target_up + transformed_height - 1 170 target_left = round(smooth_x_min + x_offset) 171 target_right = target_left + transformed_width - 1 172 173 if isinstance(char_bounding_element, Box): 174 char_bounding_box = char_bounding_element 175 else: 176 char_bounding_box = char_bounding_element.bounding_box 177 178 trimmed_up = 0 179 if target_up < char_bounding_box.up: 180 trimmed_up = char_bounding_box.up - target_up 181 target_up = char_bounding_box.up 182 183 trimmed_down = transformed_height - 1 184 if target_down > char_bounding_box.down: 185 trimmed_down -= (target_down - char_bounding_box.down) 186 target_down = char_bounding_box.down 187 188 trimmed_left = 0 189 if target_left < char_bounding_box.left: 190 trimmed_left = char_bounding_box.left - target_left 191 target_left = char_bounding_box.left 192 193 trimmed_right = transformed_width - 1 194 if target_right > char_bounding_box.right: 195 trimmed_right -= (target_right - char_bounding_box.right) 196 target_right = char_bounding_box.right 197 198 target_box = Box( 199 up=target_up, 200 down=target_down, 201 left=target_left, 202 right=target_right, 203 ) 204 np_transformed_external_mask = \ 205 np_transformed_external_mask[ 206 trimmed_up:trimmed_down + 1, 207 trimmed_left:trimmed_right + 1 208 ] 209 char_mask = Mask(mat=np_transformed_external_mask, box=target_box) 210 211 if isinstance(char_bounding_element, Polygon): 212 char_inverted_mask = char_bounding_element.mask.to_inverted_mask() 213 char_inverted_mask = target_box.extract_mask(char_inverted_mask) 214 char_inverted_mask.fill_mask(char_mask, 0) 215 216 char_masks.append(char_mask) 217 218 # Fill. 219 char_mask.fill_mask(combined_chars_mask, 1, keep_max_value=True) 220 221 return CharMask( 222 combined_chars_mask=combined_chars_mask, 223 char_masks=char_masks, 224 )