Skip to content

Commit c4fbe01

Browse files
author
Olexandr Dubchak
committed
refactor and text on rectangles
1 parent bc5af1b commit c4fbe01

File tree

6 files changed

+106
-70
lines changed

6 files changed

+106
-70
lines changed

‎core/detection/yolo/detect.py

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import cv2
22
import numpy as np
33
from yolov4.tf import YOLOv4
4+
import pytesseract
45

5-
from core.logic.rectangle_graph import YoloRectangle, Point
6+
from core.logic.rectangle_graph import YoloRectangle, Point, RectangleDocumentGraph
67

78
yolo = YOLOv4()
89

@@ -11,31 +12,45 @@
1112

1213
yolo.make_model()
1314
yolo.load_weights('./weights/v2/custom_yolov4_v2.weights', weights_type="yolo")
14-
yolo.summary(summary_type="yolo")
15-
yolo.summary()
15+
#yolo.summary(summary_type="yolo")
16+
#yolo.summary()
1617

1718
#yolo.inference(media_path="piia.jpg")
1819
#yolo.get_yolo_detections()
1920

20-
frame = cv2.imread('im4.jpg')
21-
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
22-
bboxes = yolo.predict(frame_rgb, 0.25)
21+
frame = cv2.imread('im3.jpg')
22+
#frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
23+
bboxes = yolo.predict(frame, 0.25)
2324

2425

25-
image = np.copy(frame_rgb)
26+
image = np.copy(frame)
2627
height, width, _ = image.shape
2728
res_bboxes = bboxes * np.array([width, height, width, height, 1, 1])
2829

2930
print('fuck')
3031
#
3132

32-
kek = [res_bboxes[0], res_bboxes[5]]
33+
kek = [res_bboxes[10], res_bboxes[5]]
34+
fst_bbox = res_bboxes[10]
35+
yolo_recs = []
3336
for x, y, width, height, _, _ in res_bboxes:
3437
rec = YoloRectangle(center_point=Point(x, y), width=width, height=height)
3538
cv2.rectangle(image, (int(rec.point1.x), int(rec.point1.y)), (int(rec.point2.x), int(rec.point2.y)), color=(0, 255, 0), thickness=3)
39+
roi = image[rec.bottom_right_point.int_y: rec.top_left_point.int_y, rec.top_left_point.int_x: rec.bottom_right_point.int_x,:]
40+
text = pytesseract.image_to_string(roi, config='--oem 3 --psm 6', lang='ukr+eng')
41+
rec.text = text
42+
yolo_recs.append(rec)
3643

37-
cv2.imwrite('im41.jpg', image)
44+
graph = RectangleDocumentGraph(yolo_recs)
45+
graph.build_rectangle_graph()
46+
47+
#cv2.imwrite('im41.jpg', image)
3848
#cv2.rectangle(image, (10,10), (100, 100), color=(0, 255, 0), thickness=3)
49+
# rec = YoloRectangle(center_point=Point(fst_bbox[0], fst_bbox[1]), width=fst_bbox[2], height=fst_bbox[3])
50+
# roi = image[rec.bottom_right_point.int_y: rec.top_left_point.int_y, rec.top_left_point.int_x: rec.bottom_right_point.int_x,:]
51+
# img_box =0
52+
# text = pytesseract.image_to_string(roi, config='--oem 3 --psm 6', lang='ukr+eng')
53+
# cv2.imshow('ROI', roi)
3954
cv2.imshow('shit', image)
4055
k = cv2.waitKey(0)
4156
# yolo.inference(media_path="road.mp4", is_image=False)

‎core/detection/yolo/tesseract_test.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import pytesseract
33
import numpy as np
44

5-
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
5+
#pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
66

77

88

@@ -70,7 +70,7 @@ def match_template(image, template):
7070

7171

7272

73-
img = cv2.imread('im1.jpg')
73+
img = cv2.imread('im3.jpg')
7474
#img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
7575

7676
# gray = get_grayscale(img)
@@ -98,6 +98,6 @@ def match_template(image, template):
9898
img = cv2.rectangle(img, (int(b[1]), h - int(b[2])), (int(b[3]), h - int(b[4])), (0, 255, 0), 2)
9999

