Skip to content
Snippets Groups Projects
Commit 9ab6085f authored by Martin Maarand's avatar Martin Maarand
Browse files

add skew extraction modes

parent a297c651
No related branches found
No related tags found
1 merge request!19add skew extraction modes
Pipeline #74317 passed
......@@ -51,6 +51,8 @@ class Extraction(Enum):
min_area_rect: int = 2
deskew_polygon: int = 3
deskew_min_area_rect: int = 4
skew_polygon: int = 5
skew_min_area_rect: int = 6
class HTRDataGenerator:
......@@ -67,6 +69,7 @@ class HTRDataGenerator:
accepted_worker_version_ids=None,
transcription_type=TEXT_LINE,
max_deskew_angle=45,
skew_angle=0,
should_rotate=False,
scale_x=None,
scale_y_top=None,
......@@ -91,6 +94,7 @@ class HTRDataGenerator:
self.skipped_vertical_lines_count = 0
self.accepted_lines_count = 0
self.max_deskew_angle = max_deskew_angle
self.skew_angle = skew_angle
self.should_rotate = should_rotate
if scale_x or scale_y_top or scale_y_bottom:
self.should_resize_polygons = True
......@@ -397,6 +401,32 @@ class HTRDataGenerator:
trimmed_img = self.rotate_and_trim(min_rect_img, rotate_angle)
self._save_line_image(page_id, i, trimmed_img, manifest_fp, trans)
elif self.extraction_mode == Extraction.skew_polygon:
for i, trans in enumerate(sorted_lines):
rotate_angle = self.skew_angle
# get polygon image
polygon_img = extract_polygon_image(
img, polygon=trans.polygon, rect=trans.rect
)
trimmed_img = self.rotate_and_trim(polygon_img, rotate_angle)
self._save_line_image(page_id, i, trimmed_img, manifest_fp, trans)
elif self.extraction_mode == Extraction.skew_min_area_rect:
for i, trans in enumerate(sorted_lines):
rotate_angle = self.skew_angle
min_rect_img = extract_min_area_rect_image(
img, polygon=trans.polygon, rect=trans.rect
)
trimmed_img = self.rotate_and_trim(min_rect_img, rotate_angle)
self._save_line_image(page_id, i, trimmed_img, manifest_fp, trans)
else:
raise ValueError(f"Unsupported extraction mode: {self.extraction_mode}")
......@@ -628,6 +658,15 @@ def create_parser():
"then that line won't be deskewed/rotated.",
)
parser.add_argument(
"--skew_angle",
type=int,
default=0,
help="Angle by which the line image will be rotated. Useful for data augmnetation"
" - creating skewed text lines for a more robust model."
" Only used with skew_* extraction modes.",
)
parser.add_argument(
"--should_rotate",
action="store_true",
......@@ -762,6 +801,7 @@ def main():
transcription_type=args.transcription_type,
accepted_worker_version_ids=args.accepted_worker_version_ids,
max_deskew_angle=args.max_deskew_angle,
skew_angle=args.skew_angle,
should_rotate=args.should_rotate,
scale_x=args.scale_x,
scale_y_top=args.scale_y_top,
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment