forked from sightmachine/SimpleCV
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathDiffSegmentation.py
More file actions
146 lines (123 loc) · 4.16 KB
/
Copy pathDiffSegmentation.py
File metadata and controls
146 lines (123 loc) · 4.16 KB
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
from SimpleCV.base import *
from SimpleCV.Features import Feature, FeatureSet, BlobMaker
from SimpleCV.ImageClass import Image
from SimpleCV.Segmentation.SegmentationBase import SegmentationBase
class DiffSegmentation(SegmentationBase):
"""
This method will do image segmentation by looking at the difference between
two frames.
grayOnly - use only gray images.
threshold - The value at which we consider the color difference to
be significant enough to be foreground imagery.
The general usage is
>>> segmentor = DiffSegmentation()
>>> cam = Camera()
>>> while(1):
>>> segmentor.addImage(cam.getImage())
>>> if(segmentor.isReady()):
>>> img = segmentor.getSegmentedImage()
"""
mError = False
mLastImg = None
mCurrImg = None
mDiffImg = None
mColorImg = None
mGrayOnlyMode = True
mThreshold = 10
mBlobMaker = None
def __init__(self, grayOnly=False, threshold = (10,10,10) ):
self.mGrayOnlyMode = grayOnly
self.mThreshold = threshold
self.mError = False
self.mCurrImg = None
self.mLastImg = None
self.mDiffImg = None
self.mColorImg = None
self.mBlobMaker = BlobMaker()
def addImage(self, img):
"""
Add a single image to the segmentation algorithm
"""
if( img is None ):
return
if( self.mLastImg == None ):
if( self.mGrayOnlyMode ):
self.mLastImg = img.toGray()
self.mDiffImg = Image(self.mLastImg.getEmpty(1))
self.mCurrImg = None
else:
self.mLastImg = img
self.mDiffImg = Image(self.mLastImg.getEmpty(3))
self.mCurrImg = None
else:
if( self.mCurrImg is not None ): #catch the first step
self.mLastImg = self.mCurrImg
if( self.mGrayOnlyMode ):
self.mColorImg = img
self.mCurrImg = img.toGray()
else:
self.mColorImg = img
self.mCurrImg = img
cv.AbsDiff(self.mCurrImg.getBitmap(),self.mLastImg.getBitmap(),self.mDiffImg.getBitmap())
return
def isReady(self):
"""
Returns true if the camera has a segmented image ready.
"""
if( self.mDiffImg is None ):
return False
else:
return True
def isError(self):
"""
Returns true if the segmentation system has detected an error.
Eventually we'll consruct a syntax of errors so this becomes
more expressive
"""
return self.mError #need to make a generic error checker
def resetError(self):
"""
Clear the previous error.
"""
self.mError = False
return
def reset(self):
"""
Perform a reset of the segmentation systems underlying data.
"""
self.mCurrImg = None
self.mLastImg = None
self.mDiffImg = None
def getRawImage(self):
"""
Return the segmented image with white representing the foreground
and black the background.
"""
return self.mDiffImg
def getSegmentedImage(self, whiteFG=True):
"""
Return the segmented image with white representing the foreground
and black the background.
"""
retVal = None
if( whiteFG ):
retVal = self.mDiffImg.binarize(thresh=self.mThreshold)
else:
retVal = self.mDiffImg.binarize(thresh=self.mThreshold).invert()
return retVal
def getSegmentedBlobs(self):
"""
return the segmented blobs from the fg/bg image
"""
retVal = []
if( self.mColorImg is not None and self.mDiffImg is not None ):
retVal = self.mBlobMaker.extractFromBinary(self.mDiffImg.binarize(thresh=self.mThreshold),self.mColorImg)
return retVal
def __getstate__(self):
mydict = self.__dict__.copy()
self.mBlobMaker = None
del mydict['mBlobMaker']
return mydict
def __setstate__(self, mydict):
self.__dict__ = mydict
self.mBlobMaker = BlobMaker()