cam = Camera()
path = "./data/"
ext = ".png"
suit_ptr = 0
rank_ptr = 0
current_dir = ""
allDone = False
for s in SUITS:
for r in RANKS:
directory = path + s + "/" + r + "/"
if not os.path.exists(directory):
os.makedirs(directory)
print "Current Data: " + str(RANKS[rank_ptr]) + str(SUITS[suit_ptr])
while not allDone:
keys = disp.checkEvents()
img = cam.getImage()
for k in keys:
if (k == pg.K_SPACE):
directory = path + SUITS[suit_ptr] + "/" + RANKS[rank_ptr] + "/"
files = glob.glob(directory + "*.*")
count = len(files)
fname = directory + RANKS[rank_ptr] + SUITS[suit_ptr] + "-" + str(
count) + ext
img.save(fname)
print "Saved: " + fname
if (k == pg.K_n):
if rank_ptr == len(RANKS) - 1 and suit_ptr == len(SUITS) - 1:
allDone = True
elif rank_ptr == len(RANKS) - 1:
rank_ptr = 0
def turk(self,
saveOriginal=False,
disp_size=(800, 600),
showKeys=True,
font_size=16,
color=Color.RED,
spacing=10):
"""
**SUMMARY**
This function does the turning of the data. The method goes through each image,
applies the preprocessing (which can return multiple images), displays each image
with an optional display of the key mapping. The user than selects the key that describes
the class of the image. The image is then post processed and saved to the directory.
The escape key kills the turking, the space key skips an image.
**PARAMETERS**
* *saveOriginal* - if true save the original image versus the preprocessed image.
* *disp_size* - size of the display to create.
* *showKeys* - Show the key mapping for the turking. Note that on small images this may not render correctly.
* *font_size* - the font size for the turking display.
* *color* - the font color.
* *spacing* - the spacing between each line of text on the display.
**RETURNS**
Nothing but stores each image in the directory. The image sets are also available
via the getClass method.
**EXAMPLE**
>>>> def GetBlobs(img):
>>>> blobs = img.findBlobs()
>>>> return [b.mMask for b in blobs]
>>>> def ScaleIng(img):
>>>> return img.resize(100,100).invert()
>>>> turker = TurkingModule(['./data/'],['./turked/'],['apple','banana','cherry'],['a','b','c'],preProcess=GetBlobs,postProcess=ScaleInv]
>>>> turker.turk()
>>>> # ~~~ stuff ~~~
>>>> turker.save('./derp.pkl')
** TODO **
TODO: fix the display so that it renders correctly no matter what the image size.
TODO: Make it so you can stop and start turking at any given spot in the process
"""
disp = Display(disp_size)
bail = False
for img in self.srcImgs:
print img.filename
samples = self.preProcess(img)
for sample in samples:
if (showKeys):
sample = self._drawControls(sample, font_size, color,
spacing)
sample.save(disp)
gotKey = False
while (not gotKey):
keys = disp.checkEvents(True)
for k in keys:
if k in self.keyMap:
if saveOriginal:
self._saveIt(img, self.keyMap[k])
else:
self._saveIt(sample, self.keyMap[k])
gotKey = True
if k == 'space':
gotKey = True # skip
if k == 'escape':
return
class Gui:
"""
Default layer groups available to display
First element is the 'base' layer, and must be an image,
the rest must be features and will be drawn on top
"""
layersets = {
'default': ['raw', 'yellow', 'blue', 'ball'],
'yellow': ['threshY', 'yellow'],
'blue': ['threshB', 'blue'],
'ball': ['threshR', 'ball']
}
def __init__(self, size=(720, 540)):
self._layers = {
# Base layers
'raw': None,
'threshY': None,
'threshB': None,
'threshR': None,
# Overlay layers
'yellow': None,
'blue': None,
'ball' : None,
}
# These layers are drawn regardless of the current layerset
self._persistentLayers = {
'mouse': None
}
self._currentLayerset = self.layersets['default']
self._display = Display(size)
self._eventHandler = Gui.EventHandler()
self._lastMouseState = 0
self._showMouse = True
self._lastFrame = None
self._lastFrameTime = time.time()
def __draw(self):
iterator = iter(self._currentLayerset)
# First element is the base layer
baseLayer = self._layers[iterator.next()]
if baseLayer is None:
return
size = baseLayer.size()
# Draw all entities to one layer (for speed)
entityLayer = baseLayer.dl()
for key in iterator:
toDraw = self._layers[key]
if toDraw is None:
continue
elif isinstance(toDraw, DrawingLayer):
baseLayer.addDrawingLayer(toDraw)
else:
toDraw.draw(entityLayer)
for layer in self._persistentLayers.itervalues():
if layer is not None:
baseLayer.addDrawingLayer(layer)
finalImage = baseLayer.applyLayers()
self._display.writeFrame(finalImage, fit=False)
def __updateFps(self):
smoothConst = 0.1
thisFrameTime = time.time()
thisFrame = thisFrameTime - self._lastFrameTime
if self._lastFrame is not None:
# Smooth values
thisFrame = thisFrame * (1 - smoothConst) + smoothConst * self._lastFrame
fps = 1.0 / thisFrame
self._lastFrame = thisFrame
self._lastFrameTime = thisFrameTime
layer = self._layers['raw'].dl()
layer.ezViewText('{0:.1f} fps'.format(fps), (10, 10))
def drawCrosshair(self, pos, layerName = None):
size = self._layers['raw'].size()
if layerName is not None:
layer = self.getDrawingLayer()
else:
layer = self._layers['raw'].dl()
layer.line((0, pos[1]), (size[0], pos[1]), color=(0, 0, 255))
layer.line((pos[0], 0), (pos[0], size[1]), color=(0, 0, 255))
if layerName is not None:
self.updateLayer(layerName, layer)
def loop(self):
"""
Draw the image to the display, and process any events
"""
for event in pygame.event.get(pygame.KEYDOWN):
self._eventHandler.processKey(chr(event.key % 0x100))
self._display.checkEvents()
mouseX = self._display.mouseX
mouseY = self._display.mouseY
if self._showMouse:
self.drawCrosshair((mouseX, mouseY), 'mouse')
mouseLeft = self._display.mouseLeft
# Only fire click event once for each click
if mouseLeft == 1 and self._lastMouseState == 0:
self._eventHandler.processClick((mouseX, mouseY))
self._lastMouseState = mouseLeft
# Processing OpenCV events requires calling cv.WaitKey() with a reasonable timeout,
# which hits our framerate hard (NOTE: Need to confirm this on DICE), so only do
# this if the focus isn't on the pygame (image) window`
if not pygame.key.get_focused():
c = cv.WaitKey(2)
self._eventHandler.processKey(chr(c % 0x100))
self.__updateFps()
self.__draw()
def getEventHandler(self):
return self._eventHandler
def getDrawingLayer(self):
return DrawingLayer(self._layers['raw'].size())
def updateLayer(self, name, layer):
"""
Update the layer specified by 'name'
If the layer name is not in the known list of layers,
then it will be drawn regardless of the current view setting
"""
if name in self._layers.keys():
self._layers[name] = layer
else:
self._persistentLayers[name] = layer
def switchLayerset(self, name):
assert name in self.layersets.keys(), 'Unknown layerset ' + name + '!'
self._currentLayerset = self.layersets[name]
def setShowMouse(self, showMouse):
if not showMouse:
self.updateLayer('mouse', None)
self._showMouse = showMouse
class EventHandler:
def __init__(self):
self._listeners = {}
self._clickListener = None
def processKey(self, key):
if key in self._listeners.keys():
self._listeners[key]()
def processClick(self, where):
if self._clickListener is not None:
self._clickListener(where)
def addListener(self, key, callback):
"""
Adds a function callback for a key.
"""
assert callable(callback), '"callback" must be callable'
self._listeners[key] = callback
def setClickListener(self, callback):
"""
Sets a function to be called on clicking on the image.
The function will be passed a tuple with the (x,y) of the click.
Setting a new callback will override the last one (or pass None to clear)
"""
assert callback is None or callable(callback), '"callback" must be callable'
self._clickListener = callback
def turk(self,saveOriginal=False,disp_size=(800,600),showKeys=True,font_size=16,color=Color.RED,spacing=10 ):
"""
**SUMMARY**
This function does the turning of the data. The method goes through each image,
applies the preprocessing (which can return multiple images), displays each image
with an optional display of the key mapping. The user than selects the key that describes
the class of the image. The image is then post processed and saved to the directory.
The escape key kills the turking, the space key skips an image.
**PARAMETERS**
* *saveOriginal* - if true save the original image versus the preprocessed image.
* *disp_size* - size of the display to create.
* *showKeys* - Show the key mapping for the turking. Note that on small images this may not render correctly.
* *font_size* - the font size for the turking display.
* *color* - the font color.
* *spacing* - the spacing between each line of text on the display.
**RETURNS**
Nothing but stores each image in the directory. The image sets are also available
via the getClass method.
**EXAMPLE**
>>>> def GetBlobs(img):
>>>> blobs = img.findBlobs()
>>>> return [b.mMask for b in blobs]
>>>> def ScaleIng(img):
>>>> return img.resize(100,100).invert()
>>>> turker = TurkingModule(['./data/'],['./turked/'],['apple','banana','cherry'],['a','b','c'],preProcess=GetBlobs,postProcess=ScaleInv]
>>>> turker.turk()
>>>> # ~~~ stuff ~~~
>>>> turker.save('./derp.pkl')
** TODO **
TODO: fix the display so that it renders correctly no matter what the image size.
TODO: Make it so you can stop and start turking at any given spot in the process
"""
disp = Display(disp_size)
bail = False
for img in self.srcImgs:
print img.filename
samples = self.preProcess(img)
for sample in samples:
if( showKeys ):
sample = self._drawControls(sample,font_size,color,spacing )
sample.save(disp)
gotKey = False
while( not gotKey ):
keys = disp.checkEvents(True)
for k in keys:
if k in self.keyMap:
if saveOriginal:
self._saveIt(img,self.keyMap[k])
else:
self._saveIt(sample,self.keyMap[k])
gotKey = True
if k == 'space':
gotKey = True # skip
if k == 'escape':
return
class Gui:
_layer_sets = { 'default': ['raw', 'robot0', 'robot1', 'robot2', 'robot3', 'ball'],
'squares1': ['squares', 'robot0', 'robot1', 'robot2', 'robot3', 'ball'],
'squares2': ['squares', 'robot0', 'robot1', 'robot2', 'robot3', 'ball'],
'squares3': ['squares', 'robot0', 'robot1', 'robot2', 'robot3', 'ball'],
'squares4': ['squares', 'robot0', 'robot1', 'robot2', 'robot3', 'ball'],
'ball': ['threshR', 'ball'] }
_layers = { 'raw': None,
'threshR': None,
'ball': None,
'robot0': None,
'robot1': None,
'robot2': None,
'robot3': None,
'squares': None }
_persistent_layers = { 'mouse': None }
def __init__(self, size=(720, 540)):
self._current_layer_set = self._layer_sets['default']
self._display = Display(size)
self._event_handler = Gui.EventHandler()
self._last_mouse_state = 0
self._show_mouse = True
self._last_frame = None
self._last_frame_time = time.time()
def __draw(self):
iterator = iter(self._current_layer_set)
base_layer = self._layers[iterator.next()]
if base_layer is None:
return
size = base_layer.size()
entity_layer = base_layer.dl()
for key in iterator:
to_draw = self._layers[key]
if isinstance(to_draw, DrawingLayer):
base_layer.addDrawingLayer(to_draw)
elif not to_draw is None:
to_draw.draw(entity_layer)
for layer in self._persistent_layers.itervalues():
if layer is not None:
base_layer.addDrawingLayer(layer)
final_image = base_layer.applyLayers()
self._display.writeFrame(final_image, fit=False)
def __update_fps(self):
this_frame_time = time.time()
this_frame = this_frame_time - self._last_frame_time
fps = 1.0 / this_frame
self._lastFrame = this_frame
self._last_frame_time = this_frame_time
layer = self._layers[self._current_layer_set[0]].dl()
layer.ezViewText('{0:.1f} fps'.format(fps), (10, 10))
def draw_crosshair(self, pos, layer_name = None):
size = self._layers['raw'].size()
if layer_name is not None:
layer = DrawingLayer(self._layers['raw'].size())
else:
layer = self._layers['raw'].dl()
layer.line((0, pos[1]), (size[0], pos[1]), color=(0, 0, 255))
layer.line((pos[0], 0), (pos[0], size[1]), color=(0, 0, 255))
if layer_name is not None:
self.update_layer(layer_name, layer)
def process_update(self):
"""Draw the image to the display, and process any events
"""
for event in pygame.event.get(pygame.KEYDOWN):
self._event_handler.process_key(chr(event.key % 0x100))
self._display.checkEvents()
mouseX = self._display.mouseX
mouseY = self._display.mouseY
if self._show_mouse:
self.draw_crosshair((mouseX, mouseY), 'mouse')
mouse_left = self._display.mouseLeft
# Only fire click event once for each click
if mouse_left == 1 and self._last_mouse_state == 0:
self._event_handler.process_click((mouseX, mouseY))
self._last_mouse_state = mouse_left
# Processing OpenCV events requires calling cv.WaitKey() with a reasonable timeout,
# which hits our framerate hard (NOTE: Need to confirm this on DICE), so only do
# this if the focus isn't on the pygame (image) window
if not pygame.key.get_focused():
c = cv.WaitKey(2)
self._event_handler.process_key(chr(c % 0x100))
self.__update_fps()
self.__draw()
def get_event_handler(self):
return self._event_handler
def update_layer(self, name, layer):
"""Update the layer specified by 'name'
If the layer name is not in the known list of layers,
then it will be drawn regardless of the current view setting
"""
if name in self._layers.keys():
self._layers[name] = layer
else:
self._persistent_layers[name] = layer
def switch_layer_set(self, name):
assert name in self._layer_sets.keys(), 'Unknown layerset ' + name + '!'
self._current_layer_set = self._layer_sets[name]
def set_show_mouse(self, show_mouse):
if not show_mouse:
self.update_layer('mouse', None)
self._show_mouse = show_mouse
class EventHandler:
def __init__(self):
self._listeners = {}
self._click_listener = None
def process_key(self, key):
if key in self._listeners.keys():
self._listeners[key]()
def process_click(self, where):
if self._click_listener is not None:
self._click_listener(where)
def add_listener(self, key, callback):
"""Adds a function callback for a key.
"""
assert callable(callback), '"callback" must be callable'
self._listeners[key] = callback
def set_click_listener(self, callback):
"""Sets a function to be called on clicking on the image.
The function will be passed a tuple with the (x, y) of the click.
Setting a new callback will override the last one (or pass None to clear)
"""
assert callback is None or callable(callback), '"callback" must be callable'
self._click_listener = callback
cam = Camera()
path = "./data/"
ext = ".png"
suit_ptr = 0
rank_ptr = 0
current_dir = ""
allDone = False
for s in SUITS:
for r in RANKS:
directory = path+s+"/"+r+"/"
if not os.path.exists(directory):
os.makedirs(directory)
print "Current Data: " + str(RANKS[rank_ptr])+str(SUITS[suit_ptr])
while not allDone :
keys = disp.checkEvents()
img = cam.getImage()
for k in keys:
if( k == pg.K_SPACE ):
directory = path+SUITS[suit_ptr]+"/"+RANKS[rank_ptr]+"/"
files = glob.glob(directory+"*.*")
count = len(files)
fname = directory+RANKS[rank_ptr]+SUITS[suit_ptr]+"-"+str(count)+ext
img.save(fname)
print "Saved: " + fname
if( k == pg.K_n ):
if rank_ptr == len(RANKS)-1 and suit_ptr == len(SUITS)-1:
allDone = True
elif rank_ptr == len(RANKS)-1:
rank_ptr = 0
suit_ptr = suit_ptr + 1
class Gui:
"""
Default layer groups available to display
First element is the 'base' layer, and must be an image,
the rest must be features and will be drawn on top
"""
layersets = {
'default': ['raw', 'yellow', 'blue', 'ball'],
'yellow': ['threshY', 'yellow'],
'blue': ['threshB', 'blue'],
'ball': ['threshR', 'ball']
}
h_min = s_min = v_min = 255
h_max = s_max = v_max = 0
# Show pixel mode: turn on/off by 'p', show the HSV of the pixel at the mouse
showPixel = False
# Record pixel mode: use the current pixel to update threshold
recordPixel = False
def __init__(self, noGui):
self._layers = {
# Base layers
'raw': None,
'threshY': None,
'threshB': None,
'threshR': None,
# Overlay layers
'yellow': None,
'blue': None,
'ball': None,
}
# These layers are drawn regardless of the current layerset
self._persistentLayers = {
'mouse': None
}
self._currentLayerset = self.layersets['default']
self._display = Display()
self._eventHandler = Gui.EventHandler()
self._lastMouseState = 0
self._showMouse = True
self._lastFrame = None
self._lastFrameTime = time.time()
def __draw(self):
iterator = iter(self._currentLayerset)
# First element is the base layer
baseLayer = self._layers[iterator.next()]
if baseLayer is None:
return
# Draw all entities to one layer (for speed)
entityLayer = baseLayer.dl()
for key in iterator:
toDraw = self._layers[key]
if toDraw is None:
continue
elif isinstance(toDraw, DrawingLayer):
baseLayer.addDrawingLayer(toDraw)
else:
toDraw.draw(entityLayer)
for layer in self._persistentLayers.itervalues():
if layer is not None:
baseLayer.addDrawingLayer(layer)
finalImage = baseLayer.applyLayers()
self._display.writeFrame(finalImage, fit=False)
def __updateFps(self):
smoothConst = 0.1
thisFrameTime = time.time()
thisFrame = thisFrameTime - self._lastFrameTime
if self._lastFrame is not None:
# Smooth values
thisFrame = thisFrame * (1 - smoothConst) + smoothConst * self._lastFrame
self._lastFrame = thisFrame
self._lastFrameTime = thisFrameTime
def drawCrosshair(self, pos, layerName=None):
size = self._layers['raw'].size()
if layerName is not None:
layer = self.getDrawingLayer()
else:
layer = self._layers['raw'].dl()
layer.line((0, pos[1]), (size[0], pos[1]), color=(0, 0, 255))
layer.line((pos[0], 0), (pos[0], size[1]), color=(0, 0, 255))
if layerName is not None:
self.updateLayer(layerName, layer)
def loop(self):
"""
Draw the image to the display, and process any events
"""
for event in pygame.event.get(pygame.KEYDOWN):
self._eventHandler.processKey(chr(event.key % 0x100))
self._display.checkEvents()
mouseX = self._display.mouseX
mouseY = self._display.mouseY
if self._showMouse:
self.drawCrosshair((mouseX, mouseY), 'mouse')
mouseLeft = self._display.mouseLeft
# Only fire click event once for each click
if mouseLeft == 1 and self._lastMouseState == 0:
self._eventHandler.processClick((mouseX, mouseY))
self._lastMouseState = mouseLeft
# Processing OpenCV events requires calling cv.WaitKey() with a reasonable timeout,
# which hits our framerate hard (NOTE: Need to confirm this on DICE), so only do
# this if the focus isn't on the pygame (image) window`
if not pygame.key.get_focused():
c = cv.WaitKey(2)
self._eventHandler.processKey(chr(c % 0x100))
if self.showPixel:
self.__updatePixel()
self.__updateFps()
self.__draw()
def getEventHandler(self):
return self._eventHandler
def getDrawingLayer(self):
return DrawingLayer(self._layers['raw'].size())
def updateLayer(self, name, layer):
"""
Update the layer specified by 'name'
If the layer name is not in the known list of layers,
then it will be drawn regardless of the current view setting
"""
if name in self._layers.keys():
self._layers[name] = layer
else:
self._persistentLayers[name] = layer
def switchLayerset(self, name):
assert name in self.layersets.keys(), 'Unknown layerset ' + name + '!'
self._currentLayerset = self.layersets[name]
def setShowMouse(self, showMouse):
if not showMouse:
self.updateLayer('mouse', None)
self._showMouse = showMouse
def toggleShowPixel(self):
self.showPixel = not self.showPixel
def toggleRecordPixel(self):
self.recordPixel = not self.recordPixel
def getMousePos(self):
return (self._display.mouseX, self._display.mouseY)
def getPixelHSV(self, x, y):
return self._layers['raw'].crop(x, y, 1, 1).toHSV()[0, 0]
def getPixel(self, x, y):
return self._layers['raw'][x, y]
#check (x,y) is in the image range
def inRange(self, x, y):
h = self._layers['raw'].height
w = self._layers['raw'].width
return (0 < x and x < w and 0 < y and y < h)
# Display the pixel HSV
def __updatePixel(self):
(x, y) = self.getMousePos()
if self.inRange(x, y):
(h, s, v) = self.getPixelHSV(x, y)
else:
return
if self.recordPixel:
self.updateMinMax(h, s, v)
# drawingLayer = self.getDrawingLayer()
# drawingLayer.text('Pixel ({0}, {1}) HSV = ({2}, {3}, {4})'.format(x,y,h,s,v),(10,10), fgcolor(255, 255, 255))
# drawingLayer.ezViewText('HSV_min = ({0}, {1}, {2}) '.format(self.h_min, self.s_min, self.v_min),(10,30))
# drawingLayer.ezViewText('HSV_max = ({0}, {1}, {2}) '.format(self.h_max, self.s_max, self.v_max),(10,50))
print ('Pixel ({0}, {1}) HSV = ({2}, {3}, {4})'.format(x, y, h, s, v))
print ('HSV_min = ({0}, {1}, {2}) '.format(self.h_min, self.s_min, self.v_min))
print ('HSV_max = ({0}, {1}, {2}) '.format(self.h_max, self.s_max, self.v_max))
if self.recordPixel:
print 'Recording Pixel'
#drawingLayer.ezViewText('Recording Pixel',(10,70))
#print 'pixel ',(x,y), ' HSV = ', (h,s,v), ' RGB = ', (r,g,b),
#print 'HSV_min = ', (self.h_min, self.s_min, self.v_min),
#print 'HSV_max = ', (self.h_max, self.s_max, self.v_max)
def resetMinMax(self):
self.h_min = self.s_min = self.v_min = 255
self.h_max = self.s_max = self.v_max = 0
# use the given HSV to update threshold
def updateMinMax(self, h, s, v):
self.h_min = min(self.h_min, h)
self.s_min = min(self.s_min, s)
self.v_min = min(self.v_min, v)
self.h_max = max(self.h_max, h)
self.s_max = max(self.s_max, s)
self.v_max = max(self.v_max, v)
class EventHandler:
def __init__(self):
self._listeners = {}
self._clickListener = None
def processKey(self, key):
if key in self._listeners.keys():
self._listeners[key]()
def processClick(self, where):
if self._clickListener is not None:
self._clickListener(where)
def addListener(self, key, callback):
"""
Adds a function callback for a key.
"""
assert callable(callback), '"callback" must be callable'
self._listeners[key] = callback
def setClickListener(self, callback):
"""
Sets a function to be called on clicking on the image.
The function will be passed a tuple with the (x,y) of the click.
Setting a new callback will override the last one (or pass None to clear)
"""
assert callback is None or callable(callback), '"callback" must be callable'
self._clickListener = callback