100100
#print(boxes)
101-
cv2.imwrite('tes_res_rec.jpg', img)
102-
#cv2.imshow('img', img)
103-
#cv2.waitKey(0)
101+
#cv2.imwrite('tes_res_rec.jpg', img)
102+
cv2.imshow('img', img)
103+
cv2.waitKey(0)

‎core/logic/graph.py

Lines changed: 45 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,16 @@
33

44
class Node:
55
newid = itertools.count()
6-
TOP = 'TOP'
7-
BOTTOM = 'BOTTOM'
8-
RIGHT = 'RIGHT'
9-
LEFT = 'LEFT'
10-
TOP_RIGHT = 'TOP_RIGHT'
11-
TOP_LEFT = 'TOP_LEFT'
12-
BOTTOM_RIGHT = 'BOTTOM_RIGHT'
13-
BOTTOM_LEFT = 'BOTTOM_LEFT'
6+
7+
class NodePosition:
8+
TOP = 'TOP'
9+
BOTTOM = 'BOTTOM'
10+
RIGHT = 'RIGHT'
11+
LEFT = 'LEFT'
12+
TOP_RIGHT = 'TOP_RIGHT'
13+
TOP_LEFT = 'TOP_LEFT'
14+
BOTTOM_RIGHT = 'BOTTOM_RIGHT'
15+
BOTTOM_LEFT = 'BOTTOM_LEFT'
1416

