Spaces:
Build error
Build error
| # Crop by word bounding box | |
| # Locate script with gt.mat | |
| # $ python crop_by_word_bb.py | |
| import os | |
| import re | |
| import cv2 | |
| import scipy.io as sio | |
| from itertools import chain | |
| import numpy as np | |
| import math | |
| mat_contents = sio.loadmat('gt.mat') | |
| image_names = mat_contents['imnames'][0] | |
| cropped_indx = 0 | |
| start_img_indx = 0 | |
| gt_file = open('gt_oabc.txt', 'a') | |
| err_file = open('err_oabc.txt', 'a') | |
| for img_indx in range(start_img_indx, len(image_names)): | |
| # Get image name | |
| image_name_new = image_names[img_indx][0] | |
| # print(image_name_new) | |
| image_name = '/home/yxwang/pytorch/dataset/SynthText/img/'+ image_name_new | |
| # print('IMAGE : {}.{}'.format(img_indx, image_name)) | |
| print('evaluating {} image'.format(img_indx), end='\r') | |
| # Get text in image | |
| txt = mat_contents['txt'][0][img_indx] | |
| txt = [re.split(' \n|\n |\n| ', t.strip()) for t in txt] | |
| txt = list(chain(*txt)) | |
| txt = [t for t in txt if len(t) > 0 ] | |
| # print(txt) # ['Lines:', 'I', 'lost', 'Kevin', 'will', 'line', 'and', 'and', 'the', '(and', 'the', 'out', 'you', "don't", 'pkg'] | |
| # assert 1<0 | |
| # Open image | |
| #img = Image.open(image_name) | |
| img = cv2.imread(image_name, cv2.IMREAD_COLOR) | |
| img_height, img_width, _ = img.shape | |
| # Validation | |
| if len(np.shape(mat_contents['wordBB'][0][img_indx])) == 2: | |
| wordBBlen = 1 | |
| else: | |
| wordBBlen = mat_contents['wordBB'][0][img_indx].shape[-1] | |
| if wordBBlen == len(txt): | |
| # Crop image and save | |
| for word_indx in range(len(txt)): | |
| # print('txt--',txt) | |
| txt_temp = txt[word_indx] | |
| len_now = len(txt_temp) | |
| # txt_temp = re.sub('[^0-9a-zA-Z]+', '', txt_temp) | |
| # print('txt_temp-1-',txt_temp) | |
| txt_temp = re.sub('[^a-zA-Z]+', '', txt_temp) | |
| # print('txt_temp-2-',txt_temp) | |
| if len_now - len(txt_temp) != 0: | |
| print('txt_temp-2-', txt_temp) | |
| if len(np.shape(mat_contents['wordBB'][0][img_indx])) == 2: # only one word (2,4) | |
| wordBB = mat_contents['wordBB'][0][img_indx] | |
| else: # many words (2,4,num_words) | |
| wordBB = mat_contents['wordBB'][0][img_indx][:, :, word_indx] | |
| if np.shape(wordBB) != (2, 4): | |
| err_log = 'malformed box index: {}\t{}\t{}\n'.format(image_name, txt[word_indx], wordBB) | |
| err_file.write(err_log) | |
| # print(err_log) | |
| continue | |
| pts1 = np.float32([[wordBB[0][0], wordBB[1][0]], | |
| [wordBB[0][3], wordBB[1][3]], | |
| [wordBB[0][1], wordBB[1][1]], | |
| [wordBB[0][2], wordBB[1][2]]]) | |
| height = math.sqrt((wordBB[0][0] - wordBB[0][3])**2 + (wordBB[1][0] - wordBB[1][3])**2) | |
| width = math.sqrt((wordBB[0][0] - wordBB[0][1])**2 + (wordBB[1][0] - wordBB[1][1])**2) | |
| # Coord validation check | |
| if (height * width) <= 0: | |
| err_log = 'empty file : {}\t{}\t{}\n'.format(image_name, txt[word_indx], wordBB) | |
| err_file.write(err_log) | |
| # print(err_log) | |
| continue | |
| elif (height * width) > (img_height * img_width): | |
| err_log = 'too big box : {}\t{}\t{}\n'.format(image_name, txt[word_indx], wordBB) | |
| err_file.write(err_log) | |
| # print(err_log) | |
| continue | |
| else: | |
| valid = True | |
| for i in range(2): | |
| for j in range(4): | |
| if wordBB[i][j] < 0 or wordBB[i][j] > img.shape[1 - i]: | |
| valid = False | |
| break | |
| if not valid: | |
| break | |
| if not valid: | |
| err_log = 'invalid coord : {}\t{}\t{}\t{}\t{}\n'.format( | |
| image_name, txt[word_indx], wordBB, (width, height), (img_width, img_height)) | |
| err_file.write(err_log) | |
| # print(err_log) | |
| continue | |
| pts2 = np.float32([[0, 0], | |
| [0, height], | |
| [width, 0], | |
| [width, height]]) | |
| x_min = np.int(round(min(wordBB[0][0], wordBB[0][1], wordBB[0][2], wordBB[0][3]))) | |
| x_max = np.int(round(max(wordBB[0][0], wordBB[0][1], wordBB[0][2], wordBB[0][3]))) | |
| y_min = np.int(round(min(wordBB[1][0], wordBB[1][1], wordBB[1][2], wordBB[1][3]))) | |
| y_max = np.int(round(max(wordBB[1][0], wordBB[1][1], wordBB[1][2], wordBB[1][3]))) | |
| # print(x_min, x_max, y_min, y_max) | |
| # print(img.shape) | |
| # assert 1<0 | |
| if len(img.shape) == 3: | |
| img_cropped = img[ y_min:y_max:1, x_min:x_max:1, :] | |
| else: | |
| img_cropped = img[ y_min:y_max:1, x_min:x_max:1] | |
| dir_name = '/home/yxwang/pytorch/dataset/SynthText/cropped-oabc/{}'.format(image_name_new.split('/')[0]) | |
| # print('dir_name--',dir_name) | |
| if not os.path.exists(dir_name): | |
| os.mkdir(dir_name) | |
| cropped_file_name = "{}/{}_{}_{}.jpg".format(dir_name, cropped_indx, | |
| image_name.split('/')[-1][:-len('.jpg')], word_indx) | |
| # print('cropped_file_name--',cropped_file_name) | |
| # print('img_cropped--',img_cropped.shape) | |
| if img_cropped.shape[0] == 0 or img_cropped.shape[1] == 0: | |
| err_log = 'word_box_mismatch : {}\t{}\t{}\n'.format(image_name, mat_contents['txt'][0][ | |
| img_indx], mat_contents['wordBB'][0][img_indx]) | |
| err_file.write(err_log) | |
| # print(err_log) | |
| continue | |
| # print('img_cropped--',img_cropped) | |
| # img_cropped.save(cropped_file_name) | |
| cv2.imwrite(cropped_file_name, img_cropped) | |
| cropped_indx += 1 | |
| gt_file.write('%s\t%s\n' % (cropped_file_name, txt[word_indx])) | |
| # if cropped_indx>10: | |
| # assert 1<0 | |
| # assert 1 < 0 | |
| else: | |
| err_log = 'word_box_mismatch : {}\t{}\t{}\n'.format(image_name, mat_contents['txt'][0][ | |
| img_indx], mat_contents['wordBB'][0][img_indx]) | |
| err_file.write(err_log) | |
| # print(err_log) | |
| gt_file.close() | |
| err_file.close() | |