-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathsegmentationDataset.py
148 lines (134 loc) · 6.9 KB
/
segmentationDataset.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
import numpy as np
import cv2
from matplotlib import pyplot as plt
import preObj as prePro
import detectorDescriptor2 as detDes
import glob
import os
import time
rootInputName = "FinalScenarioDataset/"
rootOutputName = "FinalScenarioDataset/BoundingBoxes/"
formatName = "*.jpg"
listDir = [];
#listDir.append('kettle/')
listDir.append("basket/")
listDir.append("kettle/")
#listDir.append("TrainingSetOld/apple/green/")
#listDir.append("TrainingSetOld/apple/red/")
#listDir.append("TrainingSetOld/banana/green/")
#listDir.append("TrainingSetOld/banana/yellow/")
#listDir.append("TrainingSetOld/cube/blue/")
#listDir.append("TrainingSetOld/cube/yellow/")
#listDir.append("TrainingSetOld/phone/green/")
#listDir.append("TrainingSetOld/phone/yellow/")
#listDir.append("TrainingSetNao/banana/red/")
#listDir.append("TrainingSetNao/banana/yellow/")
#listDir.append("TrainingSetNao/cube/yellow/")
#listDir.append("TrainingSetNao/cube/red/")
#listDir.append("TrainingSetNao/phone/green/")
#listDir.append("TrainingSetNao/phone/red_blinds_closed/")
#listDir.append("TrainingSetNao/phone/red_blinds_open/")
#listDir.append("TrainingSetNew/apple/")
#listDir.append("TrainingSetNew/banana/")
#listDir.append("TrainingSetNew/cube/")
#listDir.append("TrainingSetOccMul/")
dirCount = 0
for direct in listDir:
fileList = glob.glob(rootInputName + listDir[dirCount] + formatName)
print ("Current Directory Name:" + rootInputName + listDir[dirCount])
count = 0
for files in fileList:
frame=cv2.imread(fileList[count])
base = os.path.basename(files)
fileName = os.path.splitext(base)[0]
print ("Processing Image " + fileList[count])
time_whole_begin = time.time()
grayScaleInput = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# Apply meanshift on the RGB image
time_meanshift_begin = time.time()
meanShiftResult = prePro.meanShift(frame)
time_meanshift_end = time.time()
# Convert the result of the Mean shifted image into Grayscale
meanShiftGray = cv2.cvtColor(meanShiftResult, cv2.COLOR_BGR2GRAY)
# Apply adaptive thresholding on the resulting greyscale image
meanShiftAdapResult = prePro.adapThresh(meanShiftGray)
'''########################################################################################################################
Uncomment the following lines to add morphology
Comment the following lines if morphology need not be used
############################################## Morphology begins ##########################################################################'''
# kernel = np.ones((5,5),np.uint8)
# opening = cv2.dilate(meanShiftAdapResult,kernel,iterations=1)
# # Draw Contours on the Input image with results from the meanshift
# contourPlot = prePro.contourDraw(frame, opening)
# cv2.imwrite('contourplot.png',contourPlot)
# # Find the contours on the mean shifted image
# contours, hierarchy = prePro.contourFindFull(opening)
# # Draw Contours on the Input image with results from the meanshift
'''##############################################Morphology ends ##########################################################################
Comment the following lines to add morphology
Uncomment them if no morphology needs to be used
######################################### No Morphology begins ########################################################################'''
contourPlot = prePro.contourDraw(frame, meanShiftAdapResult)
#cv2.imshow('contourplot',contourPlot)
# Find the contours on the mean shifted image
contours, hierarchy = prePro.contourFindFull(meanShiftAdapResult)
'''######################################## No Morphology ends ########################################################################'''
## Use Histogram equalization
boundBoxContour = grayScaleInput.copy()
counter = 0
for cnt in contours:
if cv2.contourArea(cnt)>500:
print("Processing Contour no.", str(counter))
[x, y, w, h] = cv2.boundingRect(cnt)
extendBBox20 = 20
extendBBox10 = 10
left = 0
right = 0
top = 0
bottom = 0
#
#Extend it by 10 pixels to avoid missing the key points on the edges
if x-extendBBox20 > 0:
left = x-extendBBox20
elif x-extendBBox10 > 0:
left = x-extendBBox10
else:
left = x
if y-extendBBox20 > 0:
top = y-extendBBox20
elif y-extendBBox10 > 0:
top = y-extendBBox10
else:
top = y
if x+w+extendBBox20 < boundBoxContour.shape[0]:
right = x+w+extendBBox20
elif x+w+extendBBox10 < boundBoxContour.shape[0]:
right = x+w+extendBBox10
else:
right = x+w
if y+h+extendBBox20 < boundBoxContour.shape[1]:
bottom = y+h+extendBBox20
elif y+h+extendBBox10 < boundBoxContour.shape[1]:
bottom = y+h+extendBBox10
else:
bottom = y+h
roiImage = boundBoxContour[top:bottom,left:right]
#roiImage = boundBoxContour[y-extendBBox:y+h+extendBBox, x-extendBBox:x+w+extendBBox]
roiImageFiltered = roiImage #cv2.medianBlur(roiImage, 3)
kp, roiKeyPointImage = detDes.featureDetectCorner(roiImageFiltered)
kp, des, roiKeyPointImage = detDes.featureDescriptorORB(roiImageFiltered, kp)
if np.size(kp)>0:
# print 'directory number ' + str(dirCount*10000)
# print 'file number ' + str(count*1000)
# print 'contour number ' + str(counter)
# print 'resulting image number ' + str(dirCount*10000+ count*1000 + counter)
cv2.imwrite(rootOutputName+listDir[dirCount]+fileName + str(counter) + ".png", roiImage)
cv2.imwrite(rootOutputName+listDir[dirCount]+fileName + str(counter) + "kp" + ".png", roiKeyPointImage)
print ("Path: " +rootOutputName + str(dirCount*1000+ count*100 + counter*10) + ".png")
print ("Found some non-zero keypoints for the countours.")
counter = counter + 1
count = count + 1
time_whole_image_end = time.time()
print ("Time spent on meanshift for this image is" + str(time_meanshift_end - time_meanshift_begin) + "secs" )
print ("Time spent on the entire processing is" + str(time_whole_image_end - time_whole_begin) + "secs")
dirCount = dirCount + 1