1517
class NodeType:
1618
ANCHOR = 'anchor'
@@ -35,52 +37,52 @@ def __init__(self, value='', name='unkown', node_type=NodeType.DEFAULT, node_id:
3537
@classmethod
3638
def get_reversed_position(cls, position):
3739
reversed = {
38-
cls.TOP: cls.BOTTOM,
39-
cls.BOTTOM: cls.TOP,
40-
cls.RIGHT: cls.LEFT,
41-
cls.LEFT: cls.RIGHT,
42-
cls.TOP_RIGHT: cls.BOTTOM_LEFT,
43-
cls.TOP_LEFT: cls.BOTTOM_RIGHT,
44-
cls.BOTTOM_LEFT: cls.TOP_RIGHT,
45-
cls.BOTTOM_RIGHT: cls.TOP_LEFT,
40+
cls.NodePosition.TOP: cls.NodePosition.BOTTOM,
41+
cls.NodePosition.BOTTOM: cls.NodePosition.TOP,
42+
cls.NodePosition.RIGHT: cls.NodePosition.LEFT,
43+
cls.NodePosition.LEFT: cls.NodePosition.RIGHT,
44+
cls.NodePosition.TOP_RIGHT: cls.NodePosition.BOTTOM_LEFT,
45+
cls.NodePosition.TOP_LEFT: cls.NodePosition.BOTTOM_RIGHT,
46+
cls.NodePosition.BOTTOM_LEFT: cls.NodePosition.TOP_RIGHT,
47+
cls.NodePosition.BOTTOM_RIGHT: cls.NodePosition.TOP_LEFT,
4648
}
4749
return reversed.get(position)
4850

4951

5052
def get_corresponding_attr(self, position: str):
51-
if position == self.RIGHT:
53+
if position == self.NodePosition.RIGHT:
5254
return self.right_node
53-
elif position == self.TOP_RIGHT:
55+
elif position == self.NodePosition.TOP_RIGHT:
5456
return self.top_right_node
55-
elif position == self.TOP:
57+
elif position == self.NodePosition.TOP:
5658
return self.top_node
57-
elif position == self.TOP_LEFT:
59+
elif position == self.NodePosition.TOP_LEFT:
5860
return self.top_left_node
59-
elif position == self.LEFT:
61+
elif position == self.NodePosition.LEFT:
6062
return self.left_node
61-
elif position == self.BOTTOM_LEFT:
63+
elif position == self.NodePosition.BOTTOM_LEFT:
6264
return self.bottom_left_node
63-
elif position == self.BOTTOM:
65+
elif position == self.NodePosition.BOTTOM:
6466
return self.bottom_node
65-
elif position == self.BOTTOM_RIGHT:
67+
elif position == self.NodePosition.BOTTOM_RIGHT:
6668
return self.bottom_right_node
6769

6870
def set_to_corresponding_attr(self, position: str, node):
69-
if position == self.RIGHT:
71+
if position == self.NodePosition.RIGHT:
7072
self.right_node = node
71-
elif position == self.TOP_RIGHT:
73+
elif position == self.NodePosition.TOP_RIGHT:
7274
self.top_right_node = node
73-
elif position == self.TOP:
75+
elif position == self.NodePosition.TOP:
7476
self.top_node = node
75-
elif position == self.TOP_LEFT:
77+
elif position == self.NodePosition.TOP_LEFT:
7678
self.top_left_node = node
77-
elif position == self.LEFT:
79+
elif position == self.NodePosition.LEFT:
7880
self.left_node = node
79-
elif position == self.BOTTOM_LEFT:
81+
elif position == self.NodePosition.BOTTOM_LEFT:
8082
self.bottom_left_node = node
81-
elif position == self.BOTTOM:
83+
elif position == self.NodePosition.BOTTOM:
8284
self.bottom_node = node
83-
elif position == self.BOTTOM_RIGHT:
85+
elif position == self.NodePosition.BOTTOM_RIGHT:
8486
self.bottom_right_node = node
8587
return self.get_corresponding_attr(position)
8688

@@ -90,6 +92,9 @@ def set_relation(self, other_node, position):
9092
reversed_position = self.get_reversed_position(position)
9193
other_node.set_to_corresponding_attr(reversed_position, self)
9294

95+
def __repr__(self):
96+
return f'{self.id} | {self.value}'
97+
9398

9499
class DocumentGraph:
95100
def __init__(self, nodes=None):
@@ -111,11 +116,15 @@ def get_node_by_id(self, id:int)->Node:
111116
return None
112117
return self.nodes[index]
113118

114-
def find_corresponding_node(self, node, graph2):
115-
pass
116119

117-
def compare_graphs(self, graph2):
120+
class PhotoDocumentGraph(DocumentGraph):
121+
def extract_features(self, nodes):
118122
pass
119123

120-
def get_related_info_from_graph(self, graph2):
124+
125+
class TemplateDocumentGraph(DocumentGraph):
126+
127+
def compare(self, graph):
128+
"""Compares template graph with graph from photo, returns corresponding nodes"""
121129
pass
130+

‎core/logic/im3.jpg

424 KB
Loading

‎core/logic/main.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
rec1, rec2 = recs[0], recs[5]
1919

2020

21+
22+
2123
graph = RectangleDocumentGraph(recs)
2224
graph.build_rectangle_graph()
2325

‎core/logic/rectangle_graph.py

Lines changed: 30 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,14 @@ def __init__(self, x, y):
1111
self.x = x
1212
self.y = y
1313

14+
@property
15+
def int_x(self):
16+
return int(self.x)
17+
18+
@property
19+
def int_y(self):
20+
return int(self.y)
21+
1422
def get_distance(self, point):
1523
return math.sqrt((self.x - point.x)**2 + (self.y - point.y)**2)
1624

@@ -44,7 +52,7 @@ class Rectangle:
4452
BOTTOM_LEFT = 'BOTTOM_LEFT'
4553

4654

47-
def __init__(self, point1, point2):
55+
def __init__(self, point1, point2, text=''):
4856
"""
4957
rectangle on the image
5058
@@ -63,6 +71,7 @@ def __init__(self, point1, point2):
6371
self.length = abs(point2.x - point1.x)
6472
self.height = abs(point2.y - point1.y)
6573
self.center = Point((point2.x + point1.x) / 2, (point2.y + point1.y) / 2)
74+
self.text = text
6675

6776
@property
6877
def top_left_point(self):
@@ -83,21 +92,21 @@ def bottom_right_point(self):
8392
def get_distance(self, rectangle):
8493
position = self.get_relative_position(rectangle)
8594
res = 0
86-
if position == Node.RIGHT:
95+
if position == Node.NodePosition.RIGHT:
8796
res = rectangle.top_left_point.x - self.top_right_point.x
88-
elif position == Node.TOP_RIGHT:
97+
elif position == Node.NodePosition.TOP_RIGHT:
8998
res = self.top_right_point.get_distance(rectangle.bottom_left_point)
90-
elif position == Node.TOP:
99+
elif position == Node.NodePosition.TOP:
91100
res = self.top_left_point.y - rectangle.bottom_left_point.y
92-
elif position == Node.TOP_LEFT:
101+
elif position == Node.NodePosition.TOP_LEFT:
93102
res = self.top_left_point.get_distance(rectangle.bottom_right_point)
94-
elif position == Node.LEFT:
103+
elif position == Node.NodePosition.LEFT:
95104
res = self.top_left_point.x - rectangle.top_right_point.x
96-
elif position == Node.BOTTOM_LEFT:
105+
elif position == Node.NodePosition.BOTTOM_LEFT:
97106
res = self.bottom_left_point.get_distance(rectangle.top_right_point)
98-
elif position == Node.BOTTOM:
107+
elif position == Node.NodePosition.BOTTOM:
99108
res = rectangle.top_left_point.y - self.bottom_left_point.y
100-
elif position == Node.BOTTOM_RIGHT:
109+
elif position == Node.NodePosition.BOTTOM_RIGHT:
101110
res = self.bottom_right_point.get_distance(rectangle.top_left_point)
102111
return abs(res)
103112

@@ -106,25 +115,25 @@ def get_relative_position(self, rect):
106115
#angle = self.center.get_angle(rect.center)
107116
angle = self.point1.get_angle(rect.point1)
108117
if angle <= 22.5 or angle >= 337.5:
109-
return Node.RIGHT
118+
return Node.NodePosition.RIGHT
110119
elif 67.5 > angle > 22.5:
111-
return Node.TOP_RIGHT
120+
return Node.NodePosition.TOP_RIGHT
112121
elif 112.5 >= angle >= 67.5:
113-
return Node.TOP
122+
return Node.NodePosition.TOP
114123
elif 157.5 > angle > 112.5:
115-
return Node.TOP_LEFT
124+
return Node.NodePosition.TOP_LEFT
116125
elif 202.5 >= angle >= 157.5:
117-
return Node.LEFT
126+
return Node.NodePosition.LEFT
118127
elif 247.5 > angle > 202.5:
119-
return Node.BOTTOM_LEFT
128+
return Node.NodePosition.BOTTOM_LEFT
120129
elif 292.5 >= angle >= 247.5:
121-
return Node.BOTTOM
130+
return Node.NodePosition.BOTTOM
122131
elif 337.5 > angle > 292.5:
123-
return Node.BOTTOM_RIGHT
132+
return Node.NodePosition.BOTTOM_RIGHT
124133

125134

126135
class YoloRectangle(Rectangle):
127-
def __init__(self, center_point, width, height):
136+
def __init__(self, center_point, width, height, text=''):
128137
half_width = int(width / 2)
129138
half_height = int(height / 2)
130139

@@ -137,7 +146,7 @@ def __init__(self, center_point, width, height):
137146

138147
top_left = Point(top, left)
139148
bottom_right = Point(center_point.x + half_width, center_point.y + half_height)
140-
super().__init__(top_left, bottom_right)
149+
super().__init__(top_left, bottom_right, text)
141150

142151

143152
class RectangleDocumentGraph:
@@ -149,7 +158,8 @@ def __init__(self, rectangles: List[Rectangle]):
149158
self.calculate_distances()
150159
self.graph = DocumentGraph()
151160
for i in range(self.n):
152-
node = Node(node_id=i, value=f'{self.rectangles[i].top_left_point} ___ {self.rectangles[i].bottom_right_point}')
161+
value = self.rectangles[i].text or f'{self.rectangles[i].top_left_point} ___ {self.rectangles[i].bottom_right_point}'
162+
node = Node(node_id=i, value=value)
153163
self.graph.add_to_graph(node)
154164

155165
def build_rectangle_graph(self):

0 commit comments

Comments
 (0)