Mobile-Agent / MobileAgent /text_localization.py
阳渠
Mobile-Agent-v2
1e96bca
raw
history blame contribute delete
No virus
1.69 kB
import cv2
import numpy as np
from MobileAgent.crop import crop_image, calculate_size
from PIL import Image
def order_point(coor):
arr = np.array(coor).reshape([4, 2])
sum_ = np.sum(arr, 0)
centroid = sum_ / arr.shape[0]
theta = np.arctan2(arr[:, 1] - centroid[1], arr[:, 0] - centroid[0])
sort_points = arr[np.argsort(theta)]
sort_points = sort_points.reshape([4, -1])
if sort_points[0][0] > centroid[0]:
sort_points = np.concatenate([sort_points[3:], sort_points[:3]])
sort_points = sort_points.reshape([4, 2]).astype('float32')
return sort_points
def longest_common_substring_length(str1, str2):
m = len(str1)
n = len(str2)
dp = [[0] * (n + 1) for _ in range(m + 1)]
for i in range(1, m + 1):
for j in range(1, n + 1):
if str1[i - 1] == str2[j - 1]:
dp[i][j] = dp[i - 1][j - 1] + 1
else:
dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
return dp[m][n]
def ocr(image_path, ocr_detection, ocr_recognition):
text_data = []
coordinate = []
image_full = cv2.imread(image_path)
det_result = ocr_detection(image_full)
det_result = det_result['polygons']
for i in range(det_result.shape[0]):
pts = order_point(det_result[i])
image_crop = crop_image(image_full, pts)
try:
result = ocr_recognition(image_crop)['text'][0]
except:
continue
box = [int(e) for e in list(pts.reshape(-1))]
box = [box[0], box[1], box[4], box[5]]
text_data.append(result)
coordinate.append(box)
else:
return text_data, coordinate