qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-07-12 06:39:31 +02:00
Code Issues Projects Releases Wiki Activity

train: make preprocess_imgs_ocr work for transformer-ocr

Browse source
This commit is contained in:
kba 2026-07-09 19:12:09 +02:00
parent b2f3a8f2d8
commit affddd6c85
2 changed files with 43 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

3
src/eynollah/training/train.py
View file

@ -724,7 +724,8 @@ def run(_config,
config, config,
dir_img, dir_img,
dir_lab, dir_lab,
char_to_num=char_to_num char_to_num=char_to_num,
processor=processor,
) )
) )

69
src/eynollah/training/utils.py
View file

@ -789,7 +789,7 @@ def preprocess_imgs(config,
lab = cv2.imread(os.path.join(dir_lab, img_name + '.png')) lab = cv2.imread(os.path.join(dir_lab, img_name + '.png'))
elif config['task'] == "enhancement": elif config['task'] == "enhancement":
lab = cv2.imread(os.path.join(dir_lab, img)) lab = cv2.imread(os.path.join(dir_lab, img))
elif config['task'] == "cnn-rnn-ocr": elif config['task'] in ["cnn-rnn-ocr", "transformer-ocr"]:
# assert lab == 'img_name + '.txt' # assert lab == 'img_name + '.txt'
with open(os.path.join(dir_lab, img_name + '.txt'), 'r') as f: with open(os.path.join(dir_lab, img_name + '.txt'), 'r') as f:
lab = f.read().split('\n')[0] lab = f.read().split('\n')[0]
@ -797,7 +797,7 @@ def preprocess_imgs(config,
lab = None lab = None
try: try:
if config['task'] == "cnn-rnn-ocr": if config['task'] in ["cnn-rnn-ocr", "transformer-ocr"]:
yield from preprocess_img_ocr(img, img_name, lab, **config) yield from preprocess_img_ocr(img, img_name, lab, **config)
continue continue
else: else:
@ -1116,14 +1116,25 @@ def preprocess_img_ocr(
number_of_backgrounds_per_image=None, number_of_backgrounds_per_image=None,
list_all_possible_background_images=None, list_all_possible_background_images=None,
list_all_possible_foreground_rgbs=None, list_all_possible_foreground_rgbs=None,
task=None,
processor=None,
**kwargs **kwargs
): ):
def scale_image(img): def scale_image(img):
return scale_padd_image_for_ocr(img, input_height, input_width).astype(np.float32) / 255. return scale_padd_image_for_ocr(img, input_height, input_width).astype(np.float32) / 255.
#lab = vectorize_label(lab, char_to_num, padding_token, max_len) #lab = vectorize_label(lab, char_to_num, padding_token, max_len)
# now padded at Dataset.padded_batch # now padded at Dataset.padded_batch
if task == 'cnn-rnn-ocr':
assert char_to_num, 'task is cnn-rnn-ocr, so preprocess_imgs_ocr should be passed "char_to_num"'
lab = char_to_num(tf.strings.unicode_split(lab, input_encoding="UTF-8")) lab = char_to_num(tf.strings.unicode_split(lab, input_encoding="UTF-8"))
yield scale_image(img), lab yield_encoder = lambda x: x
elif task == 'transformer-ocr':
assert processor, 'task is transformer-ocr, so preprocess_imgs_ocr should be passed "processor"'
# TODO make max_length configurable again, if deemed sensible
lab = [l if l != self.processor.tokenizer.pad_token_id else -100
for l in processor.tokenizer(lab, padding="max_length", max_length=128).input_ids]
yield_encoder = lambda img_, lab_: {"pixel_values": processor(Image.fromarray(img_), return_tensors="pt").pixel_values.squeeze(), "labels": torch.tensor(lab_)}
yield yield_encoder(scale_image(img), lab)
#to_yield = {"image": ret_x, "label": ret_y} #to_yield = {"image": ret_x, "label": ret_y}
if dir_img_bin: if dir_img_bin:
@ -1139,32 +1150,32 @@ def preprocess_img_ocr(
for padd_col in padd_colors: for padd_col in padd_colors:
img_pad = do_padding_for_ocr(img, 1.2, padd_col) img_pad = do_padding_for_ocr(img, 1.2, padd_col)
img_rot = rotation_not_90_func_single_image(img_pad, thetha_ind) img_rot = rotation_not_90_func_single_image(img_pad, thetha_ind)
yield scale_image(img_rot), lab yield yield_encoder(scale_image(img_rot), lab)
if rotation_not_90: if rotation_not_90:
for thetha_ind in thetha: for thetha_ind in thetha:
img_rot = rotation_not_90_func_single_image(img, thetha_ind) img_rot = rotation_not_90_func_single_image(img, thetha_ind)
yield scale_image(img_rot), lab yield yield_encoder(scale_image(img_rot), lab)
if blur_aug: if blur_aug:
for blur_type in blur_k: for blur_type in blur_k:
img_blur = bluring(img, blur_type) img_blur = bluring(img, blur_type)
yield scale_image(img_blur), lab yield yield_encoder(scale_image(img_blur), lab)
if degrading: if degrading:
for deg_scale_ind in degrade_scales: for deg_scale_ind in degrade_scales:
img_deg = do_degrading(img, deg_scale_ind) img_deg = do_degrading(img, deg_scale_ind)
yield scale_image(img_deg), lab yield yield_encoder(scale_image(img_deg), lab)
if bin_deg: if bin_deg:
for deg_scale_ind in degrade_scales: for deg_scale_ind in degrade_scales:
img_deg = do_degrading(img_bin_corr, deg_scale_ind) img_deg = do_degrading(img_bin_corr, deg_scale_ind)
yield scale_image(img_deg), lab yield yield_encoder(scale_image(img_deg), lab)
if brightening: if brightening:
for bright_scale_ind in brightness: for bright_scale_ind in brightness:
img_bright = do_brightening(img, bright_scale_ind) img_bright = do_brightening(img, bright_scale_ind)
yield scale_image(img_bright), lab yield yield_encoder(scale_image(img_bright), lab)
if padding_white: if padding_white:
for padding_size in white_padds: for padding_size in white_padds:
for padd_col in padd_colors: for padd_col in padd_colors:
img_pad = do_padding_for_ocr(img, padding_size, padd_col) img_pad = do_padding_for_ocr(img, padding_size, padd_col)
yield scale_image(img_pad), lab yield yield_encoder(scale_image(img_pad), lab)
if adding_rgb_foreground: if adding_rgb_foreground:
for i_n in range(number_of_backgrounds_per_image): for i_n in range(number_of_backgrounds_per_image):
background_image_chosen_name = random.choice(list_all_possible_background_images) background_image_chosen_name = random.choice(list_all_possible_background_images)
@ -1178,7 +1189,7 @@ def preprocess_img_ocr(
img_fg = \ img_fg = \
return_binary_image_with_given_rgb_background_and_given_foreground_rgb( return_binary_image_with_given_rgb_background_and_given_foreground_rgb(
img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen)
yield scale_image(img_fg), lab yield yield_encoder(scale_image(img_fg), lab)
if adding_rgb_background: if adding_rgb_background:
for i_n in range(number_of_backgrounds_per_image): for i_n in range(number_of_backgrounds_per_image):
background_image_chosen_name = random.choice(list_all_possible_background_images) background_image_chosen_name = random.choice(list_all_possible_background_images)
@ -1186,59 +1197,59 @@ def preprocess_img_ocr(
cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name)
img_bg = \ img_bg = \
return_binary_image_with_given_rgb_background(img_bin_corr, img_rgb_background_chosen) return_binary_image_with_given_rgb_background(img_bin_corr, img_rgb_background_chosen)
yield scale_image(img_bg), lab yield yield_encoder(scale_image(img_bg), lab)
if binarization: if binarization:
yield scale_image(img_bin_corr), lab yield yield_encoder(scale_image(img_bin_corr), lab)
if image_inversion: if image_inversion:
img_inv = invert_image(img_bin_corr) img_inv = invert_image(img_bin_corr)
yield scale_image(img_inv), lab yield yield_encoder(scale_image(img_inv), lab)
if channels_shuffling: if channels_shuffling:
for shuffle_index in shuffle_indexes: for shuffle_index in shuffle_indexes:
img_shuf = return_shuffled_channels(img, shuffle_index) img_shuf = return_shuffled_channels(img, shuffle_index)
yield scale_image(img_shuf), lab yield yield_encoder(scale_image(img_shuf), lab)
if add_red_textlines: if add_red_textlines:
img_red = return_image_with_red_elements(img, img_bin_corr) img_red = return_image_with_red_elements(img, img_bin_corr)
yield scale_image(img_red), lab yield yield_encoder(scale_image(img_red), lab)
if white_noise_strap: if white_noise_strap:
img_noisy = return_image_with_strapped_white_noises(img) img_noisy = return_image_with_strapped_white_noises(img)
yield scale_image(img_noisy), lab yield yield_encoder(scale_image(img_noisy), lab)
if textline_skewing: if textline_skewing:
for des_scale_ind in skewing_amplitudes: for des_scale_ind in skewing_amplitudes:
img_rot = do_deskewing(img, des_scale_ind) img_rot = do_deskewing(img, des_scale_ind)
yield scale_image(img_rot), lab yield yield_encoder(scale_image(img_rot), lab)
if textline_skewing_bin: if textline_skewing_bin:
for des_scale_ind in skewing_amplitudes: for des_scale_ind in skewing_amplitudes:
img_rot = do_deskewing(img_bin_corr, des_scale_ind) img_rot = do_deskewing(img_bin_corr, des_scale_ind)
yield scale_image(img_rot), lab yield yield_encoder(scale_image(img_rot), lab)
if textline_left_in_depth: if textline_left_in_depth:
img_warp = do_direction_in_depth(img, 'left') img_warp = do_direction_in_depth(img, 'left')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_left_in_depth_bin: if textline_left_in_depth_bin:
img_warp = do_direction_in_depth(img_bin_corr, 'left') img_warp = do_direction_in_depth(img_bin_corr, 'left')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_right_in_depth: if textline_right_in_depth:
img_warp = do_direction_in_depth(img, 'right') img_warp = do_direction_in_depth(img, 'right')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_right_in_depth_bin: if textline_right_in_depth_bin:
img_warp = do_direction_in_depth(img_bin_corr, 'right') img_warp = do_direction_in_depth(img_bin_corr, 'right')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_up_in_depth: if textline_up_in_depth:
img_warp = do_direction_in_depth(img, 'up') img_warp = do_direction_in_depth(img, 'up')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_up_in_depth_bin: if textline_up_in_depth_bin:
img_warp = do_direction_in_depth(img_bin_corr, 'up') img_warp = do_direction_in_depth(img_bin_corr, 'up')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_down_in_depth: if textline_down_in_depth:
img_warp = do_direction_in_depth(img, 'down') img_warp = do_direction_in_depth(img, 'down')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if textline_down_in_depth_bin: if textline_down_in_depth_bin:
img_warp = do_direction_in_depth(img_bin_corr, 'down') img_warp = do_direction_in_depth(img_bin_corr, 'down')
yield scale_image(img_warp), lab yield yield_encoder(scale_image(img_warp), lab)
if pepper_aug: if pepper_aug:
for pepper_ind in pepper_indexes: for pepper_ind in pepper_indexes:
img_noisy = add_salt_and_pepper_noise(img, pepper_ind, pepper_ind) img_noisy = add_salt_and_pepper_noise(img, pepper_ind, pepper_ind)
yield scale_image(img_noisy), lab yield yield_encoder(scale_image(img_noisy), lab)
if pepper_bin_aug: if pepper_bin_aug:
for pepper_ind in pepper_indexes: for pepper_ind in pepper_indexes:
img_noisy = add_salt_and_pepper_noise(img_bin_corr, pepper_ind, pepper_ind) img_noisy = add_salt_and_pepper_noise(img_bin_corr, pepper_ind, pepper_ind)
yield scale_image(img_noisy), lab yield yield_encoder(scale_image(img_noisy), lab)