pipeline.py
1.91 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import time
import cv2
from bank_ocr_inference import bill_ocr, extract_bank_info
from inference import Yolov5
from models.yolov5_config import config
def enlarge_position(box):
x1, y1, x2, y2 = box
w, h = abs(x2 - x1), abs(y2 - y1)
y1, y2 = max(y1 - h // 3, 0), y2 + h // 3
x1, x2 = max(x1 - w // 8, 0), x2 + w // 8
return [x1, y1, x2, y2]
def tamper_detect(image):
st = time.time()
ocr_results = bill_ocr(image)
et1 = time.time()
info_results = extract_bank_info(ocr_results)
et2 = time.time()
tamper_results = []
if len(info_results) != 0:
for info_result in info_results:
box = [info_result[1][0], info_result[1][1], info_result[1][4], info_result[1][5]]
x1, y1, x2, y2 = enlarge_position(box)
# x1, y1, x2, y2 = box
info_image = image[y1:y2, x1:x2, :]
cv2.imshow('info_image', info_image)
results = detector.detect(info_image)
print(results)
if len(results['result']) != 0:
for res in results['result']:
left = int(res['left'])
top = int(res['top'])
width = int(res['width'])
height = int(res['height'])
absolute_position = [x1 + left, y1 + top, x1 + left + width, y1 + top + height]
tamper_results .append(absolute_position)
print(tamper_results)
et3 = time.time()
print(f'all:{et3 - st} ocr:{et1 - st} extract:{et2 - et1} yolo:{et3 - et2}')
for i in tamper_results:
cv2.rectangle(image, tuple(i[:2]), tuple(i[2:]), (0, 0, 255), 2)
cv2.imshow('info', image)
cv2.waitKey(0)
if __name__ == '__main__':
detector = Yolov5(config)
image = cv2.imread(
"/data/situ_invoice_bill_data/new_data/qfs_bank_bill_person_ps/gongshang/tampered/images/val/ps3/CH-B006369332_page_67_img_0.jpg")
tamper_detect(image)