-
Notifications
You must be signed in to change notification settings - Fork 2
/
polygon_analysis.py
executable file
·178 lines (125 loc) · 5.81 KB
/
polygon_analysis.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
#!/usr/bin/env python
import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
from lib.model import Polygon, Point, Image, SurfaceResult
from lib.utils import get_image, main_process, Arg, Campaign, Phase
WIDTH=400
HEIGHT=400
DEFAULT_THRESHOLD=0.45
THRESHOLD_AREA=100
IMAGE_TYPE_THRES="threshold"
IMAGE_TYPE_POLY="polygon"
IMAGE_TYPE_ALL="all"
def draw_polys(polys, label, selected_idx=None, color="red") :
first=True
for idx, poly in enumerate(polys):
points = poly.points
color = "red" if (idx == selected_idx) else color
for i in range(-1, len(points) - 1):
x = [points[i].x, points[i + 1].x]
y = [points[i].y, points[i + 1].y]
if first :
first = False
else:
label = None
plt.plot(x, y, marker="+", color=color, label=label)
def points2cv2(points) :
"""Transform list of points to be used by opencv2"""
return np.array([[pt.x, pt.y] for pt in points]).astype(np.int32)
def process_img(
img : Image, threshold=DEFAULT_THRESHOLD, out=None,
image_type=IMAGE_TYPE_THRES, display=False, campaign=Campaign.GOOGLE,
polys_to_draw=None,
selected_idx=None, **kwargs):
matrix = np.zeros((WIDTH, HEIGHT))
res = SurfaceResult(img.id)
# No input polygon ? Not part of phase 2 : skip.
if len(img.polygons) == 0 :
return None
# Draw polygons
for poly in img.polygons :
pts = list((pt.x, pt.y) for pt in poly.points)
poly_img = np.zeros((WIDTH, HEIGHT), np.uint8)
cv2.fillPoly(poly_img, np.array([pts]), 1)
matrix += poly_img
# Blur image to prevent noise
matrix = cv2.blur(matrix, (3, 3))
nb_actors = len(set(poly.action.actorId for poly in img.polygons))
# Threshold < 1 is ratio of number of actors
if threshold < 1 :
threshold = threshold * nb_actors
# Apply threshold
ret, thres = cv2.threshold(matrix, threshold, 255, 0)
thres = thres.astype(np.uint8)
# Find polygons in the resulting threshold mask
contours, _ = cv2.findContours(thres, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours :
area = cv2.contourArea(contour)
if area < THRESHOLD_AREA :
continue
epsilon = 0.01 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
points = list(Point(pt[0][0], pt[0][1]) for pt in approx)
# Compute score as mean of values within polygon
poly_mask = np.zeros((WIDTH, HEIGHT), np.int8)
pts = points2cv2(points)
cv2.fillPoly(poly_mask, [pts], 1)
# Mean value
score = np.sum(matrix * poly_mask) / np.sum(poly_mask)
# Update output
poly = Polygon(score=score, area=area)
poly.points = points
res.polygons.append(poly)
if out :
if image_type == IMAGE_TYPE_THRES :
out_img = thres
elif image_type == IMAGE_TYPE_ALL :
out_matrix = (matrix * 255 / nb_actors).astype(np.uint8)
out_matrix[thres == 0] = 0
out_img = cv2.applyColorMap(out_matrix, cv2.COLORMAP_OCEAN)
for poly in res.polygons :
cv2.polylines(out_img, [points2cv2(poly.points)], True, (0, 0, 255))
elif image_type == IMAGE_TYPE_POLY :
out_img = np.zeros((WIDTH, HEIGHT), np.uint8)
for poly in res.polygons:
cv2.fillPoly(out_img, [points2cv2(poly.points)], (255,))
cv2.imwrite(os.path.join(out, "%s.png" % img.id), out_img)
# Plots
if display :
fig = plt.gcf()
path = get_image(img.id, phase=Phase.SURF, campaign=campaign)
image = plt.imread(path)
heat_ax = plt.subplot2grid((1, 2), (0, 1))
matrix[matrix < 0.1] = np.nan
heat_plt = plt.imshow(matrix)
# draw_polys(polys, selected_idx)
plt.subplot2grid((1, 2), (0, 0))
plt.imshow(image)
# Annotations in red
draw_polys(img.polygons, color="red", label="annotations")
# Selected poly in blue
selected_polys = polys_to_draw if polys_to_draw else res.polygons
draw_polys(selected_polys, color="blue", label="selected polygons (> %0.1f)" % threshold, selected_idx=selected_idx)
plt.legend(loc='upper left')
# Scale
cax = fig.add_axes(
[heat_ax.get_position().x1 + 0.01, heat_ax.get_position().y0, 0.02, heat_ax.get_position().height])
cbar = plt.colorbar(heat_plt, cax=cax) # Similar to fig.colorbar(im, cax = cax)
cbar.set_label(r'polygon annotation consensus (-)', fontsize=12)
plt.show()
# Only return matches
if len(res.polygons) == 0:
return None
else:
return res
if __name__ == '__main__':
threshold_arg = Arg('--threshold', '-t', type=float, default=DEFAULT_THRESHOLD,
help="Threshold value as fraction of number of actors. %.2f by default" % DEFAULT_THRESHOLD)
image_type = Arg('--image-type', '-it', choices=[IMAGE_TYPE_POLY, IMAGE_TYPE_THRES, IMAGE_TYPE_ALL], default=IMAGE_TYPE_THRES,
help="Type of output images. "
"'polygon' : Outputs binary image of best polygon. "
"'threshold (default)' : Outputs binary image of threshold (before detection of polygon). "
"'all' : Outputs both raw level of detection in gray and final polygon in red.")
main_process(process_img, [threshold_arg, image_type])