pytorch、TensorFlow、深度学习实现基于EAST自然文本检测系统.zip

import os import math import random import torch import torchvision.transforms as transforms from torch.utils import data import glob as gb import numpy as np import cv2 import csv import sys import matplotlib.pyplot as plt import matplotlib.patches as Patches from shapely.geometry import Polygon from PIL import Image import warnings from geo_map_cython_lib import gen_geo_map def get_images(root): ''' get images's path and name return: files -- a list of img path name -- a list of img name ''' assert os.path.isdir(root) == True, 'get_img get a wrong path:{} to imgs'.format(root) files = [] for ext in ['jpg', 'png', 'jpeg', 'JPG']: files.extend(gb.glob(os.path.join(root, '*.{}'.format(ext)))) name = [] for i in range(len(files)): name.append(files[i].split('/')[-1]) # check for i in range(len(files)): assert os.path.basename(files[i]) == name[i], 'img path cant corresponding to img name' print('EAST <==> Prepare <==> Total:{} imgs for train'.format(len(files))) files = sorted(files) name = sorted(name) return files, name def load_annoataion(p): ''' load annotation from the text file Note: modified 1. top left vertice 2. clockwise :param p: :return: ''' text_polys = [] text_tags = [] if not os.path.exists(p): return np.array(text_polys, dtype=np.float32) with open(p, 'r') as f: reader = csv.reader(f) for line in reader: label = line[-1]# strip BOM. \ufeff for python3, \xef\xbb\bf for python2 line = [i.strip('\ufeff').strip('\xef\xbb\xbf') for i in line] x1, y1, x2, y2, x3, y3, x4, y4 = list(map(int, line[:8])) text_polys.append([[x1, y1], [x2, y2], [x3, y3], [x4, y4]]) if label == '*' or label == '###': text_tags.append(True) else: text_tags.append(False) return np.array(text_polys, dtype=np.float32), np.array(text_tags, dtype=np.bool) #return text_polys, text_tags def polygon_area(poly): ''' compute area of a polygon :param poly: :return: ''' poly_ = np.array(poly) assert poly_.shape == (4,2), 'poly shape should be 4,2' edge = [ (poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]), (poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]), (poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]), (poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1]) ] return np.sum(edge)/2. def calculate_distance(c1, c2): return math.sqrt(math.pow(c1[0]-c2[0], 2) + math.pow(c1[1]-c2[1], 2)) def choose_best_begin_point(pre_result): """ find top-left vertice and resort """ final_result = [] for coordinate in pre_result: x1 = coordinate[0][0] y1 = coordinate[0][1] x2 = coordinate[1][0] y2 = coordinate[1][1] x3 = coordinate[2][0] y3 = coordinate[2][1] x4 = coordinate[3][0] y4 = coordinate[3][1] xmin = min(x1, x2, x3, x4) ymin = min(y1, y2, y3, y4) xmax = max(x1, x2, x3, x4) ymax = max(y1, y2, y3, y4) combinate = [[[x1, y1], [x2, y2], [x3, y3], [x4, y4]], [[x2, y2], [x3, y3], [x4, y4], [x1, y1]], [[x3, y3], [x4, y4], [x1, y1], [x2, y2]], [[x4, y4], [x1, y1], [x2, y2], [x3, y3]]] dst_coordinate = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]] force = 100000000.0 force_flag = 0 for i in range(4): temp_force = calculate_distance(combinate[i][0], dst_coordinate[0]) + calculate_distance(combinate[i][1], dst_coordinate[1]) + calculate_distance(combinate[i][2], dst_coordinate[2]) + calculate_distance(combinate[i][3], dst_coordinate[3]) if temp_force < force: force = temp_force force_flag = i #if force_flag != 0: # print("choose one direction!") final_result.append(combinate[force_flag]) return final_result def check_and_validate_polys(polys, tags, xxx_todo_changeme): ''' check so that the text poly is in the same direction, and also filter some invalid polygons :param polys: :param tags: :return: ''' (h, w) = xxx_todo_changeme if polys.shape[0] == 0: return polys polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w-1) polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h-1) validated_polys = [] validated_tags = [] # find top-left and clockwise polys = choose_best_begin_point(polys) for poly, tag in zip(polys, tags): p_area = polygon_area(poly) if abs(p_area) < 1: # print poly #print('invalid poly') continue if p_area > 0: #print('poly in wrong direction') poly = poly[(0, 3, 2, 1), :] validated_polys.append(poly) validated_tags.append(tag) return np.array(validated_polys), np.array(validated_tags) def crop_area(im, polys, tags, crop_background=False, max_tries=5000, vis = False, img_name = None): ''' make random crop from the input image :param im: :param polys: :param tags: :param crop_background: :param max_tries: :return: ''' h, w, _ = im.shape pad_h = h//10 pad_w = w//10 h_array = np.zeros((h + pad_h*2), dtype=np.int32) w_array = np.zeros((w + pad_w*2), dtype=np.int32) if polys.shape[0] == 0: return im, [], [] for poly in polys: poly = np.round(poly, decimals=0).astype(np.int32) minx = np.min(poly[:, 0]) maxx = np.max(poly[:, 0]) w_array[minx+pad_w:maxx+pad_w] = 1 miny = np.min(poly[:, 1]) maxy = np.max(poly[:, 1]) h_array[miny+pad_h:maxy+pad_h] = 1 # ensure the cropped area not across a text h_axis = np.where(h_array == 0)[0] w_axis = np.where(w_array == 0)[0] if len(h_axis) == 0 or len(w_axis) == 0: return im, polys, tags for i in range(max_tries): #print('we have try {} times'.format(i)) xx = np.random.choice(w_axis, size=2) xmin = np.min(xx) - pad_w xmax = np.max(xx) - pad_w xmin = np.clip(xmin, 0, w-1) xmax = np.clip(xmax, 0, w-1) yy = np.random.choice(h_axis, size=2) ymin = np.min(yy) - pad_h ymax = np.max(yy) - pad_h ymin = np.clip(ymin, 0, h-1) ymax = np.clip(ymax, 0, h-1) # if xmax - xmin < FLAGS.min_crop_side_ratio*w or ymax - ymin < FLAGS.min_crop_side_ratio*h: if xmax - xmin < 0.1*w or ymax - ymin < 0.1*h: # area too small continue if polys.shape[0] != 0: poly_axis_in_area = (polys[:, :, 0] >= xmin) & (polys[:, :, 0] <= xmax) \ & (polys[:, :, 1] >= ymin) & (polys[:, :, 1] <= ymax) selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0] else: selected_polys = [] if len(selected_polys) == 0: # no text in this area if crop_background == True: im = im[ymin:ymax+1, xmin:xmax+1, :] polys = [] tags = [] return im, polys, tags else: continue else: if crop_background == False: im = im[ymin:ymax+1, xmin:xmax+1, :] polys = polys.tolist() polys = [polys[i] for i in selected_polys] polys = np.array(polys) polys[:, :, 0] -= xmin #ndarray polys[:, :, 1] -= ymin polys = polys.astype(np.int32) polys = polys.tolist() tags = tags.tolist() tags = [tags[i] for i in selected_polys] return im, polys, tags