AdvancedEAST-master.rar

import numpy as np import os from PIL import Image, ImageDraw from tqdm import tqdm import cfg #判断点是否在四边形内部 def point_inside_of_quad(px, py, quad_xy_list, p_min, p_max): if (p_min[0] <= px <= p_max[0]) and (p_min[1] <= py <= p_max[1]): xy_list = np.zeros((4, 2)) xy_list[:3, :] = quad_xy_list[1:4, :] - quad_xy_list[:3, :] xy_list[3] = quad_xy_list[0, :] - quad_xy_list[3, :] yx_list = np.zeros((4, 2)) yx_list[:, :] = quad_xy_list[:, -1:-3:-1] a = xy_list * ([py, px] - yx_list) b = a[:, 0] - a[:, 1] if np.amin(b) >= 0 or np.amax(b) <= 0: return True else: return False else: return False #点属于哪类标签（-1中间区域红色，0是头部黄色，1是尾部绿色） def point_inside_of_nth_quad(px, py, xy_list, shrink_1, long_edge): nth = -1 vs = [[[0, 0, 3, 3, 0], [1, 1, 2, 2, 1]], [[0, 0, 1, 1, 0], [2, 2, 3, 3, 2]]] for ith in range(2): quad_xy_list = np.concatenate(( np.reshape(xy_list[vs[long_edge][ith][0]], (1, 2)), np.reshape(shrink_1[vs[long_edge][ith][1]], (1, 2)), np.reshape(shrink_1[vs[long_edge][ith][2]], (1, 2)), np.reshape(xy_list[vs[long_edge][ith][3]], (1, 2))), axis=0) p_min = np.amin(quad_xy_list, axis=0) p_max = np.amax(quad_xy_list, axis=0) if point_inside_of_quad(px, py, quad_xy_list, p_min, p_max): if nth == -1: nth = ith else: nth = -1 break return nth ''' #对标签进行缩小，按照最短边进行缩放比例 @input --xy_list : 四个点坐标 --ratio ：缩放比例 @:return --temp_new_xy_list : 只裁长边的结果 --new_xy_list ：全部裁剪的结果 ''' def shrink(xy_list, ratio=cfg.shrink_ratio): if ratio == 0.0: return xy_list, xy_list diff_1to3 = xy_list[:3, :] - xy_list[1:4, :] diff_4 = xy_list[3:4, :] - xy_list[0:1, :] diff = np.concatenate((diff_1to3, diff_4), axis=0) dis = np.sqrt(np.sum(np.square(diff), axis=-1)) # determine which are long or short edges long_edge = int(np.argmax(np.sum(np.reshape(dis, (2, 2)), axis=0))) short_edge = 1 - long_edge # cal r length array #r = [np.minimum(dis[i], dis[(i + 1) % 4]) for i in range(4)] r = [np.minimum(dis[3], dis[(3 + 1) % 4])] for i in range(3): r.append(np.minimum(dis[i], dis[(i + 1) % 4])) # cal theta array diff_abs = np.abs(diff).astype(np.float32) diff_abs[:, 0] += cfg.epsilon theta = np.arctan(diff_abs[:, 1] / diff_abs[:, 0]) # shrink two long edges temp_new_xy_list = np.copy(xy_list) shrink_edge(xy_list, temp_new_xy_list, long_edge, r, theta, ratio) shrink_edge(xy_list, temp_new_xy_list, long_edge + 2, r, theta, ratio) # shrink two short edges new_xy_list = np.copy(temp_new_xy_list) shrink_edge(temp_new_xy_list, new_xy_list, short_edge, r, theta, ratio) shrink_edge(temp_new_xy_list, new_xy_list, short_edge + 2, r, theta, ratio) return temp_new_xy_list, new_xy_list, long_edge ''' #对其中的一个边进行缩放操作（其实就是对四边形两个点） --xy_list : 四个顶点坐标 --new_xy_list ：xy_list的复制 --edge ：需要缩放的边（下标） --r : 顶点之间的X、Y差值 --theta ：每条边与X轴的夹角 --ratio : 缩放比例 ''' def shrink_edge(xy_list, new_xy_list, edge, r, theta, ratio=cfg.shrink_ratio): if ratio == 0.0: return start_point = edge end_point = (edge + 1) % 4 long_start_sign_x = np.sign( xy_list[end_point, 0] - xy_list[start_point, 0])#顶点移动方向（因为end和start的移动相反，且某个轴不移动为0） new_xy_list[start_point, 0] = \ xy_list[start_point, 0] + \ long_start_sign_x * ratio * r[start_point] * np.cos(theta[start_point])#先把缩放按照theta投影到X方向，之后进行缩放 long_start_sign_y = np.sign( xy_list[end_point, 1] - xy_list[start_point, 1]) new_xy_list[start_point, 1] = \ xy_list[start_point, 1] + \ long_start_sign_y * ratio * r[start_point] * np.sin(theta[start_point]) # long edge one, end point long_end_sign_x = -1 * long_start_sign_x new_xy_list[end_point, 0] = \ xy_list[end_point, 0] + \ long_end_sign_x * ratio * r[end_point] * np.cos(theta[start_point]) long_end_sign_y = -1 * long_start_sign_y new_xy_list[end_point, 1] = \ xy_list[end_point, 1] + \ long_end_sign_y * ratio * r[end_point] * np.sin(theta[start_point]) ''' @生成标签 @output batch*width*hight*7 7含义 ==> 1位是否是目标点、1位是否是头尾点、1位是头或者尾点(头0黄色+尾1绿色)、2位(左上+左下)、2位(右上+右下) 其中后面的两个顺序有可能不一样，取决于长边所在的位置 ''' def process_label(data_dir=cfg.data_dir): with open(os.path.join(data_dir, cfg.val_fname), 'r') as f_val: f_list = f_val.readlines() with open(os.path.join(data_dir, cfg.train_fname), 'r') as f_train: f_list.extend(f_train.readlines()) for line, _ in zip(f_list, tqdm(range(len(f_list)))): line_cols = str(line).strip().split(',') img_name, width, height = \ line_cols[0].strip(), int(line_cols[1].strip()), \ int(line_cols[2].strip()) gt = np.zeros((height // cfg.pixel_size, width // cfg.pixel_size, 7)) train_label_dir = os.path.join(data_dir, cfg.train_label_dir_name) xy_list_array = np.load(os.path.join(train_label_dir, img_name[:-4] + '.npy')) train_image_dir = os.path.join(data_dir, cfg.train_image_dir_name) with Image.open(os.path.join(train_image_dir, img_name)) as im: draw = ImageDraw.Draw(im) for xy_list in xy_list_array: #xy_list = np.array([[10, 10], [20, 10], [10, 50], [20, 50]]) _, shrink_xy_list, _ = shrink(xy_list, cfg.shrink_ratio) shrink_1, _, long_edge = shrink(xy_list, cfg.shrink_side_ratio) p_min = np.amin(shrink_xy_list, axis=0) p_max = np.amax(shrink_xy_list, axis=0) # floor of the float ji_min = (p_min / cfg.pixel_size - 0.5).astype(int) - 1 # +1 for ceil of the float and +1 for include the end ji_max = (p_max / cfg.pixel_size - 0.5).astype(int) + 3 imin = np.maximum(0, ji_min[1]) imax = np.minimum(height // cfg.pixel_size, ji_max[1]) jmin = np.maximum(0, ji_min[0]) jmax = np.minimum(width // cfg.pixel_size, ji_max[0]) for i in range(imin, imax): for j in range(jmin, jmax): px = (j + 0.5) * cfg.pixel_size py = (i + 0.5) * cfg.pixel_size if point_inside_of_quad(px, py, shrink_xy_list, p_min, p_max): gt[i, j, 0] = 1 line_width, line_color = 1, 'red' ith = point_inside_of_nth_quad(px, py, xy_list, shrink_1, long_edge) vs = [[[3, 0], [1, 2]], [[0, 1], [2, 3]]] if ith in range(2): gt[i, j, 1] = 1 if ith == 0: line_width, line_color = 2, 'yellow' else: line_width, line_color = 2, 'green' gt[i, j, 2:3] = ith gt[i,