def seal_detect(img, model, ratio_w, ratio_h): boxes = detect.detect(img, model, device) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) orig_vertices = [] theta = 0 if boxes is not None and boxes.size: for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices) theta /= len(boxes) tmp_img, vertices = crop_mode.rotate_allimg(img, orig_vertices, -theta / math.pi * 180) dict_centers = {} for i, vertice in enumerate(vertices): avg_x = int(crop_mode.averagenum(vertice[::2])) avg_y = int(crop_mode.averagenum(vertice[1::2])) dict_centers[str(avg_x) + ',' + str(avg_y)] = i centers = crop_mode.sort_centers(dict_centers, 1) xcenters = [] for center in centers: xcenters.append([center]) shape = [] for i, xcenter in enumerate(xcenters): for center in xcenter: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'seal' shape.append(anno) return shape, boxes
def predict_mode3(img, img_name, rotate_model, seal_detect_model, detect_model, handword_model, handnum_model, word_model, num_model, char_model, seal_model): rotate_flag = rotate_det(img, rotate_model) if rotate_flag == 0: rotate_img = img.rotate(180, Image.BILINEAR) else: rotate_img = img rotate_img = rotate_img.convert("RGB") w, h = rotate_img.size ratio_w = 512 / w ratio_h = 512 / h img_tmp = rotate_img.resize((512, 512)) shape, seal_boxes = seal_detect(img_tmp, seal_detect_model, ratio_w, ratio_h) boxes = detect.detect(img_tmp, detect_model) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) plot_img = detect.plot_boxes(rotate_img, boxes) plot_img = detect.plot_boxes(plot_img, seal_boxes) plot_img.save( os.path.join('/home/flask_web/uploads', img_name[:-4] + "_result" + ".jpg")) print('detection result saved') orig_vertices = [] theta = 0 for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices) theta /= len(boxes) tmp_img, vertices = crop_mode.rotate_allimg(rotate_img, orig_vertices, -theta / math.pi * 180) dict_centers = {} for i, vertice in enumerate(vertices): avg_x = int(crop_mode.averagenum(vertice[::2])) avg_y = int(crop_mode.averagenum(vertice[1::2])) dict_centers[str(avg_x) + ',' + str(avg_y)] = i centers = crop_mode.sort_centers(dict_centers, 1) k = 0 xcenters = [] index = [3, 3, 1, 2, 6, 3] for i, j in enumerate(index): xcenter = crop_mode.sort_xcenters(centers[k:k + j], 0) if i == 4: if int(xcenter[0][0].split(',')[1]) > int( xcenter[1][0].split(',')[1]): tmp = xcenter[0] xcenter[0] = xcenter[1] xcenter[1] = tmp if int(xcenter[2][0].split(',')[1]) > int( xcenter[3][0].split(',')[1]): tmp = xcenter[2] xcenter[2] = xcenter[3] xcenter[3] = tmp if int(xcenter[3][0].split(',')[1]) > int( xcenter[4][0].split(',')[1]): tmp = xcenter[3] xcenter[3] = xcenter[4] xcenter[4] = tmp k += j xcenters.append(xcenter) for i, xcenter in enumerate(xcenters): if i == 0: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '年' shape.append(anno) if j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '月' shape.append(anno) if j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '日' shape.append(anno) else: break if i == 1: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '收款人' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '代理付款行' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '代理付款行号' shape.append(anno) else: break if i == 2: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票金额' shape.append(anno) else: break if i == 3: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '实际结算金额(大写)' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handnum' anno['name'] = '实际结算金额(小写)' shape.append(anno) else: break if i == 4: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '账号' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '密押' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '申请人' shape.append(anno) elif j == 3: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '出票行' shape.append(anno) elif j == 4: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '备注' shape.append(anno) elif j == 5: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '出票行号' shape.append(anno) else: break if i == 5: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handnum' anno['name'] = '多余金额' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '复核' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '记账' shape.append(anno) else: break results = recognize(shape, rotate_img, rotate_flag, handword_model, handnum_model, word_model, num_model, char_model, seal_model) return results
def do_predict(app, img_path, origin_name): img_name = img_path[-36:] db_image = Image.query.filter_by(img_name=img_name).first() img_mode = int(origin_name[14]) print('image mode is : ', img_mode) img = PIL_Image.open(img_path).convert('L') if img_mode == 1: db_image.img_mode = 1 # results = predict.predict_mode1(img, img_name, rotate_model, seal_model_mode1, detect_model_mode1, # handword_model, handnum_model, word_model, num_model, char_model, seal_model) rotate_flag = rotate_det(img, rotate_model) if rotate_flag == 0: rotate_img = img.rotate(180, PIL_Image.BILINEAR) else: rotate_img = img rotate_img = rotate_img.convert("RGB") w, h = rotate_img.size ratio_w = 512 / w ratio_h = 512 / h img_tmp = rotate_img.resize((512, 512)) seal_model_mode1 = EAST(pretrained=False).to(device) seal_model_mode1.load_state_dict(torch.load(stamp_model_one, map_location=torch.device('cpu'))) seal_model_mode1.eval() shape, seal_boxes = seal_detect(img_tmp, seal_model_mode1, ratio_w, ratio_h) detect_model_mode1 = EAST(pretrained=False).to(device) detect_model_mode1.load_state_dict(torch.load(detect_model_one, map_location=torch.device('cpu'))) detect_model_mode1.eval() boxes = detect.detect(img_tmp, detect_model_mode1, device) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) plot_img = detect.plot_boxes(rotate_img, boxes) plot_img = detect.plot_boxes(plot_img, seal_boxes) plot_img.save(os.path.join(app.config['UPLOAD_PATH'], img_name[:-4] + "_result" + ".jpg")) print('detection result saved') orig_vertices = [] theta = 0 for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices) theta /= len(boxes) tmp_img, vertices = crop_mode.rotate_allimg(rotate_img, orig_vertices, - theta / math.pi * 180) dict_centers = {} for i, vertice in enumerate(vertices): avg_x = int(crop_mode.averagenum(vertice[::2])) avg_y = int(crop_mode.averagenum(vertice[1::2])) dict_centers[str(avg_x) + ',' + str(avg_y)] = i centers = crop_mode.sort_centers(dict_centers, 1) k = 0 xcenters = [] index = [6, 2, 2, 1, 2] for i, j in enumerate(index): xcenter = crop_mode.sort_xcenters(centers[k:k + j], 0) if i == 0: if int(xcenter[0][0].split(',')[1]) < int(xcenter[1][0].split(',')[1]): tmp = xcenter[0] xcenter[0] = xcenter[1] xcenter[1] = tmp if int(xcenter[5][0].split(',')[1]) < int(xcenter[4][0].split(',')[1]): tmp = xcenter[5] xcenter[5] = xcenter[4] xcenter[4] = tmp k += j xcenters.append(xcenter) for i, xcenter in enumerate(xcenters): if i == 0: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '收款人' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(年)' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(月)' shape.append(anno) elif j == 3: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(日)' shape.append(anno) elif j == 4: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '出票人账号' shape.append(anno) elif j == 5: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '付款行名称' shape.append(anno) else: break if i == 1: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '人民币(大写)' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handnum' anno['name'] = '人民币(小写)' shape.append(anno) else: break if i == 2: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '行号' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '用途' shape.append(anno) else: break if i == 3: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '密码' shape.append(anno) else: break if i == 4: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '复核' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '记账' shape.append(anno) else: break results = recognize(shape, rotate_img, rotate_flag) elif img_mode == 2: db_image.img_mode = 2 # results = predict.predict_mode2(img, img_name, rotate_model, seal_model_mode2, detect_model_mode2, # handword_model, handnum_model, word_model, num_model, char_model, seal_model) rotate_flag = rotate_det(img, rotate_model) if rotate_flag == 0: rotate_img = img.rotate(180, PIL_Image.BILINEAR) else: rotate_img = img rotate_img = rotate_img.convert("RGB") w, h = rotate_img.size ratio_w = 512 / w ratio_h = 512 / h img_tmp = rotate_img.resize((512, 512)) seal_model_mode2 = EAST(pretrained=False).to(device) seal_model_mode2.load_state_dict(torch.load(stamp_model_two, map_location=torch.device('cpu'))) seal_model_mode2.eval() shape, seal_boxes = seal_detect(img_tmp, seal_model_mode2, ratio_w, ratio_h) detect_model_mode2 = EAST(pretrained=False).to(device) detect_model_mode2.load_state_dict(torch.load(detect_model_two, map_location=torch.device('cpu'))) detect_model_mode2.eval() boxes = detect.detect(img_tmp, detect_model_mode2, device) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) plot_img = detect.plot_boxes(rotate_img, boxes) plot_img = detect.plot_boxes(plot_img, seal_boxes) plot_img.save(os.path.join(app.config['UPLOAD_PATH'], img_name[:-4] + "_result" + ".jpg")) print('detection result saved') orig_vertices = [] theta = 0 for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices) theta /= len(boxes) tmp_img, vertices = crop_mode.rotate_allimg(rotate_img, orig_vertices, - theta / math.pi * 180) dict_centers = {} for i, vertice in enumerate(vertices): avg_x = int(crop_mode.averagenum(vertice[::2])) avg_y = int(crop_mode.averagenum(vertice[1::2])) dict_centers[str(avg_x) + ',' + str(avg_y)] = i centers = crop_mode.sort_centers(dict_centers, 1) k = 0 xcenters = [] index = [3, 2, 2, 1, 4] for i, j in enumerate(index): xcenter = crop_mode.sort_xcenters(centers[k:k + j], 0) k += j xcenters.append(xcenter) for i, xcenter in enumerate(xcenters): if i == 0: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(年)' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(月)' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(日)' shape.append(anno) else: break if i == 1: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '人民币(大写)' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handnum' anno['name'] = '人民币(小写)' shape.append(anno) else: break if i == 2: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'char' anno['name'] = '使用现金' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '行号' shape.append(anno) else: break if i == 3: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '密押' shape.append(anno) else: break if i == 4: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '备注' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '经办' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '复核' shape.append(anno) elif j == 3: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '出纳' shape.append(anno) else: break results = recognize(shape, rotate_img, rotate_flag) elif img_mode == 3: db_image.img_mode = 3 # results = predict.predict_mode3(img, img_name, rotate_model, seal_model_mode3, detect_model_mode3, # handword_model, handnum_model, word_model, num_model, char_model, seal_model) rotate_flag = rotate_det(img, rotate_model) if rotate_flag == 0: rotate_img = img.rotate(180, PIL_Image.BILINEAR) else: rotate_img = img rotate_img = rotate_img.convert("RGB") w, h = rotate_img.size ratio_w = 512 / w ratio_h = 512 / h img_tmp = rotate_img.resize((512, 512)) seal_model_mode3 = EAST(pretrained=False).to(device) seal_model_mode3.load_state_dict(torch.load(stamp_model_three, map_location=torch.device('cpu'))) seal_model_mode3.eval() shape, seal_boxes = seal_detect(img_tmp, seal_model_mode3, ratio_w, ratio_h) detect_model_mode3 = EAST(pretrained=False).to(device) detect_model_mode3.load_state_dict(torch.load(detect_model_three, map_location=torch.device('cpu'))) detect_model_mode3.eval() boxes = detect.detect(img_tmp, detect_model_mode3, device) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) plot_img = detect.plot_boxes(rotate_img, boxes) plot_img = detect.plot_boxes(plot_img, seal_boxes) plot_img.save(os.path.join(app.config['UPLOAD_PATH'], img_name[:-4] + "_result" + ".jpg")) print('detection result saved') orig_vertices = [] theta = 0 for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices) theta /= len(boxes) tmp_img, vertices = crop_mode.rotate_allimg(rotate_img, orig_vertices, - theta / math.pi * 180) dict_centers = {} for i, vertice in enumerate(vertices): avg_x = int(crop_mode.averagenum(vertice[::2])) avg_y = int(crop_mode.averagenum(vertice[1::2])) dict_centers[str(avg_x) + ',' + str(avg_y)] = i centers = crop_mode.sort_centers(dict_centers, 1) k = 0 xcenters = [] index = [3, 3, 1, 2, 6, 3] for i, j in enumerate(index): xcenter = crop_mode.sort_xcenters(centers[k:k + j], 0) if i == 4: if int(xcenter[0][0].split(',')[1]) > int(xcenter[1][0].split(',')[1]): tmp = xcenter[0] xcenter[0] = xcenter[1] xcenter[1] = tmp if int(xcenter[2][0].split(',')[1]) > int(xcenter[3][0].split(',')[1]): tmp = xcenter[2] xcenter[2] = xcenter[3] xcenter[3] = tmp if int(xcenter[3][0].split(',')[1]) > int(xcenter[4][0].split(',')[1]): tmp = xcenter[3] xcenter[3] = xcenter[4] xcenter[4] = tmp k += j xcenters.append(xcenter) for i, xcenter in enumerate(xcenters): if i == 0: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(年)' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(月)' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票日期(日)' shape.append(anno) else: break if i == 1: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '收款人' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '代理付款行' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '代理付款行号' shape.append(anno) else: break if i == 2: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '出票金额' shape.append(anno) else: break if i == 3: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handword' anno['name'] = '实际结算金额(大写)' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handnum' anno['name'] = '实际结算金额(小写)' shape.append(anno) else: break if i == 4: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '账号' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '密押' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '申请人' shape.append(anno) elif j == 3: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '出票行' shape.append(anno) elif j == 4: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'word' anno['name'] = '备注' shape.append(anno) elif j == 5: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '出票行号' shape.append(anno) else: break if i == 5: for j, center in enumerate(xcenter): if j == 0: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'handnum' anno['name'] = '多余金额' shape.append(anno) elif j == 1: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '复核' shape.append(anno) elif j == 2: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'num' anno['name'] = '记账' shape.append(anno) else: break results = recognize(shape, rotate_img, rotate_flag) else: print('the image mode is wrong !') return with open(os.path.join(app.config['UPLOAD_PATH'], img_name[:-4] + '.txt'), 'w') as writer: for result in results: line = result[0] + ";" + result[1] + '\n' writer.write(line) db_image.result_flag = 'Done!!!' db.session.commit() print("Task is done!")