-
Notifications
You must be signed in to change notification settings - Fork 0
/
group_nodes_2_text.py
147 lines (118 loc) · 4.61 KB
/
group_nodes_2_text.py
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
import matplotlib, cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import distance
from get_nodes import le_main
try:
import Image
except:
from PIL import Image
import pytesseract
import os.path
from pytesseract import image_to_string
# %matplotlib inline
def show(img, cvt=cv2.COLOR_GRAY2RGB, do_cvt=True):
plt.figure(figsize=(10,10))
if do_cvt: plt.imshow(cv2.cvtColor(img, cvt))
else: plt.imshow(img)
plt.show()
imfile = "images/tree_image_ref.png"
#imfile = "images/Tree1.png"
#just_tree = cv2.imread("images/Tree1.png") #use k=12
#just_tree = cv2.imread("images/Tree2.png") #use k=1
just_tree = cv2.imread(imfile) #use k=9 to get only tips, k=4 gets some of the labels but is not as clean
#label_rectangles
#input: png image of the tree, parameter k
#Parameter k:
#adjust k higher so that phrases are connected, but lower such that all phrases are separate from the tree
#output: (image, list of rectangles in the format (x,y,w,h))
def label_rectangles(img, thresh=0.5, k=9, iterations=1):
img = img.copy()
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
_,thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
#set the amount to dilate to merge words in a phrase together
kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(k,k))
#print kernel
dilated = cv2.dilate(thresh,kernel,iterations)
im2, contours, hierarchy = cv2.findContours(dilated,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
#get the rectangle boundaries we want
heights = []
rectangles = []
for contour in contours:
[x,y,w,h] = cv2.boundingRect(contour)
if h < 5 or w < 5: continue
#if h > grays[i].shape[0] * 0.5 or w > grays[i].shape[1] * 0.5: continue
rectangles.append((x,y,w,h))
heights.append(h)
#remove the tree rectangle
#other way: check if square shaped
good_rectangles = []
for r in rectangles:
(x,y,w,h) = r
if h<3*np.std(heights):
good_rectangles.append(r)
cv2.rectangle(img, (x,y), (x+w,y+h), (255,0,255), 2)
return (img, good_rectangles)
result = label_rectangles(just_tree)
# show(result[0], False)
# contours = text_contours(just_tree)
# given the rectangles to a text detection program, which hopefully returns ((x,y,w,h),"Genus species")
# also assume we have node locations [(x,y)....]
# now we need to find closest leaf
# labels = [((1,1,10,10), "Genus species")]
labels = [(pos, "Label " + str(result[1].index(pos)) ) for pos in result[1] ]
nodes = le_main(imfile, 1)
associations = []
for label in labels:
r,lname = label
x,y,w,h = r
# get distances to all nodes
distances = set()
for node in nodes:
x1,y1 = node
d = distance.euclidean((x,y),node)
distances.add((d,node))
#get the node that is the minimum distance and name it
associations.append((min(distances)[1],(x,y),lname))
tree = result[0].copy()
# print list(associations)[0][0]
def print_leaf(leaf):
nodepos = tuple(map(int, leaf[0]))
labelpos = tuple(map(int, leaf[1]))
font = cv2.FONT_HERSHEY_SIMPLEX
#cv2.putText(tree, leaf[2], nodepos, font,0.5, 135)
cv2.circle(tree, nodepos, 4, (255,0,0), 2)
cv2.circle(tree, labelpos, 4, (255,0,0),2)
cv2.line(tree, nodepos, labelpos, (0,0,255), 2)
# map(print_leaf, associations)
# show(tree, False)
def labels_for_tree(just_tree, nodes):
# just_tree = cv2.imread(image_path) #use k=9 to get only tips, k=4 gets some of the labels but is not as clean
result = label_rectangles(just_tree)
labels = [(pos, get_label(Image.fromarray(just_tree), pos) ) for pos in result[1] ]
associations = []
for label in labels:
r,lname = label
x,y,w,h = r
# get distances to all nodes
distances = set()
for node in nodes:
x1,y1 = node
d = distance.euclidean((x,y),node)
distances.add((d,node))
#get the node that is the minimum distance and name it
associations.append((min(distances)[1],(x,y),lname))
tree = result[0].copy()
return associations
# map(print_leaf, associations)
# show(tree, False)
def imgToStr(im):
gray_im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
gray_image = Image.fromarray(gray_im)
str = pytesseract.image_to_string(gray_image, config="--user-words config/user-words.txt --user-patterns config/user-patterns.txt --tessedit_char_whitelist AaBbCcDdEeFfGgHhIiJjKkLlMmNnOoPpQqRrSsTtUuVvWwXxYyZz_0123456789")
return str
def get_label(img, pos):
x,y,w,h = pos
box = (x, y, x + w, y + h)
ocr_ready = img.crop(box)
return imgToStr(ocr_ready)