-
Notifications
You must be signed in to change notification settings - Fork 0
/
segmentation.py
138 lines (103 loc) · 4.27 KB
/
segmentation.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
"""
@class: segmentation.py
@author: Julian Blos
"""
from classification import classify
import numpy as np
import pyttsx3
import cv2
# Setting up the text-to-speech engine
tts_engine = pyttsx3.init()
tts_rate = tts_engine.getProperty('rate')
# Open the image from the testing dataset
image = cv2.imread('testing/IMAG111.BMP')
# Show the image from the testing dataset
cv2.imshow("Original", image)
# Convert the image to grayscale
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# Apply a distortion filter (median blur)
blurred = cv2.medianBlur(gray, 7)
# Show the blurred image
cv2.imshow("Blurred", blurred)
# Apply a manual threshold of 90
(T, thresh) = cv2.threshold(blurred, 90, 255, cv2.THRESH_BINARY)
# Show the thresholded image
cv2.imshow("Thresholded", thresh)
# Labeling the connected components within the frame
# Get the connected components of the thresholded image
components = cv2.connectedComponentsWithStats(thresh, 4, cv2.CV_32S)
# Number of labels of the connected components
num_labels = components[0]
# Labels of the connected components
labels = components[1]
# Stats of the connected components
stats = components[2]
# Centroids of the connected components
centroids = components[3]
# For every label in the image
for label in range(num_labels):
# Ignore the first label as this corresponds to the image background
if label == 0:
continue
# Dimensions for the Region of Interest (ROI)
# Get the width of the bounding box of the connected component
width = stats[label, cv2.CC_STAT_WIDTH]
# Get the height of the bounding box of the connected component
height = stats[label, cv2.CC_STAT_HEIGHT]
# Get the leftmost coordinate of the bounding box of the connected component
x = stats[label, cv2.CC_STAT_LEFT]
# Get the upmost coordinate of the bounding box of the connected component
y = stats[label, cv2.CC_STAT_TOP]
# Creating the ROI of the connected component
roi = thresh[y:y + height, x:x + width]
# Calculating the image moments of the ROI
img_moments = cv2.moments(roi)
# Calculating the Hu moment invariants of the ROI
hu = cv2.HuMoments(img_moments)
# Obtaining the first Hu moment invariant
firstHu = -np.sign(hu[0]) * np.log10(np.abs(hu[0]))
# Get the class of the corresponding connected component
obj_class = (classify(np.round(firstHu, 4)))
# Types of object classes
# Associate a color and a label to the object class of the connected component
if obj_class == 1:
classLabel = "screw"
color = (255, 0, 0) # the color for the contour of the object class
elif obj_class == 2:
classLabel = "washer"
color = (0, 255, 0)
elif obj_class == 3:
classLabel = "screw eye"
color = (0, 0, 255)
elif obj_class == 4:
classLabel = "tenterhook"
color = (0, 255, 255)
else:
classLabel = "dovetail"
color = (255, 255, 0)
# Dimensions for labeling from the coordinates of the centroid
# of the component in X and Y.
# Calculating the coordinate in X of the centroid of the component
cX = int(centroids[label, 0])
# Calculating the coordinate in Y of the centroid of the component
cY = int(centroids[label, 1])
# Add the corresponding label according to the class of the component
cv2.putText(image, classLabel, (cX, cY), cv2.FONT_HERSHEY_SIMPLEX,
0.5, color, 2)
# Finding the contours of the ROI
contours = cv2.findContours(roi.copy(), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
# Prepare the contours as required by OpenCV 4
contours = contours[0]
# Iterate according to the number of contours
for c in contours:
# Draw the contours for the objects in the image
cv2.drawContours(image, [c], -1, color, 2, offset=(x, y))
# Show the image from the testing dataset correctly classified
cv2.imshow("Result", image)
# Perform the speech by saying the corresponding label for the object
# pyttsx3 speech temporarily disabled due to a bug in MacOS, see:
# https://stackoverflow.com/questions/27338298/workaround-for-pyttsx-engine-runandwait-not-returning-on-yosemite
# tts_engine.say(classLabel)
# tts_engine.runAndWait()
cv2.waitKey(0)