class CocaineController(ArenaController.ArenaController):
def __init__(self, parent, arenaMain):
super(CocaineController, self).__init__(parent, arenaMain)
#calibration
self.arenaCamCorners = []
self.arenaMidLine = []
self.arenaSide1Sign = 1
self.arenaProjCorners = []
self.fishImg = None #image of fish
#tracking
self.arenaCvMask = None
#state
self.mutex = Lock()
self.arenaData = None
self.currState = State.OFF
self.currRun = -1 #a run is started each time Start Switches is pressed
self.currSwitch = 0
self.lastSwitchTime = 0
self.fishPosUpdate = False
self.fishPos = (0,0)
self.allFish = []
self.fishSize = None
#tracking related images
self.currCvFrame = None
#init UI
#arena info group box
self.arenaGroup = QtGui.QGroupBox(self)
self.arenaGroup.setTitle('Arena Info')
self.arenaLayout = QtGui.QGridLayout()
self.arenaLayout.setHorizontalSpacing(3)
self.arenaLayout.setVerticalSpacing(3)
self.camCalibButton = QtGui.QPushButton('Set Cam Position')
self.camCalibButton.setMaximumWidth(150)
self.camCalibButton.clicked.connect(self.getArenaCameraPosition)
self.projGroup = QtGui.QGroupBox(self.arenaGroup)
self.projGroup.setTitle('Projector Position')
self.projLayout = QtGui.QGridLayout(self.projGroup)
self.projLayout.setHorizontalSpacing(3)
self.projLayout.setVerticalSpacing(3)
self.projCalibButton = QtGui.QPushButton('Calibrate Projector Position')
self.projCalibButton.setCheckable(True)
self.projCalibButton.clicked.connect(self.updateProjectorDisplay)
self.projLayout.addWidget(self.projCalibButton,0,0,1,6)
self.proj1L = QtGui.QLabel('C1')
self.proj1X = QtGui.QSpinBox()
self.proj1X.setRange(0,1000)
self.proj1X.setValue(0)
self.proj1X.setMaximumWidth(50)
self.proj1Y = QtGui.QSpinBox()
self.proj1Y.setValue(0)
self.proj1Y.setMaximumWidth(50)
self.proj1Y.setRange(0,1000)
self.projLayout.addWidget(self.proj1L,1,0)
self.projLayout.addWidget(self.proj1X,1,1)
self.projLayout.addWidget(self.proj1Y,1,2)
self.proj1X.valueChanged.connect(self.updateProjectorDisplay)
self.proj1Y.valueChanged.connect(self.updateProjectorDisplay)
self.proj2L = QtGui.QLabel('C2')
self.proj2X = QtGui.QSpinBox()
self.proj2X.setRange(0,1000)
self.proj2X.setValue(0)
self.proj2X.setMaximumWidth(50)
self.proj2Y = QtGui.QSpinBox()
self.proj2Y.setRange(0,1000)
self.proj2Y.setValue(400)
self.proj2Y.setMaximumWidth(50)
self.proj2X.valueChanged.connect(self.updateProjectorDisplay)
self.proj2Y.valueChanged.connect(self.updateProjectorDisplay)
self.projLayout.addWidget(self.proj2L,2,0)
self.projLayout.addWidget(self.proj2X,2,1)
self.projLayout.addWidget(self.proj2Y,2,2)
self.proj3L = QtGui.QLabel('C3')
self.proj3X = QtGui.QSpinBox()
self.proj3X.setRange(0,1000)
self.proj3X.setValue(848)
self.proj3X.setMaximumWidth(50)
self.proj3Y = QtGui.QSpinBox()
self.proj3Y.setRange(0,1000)
self.proj3Y.setValue(400)
self.proj3Y.setMaximumWidth(50)
self.projLayout.addWidget(self.proj3L,2,3)
self.projLayout.addWidget(self.proj3X,2,4)
self.projLayout.addWidget(self.proj3Y,2,5)
self.proj3X.valueChanged.connect(self.updateProjectorDisplay)
self.proj3Y.valueChanged.connect(self.updateProjectorDisplay)
self.proj4L = QtGui.QLabel('C4')
self.proj4X = QtGui.QSpinBox()
self.proj4X.setRange(0,1000)
self.proj4X.setValue(848)
self.proj4X.setMaximumWidth(50)
self.proj4Y = QtGui.QSpinBox()
self.proj4Y.setRange(0,1000)
self.proj4Y.setValue(0)
self.proj4Y.setMaximumWidth(50)
self.projLayout.addWidget(self.proj4L,1,3)
self.projLayout.addWidget(self.proj4X,1,4)
self.projLayout.addWidget(self.proj4Y,1,5)
self.proj4X.valueChanged.connect(self.updateProjectorDisplay)
self.proj4Y.valueChanged.connect(self.updateProjectorDisplay)
self.projLayout.setColumnStretch(6,1)
self.projGroup.setLayout(self.projLayout)
self.arenaLayout.addWidget(self.camCalibButton, 0,0)
self.arenaLayout.addWidget(self.projGroup,1,0)
self.arenaLayout.setColumnStretch(1,1)
self.arenaGroup.setLayout(self.arenaLayout)
#tracking group box
self.trackWidget = FishTrackerWidget(self, self.arenaMain.ftDisp)
self.startButton = QtGui.QPushButton('Start Switches')
self.startButton.setMaximumWidth(150)
self.startButton.setCheckable(True)
self.startButton.clicked.connect(self.startSwitches)
#experimental parameters groupbox
self.paramGroup = QtGui.QGroupBox()
self.paramGroup.setTitle('Exprimental Parameters')
self.paramLayout = QtGui.QGridLayout()
self.paramLayout.setHorizontalSpacing(2)
self.paramLayout.setVerticalSpacing(2)
self.paramNumSwitches = LabeledSpinBox(None,'NumSwitches',0,100,1,60)
self.paramLayout.addWidget(self.paramNumSwitches,1,0,1,2)
self.paramSwitchTime = LabeledSpinBox(None,'SwitchTime (s)',1,3600,300,60)
self.paramLayout.addWidget(self.paramSwitchTime,1,2,1,2)
self.paramCond = QtGui.QComboBox()
self.paramCond.addItem('Pre Train')
self.paramCond.addItem('Cocaine')
self.paramCond.addItem('Neutral')
self.paramCond.addItem('Post Train')
self.paramCond.setCurrentIndex(0)
self.labelCond = QtGui.QLabel('Condition')
self.paramLayout.addWidget(self.paramCond,2,0)
self.paramLayout.addWidget(self.labelCond,2,1)
self.paramColor1 = QtGui.QComboBox()
self.paramColor1.addItem('White')
self.paramColor1.addItem('Red')
self.paramColor1.addItem('Blue')
self.paramColor1.addItem('Gray')
self.paramColor1.setCurrentIndex(1)
self.labelColor1 = QtGui.QLabel('Color1')
self.paramLayout.addWidget(self.paramColor1,2,2)
self.paramLayout.addWidget(self.labelColor1,2,3)
self.paramColor2 = QtGui.QComboBox()
self.paramColor2.addItem('White')
self.paramColor2.addItem('Red')
self.paramColor2.addItem('Blue')
self.paramColor2.addItem('Gray')
self.paramColor2.setCurrentIndex(2)
self.labelColor2 = QtGui.QLabel('Color2')
self.paramLayout.addWidget(self.paramColor2,3,0)
self.paramLayout.addWidget(self.labelColor2,3,1)
self.paramOffColor = QtGui.QComboBox()
self.paramOffColor.addItem('White')
self.paramOffColor.addItem('Red')
self.paramOffColor.addItem('Blue')
self.paramOffColor.addItem('Gray')
self.paramOffColor.setCurrentIndex(0)
self.labelOffColor = QtGui.QLabel('OffColor')
self.paramLayout.addWidget(self.paramOffColor,3,2)
self.paramLayout.addWidget(self.labelOffColor,3,3)
self.paramDrugColor = QtGui.QComboBox()
self.paramDrugColor.addItem('White')
self.paramDrugColor.addItem('Red')
self.paramDrugColor.addItem('Blue')
self.paramDrugColor.addItem('Gray')
self.paramDrugColor.setCurrentIndex(0)
self.labelDrugColor = QtGui.QLabel('DrugColor')
self.paramLayout.addWidget(self.paramDrugColor,4,0)
self.paramLayout.addWidget(self.labelDrugColor,4,1)
self.paramConc = QtGui.QDoubleSpinBox()
self.paramConc.setRange(0.0,1000.0)
self.paramConc.setValue(10.0)
self.labelConc = QtGui.QLabel('Conc mg/L')
self.paramLayout.addWidget(self.paramConc,5,0)
self.paramLayout.addWidget(self.labelConc,5,1)
self.paramNumFish = LabeledSpinBox(None,'NumFish',1,10,1,60)
self.paramLayout.addWidget(self.paramNumFish,5,2,1,2)
self.paramGroup.setLayout(self.paramLayout)
#Experimental info group
self.infoGroup = QtGui.QGroupBox()
self.infoGroup.setTitle('Experiment Info')
self.infoLayout = QtGui.QGridLayout()
self.infoLayout.setHorizontalSpacing(3)
self.infoLayout.setVerticalSpacing(3)
self.labelDir = QtGui.QLabel('Dir: ')
self.infoDir = PathSelectorWidget(browseCaption='Experimental Data Directory')
self.infoLayout.addWidget(self.labelDir,0,0)
self.infoLayout.addWidget(self.infoDir,0,1)
self.labelDOB = QtGui.QLabel('DOB: ')
self.infoDOB = QtGui.QDateEdit(QtCore.QDate.currentDate())
self.infoLayout.addWidget(self.labelDOB,1,0)
self.infoLayout.addWidget(self.infoDOB,1,1)
self.labelType = QtGui.QLabel('Line: ')
self.infoType = QtGui.QLineEdit('Nacre')
self.infoLayout.addWidget(self.labelType,2,0)
self.infoLayout.addWidget(self.infoType,2,1)
self.infoFish = QtGui.QPushButton('Snap Fish Image')
self.infoFish.clicked.connect(self.getFishSize)
self.infoLayout.addWidget(self.infoFish,3,0,1,2)
self.infoGroup.setLayout(self.infoLayout)
self.settingsLayout = QtGui.QGridLayout()
self.settingsLayout.addWidget(self.arenaGroup,0,0)
self.settingsLayout.addWidget(self.trackWidget,1,0)
self.settingsLayout.addWidget(self.startButton,2,0)
self.settingsLayout.addWidget(self.paramGroup,3,0)
self.settingsLayout.addWidget(self.infoGroup,4,0)
self.setLayout(self.settingsLayout)
#---------------------------------------------------
# OVERLOADED METHODS
#---------------------------------------------------
#update to overload an return image based on dispmode.
#def getArenaView(self):
# return self.arenaView
def onNewFrame(self, frame, time):
self.mutex.acquire()
try:
self.frameTime = time
self.currCvFrame = frame # may need to make a deep copy
(found, pos, view, allFish) = self.trackWidget.findFish(self.currCvFrame)
if found:
self.fishPosUpdate = found
self.fishPos = pos
self.allFish = allFish
if not self.currState == State.OFF:
self.arenaData['runs'][self.currRun]['video'].append((self.frameTime, None))
# if self.paramNumFish.value() == 1:
# self.arenaData['runs'][self.currRun]['tracking'].append((self.frameTime, self.fishPos[0], self.fishPos[1]))
# else:
d = [self.frameTime, pos[0], pos[1]]
for nFish in range(1, min(self.paramNumFish.value(),len(allFish))):
d.append(allFish[nFish][0])
d.append(allFish[nFish][1])
# print d
self.arenaData['runs'][self.currRun]['tracking'].append(tuple(d))
self.arenaView = view
except:
print 'CocaineController:onNewFrame failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def updateState(self):
self.updateExperimentalState()
def updateExperimentalState(self):
if not self.currState == State.OFF:
self.mutex.acquire()
try:
t = time.time()
#only update status bar if arena is selected.
if self.isCurrent():
self.arenaMain.statusBar().showMessage('SwitchesCompleted: %d TimeSinceSwitch: %f' % (self.currSwitch, t - self.lastSwitchTime))
#handle State Changes
if t >= self.lastSwitchTime + self.paramSwitchTime.value():
self.currSwitch += 1
self.lastSwitchTime = t
if self.currSwitch > self.paramNumSwitches.value():
self.currState = State.OFF
self.arenaData['runs'][self.currRun]['endTime'] = t
self.startButton.setText('Start Switches')
self.startButton.setChecked(False)
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
self.updateProjectorDisplay()
self.arenaData['runs'][self.currRun]['switchTimes'].append(t)
self.saveResults()
except:
print 'CocaineController:updateState failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def isReadyToStart(self):
return os.path.exists(self.infoDir.text()) and self.trackWidget.getBackgroundImage() and self.fishImg and self.arenaCamCorners
def drawProjectorDisplay(self, painter):
if self.currState == State.OFF and self.projCalibButton.isChecked():
brush = QtGui.QBrush(QtCore.Qt.blue)
pen = QtGui.QPen(QtCore.Qt.black)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
else:
if self.currState == State.OFF:
#Draw whole tank
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(self.paramOffColor.currentIndex())
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
else:
side1Color = self.paramColor1.currentIndex()
side2Color = self.paramColor2.currentIndex()
if self.currSwitch%2 == 1:
side1Color = self.paramColor2.currentIndex()
side2Color = self.paramColor1.currentIndex()
a = .5
b = 1-a
#Draw side one
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(side1Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(a*self.proj2X.value() + b*self.proj3X.value(), a*self.proj2Y.value() + b*self.proj3Y.value()))
poly.append(QtCore.QPointF(a*self.proj1X.value() + b*self.proj4X.value(), a*self.proj1Y.value() + b*self.proj4Y.value()))
painter.drawPolygon(poly)
#Draw side two
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(side2Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(a*self.proj1X.value() + b*self.proj4X.value(), a*self.proj1Y.value() + b*self.proj4Y.value()))
poly.append(QtCore.QPointF(a*self.proj2X.value() + b*self.proj3X.value(), a*self.proj2Y.value() + b*self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
def drawDisplayOverlay(self, painter):
#draw the fish position
if self.fishPos and self.fishPos[0] > 0:
brush = QtGui.QBrush(QtCore.Qt.red)
painter.setBrush(brush)
painter.setPen(QtCore.Qt.NoPen)
for nFish in range(min(len(self.allFish),self.paramNumFish.value())):
painter.drawEllipse(QtCore.QPointF(self.allFish[nFish][0],self.allFish[nFish][1]), 3,3)
#draw the arena overlay
if self.arenaCamCorners:
painter.setBrush(QtCore.Qt.NoBrush)
if self.bIsCurrentArena:
pen = QtGui.QPen(QtCore.Qt.green)
else:
pen = QtGui.QPen(QtCore.Qt.blue)
pen.setWidth(3)
painter.setPen(pen)
poly = QtGui.QPolygonF()
for p in self.arenaCamCorners:
poly.append(QtCore.QPointF(p[0],p[1]))
painter.drawPolygon(poly)
def start(self):
#Global start button not used.
pass
def stop(self):
#Global stop button not used.
pass
#---------------------------------------------------
# CALLBACK METHODS
#---------------------------------------------------
def startSwitches(self):
if self.currState == State.OFF:
if self.isReadyToStart():
self.paramGroup.setDisabled(True)
self.infoGroup.setDisabled(True)
if self.arenaData == None or not str(self.infoDir.text()) == self.saveLocation:
td = datetime.datetime.now()
self.saveLocation = str(self.infoDir.text())
[p, self.fnResults] = os.path.split(self.saveLocation)
self.fnResults = self.fnResults + '_' + td.strftime('%Y-%m-%d-%H-%M-%S')
self.jsonFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '.json'
self.arenaData = {}
self.arenaData['runs'] = list()
self.initNewRunData()
t = time.time()
self.currState = State.RUNNING
self.lastSwitchTime = t
self.arenaData['runs'][self.currRun]['startTime'] = t
self.currSwitch = 0
self.startButton.setText('Stop')
else:
self.arenaMain.statusBar().showMessage('Arena not ready to start. Information is missing.')
else:
self.mutex.acquire()
try:
t = time.time()
self.currState = State.OFF
self.arenaData['runs'][self.currRun]['endTime'] = t
self.arenaData['runs'][self.currRun]['switchTimes'].append(t)
self.saveResults()
self.startButton.setText('Start Switches')
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
except:
print 'stop switches failed'
print "Unexpected error:", sys.exc_info()[0]
finally:
self.mutex.release()
self.updateProjectorDisplay()
def getArenaCameraPosition(self):
self.arenaMain.statusBar().showMessage('Click on the corners of the arena on side 1.')
self.currArenaClick = 0
self.arenaCamCorners = []
self.arenaMain.ftDisp.clicked.connect(self.handleArenaClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleArenaClicks(self, x, y):
self.currArenaClick+=1
if self.currArenaClick<5:
self.arenaCamCorners.append((x,y))
if self.currArenaClick==1:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 1.')
elif self.currArenaClick==2:
self.arenaMain.statusBar().showMessage('Now, click on the corners of the arena on side 2.')
elif self.currArenaClick==3:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 2.')
elif self.currArenaClick==4:
self.arenaMain.ftDisp.clicked.disconnect(self.handleArenaClicks)
self.arenaMain.statusBar().showMessage('')
[self.arenaMidLine, self.arenaSide1Sign] = self.processArenaCorners(self.arenaCamCorners, .5)
self.getArenaMask()
def getFishSize(self):
self.arenaMain.statusBar().showMessage('Click on the tip of the fish tail.')
self.currFishClick = 0
self.fishSize = []
self.arenaMain.ftDisp.clicked.connect(self.handleFishClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleFishClicks(self, x, y):
self.currFishClick+=1
if self.currFishClick == 1:
self.fishSize.append((x,y))
self.arenaMain.statusBar().showMessage('Click on the tip of the fish head.')
elif self.currFishClick == 2:
self.arenaMain.ftDisp.clicked.disconnect(self.handleFishClicks)
self.fishSize.append((x,y))
self.fishImg = cv.CloneImage(self.currCvFrame)
self.arenaMain.statusBar().showMessage('')
#---------------------------------------------------
# HELPER METHODS
#---------------------------------------------------
#def getBackgroundImage(self):
# if self.currCvFrame:
# self.bcvImg = cv.CloneImage(self.currCvFrame)
# self.trackWidget.setBackgroundImage(self.bcvImg)
def processArenaCorners(self, arenaCorners, linePosition):
#return the line dividing the center of the arena, and a definition of side 1.
a = 1-linePosition
b = linePosition
ac = np.array(arenaCorners)
#arenaDivideLine = [tuple(np.mean(ac[(0,3),:],axis = 0)),tuple(np.mean(ac[(1,2),:],axis = 0))]
arenaDivideLine = [(a*ac[0,0]+b*ac[3,0], a*ac[0,1]+b*ac[3,1]),(a*ac[1,0]+b*ac[2,0], a*ac[1,1]+b*ac[2,1])]
side1Sign = 1
if not self.isOnSide(arenaCorners[1], arenaDivideLine, side1Sign):
side1Sign = -1
return (arenaDivideLine, side1Sign)
def isOnSide(self, point, line, sideSign):
#return if the fish is on side1 of the arena.
side = (line[1][0] - line[0][0]) * (point[1] - line[0][1]) - (line[1][1] - line[0][1]) * (point[0] - line[0][0])
return cmp(side,0)==sideSign
#convert the arena corners into a color mask image (arena=255, not=0)
def getArenaMask(self):
if self.arenaView:
cvImg = self.currCvFrame
self.arenaCvMask = cv.CreateImage((cvImg.width,cvImg.height), cvImg.depth, cvImg.channels)
cv.SetZero(self.arenaCvMask)
cv.FillConvexPoly(self.arenaCvMask, self.arenaCamCorners, (255,)*cvImg.channels)
self.trackWidget.setTrackMask(self.arenaCvMask)
def initNewRunData(self):
#prepare output data structure
self.arenaData['runs'].append({})
cr = self.currRun = len(self.arenaData['runs']) - 1
self.arenaData['runs'][cr]['fishbirthday'] = str(self.infoDOB.date().toPyDate())
self.arenaData['runs'][cr]['fishage'] = (datetime.date.today() - self.infoDOB.date().toPyDate()).days
self.arenaData['runs'][cr]['fishstrain'] = str(self.infoType.text())
self.arenaData['runs'][cr]['fishsize'] = self.fishSize
self.arenaData['runs'][cr]['parameters'] = { 'numSwitches':self.paramNumSwitches.value(),
'SwitchDuration':self.paramSwitchTime.value(),
'Cond':str(self.paramCond.currentText()),
'Conc':str(self.paramConc.value()),
'Color1':str(self.paramColor1.currentText()),
'Color2':str(self.paramColor2.currentText()),
'OffColor':str(self.paramOffColor.currentText()),
'DrugColor':str(self.paramDrugColor.currentText()),
'CodeVersion':None }
self.arenaData['runs'][cr]['trackingParameters'] = self.trackWidget.getParameterDictionary()
self.arenaData['runs'][cr]['trackingParameters']['arenaPoly'] = self.arenaCamCorners
self.arenaData['runs'][cr]['trackingParameters']['arenaDivideLine'] = self.arenaMidLine
self.arenaData['runs'][cr]['trackingParameters']['arenaSide1Sign'] = self.arenaSide1Sign
self.arenaData['runs'][cr]['tracking'] = list() #list of tuples (frametime, posx, posy)
self.arenaData['runs'][cr]['video'] = list() #list of tuples (frametime, filename)
self.arenaData['runs'][cr]['switchTimes'] = list() #list of times at switch stimulus flipped.
self.arenaData['runs'][cr]['startTime'] = None
self.arenaData['runs'][cr]['endTime'] = None
t = datetime.datetime.now()
#save experiment images
self.bcvImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_BackImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.bcvImgFileName, self.trackWidget.getBackgroundImage())
self.fishImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_FishImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.fishImgFileName, self.fishImg)
def saveResults(self):
f = open(name=self.jsonFileName, mode='w')
json.dump(self.arenaData,f)
f.close()
def getBrush(self, colorNdx):
if colorNdx == 0:
return QtGui.QBrush(QtCore.Qt.white)
elif colorNdx == 1:
return QtGui.QBrush(QtCore.Qt.red)
elif colorNdx == 2:
return QtGui.QBrush(QtCore.Qt.blue)
elif colorNdx == 3:
return QtGui.QBrush(QtGui.QColor(128,128,128))
else:
return QtGui.QBrush(QtCore.Qt.black)
class RealTimeShockController(ArenaController.ArenaController):
def __init__(self, parent, arenaMain):
super(RealTimeShockController, self).__init__(parent, arenaMain)
#calibration
self.arenaCamCorners = []
self.arenaMidLine = []
self.arenaSide1Sign = 1
self.arenaProjCorners = []
self.fishImg = None #image of fish
#tracking
self.arenaCvMask = None
#state
self.mutex = Lock()
self.arenaData = None
self.currState = State.OFF
self.fishPosUpdate = False
self.fishPos = (0,0)
self.allFish = []
self.fishSize = None
#tracking related images
self.currCvFrame = None
#init UI
#arena info group box
self.arenaGroup = QtGui.QGroupBox(self)
self.arenaGroup.setTitle('Arena Info')
self.arenaLayout = QtGui.QGridLayout()
self.arenaLayout.setHorizontalSpacing(3)
self.arenaLayout.setVerticalSpacing(3)
self.camCalibButton = QtGui.QPushButton('Set Cam Position')
self.camCalibButton.setMaximumWidth(150)
self.camCalibButton.clicked.connect(self.getArenaCameraPosition)
self.projGroup = QtGui.QGroupBox(self.arenaGroup)
self.projGroup.setTitle('Projector Position')
self.projLayout = QtGui.QGridLayout(self.projGroup)
self.projLayout.setHorizontalSpacing(3)
self.projLayout.setVerticalSpacing(3)
self.projCalibButton = QtGui.QPushButton('Calibrate Projector Position')
self.projCalibButton.setCheckable(True)
self.projCalibButton.clicked.connect(self.updateProjectorDisplay)
self.projLayout.addWidget(self.projCalibButton,0,0,1,6)
self.proj1L = QtGui.QLabel('C1')
self.proj1X = QtGui.QSpinBox()
self.proj1X.setRange(0,1000)
self.proj1X.setValue(0)
self.proj1X.setMaximumWidth(50)
self.proj1Y = QtGui.QSpinBox()
self.proj1Y.setValue(0)
self.proj1Y.setMaximumWidth(50)
self.proj1Y.setRange(0,1000)
self.projLayout.addWidget(self.proj1L,1,0)
self.projLayout.addWidget(self.proj1X,1,1)
self.projLayout.addWidget(self.proj1Y,1,2)
self.proj1X.valueChanged.connect(self.updateProjectorDisplay)
self.proj1Y.valueChanged.connect(self.updateProjectorDisplay)
self.proj2L = QtGui.QLabel('C2')
self.proj2X = QtGui.QSpinBox()
self.proj2X.setRange(0,1000)
self.proj2X.setValue(0)
self.proj2X.setMaximumWidth(50)
self.proj2Y = QtGui.QSpinBox()
self.proj2Y.setRange(0,1000)
self.proj2Y.setValue(400)
self.proj2Y.setMaximumWidth(50)
self.proj2X.valueChanged.connect(self.updateProjectorDisplay)
self.proj2Y.valueChanged.connect(self.updateProjectorDisplay)
self.projLayout.addWidget(self.proj2L,2,0)
self.projLayout.addWidget(self.proj2X,2,1)
self.projLayout.addWidget(self.proj2Y,2,2)
self.proj3L = QtGui.QLabel('C3')
self.proj3X = QtGui.QSpinBox()
self.proj3X.setRange(0,1000)
self.proj3X.setValue(848)
self.proj3X.setMaximumWidth(50)
self.proj3Y = QtGui.QSpinBox()
self.proj3Y.setRange(0,1000)
self.proj3Y.setValue(400)
self.proj3Y.setMaximumWidth(50)
self.projLayout.addWidget(self.proj3L,2,3)
self.projLayout.addWidget(self.proj3X,2,4)
self.projLayout.addWidget(self.proj3Y,2,5)
self.proj3X.valueChanged.connect(self.updateProjectorDisplay)
self.proj3Y.valueChanged.connect(self.updateProjectorDisplay)
self.proj4L = QtGui.QLabel('C4')
self.proj4X = QtGui.QSpinBox()
self.proj4X.setRange(0,1000)
self.proj4X.setValue(848)
self.proj4X.setMaximumWidth(50)
self.proj4Y = QtGui.QSpinBox()
self.proj4Y.setRange(0,1000)
self.proj4Y.setValue(0)
self.proj4Y.setMaximumWidth(50)
self.projLayout.addWidget(self.proj4L,1,3)
self.projLayout.addWidget(self.proj4X,1,4)
self.projLayout.addWidget(self.proj4Y,1,5)
self.proj4X.valueChanged.connect(self.updateProjectorDisplay)
self.proj4Y.valueChanged.connect(self.updateProjectorDisplay)
self.projLayout.setColumnStretch(6,1)
self.projGroup.setLayout(self.projLayout)
self.arenaLayout.addWidget(self.camCalibButton, 0,0)
self.arenaLayout.addWidget(self.projGroup,1,0)
self.arenaLayout.setColumnStretch(1,1)
self.arenaGroup.setLayout(self.arenaLayout)
#tracking group box
self.trackWidget = FishTrackerWidget(self, self.arenaMain.ftDisp)
self.startButton = QtGui.QPushButton('Start')
self.startButton.setMaximumWidth(150)
self.startButton.setCheckable(True)
self.startButton.clicked.connect(self.startSwitches)
#experimental parameters groupbox
self.paramGroup = QtGui.QGroupBox()
self.paramGroup.setTitle('Exprimental Parameters')
self.paramLayout = QtGui.QGridLayout()
self.paramLayout.setHorizontalSpacing(2)
self.paramLayout.setVerticalSpacing(2)
#parameters
self.paramAcclimate = LabeledSpinBox(None, 'Acclimate (m)', 0, 180, 30, 60)
self.paramLayout.addWidget(self.paramAcclimate, 0,0,1,2)
self.paramNumShockBlocks = LabeledSpinBox(None,'NumShockBlocks',0,100,4,60) #number of side switches pre test
self.paramLayout.addWidget(self.paramNumShockBlocks,0,2,1,2)
self.paramShockingTime = LabeledSpinBox(None,'ShockingTime (s)',1,3600,60,60) #duration before two sides swap colors during pre and post
self.paramLayout.addWidget(self.paramShockingTime,1,0,1,2)
self.paramBetweenTime = LabeledSpinBox(None,'BetweenTime (s)',0,3600,30,60) #time between pre , train and post periods
self.paramLayout.addWidget(self.paramBetweenTime,1,2,1,2)
self.paramShockPeriod = LabeledSpinBox(None, 'Shock Period (ms)', 0,5000,1000,60)
self.paramLayout.addWidget(self.paramShockPeriod, 2,0,1,2)
self.paramShockDuration = LabeledSpinBox(None, 'ShockDuration (ms)', 0,1000,50,60)
self.paramLayout.addWidget(self.paramShockDuration, 2,2,1,2)
self.paramShockChan1 = LabeledSpinBox(None, 'ShockChan1', 0,10000,12,60)
self.paramLayout.addWidget(self.paramShockChan1, 3,0,1,2)
self.paramShockChan2 = LabeledSpinBox(None, 'ShockChan2', 0,10000,13,60)
self.paramLayout.addWidget(self.paramShockChan2, 3,2,1,2)
self.paramCurrChan1 = LabeledSpinBox(None, 'CurrChan Side 1', 0,16,0,60)
self.paramLayout.addWidget(self.paramCurrChan1,4,0,1,2)
self.paramCurrChan2 = LabeledSpinBox(None, 'CurrChan Side 2', 0,16,0,60)
self.paramLayout.addWidget(self.paramCurrChan2,4,2,1,2)
self.paramShockV = LabeledSpinBox(None, 'ShockV', 0,100, 10,60)
self.paramLayout.addWidget(self.paramShockV,5,0,1,2)
self.paramNumFish = LabeledSpinBox(None,'NumFish',1,10,1,60)
self.paramLayout.addWidget(self.paramNumFish,5,2,1,2)
self.paramGroup.setLayout(self.paramLayout)
#Experimental info group
self.infoGroup = QtGui.QGroupBox()
self.infoGroup.setTitle('Experiment Info')
self.infoLayout = QtGui.QGridLayout()
self.infoLayout.setHorizontalSpacing(3)
self.infoLayout.setVerticalSpacing(3)
self.labelDir = QtGui.QLabel('Dir: ')
self.infoDir = PathSelectorWidget(browseCaption='Experimental Data Directory')
self.infoLayout.addWidget(self.labelDir,0,0)
self.infoLayout.addWidget(self.infoDir,0,1)
self.labelDOB = QtGui.QLabel('DOB: ')
self.infoDOB = QtGui.QDateEdit(QtCore.QDate.currentDate())
self.infoLayout.addWidget(self.labelDOB,1,0)
self.infoLayout.addWidget(self.infoDOB,1,1)
self.labelType = QtGui.QLabel('Line: ')
self.infoType = QtGui.QLineEdit('Nacre')
self.infoLayout.addWidget(self.labelType,2,0)
self.infoLayout.addWidget(self.infoType,2,1)
self.infoFish = QtGui.QPushButton('Snap Fish Image')
self.infoFish.clicked.connect(self.getFishSize)
self.infoLayout.addWidget(self.infoFish,3,0,1,2)
self.infoGroup.setLayout(self.infoLayout)
self.settingsLayout = QtGui.QGridLayout()
self.settingsLayout.addWidget(self.startButton,0,0,1,1)
self.settingsLayout.addWidget(self.arenaGroup,2,0,1,2)
self.settingsLayout.addWidget(self.trackWidget,3,0,1,2)
self.settingsLayout.addWidget(self.paramGroup,4,0,1,2)
self.settingsLayout.addWidget(self.infoGroup,5,0,1,2)
self.setLayout(self.settingsLayout)
#---------------------------------------------------
# OVERLOADED METHODS
#---------------------------------------------------
#update to overload an return image based on dispmode.
#def getArenaView(self):
# return self.arenaView
def onNewFrame(self, frame, time):
self.mutex.acquire()
try:
self.frameTime = time
self.currCvFrame = frame # may need to make a deep copy
(found, pos, view, allFish) = self.trackWidget.findFish(self.currCvFrame)
if found:
self.fishPosUpdate = found
self.fishPos = pos
self.allFish = allFish
if not self.currState == State.OFF:
self.arenaData['video'].append((self.frameTime, None))
d = [self.frameTime, pos[0], pos[1]]
for nFish in range(1, min(self.paramNumFish.value(),len(allFish))):
d.append(allFish[nFish][0])
d.append(allFish[nFish][1])
self.arenaData['tracking'].append(tuple(d))
self.arenaView = view
except:
print 'onNewFrame failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def updateState(self):
self.updateExperimentalState()
def updateExperimentalState(self):
if not self.currState == State.OFF:
self.mutex.acquire()
try:
t = time.time()
#only update status bar if arena is selected.
if self.isCurrent():
self.arenaMain.statusBar().showMessage('Running: Block:%d/%d, CurrState=%d, TimeToNextState=%f, Side=%d'%(self.numBlocks,self.paramNumShockBlocks.value(),self.currState,self.nextStateTime-t, self.currShockSide))
if t > self.nextStateTime:
self.setShockState(False, False)
if (self.currState == State.ACCLIMATE or
self.currState == State.BETWEEN or
self.paramBetweenTime.value() == 0):
if self.numBlocks < self.paramNumShockBlocks.value():
if self.currShockSide == 2:
self.currState = State.SHOCKING
self.setShockState(True,False)
self.arenaData['stateinfo'].append((t, self.currState, 'On','Off'))
self.currShockSide = 1
print 'SHOCKING SIDE1'
else:
self.currState = State.SHOCKING
self.setShockState(False,True)
self.arenaData['stateinfo'].append((t, self.currState, 'Off','On'))
self.currShockSide = 2
print 'SHOCKING SIDE2'
self.nextStateTime = t + self.paramShockingTime.value()
self.numBlocks+=1
else:
self.currState = State.OFF
self.arenaData['stateinfo'].append((t, self.currState, 'Off','Off'))
self.startButton.setText('Start Switches')
self.startButton.setChecked(False)
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
self.saveResults()
print 'DONE'
elif self.paramBetweenTime.value() > 0:
self.currState = State.BETWEEN
self.nextStateTime = t + self.paramBetweenTime.value()
self.arenaData['stateinfo'].append((t, self.currState, 'Off','Off'))
print 'BETWEEN'
self.updateProjectorDisplay()
self.saveResults()
except:
print ':updateState failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def isReadyToStart(self):
return os.path.exists(self.infoDir.text()) and self.trackWidget.getBackgroundImage() and self.fishImg and self.arenaCamCorners
def drawProjectorDisplay(self, painter):
if self.currState == State.OFF and self.projCalibButton.isChecked() and self.isCurrent():
a = .5
b = 1-a
#Draw side one
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(QtCore.Qt.red)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(a*self.proj2X.value() + b*self.proj3X.value(), a*self.proj2Y.value() + b*self.proj3Y.value()))
poly.append(QtCore.QPointF(a*self.proj1X.value() + b*self.proj4X.value(), a*self.proj1Y.value() + b*self.proj4Y.value()))
painter.drawPolygon(poly)
painter.drawText(self.proj1X.value(),self.proj1Y.value(),'1')
#Draw side two
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(QtCore.Qt.blue)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(a*self.proj1X.value() + b*self.proj4X.value(), a*self.proj1Y.value() + b*self.proj4Y.value()))
poly.append(QtCore.QPointF(a*self.proj2X.value() + b*self.proj3X.value(), a*self.proj2Y.value() + b*self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
def drawDisplayOverlay(self, painter):
#draw the fish position
if self.fishPos and self.fishPos[0] > 0:
brush = QtGui.QBrush(QtCore.Qt.red)
painter.setBrush(brush)
painter.setPen(QtCore.Qt.NoPen)
for nFish in range(min(len(self.allFish),self.paramNumFish.value())):
painter.drawEllipse(QtCore.QPointF(self.allFish[nFish][0],self.allFish[nFish][1]), 3,3)
#draw the arena overlay
if self.arenaCamCorners:
painter.setBrush(QtCore.Qt.NoBrush)
if self.bIsCurrentArena:
pen = QtGui.QPen(QtCore.Qt.green)
else:
pen = QtGui.QPen(QtCore.Qt.blue)
pen.setWidth(3)
painter.setPen(pen)
poly = QtGui.QPolygonF()
for p in self.arenaCamCorners:
poly.append(QtCore.QPointF(p[0],p[1]))
painter.drawPolygon(poly)
if len(self.arenaCamCorners) >= 2:
pen = QtGui.QPen(QtCore.Qt.red)
pen.setWidth(2)
painter.setPen(pen)
painter.drawLine(self.arenaCamCorners[0][0],self.arenaCamCorners[0][1],
self.arenaCamCorners[1][0],self.arenaCamCorners[1][1])
def start(self):
#Global start button not used.
pass
def stop(self):
#Global stop button not used.
pass
#---------------------------------------------------
# CALLBACK METHODS
#---------------------------------------------------
def startSwitches(self):
self.mutex.acquire()
try:
if self.currState == State.OFF:
if self.isReadyToStart():
self.paramGroup.setDisabled(True)
self.infoGroup.setDisabled(True)
td = datetime.datetime.now()
self.saveLocation = str(self.infoDir.text())
[p, self.fnResults] = os.path.split(self.saveLocation)
self.fnResults = self.fnResults + '_' + td.strftime('%Y-%m-%d-%H-%M-%S')
self.jsonFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '.json'
self.initArenaData()
t = time.time()
self.numBlocks = 0
self.currState = State.ACCLIMATE
self.nextStateTime = t + self.paramAcclimate.value()*60
self.setShockState(False,False)
self.currShockSide = random.choice([1,2])
self.arenaData['stateinfo'].append((t, self.currState, 'Off', 'Off'))
self.updateProjectorDisplay()
self.startButton.setText('Stop')
print 'ACCLIMATING'
else:
self.startButton.setChecked(False)
self.arenaMain.statusBar().showMessage('Arena not ready to start. Information is missing.')
else:
t = time.time()
self.currState = State.OFF
self.setShockState(False,False)
self.startButton.setText('Start Switches')
self.arenaData['stateinfo'].append((t, self.currState, 'Off', 'Off'))
self.updateProjectorDisplay()
self.saveResults()
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
except:
print ':startSwitches failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def getArenaCameraPosition(self):
self.arenaMain.statusBar().showMessage('Click on the corners of the arena on side 1.')
self.currArenaClick = 0
self.arenaCamCorners = []
self.arenaMain.ftDisp.clicked.connect(self.handleArenaClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleArenaClicks(self, x, y):
self.currArenaClick+=1
if self.currArenaClick<5:
self.arenaCamCorners.append((x,y))
if self.currArenaClick==1:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 1.')
elif self.currArenaClick==2:
self.arenaMain.statusBar().showMessage('Now, click on the corners of the arena on side 2.')
elif self.currArenaClick==3:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 2.')
elif self.currArenaClick==4:
self.arenaMain.ftDisp.clicked.disconnect(self.handleArenaClicks)
self.arenaMain.statusBar().showMessage('')
[self.arenaMidLine, self.arenaSide1Sign] = self.processArenaCorners(self.arenaCamCorners, .5)
self.getArenaMask()
def getFishSize(self):
self.arenaMain.statusBar().showMessage('Click on the tip of the fish tail.')
self.currFishClick = 0
self.fishSize = []
self.arenaMain.ftDisp.clicked.connect(self.handleFishClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleFishClicks(self, x, y):
self.currFishClick+=1
if self.currFishClick == 1:
self.fishSize.append((x,y))
self.arenaMain.statusBar().showMessage('Click on the tip of the fish head.')
elif self.currFishClick == 2:
self.arenaMain.ftDisp.clicked.disconnect(self.handleFishClicks)
self.fishSize.append((x,y))
self.fishImg = cv.CloneImage(self.currCvFrame)
self.arenaMain.statusBar().showMessage('')
#---------------------------------------------------
# HELPER METHODS
#---------------------------------------------------
#def getBackgroundImage(self):
# if self.currCvFrame:
# self.bcvImg = cv.CloneImage(self.currCvFrame)
# self.trackWidget.setBackgroundImage(self.bcvImg)
def processArenaCorners(self, arenaCorners, linePosition):
#return the line dividing the center of the arena, and a definition of side 1.
a = 1-linePosition
b = linePosition
ac = np.array(arenaCorners)
#arenaDivideLine = [tuple(np.mean(ac[(0,3),:],axis = 0)),tuple(np.mean(ac[(1,2),:],axis = 0))]
arenaDivideLine = [(a*ac[0,0]+b*ac[3,0], a*ac[0,1]+b*ac[3,1]),(a*ac[1,0]+b*ac[2,0], a*ac[1,1]+b*ac[2,1])]
side1Sign = 1
if not self.isOnSide(arenaCorners[1], arenaDivideLine, side1Sign):
side1Sign = -1
return (arenaDivideLine, side1Sign)
def isOnSide(self, point, line, sideSign):
#return if the fish is on side1 of the arena.
side = (line[1][0] - line[0][0]) * (point[1] - line[0][1]) - (line[1][1] - line[0][1]) * (point[0] - line[0][0])
return cmp(side,0)==sideSign
#convert the arena corners into a color mask image (arena=255, not=0)
def getArenaMask(self):
if self.arenaView:
cvImg = self.currCvFrame
self.arenaCvMask = cv.CreateImage((cvImg.width,cvImg.height), cvImg.depth, cvImg.channels)
cv.SetZero(self.arenaCvMask)
cv.FillConvexPoly(self.arenaCvMask, self.arenaCamCorners, (255,)*cvImg.channels)
self.trackWidget.setTrackMask(self.arenaCvMask)
def initArenaData(self):
#prepare output data structure
self.arenaData = {}
self.arenaData['fishbirthday'] = str(self.infoDOB.date().toPyDate())
self.arenaData['fishage'] = (datetime.date.today() - self.infoDOB.date().toPyDate()).days
self.arenaData['fishstrain'] = str(self.infoType.text())
self.arenaData['fishsize'] = self.fishSize
self.arenaData['parameters'] = { 'numShockBlocks':self.paramNumShockBlocks.value(),
'acclimate (m)':self.paramAcclimate.value(),
'between': self.paramBetweenTime.value(),
'shockingTime': self.paramShockingTime.value(),
'shock period (ms)':self.paramShockPeriod.value(),
'shock dura (ms)':self.paramShockDuration.value(),
'shock chan 1':self.paramShockChan1.value(),
'shock chan 2':self.paramShockChan2.value(),
'shock V':self.paramShockV.value(),
'numFish': self.paramNumFish.value(),
'CodeVersion':None }
self.arenaData['trackingParameters'] = self.trackWidget.getParameterDictionary()
self.arenaData['trackingParameters']['arenaPoly'] = self.arenaCamCorners
self.arenaData['trackingParameters']['arenaDivideLine'] = self.arenaMidLine
self.arenaData['trackingParameters']['arenaSide1Sign'] = self.arenaSide1Sign
self.arenaData['tracking'] = list() #list of tuples (frametime, posx, posy)
self.arenaData['video'] = list() #list of tuples (frametime, filename)
self.arenaData['stateinfo'] = list() #list of times at switch stimulus flipped.
self.arenaData['shockinfo'] = list()
t = datetime.datetime.now()
#save experiment images
self.bcvImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_BackImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.bcvImgFileName, self.trackWidget.getBackgroundImage())
self.fishImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_FishImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.fishImgFileName, self.fishImg)
def saveResults(self):
f = open(name=self.jsonFileName, mode='w')
json.dump(self.arenaData,f)
f.close()
def setShockState(self, bSide1, bSide2):
if not self.arenaMain.ard:
print 'WARNING: Arduino not connected'
return
if bSide1:
curr1 = self.arenaMain.ard.pinPulse(self.paramShockChan1.value(),
self.paramShockPeriod.value(),
self.paramShockDuration.value(),
feedbackPin = self.paramCurrChan1.value())
else:
self.arenaMain.ard.pinLow(self.paramShockChan1.value())
curr1 = 0
if bSide2:
curr2 = self.arenaMain.ard.pinPulse(self.paramShockChan2.value(),
self.paramShockPeriod.value(),
self.paramShockDuration.value(),
feedbackPin = self.paramCurrChan2.value())
else:
self.arenaMain.ard.pinLow(self.paramShockChan2.value())
curr2 = 0
self.arenaData['shockinfo'].append((time.time(), bSide1, bSide2, curr1, curr2))
class ContextualHelplessnessController(ArenaController.ArenaController):
def __init__(self, parent, arenaMain):
super(ContextualHelplessnessController, self).__init__(parent, arenaMain)
#calibration
self.arenaCamCorners = []
self.arenaMidLine = []
self.arenaSide1Sign = 1
self.arenaProjCorners = []
self.fishImg = None #image of fish
#tracking
self.arenaCvMask = None
#state
self.mutex = Lock()
self.arenaData = None
self.currState = State.OFF
self.fishPosUpdate = False
self.fishPos = (0,0)
self.allFish = []
self.fishSize = None
#tracking related images
self.currCvFrame = None
#init UI
#arena info group box
self.arenaGroup = QtGui.QGroupBox(self)
self.arenaGroup.setTitle('Arena Info')
self.arenaLayout = QtGui.QGridLayout()
self.arenaLayout.setHorizontalSpacing(3)
self.arenaLayout.setVerticalSpacing(3)
self.camCalibButton = QtGui.QPushButton('Set Cam Position')
self.camCalibButton.setMaximumWidth(150)
self.camCalibButton.clicked.connect(self.getArenaCameraPosition)
self.projGroup = QtGui.QGroupBox(self.arenaGroup)
self.projGroup.setTitle('Projector Position')
self.projLayout = QtGui.QGridLayout(self.projGroup)
self.projLayout.setHorizontalSpacing(3)
self.projLayout.setVerticalSpacing(3)
self.projCalibButton = QtGui.QPushButton('Calibrate Projector Position')
self.projCalibButton.setCheckable(True)
self.projCalibButton.clicked.connect(self.projectorPositionChanged)
self.projLayout.addWidget(self.projCalibButton,0,0,1,6)
self.projPosL = QtGui.QLabel('Pos')
self.projX = QtGui.QSpinBox()
self.projX.setRange(0,1000)
self.projX.setValue(250)
self.projX.setMaximumWidth(50)
self.projY = QtGui.QSpinBox()
self.projY.setRange(0,1000)
self.projY.setValue(250)
self.projY.setMaximumWidth(50)
self.projLayout.addWidget(self.projPosL,1,0)
self.projLayout.addWidget(self.projX,1,1)
self.projLayout.addWidget(self.projY,1,2)
self.projX.valueChanged.connect(self.projectorPositionChanged)
self.projY.valueChanged.connect(self.projectorPositionChanged)
self.projSizeL = QtGui.QLabel('Size L,W')
self.projLen = QtGui.QSpinBox()
self.projLen.setRange(0,1000)
self.projLen.setValue(220)
self.projLen.setMaximumWidth(50)
self.projWid = QtGui.QSpinBox()
self.projWid.setRange(0,1000)
self.projWid.setValue(115)
self.projWid.setMaximumWidth(50)
self.projLen.valueChanged.connect(self.projectorPositionChanged)
self.projWid.valueChanged.connect(self.projectorPositionChanged)
self.projLayout.addWidget(self.projSizeL,2,0)
self.projLayout.addWidget(self.projLen,2,1)
self.projLayout.addWidget(self.projWid,2,2)
self.projRotL = QtGui.QLabel('Rotation')
self.projRot = QtGui.QSpinBox()
self.projRot.setRange(0,360)
self.projRot.setValue(270)
self.projRot.setMaximumWidth(50)
self.projLayout.addWidget(self.projRotL,2,3)
self.projLayout.addWidget(self.projRot,2,4)
self.projRot.valueChanged.connect(self.projectorPositionChanged)
self.tankLength = LabeledSpinBox(None, 'Tank Len (mm)', 0,100,46,60)
self.projLayout.addWidget(self.tankLength, 3,0,1,2)
self.tankWidth = LabeledSpinBox(None, 'Tank Wid (mm)', 0,100,21,60)
self.projLayout.addWidget(self.tankWidth, 3,2,1,2)
self.projLayout.setColumnStretch(6,1)
self.projGroup.setLayout(self.projLayout)
self.arenaLayout.addWidget(self.camCalibButton, 0,0)
self.arenaLayout.addWidget(self.projGroup,1,0)
self.arenaLayout.setColumnStretch(1,1)
self.arenaGroup.setLayout(self.arenaLayout)
#tracking group box
self.trackWidget = FishTrackerWidget(self, self.arenaMain.ftDisp)
self.startButton = QtGui.QPushButton('Start Switches')
self.startButton.setMaximumWidth(150)
self.startButton.setCheckable(True)
self.startButton.clicked.connect(self.startSwitches)
self.pauseButton = QtGui.QPushButton('Pause')
self.pauseButton.setMaximumWidth(150)
self.pauseButton.setCheckable(True)
self.pauseButton.setDisabled(True)
self.pauseButton.clicked.connect(self.pause)
self.autoPauseCheckBox = QtGui.QCheckBox('AutoPause')
#experimental parameters groupbox
self.paramGroup = QtGui.QGroupBox()
self.paramGroup.setTitle('Exprimental Parameters')
self.paramLayout = QtGui.QGridLayout()
self.paramLayout.setHorizontalSpacing(2)
self.paramLayout.setVerticalSpacing(2)
#high level flow - acclimate -> pretest -> [b/t] -> train -> [b/t] -> posttest
self.paramAcclimate = LabeledSpinBox(None, 'Acclimate (m)', 0, 240, 0, 60)
self.paramLayout.addWidget(self.paramAcclimate, 0,0,1,2)
self.paramBetweenTime = LabeledSpinBox(None,'BetweenTime (m)',0, 1440, 0, 60) #time between pre , train and post periods
self.paramLayout.addWidget(self.paramBetweenTime,0,2,1,2)
self.paramPreDuration = LabeledSpinBox(None,'PreTest (m)',0,120,15,60) #number of side switches pre test
self.paramLayout.addWidget(self.paramPreDuration,1,0,1,2)
self.paramPreOMR = QtGui.QCheckBox('PreOMR')
self.paramPreOMR.setCheckState(2)
self.paramLayout.addWidget(self.paramPreOMR,1,2,1,2)
self.paramNumTrain = LabeledSpinBox(None,'NumTrain',0,500,30,60)
self.paramLayout.addWidget(self.paramNumTrain,2,0,1,2)
self.paramTrainOMR = QtGui.QCheckBox('TrainOMR')
self.paramTrainOMR.setCheckState(2)
self.paramLayout.addWidget(self.paramTrainOMR,2,2,1,2)
self.paramPostDuration = LabeledSpinBox(None,'PostTest (m)',0,120,5,60) #number of side switches post teset
self.paramLayout.addWidget(self.paramPostDuration,3,0,1,2)
self.paramPostOMR = QtGui.QCheckBox('PostOMR')
self.paramPostOMR.setCheckState(2)
self.paramLayout.addWidget(self.paramPostOMR,3,2,1,2)
#omr parameters
self.paramOMRInterval = LabeledSpinBox(None,'OMR Interval (s)',0,600,48,60) #time between OMR tests
self.paramOMRDuration = LabeledSpinBox(None,'OMR Duration (s)',0,600,12,60) #duration of OMR tests
self.paramOMRPeriod = LabeledSpinBox(None,'OMR Grating Period (mm)',0,50,5,60) #spacing between grating bars
self.paramOMRDutyCycle = LabeledSpinBox(None, 'OMR DutyCycle %',0,100,50,60) #width of grating bars
self.paramOMRVelocity = LabeledSpinBox(None, 'OMR Speed (mm/s)',0,50,5,60) #velocity of moving gratings.
self.paramLayout.addWidget(self.paramOMRInterval,4,0,1,2)
self.paramLayout.addWidget(self.paramOMRDuration,4,2,1,2)
self.paramLayout.addWidget(self.paramOMRPeriod,5,0,1,2)
self.paramLayout.addWidget(self.paramOMRDutyCycle,5,2,1,2)
self.paramLayout.addWidget(self.paramOMRVelocity,6,0,1,2)
#training parameters
self.paramUSTime = LabeledSpinBox(None, 'US Time (s)', 1,3600,60,60)
self.paramLayout.addWidget(self.paramUSTime,7,0,1,2)
self.paramNeutralMinTime = LabeledSpinBox(None, 'Neutral Min (s)',1,3600,2,60)
self.paramLayout.addWidget(self.paramNeutralMinTime, 8,0,1,2)
self.paramNeutralMaxTime = LabeledSpinBox(None, 'Neutral Max (s)',1,3600,5,60)
self.paramLayout.addWidget(self.paramNeutralMaxTime, 8,2,1,2)
self.paramShockPeriod = LabeledSpinBox(None, 'Shock Period (ms)', 0,5000,1000,60)
self.paramLayout.addWidget(self.paramShockPeriod, 9,0,1,2)
self.paramShockDuration = LabeledSpinBox(None, 'ShockDuration (ms)', 0,1000,50,60)
self.paramLayout.addWidget(self.paramShockDuration, 9,2,1,2)
self.paramShockChan1 = LabeledSpinBox(None, 'ShockChan1', 0,10000,53,60)
self.paramLayout.addWidget(self.paramShockChan1, 10,0,1,2)
self.paramShockChan2 = LabeledSpinBox(None, 'ShockChan2', 0,10000,52,60)
self.paramLayout.addWidget(self.paramShockChan2, 10,2,1,2)
self.paramCurrChan1 = LabeledSpinBox(None, 'CurrChan Side 1', 0,16,15,60)
self.paramLayout.addWidget(self.paramCurrChan1,11,0,1,2)
self.paramCurrChan2 = LabeledSpinBox(None, 'CurrChan Side 2', 0,16,14,60)
self.paramLayout.addWidget(self.paramCurrChan2,11,2,1,2)
self.paramShockV = LabeledSpinBox(None, 'ShockV', 0,100, 5,60)
self.paramLayout.addWidget(self.paramShockV, 12,0,1,2)
#Colors
(self.paramColorAcclimate,self.labelColorAcclimate) = self.getColorComboBox('AcclimateColor', 0)
self.paramLayout.addWidget(self.paramColorAcclimate,13,0)
self.paramLayout.addWidget(self.labelColorAcclimate,13,1)
(self.paramColorPreTest,self.labelColorPreTest) = self.getColorComboBox('PreTestColor', 0)
self.paramLayout.addWidget(self.paramColorPreTest,13,2)
self.paramLayout.addWidget(self.labelColorPreTest,13,3)
(self.paramColorTrain,self.labelColorTrain) = self.getColorComboBox('TrainColor', 0)
self.paramLayout.addWidget(self.paramColorTrain,14,0)
self.paramLayout.addWidget(self.labelColorTrain,14,1)
(self.paramColorPostTest,self.labelColorPostTest) = self.getColorComboBox('PostTestColor', 0)
self.paramLayout.addWidget(self.paramColorPostTest,14,2)
self.paramLayout.addWidget(self.labelColorPostTest,14,3)
(self.paramColorBetween,self.labelColorBetween) = self.getColorComboBox('BetweenColor', 0)
self.paramLayout.addWidget(self.paramColorBetween,15,0)
self.paramLayout.addWidget(self.labelColorBetween,15,1)
(self.paramColorOMR,self.labelColorOMR) = self.getColorComboBox('OMRColor', 5)
self.paramLayout.addWidget(self.paramColorOMR,15,2)
self.paramLayout.addWidget(self.labelColorOMR,15,3)
self.paramNumFish = LabeledSpinBox(None,'NumFish',1,10,1,60)
self.paramLayout.addWidget(self.paramNumFish,16,0,1,2)
self.paramGroup.setLayout(self.paramLayout)
#Experimental info group
self.infoGroup = QtGui.QGroupBox()
self.infoGroup.setTitle('Experiment Info')
self.infoLayout = QtGui.QGridLayout()
self.infoLayout.setHorizontalSpacing(3)
self.infoLayout.setVerticalSpacing(3)
self.labelDir = QtGui.QLabel('Dir: ')
self.infoDir = PathSelectorWidget(browseCaption='Experimental Data Directory')
self.infoLayout.addWidget(self.labelDir,0,0)
self.infoLayout.addWidget(self.infoDir,0,1)
self.labelDOB = QtGui.QLabel('DOB: ')
self.infoDOB = QtGui.QDateEdit(QtCore.QDate.currentDate())
self.infoLayout.addWidget(self.labelDOB,1,0)
self.infoLayout.addWidget(self.infoDOB,1,1)
self.labelType = QtGui.QLabel('Line: ')
self.infoType = QtGui.QLineEdit('Nacre')
self.infoLayout.addWidget(self.labelType,2,0)
self.infoLayout.addWidget(self.infoType,2,1)
self.infoFish = QtGui.QPushButton('Snap Fish Image')
self.infoFish.clicked.connect(self.getFishSize)
self.infoLayout.addWidget(self.infoFish,3,0,1,2)
self.infoGroup.setLayout(self.infoLayout)
self.settingsLayout = QtGui.QGridLayout()
self.settingsLayout.addWidget(self.arenaGroup,0,0,1,2)
self.settingsLayout.addWidget(self.trackWidget,1,0,1,2)
self.settingsLayout.addWidget(self.infoGroup,3,0,1,2)
self.settingsLayout.addWidget(self.startButton,4,0,1,1)
self.settingsLayout.addWidget(self.pauseButton,5,0,1,1)
self.settingsLayout.addWidget(self.autoPauseCheckBox,5,1,1,1)
self.settingsLayout.addWidget(self.paramGroup,6,0,1,2)
self.setLayout(self.settingsLayout)
self.projectorPositionChanged()
self.t = 0
# self.omrPhase = 0
# self.omrDirection = [1,0]
def getColorComboBox(self, labelName, defaultNdx=1):
colorBox = QtGui.QComboBox()
colorBox.addItem('White')
colorBox.addItem('Red')
colorBox.addItem('Blue')
colorBox.addItem('Gray')
colorBox.addItem('Magenta')
colorBox.addItem('Black')
colorBox.setCurrentIndex(defaultNdx)
colorLabel = QtGui.QLabel(labelName)
return (colorBox,colorLabel)
#---------------------------------------------------
# OVERLOADED METHODS
#---------------------------------------------------
#update to overload an return image based on dispmode.
#def getArenaView(self):
# return self.arenaView
def onNewFrame(self, frame, time):
self.mutex.acquire()
try:
self.frameTime = time
self.currCvFrame = frame # may need to make a deep copy
(found, pos, view, allFish) = self.trackWidget.findFish(self.currCvFrame)
if found:
self.fishPosUpdate = found
self.fishPos = pos
self.allFish = allFish
if not self.currState == State.OFF:
self.arenaData['video'].append((self.frameTime, None))
d = [self.frameTime, pos[0], pos[1]]
for nFish in range(1, min(self.paramNumFish.value(),len(allFish))):
d.append(allFish[nFish][0])
d.append(allFish[nFish][1])
self.arenaData['tracking'].append(tuple(d))
self.arenaView = view
except:
print 'ClassicalConditioningController:onNewFrame failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def updateState(self):
self.updateExperimentalState()
def startOMR(self):
self.nextOMRTime = self.t + self.paramOMRInterval.value()
self.bOMR = False
self.omrDirection = (0,0)
self.arenaData['OMRinfo'].append((self.t, self.bOMR, self.omrDirection))
def updateOMR(self):
if self.t > self.nextOMRTime :
if self.bOMR:
self.nextOMRTime = self.t + self.paramOMRInterval.value()
self.bOMR = False
self.omrDirection = (0,0)
else:
self.nextOMRTime = self.t + self.paramOMRDuration.value()
self.bOMR = True
self.omrPhase = 0
self.omrLastUpdate = self.t
if self.isOnSide(self.fishPos, self.arenaMidLine, self.arenaSide1Sign):
self.omrDirection = (1,0)
else:
self.omrDirection = (-1,0)
self.arenaData['OMRinfo'].append((self.t, self.bOMR, self.omrDirection))
self.updateProjectorDisplay()
def stopOMR(self):
self.nextOMRTime = float('inf')
self.bOMR = False
self.omrDirection = (0,0)
self.arenaData['OMRinfo'].append((self.t, self.bOMR, self.omrDirection))
def updateExperimentalState(self):
# if self.paramPreOMR.isChecked():
# t = time.time()
# self.lt = self.t
# self.t = t
# print self.t - self.lt
# self.omrPhase += self.omrDirection[0] * (float(self.paramOMRVelocity.value())/float(self.paramOMRPeriod.value())) * 360.0 * (self.t-self.lt)
# print self.omrPhase
# if t - self.omrLastUpdate > 0.010:
# #print t - self.omrLastUpdate
# self.updateProjectorDisplay()
if not self.currState == State.OFF:
self.mutex.acquire()
try:
t = time.time()
self.lt = self.t
self.t = t
#only update status bar if arena is selected.
if self.isCurrent():
self.arenaMain.statusBar().showMessage('State %d'%(self.currState))
if t > self.nextStateTime:
self.setShockState(False, False)
self.stopOMR()
self.bInShockBout = False
self.nextShockBoutTime = float('inf')
self.pauseButton.setDisabled(True)
if (self.currState <= State.ACCLIMATE
and self.paramPreDuration.value() > 0):
self.currState = State.PRETEST
self.nextStateTime = t + (self.paramPreDuration.value()*60)
if self.paramPreOMR.isChecked():
self.startOMR()
elif (self.currState <= State.PRETEST
and self.paramBetweenTime.value() > 0):
self.currState = State.PREDONE
self.nextStateTime = t + self.paramBetweenTime.value()*60
self.pauseButton.setDisabled(False)
if self.autoPauseCheckBox.isChecked():
self.pauseButton.setChecked(True)
self.doPause()
elif (self.currState <= State.PREDONE
and self.paramNumTrain.value() > 0):
self.currState = State.TRAIN
self.nextStateTime = float("inf")
self.numShock = 0;
self.bInShockBout = False
self.nextShockBoutTime = t + random.uniform(self.paramNeutralMinTime.value(), self.paramNeutralMaxTime.value())
if self.paramTrainOMR.isChecked():
self.startOMR()
elif (self.currState <= State.TRAIN
and self.paramBetweenTime.value() > 0):
self.currState = State.TRAINDONE
self.nextStateTime = t + self.paramBetweenTime.value()*60
self.pauseButton.setDisabled(False)
if self.autoPauseCheckBox.isChecked():
self.pauseButton.setChecked(True)
self.doPause()
elif (self.currState <= State.TRAINDONE
and self.paramPostDuration.value() > 0):
self.currState = State.POSTTEST
self.nextStateTime = t + (self.paramPostDuration.value()*60)
if self.paramPostOMR.isChecked():
self.startOMR()
elif self.currState <= State.POSTTEST:
self.currState = State.OFF
self.startButton.setText('Start Switches')
self.startButton.setChecked(False)
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
self.saveResults()
self.arenaData['stateinfo'].append((t, self.currState, self.getSide1ColorName(), self.getSide2ColorName()))
self.updateProjectorDisplay()
self.saveResults()
#handle omr state changes and projector updates
self.updateOMR()
if self.bOMR:
self.omrPhase += self.omrDirection[0] * (float(self.paramOMRVelocity.value())/float(self.paramOMRPeriod.value())) * 360.0 * (self.t-self.lt)
self.omrPhase = self.omrPhase%360.0
if t - self.omrLastUpdate > 0.010:
#if t - self.omrLastUpdate > 0.040:
# print 'WARNING: projector slow to update', t - self.omrLastUpdate
self.updateProjectorDisplay()
#handle training shock bout state changes
if t > self.nextShockBoutTime:
if self.bInShockBout:
self.bInShockBout = False
self.setShockState(False,False)
self.numShock += 1
self.nextShockBoutTime = t + random.uniform(self.paramNeutralMinTime.value(),
self.paramNeutralMaxTime.value())
if self.numShock >= self.paramNumTrain.value():
self.nextStateTime = t
else:
self.bInShockBout = True
self.setShockState(True,True)
self.nextShockBoutTime = t + self.paramUSTime.value()
self.updateProjectorDisplay()
self.arenaData['stateinfo'].append((t, self.currState, self.getSide1ColorName(), self.getSide2ColorName()))
except:
print 'ContextualHelplessnessController:updateState failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def isReadyToStart(self):
return os.path.exists(self.infoDir.text()) and self.trackWidget.getBackgroundImage() and self.fishImg and self.arenaCamCorners
def projectorPositionChanged(self):
#the point defining the tank are such that:
#points 1 and 2 define side1
#points 2 and 3 define the 'near' connecting edge
#points 3 and 4 define side2
#points 4 and 1 define the 'far' connecting edge.
#move this code elsewhere to avoid redundant computation...
self.pts = np.array([[0 , 0],
[0 , self.projWid.value()],
[self.projLen.value(), self.projWid.value()],
[self.projLen.value(), 0]])
theta = self.projRot.value()
self.tankRotM = np.array([[cos(radians(theta)), -sin(radians(theta))],
[sin(radians(theta)), cos(radians(theta))]]);
self.pts = np.dot(self.tankRotM,self.pts.T).T
self.pts += [self.projX.value(), self.projY.value()]
self.updateProjectorDisplay()
def drawProjectorDisplay(self, painter):
if self.currState == State.OFF and self.projCalibButton.isChecked() and self.isCurrent():
#DRAW A CALIBRATION IMAGE THAT INDICATES THE PROJECTOR LOCATION AND SIDE1
side1Color = QtCore.Qt.red
side2Color = QtCore.Qt.blue
a = .5
b = 1-a
#Draw side one
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(side1Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.pts[0,0], self.pts[0,1]))
poly.append(QtCore.QPointF(self.pts[1,0], self.pts[1,1]))
poly.append(QtCore.QPointF(a*self.pts[1,0] + b*self.pts[2,0], a*self.pts[1,1] + b*self.pts[2,1]))
poly.append(QtCore.QPointF(a*self.pts[0,0] + b*self.pts[3,0], a*self.pts[0,1] + b*self.pts[3,1]))
painter.drawPolygon(poly)
painter.setPen(QtGui.QColor(168, 34, 3))
painter.setFont(QtGui.QFont('Decorative',24))
painter.drawText(self.pts[0,0],self.pts[0,1],'1')
painter.drawText(self.pts[1,0],self.pts[1,1],'2')
painter.drawText(self.pts[2,0],self.pts[2,1],'3')
painter.drawText(self.pts[3,0],self.pts[3,1],'4')
#Draw side two
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(side2Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(a*self.pts[0,0] + b*self.pts[3,0], a*self.pts[0,1] + b*self.pts[3,1]))
poly.append(QtCore.QPointF(a*self.pts[1,0] + b*self.pts[2,0], a*self.pts[1,1] + b*self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[2,0], self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[3,0], self.pts[3,1]))
painter.drawPolygon(poly)
else:
if self.currState == State.OFF:
#Draw whole tank
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(QtCore.Qt.white)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.pts[0,0], self.pts[0,1]))
poly.append(QtCore.QPointF(self.pts[1,0], self.pts[1,1]))
poly.append(QtCore.QPointF(self.pts[2,0], self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[3,0], self.pts[3,1]))
painter.drawPolygon(poly)
else:
side1Color = self.getSide1ColorNdx()
side2Color = self.getSide2ColorNdx()
a = .5
b = 1-a
#Draw side one
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(side1Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.pts[0,0], self.pts[0,1]))
poly.append(QtCore.QPointF(self.pts[1,0], self.pts[1,1]))
poly.append(QtCore.QPointF(a*self.pts[1,0] + b*self.pts[2,0], a*self.pts[1,1] + b*self.pts[2,1]))
poly.append(QtCore.QPointF(a*self.pts[0,0] + b*self.pts[3,0], a*self.pts[0,1] + b*self.pts[3,1]))
painter.drawPolygon(poly)
#Draw side two
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(side2Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(a*self.pts[0,0] + b*self.pts[3,0], a*self.pts[0,1] + b*self.pts[3,1]))
poly.append(QtCore.QPointF(a*self.pts[1,0] + b*self.pts[2,0], a*self.pts[1,1] + b*self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[2,0], self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[3,0], self.pts[3,1]))
painter.drawPolygon(poly)
#Draw OMR grating
if self.bOMR:
self.omrLastUpdate = self.t
#This stuff doesn't need to computed on every frame
mm2pix = math.hypot(self.pts[2,0]-self.pts[1,0], self.pts[2,1] - self.pts[1,1]) / self.tankLength.value()
period = self.paramOMRPeriod.value() * mm2pix
numBars = int(math.ceil(math.hypot(self.pts[2,0]-self.pts[1,0], self.pts[2,1] - self.pts[1,1]) / period))
barWidth = period * (self.paramOMRDutyCycle.value()/100.0)
#Draw the bar
phaseShift = self.omrPhase/360.0 * period
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(self.paramColorOMR.currentIndex())
painter.setBrush(brush)
painter.setPen(pen)
for nBar in range(-1,numBars+1):
#define bar in unrotated untranslated space.
bl = min(max(phaseShift + nBar * period, 0), self.projLen.value())
br = min(max(phaseShift + nBar * period + barWidth, 0), self.projLen.value())
bar = np.array([[br,0],
[bl,0],
[bl,self.projWid.value()],
[br,self.projWid.value()]])
#rotate and tranlate.
bar = np.dot(self.tankRotM, bar.T).T + [self.projX.value(), self.projY.value()]
#draw
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(bar[0,0],bar[0,1]))
poly.append(QtCore.QPointF(bar[1,0],bar[1,1]))
poly.append(QtCore.QPointF(bar[2,0],bar[2,1]))
poly.append(QtCore.QPointF(bar[3,0],bar[3,1]))
painter.drawPolygon(poly)
def drawDisplayOverlay(self, painter):
#draw the fish position
if self.fishPos and self.fishPos[0] > 0:
brush = QtGui.QBrush(QtCore.Qt.red)
painter.setBrush(brush)
painter.setPen(QtCore.Qt.NoPen)
for nFish in range(min(len(self.allFish),self.paramNumFish.value())):
painter.drawEllipse(QtCore.QPointF(self.allFish[nFish][0],self.allFish[nFish][1]), 3,3)
#draw the arena overlay
if self.arenaCamCorners:
painter.setBrush(QtCore.Qt.NoBrush)
if self.bIsCurrentArena:
pen = QtGui.QPen(QtCore.Qt.green)
else:
pen = QtGui.QPen(QtCore.Qt.blue)
pen.setWidth(3)
painter.setPen(pen)
poly = QtGui.QPolygonF()
for p in self.arenaCamCorners:
poly.append(QtCore.QPointF(p[0],p[1]))
painter.drawPolygon(poly)
if len(self.arenaCamCorners) >= 2:
pen = QtGui.QPen(QtCore.Qt.red)
pen.setWidth(2)
painter.setPen(pen)
painter.drawLine(self.arenaCamCorners[0][0],self.arenaCamCorners[0][1],
self.arenaCamCorners[1][0],self.arenaCamCorners[1][1])
painter.setPen(QtGui.QColor(168, 34, 3))
painter.setFont(QtGui.QFont('Decorative',14))
painter.drawText(self.arenaCamCorners[0][0],self.arenaCamCorners[0][1],str(self.currState))
def start(self):
#Global start button not used.
pass
def stop(self):
#Global stop button not used.
pass
#---------------------------------------------------
# CALLBACK METHODS
#---------------------------------------------------
def startSwitches(self):
self.mutex.acquire()
try:
if self.currState == State.OFF:
if self.isReadyToStart():
self.paramGroup.setDisabled(True)
self.infoGroup.setDisabled(True)
td = datetime.datetime.now()
self.saveLocation = str(self.infoDir.text())
[p, self.fnResults] = os.path.split(self.saveLocation)
self.fnResults = self.fnResults + '_' + td.strftime('%Y-%m-%d-%H-%M-%S')
self.jsonFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '.json'
self.initArenaData()
t = time.time()
self.currState = State.ACCLIMATE
self.nextStateTime = t + self.paramAcclimate.value()*60
self.bInShockBout = False
self.nextShockBoutTime = float('inf')
self.setShockState(False,False)
self.stopOMR()
self.arenaData['stateinfo'].append((t, self.currState, self.getSide1ColorName(), self.getSide2ColorName()))
self.updateProjectorDisplay()
self.startButton.setText('Stop')
else:
self.startButton.setChecked(False)
self.arenaMain.statusBar().showMessage('Arena not ready to start. Information is missing.')
else:
t = time.time()
self.currState = State.OFF
self.nextStateTime = float('inf')
self.bInShockBout = False
self.nextShockBoutTime = float('inf')
self.setShockState(False,False)
self.stopOMR()
self.startButton.setText('Start Switches')
self.arenaData['stateinfo'].append((t, self.currState, self.getSide1ColorName(), self.getSide2ColorName()))
self.updateProjectorDisplay()
self.saveResults()
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
self.pauseButton.setChecked(False)
except:
print 'ContextualHelplessnessController:startSwitches failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def doPause(self):
t = time.time()
if self.pauseButton.isChecked():
self.arenaData['pauseinfo'].append((True, t))
self.pausedRemainingTime = self.nextStateTime - t
self.nextStateTime = float('inf')
else:
self.arenaData['pauseinfo'].append((False, t))
self.nextStateTime = t + self.pausedRemainingTime
def pause(self):
self.mutex.acquire()
try:
self.doPause()
except:
print 'ContextualHelplessnessController: pause failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def getArenaCameraPosition(self):
self.arenaMain.statusBar().showMessage('Click on the corners of the arena on side 1.')
self.currArenaClick = 0
self.arenaCamCorners = []
self.arenaMain.ftDisp.clicked.connect(self.handleArenaClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleArenaClicks(self, x, y):
self.currArenaClick+=1
if self.currArenaClick<5:
self.arenaCamCorners.append((x,y))
if self.currArenaClick==1:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 1.')
elif self.currArenaClick==2:
self.arenaMain.statusBar().showMessage('Now, click on the corners of the arena on side 2.')
elif self.currArenaClick==3:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 2.')
elif self.currArenaClick==4:
self.arenaMain.ftDisp.clicked.disconnect(self.handleArenaClicks)
self.arenaMain.statusBar().showMessage('')
[self.arenaMidLine, self.arenaSide1Sign] = self.processArenaCorners(self.arenaCamCorners, .5)
self.getArenaMask()
def getFishSize(self):
self.arenaMain.statusBar().showMessage('Click on the tip of the fish tail.')
self.currFishClick = 0
self.fishSize = []
self.arenaMain.ftDisp.clicked.connect(self.handleFishClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleFishClicks(self, x, y):
self.currFishClick+=1
if self.currFishClick == 1:
self.fishSize.append((x,y))
self.arenaMain.statusBar().showMessage('Click on the tip of the fish head.')
elif self.currFishClick == 2:
self.arenaMain.ftDisp.clicked.disconnect(self.handleFishClicks)
self.fishSize.append((x,y))
self.fishImg = cv.CloneImage(self.currCvFrame)
self.arenaMain.statusBar().showMessage('')
#---------------------------------------------------
# HELPER METHODS
#---------------------------------------------------
#def getBackgroundImage(self):
# if self.currCvFrame:
# self.bcvImg = cv.CloneImage(self.currCvFrame)
# self.trackWidget.setBackgroundImage(self.bcvImg)
def processArenaCorners(self, arenaCorners, linePosition):
#return the line dividing the center of the arena, and a definition of side 1.
a = 1-linePosition
b = linePosition
ac = np.array(arenaCorners)
#arenaDivideLine = [tuple(np.mean(ac[(0,3),:],axis = 0)),tuple(np.mean(ac[(1,2),:],axis = 0))]
arenaDivideLine = [(a*ac[0,0]+b*ac[3,0], a*ac[0,1]+b*ac[3,1]),(a*ac[1,0]+b*ac[2,0], a*ac[1,1]+b*ac[2,1])]
side1Sign = 1
if not self.isOnSide(arenaCorners[1], arenaDivideLine, side1Sign):
side1Sign = -1
return (arenaDivideLine, side1Sign)
def isOnSide(self, point, line, sideSign):
#return if the fish is on side1 of the arena.
side = (line[1][0] - line[0][0]) * (point[1] - line[0][1]) - (line[1][1] - line[0][1]) * (point[0] - line[0][0])
return cmp(side,0)==sideSign
#convert the arena corners into a color mask image (arena=255, not=0)
def getArenaMask(self):
if self.arenaView:
cvImg = self.currCvFrame
self.arenaCvMask = cv.CreateImage((cvImg.width,cvImg.height), cvImg.depth, cvImg.channels)
cv.SetZero(self.arenaCvMask)
cv.FillConvexPoly(self.arenaCvMask, self.arenaCamCorners, (255,)*cvImg.channels)
self.trackWidget.setTrackMask(self.arenaCvMask)
def initArenaData(self):
#prepare output data structure
self.arenaData = {}
self.arenaData['fishbirthday'] = str(self.infoDOB.date().toPyDate())
self.arenaData['fishage'] = (datetime.date.today() - self.infoDOB.date().toPyDate()).days
self.arenaData['fishstrain'] = str(self.infoType.text())
self.arenaData['fishsize'] = self.fishSize
self.arenaData['parameters'] = { 'preDuration':self.paramPreDuration.value(),
'numTrain':self.paramNumTrain.value(),
'postDuration':self.paramPostDuration.value(),
'acclimate (m)':self.paramAcclimate.value(),
'between (m)': self.paramBetweenTime.value(),
'trainUStime': self.paramUSTime.value(),
'trainNeutralMinTime':self.paramNeutralMinTime.value(),
'trainNeutralMaxTime':self.paramNeutralMaxTime.value(),
'shock period (ms)':self.paramShockPeriod.value(),
'shock dura (ms)':self.paramShockDuration.value(),
'shock chan 1':self.paramShockChan1.value(),
'shock chan 2':self.paramShockChan2.value(),
'shock V':self.paramShockV.value(),
'OMR_pre':self.paramPreOMR.isChecked(),
'OMR_train':self.paramTrainOMR.isChecked(),
'OMR_post':self.paramPostOMR.isChecked(),
'OMR_period':self.paramOMRPeriod.value(),
'OMR_dutycycle':self.paramOMRDutyCycle.value(),
'OMR_velocity':self.paramOMRVelocity.value(),
'OMR_interval':self.paramOMRInterval.value(),
'OMR_duration':self.paramOMRDuration.value(),
'OMR_color':str(self.paramColorOMR.currentText()),
'AcclimateColor':str(self.paramColorAcclimate.currentText()),
'BetweenColor':str(self.paramColorBetween.currentText()),
'PreTestColor':str(self.paramColorPreTest.currentText()),
'TrainColor':str(self.paramColorTrain.currentText()),
'PostTestColor':str(self.paramColorPostTest.currentText()),
'numFish': self.paramNumFish.value(),
'CodeVersion':None }
self.arenaData['trackingParameters'] = self.trackWidget.getParameterDictionary()
self.arenaData['trackingParameters']['arenaPoly'] = self.arenaCamCorners
self.arenaData['trackingParameters']['arenaDivideLine'] = self.arenaMidLine
self.arenaData['trackingParameters']['arenaSide1Sign'] = self.arenaSide1Sign
self.arenaData['projectorParameters'] = {'position':[self.projX.value(), self.projY.value()],
'size':[self.projLen.value(), self.projWid.value()],
'rotation':self.projRot.value()}
self.arenaData['tankSize_mm'] = [self.tankLength.value(), self.tankWidth.value()]
self.arenaData['tracking'] = list() #list of tuples (frametime, posx, posy)
self.arenaData['video'] = list() #list of tuples (frametime, filename)
self.arenaData['stateinfo'] = list() #list of times at switch stimulus flipped.
self.arenaData['shockinfo'] = list()
self.arenaData['OMRinfo'] = list()
self.arenaData['pauseinfo'] = list()
t = datetime.datetime.now()
#save experiment images
self.bcvImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_BackImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.bcvImgFileName, self.trackWidget.getBackgroundImage())
self.fishImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_FishImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.fishImgFileName, self.fishImg)
def saveResults(self):
f = open(name=self.jsonFileName, mode='w')
json.dump(self.arenaData,f)
f.close()
def getSideColors(self):
if self.currState == State.ACCLIMATE:
return (self.paramColorAcclimate.currentIndex(), self.paramColorAcclimate.currentIndex(),
str(self.paramColorAcclimate.currentText()), str(self.paramColorAcclimate.currentText()))
if self.currState == State.PREDONE or self.currState == State.TRAINDONE:
return (self.paramColorBetween.currentIndex(), self.paramColorBetween.currentIndex(),
str(self.paramColorBetween.currentText()), str(self.paramColorBetween.currentText()))
elif self.currState == State.PRETEST:
return (self.paramColorPreTest.currentIndex(), self.paramColorPreTest.currentIndex(),
str(self.paramColorPreTest.currentText()), str(self.paramColorPreTest.currentText()))
elif self.currState == State.TRAIN:
return (self.paramColorTrain.currentIndex(), self.paramColorTrain.currentIndex(),
str(self.paramColorTrain.currentText()), str(self.paramColorTrain.currentText()))
elif self.currState == State.POSTTEST:
return (self.paramColorPostTest.currentIndex(), self.paramColorPostTest.currentIndex(),
str(self.paramColorPostTest.currentText()), str(self.paramColorPostTest.currentText()))
else:
return (-1,-1,'Black','Black')
def getSide1ColorName(self):
return self.getSideColors()[2]
def getSide2ColorName(self):
return self.getSideColors()[3]
def getSide1ColorNdx(self):
return self.getSideColors()[0]
def getSide2ColorNdx(self):
return self.getSideColors()[1]
def setShockState(self, bSide1, bSide2):
if not self.arenaMain.ard:
print 'WARNING: Arduino not connected'
return
if bSide1:
curr1 = self.arenaMain.ard.pinPulse(self.paramShockChan1.value(),
self.paramShockPeriod.value(),
self.paramShockDuration.value(),
feedbackPin = self.paramCurrChan1.value())
print 'Shocking!!!'
else:
self.arenaMain.ard.pinLow(self.paramShockChan1.value())
curr1 = 0
if bSide2:
curr2 = self.arenaMain.ard.pinPulse(self.paramShockChan2.value(),
self.paramShockPeriod.value(),
self.paramShockDuration.value(),
feedbackPin = self.paramCurrChan2.value())
print 'Shocking!!!'
else:
self.arenaMain.ard.pinLow(self.paramShockChan2.value())
curr2 = 0
self.arenaData['shockinfo'].append((time.time(), bSide1, bSide2, curr1, curr2))
def getBrush(self, colorNdx):
if colorNdx == 0:
return QtGui.QBrush(QtCore.Qt.white)
elif colorNdx == 1:
return QtGui.QBrush(QtCore.Qt.red)
elif colorNdx == 2:
return QtGui.QBrush(QtCore.Qt.blue)
elif colorNdx == 3:
return QtGui.QBrush(QtGui.QColor(128,128,128))
elif colorNdx == 4:
return QtGui.QBrush(QtCore.Qt.magenta)
elif colorNdx == 5:
return QtGui.QBrush(QtCore.Qt.black)
else:
return QtGui.QBrush(QtCore.Qt.black)
class AvoidanceController(ArenaController.ArenaController):
def __init__(self, parent, arenaMain):
super(AvoidanceController, self).__init__(parent, arenaMain)
#calibration
self.arenaCamCorners = []
self.arenaMidLine = []
self.arenaSide1Sign = 1
self.arenaProjCorners = []
self.fishImg = None #image of fish
#tracking
self.arenaCvMask = None
#state
self.mutex = Lock()
self.pulseMutex = Lock()
self.nDispMode = 0
self.bRunning = False
self.currState = State.BETWEEN
self.fishPosUpdate = False
self.fishPos = (0,0)
self.fishSize = None
self.escapeLine= []
self.escapeSign = []
#shock pulse state - only if using labjack
self.shockSide1_t = 0 #0 if not shocking, otherwise time of next shock
self.bPulseHigh1 = False #true if shock_pulse in progress
self.shockSide1_ndx = None
self.shockSide2_t = 0 #0 if not shocking, otherwise time of next shock
self.bPulseHigh2 = False #true if shock_pulse in progress
self.shockSide2_ndx = None
#tracking related images
self.currCvFrame = None
#init UI
#arena info group box
self.arenaGroup = QtGui.QGroupBox(self)
self.arenaGroup.setTitle('Arena Info')
self.arenaLayout = QtGui.QGridLayout()
self.arenaLayout.setHorizontalSpacing(3)
self.arenaLayout.setVerticalSpacing(3)
self.cntrUI = QtGui.QWidget(self)
self.cntrLayout = QtGui.QGridLayout()
self.cntrLayout.setHorizontalSpacing(3)
self.cntrLayout.setVerticalSpacing(3)
l1 = QtGui.QLabel('Side1 Control Channel:')
l2 = QtGui.QLabel('Side2 Control Channel:')
l3 = QtGui.QLabel('Current Input Channel:')
self.cntrSide1Channel = QtGui.QSpinBox()
self.cntrSide1Channel.setValue(8)
self.cntrSide1Channel.setMaximumWidth(50)
self.cntrSide1Channel.setRange(0,8000)
self.cntrSide2Channel = QtGui.QSpinBox()
self.cntrSide2Channel.setValue(9)
self.cntrSide2Channel.setMaximumWidth(50)
self.cntrSide2Channel.setRange(0,8000)
self.currInputChannel = QtGui.QSpinBox()
self.currInputChannel.setValue(10)
self.currInputChannel.setMaximumWidth(50)
self.currInputChannel.setRange(0,8000)
self.cntrLayout.addWidget(l1,0,0)
self.cntrLayout.addWidget(self.cntrSide1Channel,0,1)
self.cntrLayout.addWidget(l2,1,0)
self.cntrLayout.addWidget(self.cntrSide2Channel,1,1)
self.cntrLayout.addWidget(l3,2,0)
self.cntrLayout.addWidget(self.currInputChannel,2,1)
self.cntrLayout.setColumnStretch(2,1)
self.cntrUI.setLayout(self.cntrLayout)
self.camCalibButton = QtGui.QPushButton('Set Cam Position')
self.camCalibButton.setMaximumWidth(150)
self.camCalibButton.clicked.connect(self.getArenaCameraPosition)
self.projGroup = QtGui.QGroupBox(self.arenaGroup)
self.projGroup.setTitle('Projector Position')
self.projLayout = QtGui.QGridLayout(self.projGroup)
self.projLayout.setHorizontalSpacing(3)
self.projLayout.setVerticalSpacing(3)
self.projCalibButton = QtGui.QPushButton('Calibrate Projector Position')
self.projCalibButton.setCheckable(True)
self.projCalibButton.clicked.connect(self.updateProjectorDisplay)
self.projLayout.addWidget(self.projCalibButton,0,0,1,6)
self.proj1L = QtGui.QLabel('C1')
self.proj1X = QtGui.QSpinBox()
self.proj1X.setValue(0)
self.proj1X.setMaximumWidth(50)
self.proj1X.setRange(0,2000)
self.proj1Y = QtGui.QSpinBox()
self.proj1Y.setValue(0)
self.proj1Y.setMaximumWidth(50)
self.proj1Y.setRange(0,2000)
self.projLayout.addWidget(self.proj1L,1,0)
self.projLayout.addWidget(self.proj1X,1,1)
self.projLayout.addWidget(self.proj1Y,1,2)
self.proj1X.valueChanged.connect(self.updateProjectorDisplay)
self.proj1Y.valueChanged.connect(self.updateProjectorDisplay)
self.proj2L = QtGui.QLabel('C2')
self.proj2X = QtGui.QSpinBox()
self.proj2X.setValue(0)
self.proj2X.setMaximumWidth(50)
self.proj2X.setRange(0,2000)
self.proj2Y = QtGui.QSpinBox()
self.proj2Y.setValue(100)
self.proj2Y.setMaximumWidth(50)
self.proj2Y.setRange(0,2000)
self.proj2X.valueChanged.connect(self.updateProjectorDisplay)
self.proj2Y.valueChanged.connect(self.updateProjectorDisplay)
self.projLayout.addWidget(self.proj2L,2,0)
self.projLayout.addWidget(self.proj2X,2,1)
self.projLayout.addWidget(self.proj2Y,2,2)
self.proj3L = QtGui.QLabel('C3')
self.proj3X = QtGui.QSpinBox()
self.proj3X.setValue(200)
self.proj3X.setMaximumWidth(50)
self.proj3X.setRange(0,2000)
self.proj3Y = QtGui.QSpinBox()
self.proj3Y.setValue(100)
self.proj3Y.setMaximumWidth(50)
self.proj3Y.setRange(0,2000)
self.projLayout.addWidget(self.proj3L,2,3)
self.projLayout.addWidget(self.proj3X,2,4)
self.projLayout.addWidget(self.proj3Y,2,5)
self.proj3X.valueChanged.connect(self.updateProjectorDisplay)
self.proj3Y.valueChanged.connect(self.updateProjectorDisplay)
self.proj4L = QtGui.QLabel('C4')
self.proj4X = QtGui.QSpinBox()
self.proj4X.setValue(200)
self.proj4X.setMaximumWidth(50)
self.proj4X.setRange(0,2000)
self.proj4Y = QtGui.QSpinBox()
self.proj4Y.setValue(0)
self.proj4Y.setMaximumWidth(50)
self.proj4Y.setRange(0,2000)
self.projLayout.addWidget(self.proj4L,1,3)
self.projLayout.addWidget(self.proj4X,1,4)
self.projLayout.addWidget(self.proj4Y,1,5)
self.proj4X.valueChanged.connect(self.updateProjectorDisplay)
self.proj4Y.valueChanged.connect(self.updateProjectorDisplay)
self.projLayout.setColumnStretch(6,1)
self.projGroup.setLayout(self.projLayout)
self.arenaLayout.addWidget(self.camCalibButton, 0,0)
self.arenaLayout.addWidget(self.cntrUI,1,0)
self.arenaLayout.addWidget(self.projGroup,2,0)
self.arenaLayout.setColumnStretch(1,1)
self.arenaGroup.setLayout(self.arenaLayout)
#tracking group box
self.trackWidget = FishTrackerWidget(self, self.arenaMain.ftDisp)
#experimental parameters groupbox
self.paramGroup = QtGui.QGroupBox()
self.paramGroup.setTitle('Exprimental Parameters')
self.paramLayout = QtGui.QGridLayout()
self.paramLayout.setHorizontalSpacing(2)
self.paramLayout.setVerticalSpacing(2)
self.paramNumTrials = LabeledSpinBox(None,'trials',1,999,40,60)
self.paramLayout.addWidget(self.paramNumTrials,1,0,1,2)
self.paramAcc = LabeledSpinBox(None,'acclimation (s)',1,3600,900,60)
self.paramLayout.addWidget(self.paramAcc,1,2,1,2)
self.paramITIMin = LabeledSpinBox(None,'iti min (s)',1,300,15,60)
self.paramLayout.addWidget(self.paramITIMin,2,0,1,2)
self.paramITIMax = LabeledSpinBox(None,'iti max (s)',1,300,30,60)
self.paramLayout.addWidget(self.paramITIMax,2,2,1,2)
self.paramEscapeT = LabeledSpinBox(None,'escape (s)',1,60,10,60)
self.paramLayout.addWidget(self.paramEscapeT,3,0,1,2)
self.paramShockT = LabeledSpinBox(None,'shock (s)',1,100,30,60)
self.paramLayout.addWidget(self.paramShockT,3,2,1,2)
self.paramV = LabeledSpinBox(None,'volts',1,50,12,60)
self.paramLayout.addWidget(self.paramV,4,0,1,2)
self.paramShockPeriod = LabeledSpinBox(None,'shock period (ms)',1,5000,1000,60)
self.paramLayout.addWidget(self.paramShockPeriod,4,2,1,2)
self.paramShockDura = LabeledSpinBox(None,'shock t (ms)',1,1000,100,60)
self.paramLayout.addWidget(self.paramShockDura,5,0,1,2)
self.paramNeutral = QtGui.QComboBox()
self.paramNeutral.addItem('White')
self.paramNeutral.addItem('Red')
self.paramNeutral.addItem('Blue')
self.paramNeutral.addItem('Gray')
self.paramNeutral.setCurrentIndex(0)
self.labelNeutral = QtGui.QLabel('Neutral')
self.paramLayout.addWidget(self.paramNeutral,5,2)
self.paramLayout.addWidget(self.labelNeutral,5,3)
self.paramAvoid = QtGui.QComboBox()
self.paramAvoid.addItem('White')
self.paramAvoid.addItem('Red')
self.paramAvoid.addItem('Blue')
self.paramAvoid.addItem('Gray')
self.paramAvoid.setCurrentIndex(1)
self.labelAvoid = QtGui.QLabel('Avoid')
self.paramLayout.addWidget(self.paramAvoid,6,0)
self.paramLayout.addWidget(self.labelAvoid,6,1)
self.paramEscape = QtGui.QComboBox()
self.paramEscape.addItem('White')
self.paramEscape.addItem('Red')
self.paramEscape.addItem('Blue')
self.paramEscape.addItem('Gray')
self.paramEscape.setCurrentIndex(2)
self.labelEscape = QtGui.QLabel('Escape')
self.paramLayout.addWidget(self.paramEscape,6,2)
self.paramLayout.addWidget(self.labelEscape,6,3)
self.paramDynamic= QtGui.QCheckBox('Dynamic Escape')
self.paramLayout.addWidget(self.paramDynamic,7,0,1,2)
self.paramEscapePos = QtGui.QDoubleSpinBox()
self.paramEscapePos.setRange(0,1)
self.paramEscapePos.setValue(.5)
self.labelEscapePos = QtGui.QLabel('escape pos (0-1)')
self.paramLayout.addWidget(self.paramEscapePos,7,2)
self.paramLayout.addWidget(self.labelEscapePos,7,3)
self.paramDynamicDraw = QtGui.QCheckBox('Dynamic Draw')
self.paramDynamicDraw.setChecked(True)
self.paramLayout.addWidget(self.paramDynamicDraw,8,0,1,2)
self.paramDebug = QtGui.QPushButton('Debug')
self.paramDebug.setMaximumWidth(150)
self.paramDebug.setCheckable(True)
self.paramLayout.addWidget(self.paramDebug,8,2,1,2)
self.paramDebug.clicked.connect(self.useDebugParams)
self.paramGroup.setLayout(self.paramLayout)
#Experimental info group
self.infoGroup = QtGui.QGroupBox()
self.infoGroup.setTitle('Experiment Info')
self.infoLayout = QtGui.QGridLayout()
self.infoLayout.setHorizontalSpacing(3)
self.infoLayout.setVerticalSpacing(3)
self.labelDir = QtGui.QLabel('Dir: ')
self.infoDir = PathSelectorWidget(browseCaption='Experimental Data Directory')
self.infoLayout.addWidget(self.labelDir,0,0)
self.infoLayout.addWidget(self.infoDir,0,1)
self.labelDOB = QtGui.QLabel('DOB: ')
self.infoDOB = QtGui.QDateEdit(QtCore.QDate.currentDate())
self.infoLayout.addWidget(self.labelDOB,1,0)
self.infoLayout.addWidget(self.infoDOB,1,1)
self.labelType = QtGui.QLabel('Line: ')
self.infoType = QtGui.QLineEdit('Nacre')
self.infoLayout.addWidget(self.labelType,2,0)
self.infoLayout.addWidget(self.infoType,2,1)
self.infoFish = QtGui.QPushButton('Snap Fish Image')
self.infoFish.clicked.connect(self.getFishSize)
self.infoLayout.addWidget(self.infoFish,3,0,1,2)
self.infoGroup.setLayout(self.infoLayout)
self.settingsLayout = QtGui.QGridLayout()
self.settingsLayout.addWidget(self.arenaGroup,0,0)
self.settingsLayout.addWidget(self.trackWidget,1,0)
self.settingsLayout.addWidget(self.paramGroup,2,0)
self.settingsLayout.addWidget(self.infoGroup,3,0)
self.setLayout(self.settingsLayout)
#---------------------------------------------------
# OVERLOADED METHODS
#---------------------------------------------------
#update to overload an return image based on dispmode.
#def getArenaView(self):
# return self.arenaView
def onNewFrame(self, frame, time):
self.mutex.acquire()
try:
self.frameTime = time
self.currCvFrame = frame # may need to make a deep copy
(found, pos, view, allFish) = self.trackWidget.findFish(self.currCvFrame)
if found:
self.fishPosUpdate = found
self.fishPos = pos
if self.bRunning:
self.avoidData['video'].append((self.frameTime, None))
self.avoidData['tracking'].append((self.frameTime, self.fishPos[0], self.fishPos[1]))
self.arenaView = view
except:
print 'AvoidanceController:onNewFrame failed'
print "Unexpected error:", sys.exc_info()[0]
finally:
self.mutex.release()
def updateState(self):
self.updateExperimentalState()
self.handlePulsedShocks()
def updateExperimentalState(self):
if self.bRunning:
self.mutex.acquire()
try:
t = time.time()
#only update status bar if arena is selected.
if self.isCurrent():
if self.nTrial>-1:
#print('NextTrial# %d TimeTilNextTrial %f TrialTime %f' % (self.nTrial+1, self.timeOfNextTrial - t, t - self.avoidData['trials'][self.nTrial]['startT']))
self.arenaMain.statusBar().showMessage('NextTrial# %d TimeTilNextTrial %f TrialTime %f' % (self.nTrial+1, self.timeOfNextTrial - t, t - self.avoidData['trials'][self.nTrial]['startT']))
else:
#print('NextTrial# %d TimeTilNextTrial %f' % (self.nTrial+1, self.timeOfNextTrial - t))
self.arenaMain.statusBar().showMessage('NextTrial# %d TimeTilNextTrial %f' % (self.nTrial+1, self.timeOfNextTrial - t))
#check if fish escaped
bDidEscape = False
if self.fishPosUpdate:
self.fishPosUpdate = False
if self.escapeLine and (self.currState == State.LED or self.currState == State.SHOCK):
bDidEscape = self.isOnSide(self.fishPos, self.escapeLine, self.escapeSign)
#handle State Changes
###BETWEEN###
if self.currState==State.BETWEEN and t >= self.timeOfNextTrial:
print 'Start Trial'
self.nTrial+=1
if self.nTrial < self.paramNumTrials.value():
self.currState = State.LED;
self.timeState = t;
#determine relative fish position (0->Side1 1->Side2) (by projecting onto 1 side)
ac = self.arenaCamCorners
lengC = ((self.fishPos[1]-ac[0][1])**2 + (self.fishPos[0] - ac[0][0])**2)**.5 #length of fish vector
lengB = ((ac[3][1]-ac[0][1])**2 + (ac[3][0]-ac[0][0])**2)**.5 #length of tank side
lengA = ((self.fishPos[1]-ac[3][1])**2 + (self.fishPos[0]-ac[3][0])**2)**.5 #length of connecting vector
#dot product = fish vector length * cos(angle between fish vector and tank side)/normalized by side length
relPos = lengC * (lengB**2 + lengC**2 - lengA**2)/(2*lengC*lengB)/lengB
#determine side of fish -- not quite the same as relpos<0.5 because corners 2 and 3 not considered in relpos
if self.isOnSide(self.fishPos, self.arenaMidLine, self.arenaSide1Sign):
self.trialSide = Side.S1
else:
self.trialSide = Side.S2
#determine relative escape position
if not self.paramDynamic.isChecked():
if self.trialSide == Side.S1:
self.escapePos = self.paramEscapePos.value()
else:
self.escapePos = 1 - self.paramEscapePos.value()
else:
if self.trialSide == Side.S1:
self.escapePos = relPos + self.paramEscapePos.value()
else:
self.escapePos = relPos - self.paramEscapePos.value()
#convert relative escape position into a line
(self.escapeLine, side1Sign) = self.processArenaCorners(self.arenaCamCorners, self.escapePos)
self.escapeSign = -1 * self.trialSide * self.arenaSide1Sign
self.updateProjectorDisplay()
self.avoidData['trials'].append({'trialNum':self.nTrial,
'startT':self.timeState,
'side':self.trialSide,
'escape': (self.escapeLine, self.escapeSign),
'endT':-1})
else:
#experiment is complete
self.stopController()
###LED###
elif self.currState== State.LED:
if bDidEscape:
print 'Avoided Shocks'
self.timeState = t
self.currState=State.BETWEEN;
self.timeOfNextTrial = self.timeState + random.randint(self.paramITIMin.value(),self.paramITIMax.value())
self.escapeLine = []
self.avoidData['trials'][self.nTrial]['endT'] = self.timeState
self.avoidData['trials'][self.nTrial]['bAvoidedShock'] = True
self.saveResults()
elif (t - self.timeState) > self.paramEscapeT.value():
print 'Starting Shocks'
self.timeState = t
self.currState= State.SHOCK;
self.avoidData['trials'][self.nTrial]['bAvoidedShock'] = False
self.setShockState(self.trialSide == Side.S1, self.trialSide == Side.S2)
self.avoidData['trials'][self.nTrial]['Shock_current_start'] = self.readCurrentInput()
self.updateProjectorDisplay()
###SHOCK###
elif self.currState == State.SHOCK:
if bDidEscape or t - self.timeState > self.paramShockT.value():
print 'Ending Shocks'
self.currState=State.BETWEEN
self.timeOfNextTrial = self.timeState + random.randint(self.paramITIMin.value(),self.paramITIMax.value())
self.timeState = t;
self.escapeLine = []
self.setShockState(False, False)
self.updateProjectorDisplay()
self.avoidData['trials'][self.nTrial]['endT'] = self.timeState
self.saveResults()
#########
except:
print 'AvoidanceController:updateState failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def isReadyToStart(self):
bReady = os.path.exists(self.infoDir.text()) and self.trackWidget.getBackgroundImage() and self.fishImg and self.arenaCamCorners
print bReady
return bReady
def drawProjectorDisplay(self, painter):
if not self.bRunning and self.projCalibButton.isChecked():
brush = QtGui.QBrush(QtCore.Qt.blue)
pen = QtGui.QPen(QtCore.Qt.black)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
else:
if self.currState == State.BETWEEN:
#Draw whole tank
pen = QtGui.QPen(QtCore.Qt.NoPen)
print self.paramNeutral.currentIndex()
brush = self.getBrush(self.paramNeutral.currentIndex())
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
else:
side1Color = side2Color = self.paramEscape.currentIndex()
if self.trialSide == Side.S1:
side1Color = self.paramAvoid.currentIndex()
else:
side2Color = self.paramAvoid.currentIndex()
if self.paramDynamicDraw.isChecked():
a = float(1-self.escapePos)
b = float(self.escapePos)
else:
a = 0.5
b = 0.5
#Draw side one
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(side1Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.proj1X.value(), self.proj1Y.value()))
poly.append(QtCore.QPointF(self.proj2X.value(), self.proj2Y.value()))
poly.append(QtCore.QPointF(a*self.proj2X.value() + b*self.proj3X.value(), a*self.proj2Y.value() + b*self.proj3Y.value()))
poly.append(QtCore.QPointF(a*self.proj1X.value() + b*self.proj4X.value(), a*self.proj1Y.value() + b*self.proj4Y.value()))
painter.drawPolygon(poly)
#Draw side two
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(side2Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(a*self.proj1X.value() + b*self.proj4X.value(), a*self.proj1Y.value() + b*self.proj4Y.value()))
poly.append(QtCore.QPointF(a*self.proj2X.value() + b*self.proj3X.value(), a*self.proj2Y.value() + b*self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj3X.value(), self.proj3Y.value()))
poly.append(QtCore.QPointF(self.proj4X.value(), self.proj4Y.value()))
painter.drawPolygon(poly)
def drawDisplayOverlay(self, painter):
#draw the fish position
if self.fishPos and self.fishPos[0] > 0:
brush = QtGui.QBrush(QtCore.Qt.red)
painter.setBrush(brush)
painter.setPen(QtCore.Qt.NoPen)
painter.drawEllipse(QtCore.QPointF(self.fishPos[0],self.fishPos[1]), 3,3)
#draw the arena overlay
if self.arenaCamCorners:
painter.setBrush(QtCore.Qt.NoBrush)
if self.bIsCurrentArena:
pen = QtGui.QPen(QtCore.Qt.green)
else:
pen = QtGui.QPen(QtCore.Qt.blue)
pen.setWidth(3)
painter.setPen(pen)
poly = QtGui.QPolygonF()
for p in self.arenaCamCorners:
poly.append(QtCore.QPointF(p[0],p[1]))
painter.drawPolygon(poly)
#draw escape line
if self.escapeLine:
pen = QtGui.QPen()
pen.setColor(QtCore.Qt.red)
pen.setWidth(1)
l = self.escapeLine
painter.drawLine(QtCore.QPointF(l[0][0],l[0][1]), QtCore.QPointF(l[1][0],l[1][1]))
def start(self):
self.mutex.acquire()
try:
self.paramGroup.setDisabled(True)
self.infoGroup.setDisabled(True)
self.initResults()
#initialize experimental state
self.nTrial = -1
self.currState = State.BETWEEN
self.setShockState(False,False)
self.updateProjectorDisplay()
self.timeState = time.time()
self.timeOfNextTrial = self.timeState + self.paramAcc.value()
self.trialSide = Side.S1
self.escapeLine = []
self.bRunning = True
finally:
self.mutex.release()
def stop(self):
self.mutex.acquire()
try:
self.stopController()
finally:
self.mutex.release()
def stopController(self):
if self.bRunning:
self.bRunning = False
self.saveResults()
print 'Experiment stopped successfully. Saved experimental data.'
self.currState = State.BETWEEN
self.setShockState(False,False)
self.nTrial = -1
self.escapeLine = []
self.updateProjectorDisplay()
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
#---------------------------------------------------
# CALLBACK METHODS
#---------------------------------------------------
def getArenaCameraPosition(self):
self.arenaMain.statusBar().showMessage('Click on the corners of the arena on side 1.')
self.currArenaClick = 0
self.arenaCamCorners = []
self.arenaMain.ftDisp.clicked.connect(self.handleArenaClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleArenaClicks(self, x, y):
self.currArenaClick+=1
if self.currArenaClick<5:
self.arenaCamCorners.append((x,y))
if self.currArenaClick==1:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 1.')
elif self.currArenaClick==2:
self.arenaMain.statusBar().showMessage('Now, click on the corners of the arena on side 2.')
elif self.currArenaClick==3:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 2.')
elif self.currArenaClick==4:
self.arenaMain.ftDisp.clicked.disconnect(self.handleArenaClicks)
self.arenaMain.statusBar().showMessage('')
[self.arenaMidLine, self.arenaSide1Sign] = self.processArenaCorners(self.arenaCamCorners, .5)
self.getArenaMask()
def getFishSize(self):
self.arenaMain.statusBar().showMessage('Click on the tip of the fish tail.')
self.currFishClick = 0
self.fishSize = []
self.arenaMain.ftDisp.clicked.connect(self.handleFishClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleFishClicks(self, x, y):
self.currFishClick+=1
if self.currFishClick == 1:
self.fishSize.append((x,y))
self.arenaMain.statusBar().showMessage('Click on the tip of the fish head.')
elif self.currFishClick == 2:
self.arenaMain.ftDisp.clicked.disconnect(self.handleFishClicks)
self.fishSize.append((x,y))
self.fishImg = cv.CloneImage(self.currCvFrame)
self.arenaMain.statusBar().showMessage('')
def useDebugParams(self):
if self.paramDebug.isChecked():
self.debugVals = {}
self.debugVals['dn'] = self.paramNumTrials.value()
self.debugVals['da'] = self.paramAcc.value()
self.debugVals['ds'] = self.paramITIMin.value()
self.debugVals['dl'] = self.paramITIMax.value()
self.debugVals['de'] = self.paramEscapeT.value()
self.debugVals['dt'] = self.paramShockT.value()
self.debugVals['dv'] = self.paramV.value()
self.paramNumTrials.setValue(5)
self.paramAcc.setValue(5)
self.paramITIMin.setValue(3)
self.paramITIMax.setValue(3)
self.paramEscapeT.setValue(5)
self.paramShockT.setValue(5)
self.paramV.setValue(0)
else:
self.paramNumTrials.setValue(self.debugVals['dn'])
self.paramAcc.setValue(self.debugVals['da'])
self.paramITIMin.setValue(self.debugVals['ds'])
self.paramITIMax.setValue(self.debugVals['dl'])
self.paramEscapeT.setValue(self.debugVals['de'])
self.paramShockT.setValue(self.debugVals['dt'])
self.paramV.setValue(self.debugVals['dv'])
#---------------------------------------------------
# HELPER METHODS
#---------------------------------------------------
# def getBackgroundImage(self):
# if self.currCvFrame:
# self.bcvImg = cv.CloneImage(self.currCvFrame)
# self.trackWidget.setBackgroundImage(self.bcvImg)
#convert the arena corners into a color mask image (arena=255, not=0)
def getArenaMask(self):
if self.arenaView:
cvImg = self.currCvFrame
self.arenaCvMask = cv.CreateImage((cvImg.width,cvImg.height), cvImg.depth, cvImg.channels)
cv.SetZero(self.arenaCvMask)
cv.FillConvexPoly(self.arenaCvMask, self.arenaCamCorners, (255,)*cvImg.channels)
self.trackWidget.setTrackMask(self.arenaCvMask)
def initResults(self):
#prepare output data structure
self.avoidData = {}
self.avoidData['fishbirthday'] = str(self.infoDOB.date().toPyDate())
self.avoidData['fishage'] = (datetime.date.today() - self.infoDOB.date().toPyDate()).days
self.avoidData['fishstrain'] = str(self.infoType.text())
self.avoidData['fishsize'] = self.fishSize
self.avoidData['parameters'] = { 'numtrials':self.paramNumTrials.value(),
'LEDTimeMS':self.paramEscapeT.value(),
'ShockTimeMS':self.paramShockT.value(),
'ShockV':self.paramV.value(),
'AcclimationTime':self.paramAcc.value(),
'MinTrialInterval':self.paramITIMin.value(),
'MaxTrialInterval':self.paramITIMax.value(),
'ShockPeriodT':self.paramShockPeriod.value(),
'ShockDurationT':self.paramShockDura.value(),
'NeuralStimulus':str(self.paramNeutral.currentText()),
'AvoidStimulus':str(self.paramAvoid.currentText()),
'EscapeStimulus':str(self.paramEscape.currentText()),
'isDynamicEscape':self.paramDynamic.isChecked(),
'inDynamicDraw':self.paramDynamicDraw.isChecked(),
'fEscapePosition':self.paramEscapePos.value(),
'CodeVersion':None }
self.avoidData['trackingParameters'] = self.trackWidget.getParameterDictionary()
self.avoidData['trackingParameters']['arenaPoly'] = self.arenaCamCorners
self.avoidData['trackingParameters']['arenaDivideLine'] = self.arenaMidLine
self.avoidData['trackingParameters']['arenaSide1Sign'] = self.arenaSide1Sign
self.avoidData['trials'] = list() #outcome on each trial
self.avoidData['tracking'] = list() #list of tuples (frametime, posx, posy)
self.avoidData['video'] = list() #list of tuples (frametime, filename)
self.avoidData['current'] = list() #list of tuples (time, currentValue)
self.avoidData['shockPulses'] = list() #list of tuples (side, target_start, actual_start, target_end, actual_end)
#get experiment filename
[p, self.fnResults] = os.path.split(str(self.infoDir.text()))
self.fnResults = self.fnResults + '_' + datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
self.jsonFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '.json'
#save experiment images
self.bcvImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_BackImg.tiff'
cv.SaveImage(self.bcvImgFileName, self.trackWidget.getBackgroundImage())
self.fishImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_FishImg.tiff'
cv.SaveImage(self.fishImgFileName, self.fishImg)
def processArenaCorners(self, arenaCorners, linePosition):
#return the line dividing the center of the arena, and a definition of side 1.
a = 1-linePosition
b = linePosition
ac = np.array(arenaCorners)
#arenaDivideLine = [tuple(np.mean(ac[(0,3),:],axis = 0)),tuple(np.mean(ac[(1,2),:],axis = 0))]
arenaDivideLine = [(a*ac[0,0]+b*ac[3,0], a*ac[0,1]+b*ac[3,1]),(a*ac[1,0]+b*ac[2,0], a*ac[1,1]+b*ac[2,1])]
side1Sign = 1
if not self.isOnSide(arenaCorners[1], arenaDivideLine, side1Sign):
side1Sign = -1
return (arenaDivideLine, side1Sign)
def isOnSide(self, point, line, sideSign):
#return if the fish is on side1 of the arena.
side = (line[1][0] - line[0][0]) * (point[1] - line[0][1]) - (line[1][1] - line[0][1]) * (point[0] - line[0][0])
return cmp(side,0)==sideSign
def saveResults(self):
f = open(name=self.jsonFileName, mode='w')
json.dump(self.avoidData,f)
f.close()
def getBrush(self, colorNdx):
if colorNdx == 0:
return QtGui.QBrush(QtCore.Qt.white)
elif colorNdx == 1:
return QtGui.QBrush(QtCore.Qt.red)
elif colorNdx == 2:
return QtGui.QBrush(QtCore.Qt.blue)
elif colorNdx == 3:
return QtGui.QBrush(QtGui.QColor(128,128,128))
else:
return QtGui.QBrush(QtCore.Qt.black)
def setShockState(self, bSide1, bSide2):
if bSide1:
self.arenaMain.ard.pinPulse(self.cntrSide1Channel.value(),
self.paramShockPeriod.value(),
self.paramShockDura.value())
else:
self.arenaMain.ard.pinLow(self.cntrSide1Channel.value())
if bSide2:
self.arenaMain.ard.pinPulse(self.cntrSide2Channel.value(),
self.paramShockPeriod.value(),
self.paramShockDura.value())
else:
self.arenaMain.ard.pinLow(self.cntrSide2Channel.value())
"""
elif bLabJack:
self.pulseMutex.acquire()
try:
if bSide1 and not self.shockSide1_t:
self.shockSide1_t = time.time()
self.bPulseHigh1 = False
elif not bSide1 and self.shockSide1_t:
if self.bPulseHigh1:
self.arenaMain.labjack.writeRegister(self.cntrSide1Channel.value(), 0)
if self.bRunning: self.avoidData['shockPulses'][self.shockSide1_ndx][3] = self.shockSide1_t + self.paramShockDura.value()/1000
if self.bRunning: self.avoidData['shockPulses'][self.shockSide1_ndx][4] = time.time()
self.bPulseHigh1 = False
self.shockSide1_ndx = None
self.shockSide1_t = 0
if bSide2 and not self.shockSide2_t:
self.shockSide2_t = time.time()
self.bPulseHigh2 = False
elif not bSide2 and self.shockSide2_t:
if self.bPulseHigh2:
self.arenaMain.labjack.writeRegister(self.cntrSide2Channel.value(), 0)
if self.bRunning: self.avoidData['shockPulses'][self.shockSide2_ndx][3] = self.shockSide2_t + self.paramShockDura.value()/1000
if self.bRunning: self.avoidData['shockPulses'][self.shockSide2_ndx][4] = time.time()
self.bPulseHigh2 = False
self.shockSide2_ndx = None
self.shockSide2_t = 0
finally:
self.pulseMutex.release()
"""
def handlePulsedShocks(self):
pass
"""
self.pulseMutex.acquire()
try:
t = time.time()
if self.shockSide1_t:
if not self.bPulseHigh1 and t > self.shockSide1_t:
self.arenaMain.labjack.writeRegister(self.cntrSide1Channel.value(), 5)
if self.bRunning: self.avoidData['shockPulses'].append([1, self.shockSide1_s, t, -1, -1])
self.bPulseHigh1 = True
self.shockSide1_ndx = len(self.avoidData['shockPulses'])-1
elif t > self.shockSide1_t + self.paramShockDura.value()/1000:
self.arenaMain.labjack.writeRegister(self.cntrSide1Channel.value(), 0)
if self.bRunning: self.avoidData['shockPulses'][self.shockSide1_ndx][3] = self.shockSide1_t + self.paramShockDura.value()/1000
if self.bRunning: self.avoidData['shockPulses'][self.shockSide1_ndx][4] = t
self.shockSide1_t = self.shockSide1_t + self.paramShockPeriod.value()/1000
self.bPulseHigh1 = False
self.shockSide1_ndx = None
if self.shockSide2_t:
if not self.bPulseHigh2 and t > self.shockSide2_t:
self.arenaMain.labjack.writeRegister(self.cntrSide2Channel.value(), 5)
if self.bRunning: self.avoidData['shockPulses'].append([2, self.shockSide2_s, t, -1, -1])
self.bPulseHigh2 = True
self.shockSide2_ndx = len(self.avoidData['shockPulses'])-1
elif t > self.shockSide2_t + self.paramShockDura.value()/1000:
self.arenaMain.labjack.writeRegister(self.cntrSide2Channel.value(), 0)
if self.bRunning: self.avoidData['shockPulses'][self.shockSide2_ndx][3] = self.shockSide2_t + self.paramShockDura.value()/1000
if self.bRunning: self.avoidData['shockPulses'][self.shockSide2_ndx][4] = t
self.shockSide2_t = self.shockSide2_t + self.paramShockPeriod.value()/1000
self.bPulseHigh2 = False
self.shockSide2_ndx = None
finally:
self.pulseMutex.release()
"""
def readCurrentInput(self):
return None
#self.arenaMain.ard.analogRead(self.currInputChannel.value())
"""
class OMRController(ArenaController.ArenaController, Machine):
def __init__(self, parent, arenaMain):
super(OMRController, self).__init__(parent, arenaMain)
print type(parent)
#calibration
self.arenaCamCorners = None
self.arenaMidLine = []
self.arenaSide1Sign = 1
self.arenaProjCorners = []
self.fishImg = None #image of fish
#tracking
self.arenaCvMask = None
#state
self.mutex = Lock()
states = ['off', 'between_session', 'omr_to_s1', 'no_omr_post_s1', 'omr_to_s2', 'no_omr_post_s2'] # off between L-N-R-N-L-N-R between LNRNLNR between off
Machine.__init__(self, states=states, initial='off', after_state_change='update_state_data')
self.add_transition('begin', 'off', 'between_session', conditions='isReadyToStart')
self.add_transition('start_session', 'between_session', 'omr_to_s1')
self.add_transition('start_omr', 'no_omr_post_s1', 'omr_to_s2')
self.add_transition('start_omr', 'no_omr_post_s2', 'omr_to_s1')
self.add_transition('stop_omr', 'omr_to_s1', 'no_omr_post_s1')
self.add_transition('stop_omr', 'omr_to_s2', 'no_omr_post_s2')
self.add_transition('stop_session', ['omr_to_s1', 'omr_to_s2', 'no_omr_post_s1', 'no_omr_post_s2'], 'between_session')
self.add_transition('quit', '*', 'off')
self.fishPosUpdate = False
self.fishPos = (0,0)
self.fishPosBuffer = []
self.fishPosBufferT = []
self.currCvFrame = None
#Note additional state variable will be created during state transitions (like nextStateTime)
#data results
self.arenaData = None
self.fishSize = None
#init UI
#arena info group box
self.arenaGroup = QtGui.QGroupBox(self)
self.arenaGroup.setTitle('Arena Info')
self.arenaLayout = QtGui.QGridLayout()
self.arenaLayout.setHorizontalSpacing(3)
self.arenaLayout.setVerticalSpacing(3)
self.camCalibButton = QtGui.QPushButton('Set Cam Position')
self.camCalibButton.setMaximumWidth(150)
self.camCalibButton.clicked.connect(self.getArenaCameraPosition)
self.resetCamCorners = None
self.projGroup = QtGui.QGroupBox(self.arenaGroup)
self.projGroup.setTitle('Projector Position')
self.projLayout = QtGui.QGridLayout(self.projGroup)
self.projLayout.setHorizontalSpacing(3)
self.projLayout.setVerticalSpacing(3)
self.projCalibButton = QtGui.QPushButton('Calibrate Projector Position')
self.projCalibButton.setCheckable(True)
self.projCalibButton.clicked.connect(self.projectorPositionChanged)
self.projLayout.addWidget(self.projCalibButton,0,0,1,6)
self.projPosL = QtGui.QLabel('Pos')
self.projX = QtGui.QSpinBox()
self.projX.setRange(0,1000)
self.projX.setValue(250)
self.projX.setMaximumWidth(50)
self.projY = QtGui.QSpinBox()
self.projY.setRange(0,1000)
self.projY.setValue(250)
self.projY.setMaximumWidth(50)
self.projLayout.addWidget(self.projPosL,1,0)
self.projLayout.addWidget(self.projX,1,1)
self.projLayout.addWidget(self.projY,1,2)
self.projX.valueChanged.connect(self.projectorPositionChanged)
self.projY.valueChanged.connect(self.projectorPositionChanged)
self.projSizeL = QtGui.QLabel('Size W,L')
self.projWid = QtGui.QSpinBox()
self.projWid.setRange(0,1000)
self.projWid.setValue(115)
self.projWid.setMaximumWidth(50)
self.projLen = QtGui.QSpinBox()
self.projLen.setRange(0,1000)
self.projLen.setValue(220)
self.projLen.setMaximumWidth(50)
self.projLen.valueChanged.connect(self.projectorPositionChanged)
self.projWid.valueChanged.connect(self.projectorPositionChanged)
self.projLayout.addWidget(self.projSizeL,2,0)
self.projLayout.addWidget(self.projWid,2,1)
self.projLayout.addWidget(self.projLen,2,2)
self.projRotL = QtGui.QLabel('Rotation')
self.projRot = QtGui.QSpinBox()
self.projRot.setRange(0,360)
self.projRot.setValue(270)
self.projRot.setMaximumWidth(50)
self.projLayout.addWidget(self.projRotL,3,0)
self.projLayout.addWidget(self.projRot,3,1)
self.projRot.valueChanged.connect(self.projectorPositionChanged)
self.tankLength = LabeledSpinBox(None, 'Tank Len (mm)', 0,100,46,60)
self.projLayout.addWidget(self.tankLength, 4,0,1,2)
self.tankWidth = LabeledSpinBox(None, 'Tank Wid (mm)', 0,100,21,60)
self.projLayout.addWidget(self.tankWidth, 4,2,1,2)
self.projLayout.setColumnStretch(6,1)
self.projGroup.setLayout(self.projLayout)
self.arenaLayout.addWidget(self.camCalibButton, 0,0)
self.arenaLayout.addWidget(self.projGroup,1,0)
self.arenaLayout.setColumnStretch(1,1)
self.arenaGroup.setLayout(self.arenaLayout)
#tracking group box
self.trackWidget = FishTrackerWidget(self, self.arenaMain.ftDisp)
#start button for individual tank
self.startButton = QtGui.QPushButton('Start')
self.startButton.setMaximumWidth(150)
self.startButton.clicked.connect(self.startstop)
#whether to save movie
self.paramSaveMovie = QtGui.QCheckBox('Save Movie')
#experimental parameters groupbox
self.paramGroup = QtGui.QGroupBox()
self.paramGroup.setTitle('Exprimental Parameters')
self.paramLayout = QtGui.QGridLayout()
self.paramLayout.setHorizontalSpacing(2)
self.paramLayout.setVerticalSpacing(2)
#high level flow - Between - OMR Session ( OMRS1 - pause - OMRS2
self.paramNumSession = LabeledSpinBox(None, 'Number of Sessions', 0, 200, 20, 60) #Number of Sessions (a cluster of OMR tests)
self.paramLayout.addWidget(self.paramNumSession, 1,0,1,2)
self.paramBetween = LabeledSpinBox(None, 'Between Sessions (s)', 0, 1200, 300, 60)
self.paramLayout.addWidget(self.paramBetween, 1,2,1,2)
self.paramOMRPerSession = LabeledSpinBox(None,'#OMR Per Session', 0 ,10,4,60) # Number of OMR tests per session
self.paramLayout.addWidget(self.paramOMRPerSession, 2,0,1,2)
self.paramOMRDuration = LabeledSpinBox(None,'OMR Duration (s)',0,60,12,60) #duration of each OMR tests
self.paramLayout.addWidget(self.paramOMRDuration, 2,2,1,2)
self.paramOMRInterval = LabeledSpinBox(None,'OMR Interval (s)',0,60 ,3 ,60) #time between OMR tests
self.paramLayout.addWidget(self.paramOMRInterval,3,0,1,2)
#Vision properities of OMR
self.paramOMRPeriod = LabeledSpinBox(None,'OMR Grating Period (mm)',0,50,5,60) #spacing between grating bars
self.paramLayout.addWidget(self.paramOMRPeriod,4,0,1,2)
self.paramOMRDutyCycle = LabeledSpinBox(None, 'OMR DutyCycle %',0,100,50,60) #width of grating bars
self.paramLayout.addWidget(self.paramOMRDutyCycle,4,2,1,2)
self.paramOMRVelocity = LabeledSpinBox(None, 'OMR Speed (mm/s)',0,50,5,60) #velocity of moving gratings.
self.paramLayout.addWidget(self.paramOMRVelocity,5,0,1,2)
(self.paramColorNeutral,self.labelColorNeutral) = self.getColorComboBox('Neutral', 0)
self.paramLayout.addWidget(self.paramColorNeutral,6,0)
self.paramLayout.addWidget(self.labelColorNeutral,6,1)
(self.paramColorOMR,self.labelColorOMR) = self.getColorComboBox('OMR', 5)
self.paramLayout.addWidget(self.paramColorOMR,6,2)
self.paramLayout.addWidget(self.labelColorOMR,6,3)
self.paramGroup.setLayout(self.paramLayout)
#Experimental info group
self.infoGroup = QtGui.QGroupBox()
self.infoGroup.setTitle('Experiment Info')
self.infoLayout = QtGui.QGridLayout()
self.infoLayout.setHorizontalSpacing(3)
self.infoLayout.setVerticalSpacing(3)
self.labelDir = QtGui.QLabel('Dir: ')
self.infoDir = PathSelectorWidget(browseCaption='Experimental Data Directory', default=os.path.expanduser('~/data/test'))
self.infoLayout.addWidget(self.labelDir,0,0)
self.infoLayout.addWidget(self.infoDir,0,1)
self.labelDOB = QtGui.QLabel('DOB: ')
self.infoDOB = QtGui.QDateEdit(QtCore.QDate.currentDate())
self.infoLayout.addWidget(self.labelDOB,1,0)
self.infoLayout.addWidget(self.infoDOB,1,1)
self.labelType = QtGui.QLabel('Line: ')
self.infoType = QtGui.QLineEdit('Nacre')
self.infoLayout.addWidget(self.labelType,2,0)
self.infoLayout.addWidget(self.infoType,2,1)
self.infoFish = QtGui.QPushButton('Snap Fish Image')
self.infoFish.clicked.connect(self.getFishSize)
self.infoLayout.addWidget(self.infoFish,3,0,1,2)
self.infoGroup.setLayout(self.infoLayout)
self.settingsLayout = QtGui.QGridLayout()
self.settingsLayout.addWidget(self.arenaGroup,0,0,1,2)
self.settingsLayout.addWidget(self.trackWidget,1,0,1,2)
self.settingsLayout.addWidget(self.infoGroup,2,0,1,2)
self.settingsLayout.addWidget(self.startButton,3,0,1,1)
self.settingsLayout.addWidget(self.paramSaveMovie,3,1,1,1)
self.settingsLayout.addWidget(self.paramGroup,4,0,1,2)
self.setLayout(self.settingsLayout)
#HACK Couldn't quick initialize arena lcation until after frist image arrives, so stuck in onNewFrame
#initialize projector
self.projectorPositionChanged()
self.t = 0
def getColorComboBox(self, labelName, defaultNdx=1):
colorBox = QtGui.QComboBox()
colorBox.addItem('White')
colorBox.addItem('Red')
colorBox.addItem('Blue')
colorBox.addItem('Gray')
colorBox.addItem('Magenta')
colorBox.addItem('Black')
colorBox.setCurrentIndex(defaultNdx)
colorLabel = QtGui.QLabel(labelName)
return (colorBox,colorLabel)
def configTank(self, ardConfig, camPos, projPos):
#set tank position in camera
self.resetCamCorners = camPos
#set tank position in projector
self.projX.setValue(projPos[0])
self.projY.setValue(projPos[1])
self.projWid.setValue(projPos[2])
self.projLen.setValue(projPos[3])
self.projRot.setValue(projPos[4])
#---------------------------------------------------
# OVERLOADED METHODS
#---------------------------------------------------
#update to overload an return image based on dispmode.
#def getArenaView(self):
# return self.arenaView
def onNewFrame(self, frame, time):
self.mutex.acquire()
try:
self.frameTime = time
self.currCvFrame = frame # may need to make a deep copy
(found, pos, view, allFish) = self.trackWidget.findFish(self.currCvFrame)
if found:
self.fishPosUpdate = found
self.fishPos = pos
#if experiment is running then keep track of the average fish speed in one minute blocks
if not self.is_off():
self.fishPosBuffer.append(pos)
self.fishPosBufferT.append(time)
if (self.fishPosBufferT[-1] - self.fishPosBufferT[0]) > 60: #averaging every 60 seconds, must match constant in drawDisplayOverlay
self.fishPosBuffer = [np.sqrt((f1[0]-f0[0])**2 + (f1[1]-f0[1])**2) for f0, f1 in pairwise(self.fishPosBuffer)]
self.averageSpeed.append(np.mean(self.fishPosBuffer)/(self.fishPosBufferT[-1] - self.fishPosBufferT[0])) #pixels/sec
self.fishPosBuffer = []
self.fishPosBufferT = []
if not self.is_off():
#record location
self.arenaData['video'].append((self.frameTime, None))
d = [self.frameTime, pos[0], pos[1]]
self.arenaData['tracking'].append(tuple(d))
if self.paramSaveMovie.isChecked():
img = np.array(self.currCvFrame[:,:]) #convert IplImage into numpy array
img = img[self.arenaBB[0][1]:self.arenaBB[1][1],self.arenaBB[0][0]:self.arenaBB[1][0]]
self.movie_logger.write_frame(img)
self.arenaView = view
#HACK: couldn't quickly initilize cam corner before currCvFrame is set, so I stuck it here.
#not a great place since conditoin will be checked frequently
if self.resetCamCorners is not None:
self.setArenaCamCorners(self.resetCamCorners)
self.resetCamCorners = None
except:
print 'ClassicalConditioningController:onNewFrame failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def updateState(self):
self.updateExperimentalState()
def updateExperimentalState(self):
if not self.is_off():
self.mutex.acquire()
try:
self.t = time.time()
#Check if it time to transition to next high level state (acclimate->baseline->trials->post->off)
if self.t > self.nextStateTime:
if self.is_between_session():
if self.nSession < self.paramNumSession.value():
self.nOMR = 0 #TODO move to callback for start session
self.nSession+=1 #TODO move to callback for start session
self.start_session()
else:
self.quit()
elif self.is_omr_to_s1() or self.is_omr_to_s2():
self.stop_omr()
elif self.is_no_omr_post_s1() or self.is_no_omr_post_s2():
if self.nOMR < self.paramOMRPerSession.value():
self.start_omr()
else:
self.stop_session()
#handle omr projector updates
if self.is_omr_to_s1() or self.is_omr_to_s2():
if self.t - self.omrLastUpdate > 0.010: #only updated display once sufficient time has passed.
if self.t - self.omrLastUpdate > 0.060:
print 'WARNING: projector slow to update', self.t - self.omrLastUpdate
self.updateProjectorDisplay()
except:
print 'ContextualHelplessnessController:updateState failed'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def projectorPositionChanged(self):
#the point defining the tank are such that:
#points 1 and 2 define side1
#points 2 and 3 define the 'near' connecting edge
#points 3 and 4 define side2
#points 4 and 1 define the 'far' connecting edge.
#move this code elsewhere to avoid redundant computation...
self.pts = np.array([[0 , 0],
[0 , self.projWid.value()],
[self.projLen.value(), self.projWid.value()],
[self.projLen.value(), 0]])
theta = self.projRot.value()
self.tankRotM = np.array([[cos(radians(theta)), -sin(radians(theta))],
[sin(radians(theta)), cos(radians(theta))]]);
self.pts = np.dot(self.tankRotM,self.pts.T).T
self.pts += [self.projX.value(), self.projY.value()]
self.updateProjectorDisplay()
def drawProjectorDisplay(self, painter):
if self.is_off() and self.projCalibButton.isChecked() and self.isCurrent():
#DRAW A CALIBRATION IMAGE THAT INDICATES THE PROJECTOR LOCATION AND SIDE1
side1Color = QtCore.Qt.red
side2Color = QtCore.Qt.blue
a = .5
b = 1-a
#Draw side one
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(side1Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.pts[0,0], self.pts[0,1]))
poly.append(QtCore.QPointF(self.pts[1,0], self.pts[1,1]))
poly.append(QtCore.QPointF(a*self.pts[1,0] + b*self.pts[2,0], a*self.pts[1,1] + b*self.pts[2,1]))
poly.append(QtCore.QPointF(a*self.pts[0,0] + b*self.pts[3,0], a*self.pts[0,1] + b*self.pts[3,1]))
painter.drawPolygon(poly)
painter.setPen(QtGui.QColor(168, 34, 3))
painter.setFont(QtGui.QFont('Decorative',24))
painter.drawText(self.pts[0,0],self.pts[0,1],'1')
painter.drawText(self.pts[1,0],self.pts[1,1],'2')
painter.drawText(self.pts[2,0],self.pts[2,1],'3')
painter.drawText(self.pts[3,0],self.pts[3,1],'4')
#Draw side two
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = QtGui.QBrush(side2Color)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(a*self.pts[0,0] + b*self.pts[3,0], a*self.pts[0,1] + b*self.pts[3,1]))
poly.append(QtCore.QPointF(a*self.pts[1,0] + b*self.pts[2,0], a*self.pts[1,1] + b*self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[2,0], self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[3,0], self.pts[3,1]))
painter.drawPolygon(poly)
else:
#Draw Background
backColor = self.paramColorNeutral.currentIndex()
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(backColor)
painter.setBrush(brush)
painter.setPen(pen)
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(self.pts[0,0], self.pts[0,1]))
poly.append(QtCore.QPointF(self.pts[1,0], self.pts[1,1]))
poly.append(QtCore.QPointF(self.pts[2,0], self.pts[2,1]))
poly.append(QtCore.QPointF(self.pts[3,0], self.pts[3,1]))
painter.drawPolygon(poly)
#Draw OMR grating
if self.is_omr_to_s1() or self.is_omr_to_s2():
#update the possition according to time since last update.
if self.is_omr_to_s1():
self.omrPhase -= (float(self.paramOMRVelocity.value())/float(self.paramOMRPeriod.value())) * 360.0 * (self.t-self.omrLastUpdate)
else:
self.omrPhase += (float(self.paramOMRVelocity.value())/float(self.paramOMRPeriod.value())) * 360.0 * (self.t-self.omrLastUpdate)
self.omrPhase = self.omrPhase%360.0
self.omrLastUpdate = self.t
#This stuff doesn't need to computed on every frame
mm2pix = math.hypot(self.pts[2,0]-self.pts[1,0], self.pts[2,1] - self.pts[1,1]) / self.tankLength.value()
period = self.paramOMRPeriod.value() * mm2pix
numBars = int(math.ceil(math.hypot(self.pts[2,0]-self.pts[1,0], self.pts[2,1] - self.pts[1,1]) / period))
barWidth = period * (self.paramOMRDutyCycle.value()/100.0)
#Draw the bar
phaseShift = self.omrPhase/360.0 * period
pen = QtGui.QPen(QtCore.Qt.NoPen)
brush = self.getBrush(self.paramColorOMR.currentIndex())
painter.setBrush(brush)
painter.setPen(pen)
for nBar in range(-1,numBars+1):
#define bar in unrotated untranslated space.
bl = min(max(phaseShift + nBar * period, 0), self.projLen.value())
br = min(max(phaseShift + nBar * period + barWidth, 0), self.projLen.value())
bar = np.array([[br,0],
[bl,0],
[bl,self.projWid.value()],
[br,self.projWid.value()]])
#rotate and tranlate.
bar = np.dot(self.tankRotM, bar.T).T + [self.projX.value(), self.projY.value()]
#draw
poly = QtGui.QPolygonF()
poly.append(QtCore.QPointF(bar[0,0],bar[0,1]))
poly.append(QtCore.QPointF(bar[1,0],bar[1,1]))
poly.append(QtCore.QPointF(bar[2,0],bar[2,1]))
poly.append(QtCore.QPointF(bar[3,0],bar[3,1]))
painter.drawPolygon(poly)
def drawDisplayOverlay(self, painter):
#draw the fish position
if self.fishPos and self.fishPos[0] > 0:
brush = QtGui.QBrush(QtCore.Qt.red)
painter.setBrush(brush)
painter.setPen(QtCore.Qt.NoPen)
painter.drawEllipse(QtCore.QPointF(self.fishPos[0],self.fishPos[1]), 3,3)
#draw the arena overlay
if self.arenaCamCorners:
painter.setBrush(QtCore.Qt.NoBrush)
if self.bIsCurrentArena:
pen = QtGui.QPen(QtCore.Qt.green)
else:
pen = QtGui.QPen(QtCore.Qt.blue)
pen.setWidth(3)
painter.setPen(pen)
poly = QtGui.QPolygonF()
for p in self.arenaCamCorners:
poly.append(QtCore.QPointF(p[0],p[1]))
painter.drawPolygon(poly)
if len(self.arenaCamCorners) >= 2:
pen = QtGui.QPen(QtCore.Qt.red)
pen.setWidth(2)
painter.setPen(pen)
painter.drawLine(self.arenaCamCorners[0][0],self.arenaCamCorners[0][1],
self.arenaCamCorners[1][0],self.arenaCamCorners[1][1])
painter.setPen(QtGui.QColor(168, 34, 3))
painter.setFont(QtGui.QFont('Decorative',14))
statustext = ""
statustext = statustext + self.state
if not self.is_off():
statustext = statustext + ": %d/%d %0.1f"%(self.nSession, self.paramNumSession.value(), self.nextStateTime-self.t)
if self.is_omr_to_s1() or self.is_omr_to_s2() or self.is_no_omr_post_s1() or self.is_no_omr_post_s2():
statustext = statustext + ": %d/%d"%(self.nOMR, self.paramOMRPerSession.value())
painter.drawText(self.arenaCamCorners[0][0],self.arenaCamCorners[0][1],statustext)
#Plot average speed of fish overtime
if not self.is_off() and len(self.averageSpeed)>2:
ac = np.array(self.arenaCamCorners)
plt_b = np.max(ac[:,1]) #Plot bottom
plt_l = np.min(ac[:,0]) #Plot left
plt_r = np.max(ac[:,0]) #Plot right
totalDuration = self.paramBetween.value() + self.paramNumSession.value()*(self.paramBetween.value()+
self.paramOMRPerSession.value() * (self.paramOMRDuration.value() + self.paramOMRInterval.value()))
totalDuration = totalDuration*(60/60.0) #averaging every N seconds (see onnewframe)
ndx = np.arange(len(self.averageSpeed))
x = (np.arange(len(self.averageSpeed))/totalDuration)
x = x * (plt_r - plt_l)
x = x + plt_l
#scale = 75 / np.max(self.averageSpeed)
scale = 3000
y = plt_b - (scale * np.array(self.averageSpeed))
for ((x1,x2),(y1,y2)) in zip(pairwise(x),pairwise(y)):
painter.drawLine(x1,y1,x2,y2)
def start(self):
self.startstop()
def stop(self):
self.startstop()
#---------------------------------------------------
# STATE MACHINE CALLBACK METHODS
#---------------------------------------------------
def update_state_data(self):
# runs on every state transition...
self.updateProjectorDisplay()
self.arenaData['stateinfo'].append((self.t, 0, 0, 0, self.state))
#self.saveResults() #just save at end to avoid long interframe intervals
def on_exit_off(self):
#experiment has started so disable parts of UI
self.paramGroup.setDisabled(True)
self.infoGroup.setDisabled(True)
self.startButton.setText('Stop')
#experiment has started so prepare result files and data structures
td = datetime.datetime.now()
self.saveLocation = str(self.infoDir.text())
[p, self.fnResults] = os.path.split(self.saveLocation)
self.fnResults = self.fnResults + '_' + td.strftime('%Y-%m-%d-%H-%M-%S')
self.jsonFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '.json'
#prepare to write movie
if self.paramSaveMovie.isChecked():
self.movieFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '.mp4'
from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter as VidWriter
self.movie_logger = VidWriter(filename=self.movieFileName,
size=(self.arenaBB[1][0] - self.arenaBB[0][0],
self.arenaBB[1][1] - self.arenaBB[0][1]),
fps=15,
codec='mpeg4',
preset='ultrafast')
self.initArenaData()
self.nSession = 0
self.averageSpeed = []
def on_enter_between_session(self):
self.nextStateTime = self.t + self.paramBetween.value()
def on_enter_omr_to_s1(self):
self.nOMR += 1
self.nextStateTime = self.t + self.paramOMRDuration.value()
self.omrPhase = 0
self.omrLastUpdate = self.t #updates everytime the projectory updates
self.updateProjectorDisplay()
def on_exit_omr_to_s1(self):
self.updateProjectorDisplay()
def on_enter_omr_to_s2(self):
self.nOMR += 1
self.nextStateTime = self.t + self.paramOMRDuration.value()
self.omrPhase = 0
self.omrLastUpdate = self.t #updates everytime the projectory updates
self.updateProjectorDisplay()
def on_exit_omr_to_s2(self):
self.updateProjectorDisplay()
def on_enter_no_omr_post_s1(self):
self.nextStateTime = self.t + self.paramOMRInterval.value()
self.updateProjectorDisplay()
def on_enter_no_omr_post_s2(self):
self.nextStateTime = self.t + self.paramOMRInterval.value()
self.updateProjectorDisplay()
def on_enter_off(self):
self.nextStateTime = None
self.updateProjectorDisplay()
self.saveResults()
#note: the projector colors will be reset on call to after_state_change
if self.paramSaveMovie.isChecked():
self.movie_logger.close()
#experiment has ended so enable parts of UI
self.paramGroup.setDisabled(False)
self.infoGroup.setDisabled(False)
self.startButton.setText('Start')
def isReadyToStart(self):
if not os.path.exists(self.infoDir.text()):
try:
os.mkdir(self.infoDir.text())
except:
self.arenaMain.statusBar().showMessage('%s arena not ready to start. Experiment directory does not exist and cannot be created.'%(self.getYokedStr()))
return False
if not self.arenaCamCorners:
self.arenaMain.statusBar().showMessage('%s arena not ready to start. The arena location has not been specified.'%(self.getYokedStr()))
return False
if not self.trackWidget.getBackgroundImage():
self.arenaMain.statusBar().showMessage('%s arena not ready to start. Background image has not been taken.'%(self.getYokedStr()))
return False
#if not self.fishImg:
# self.arenaMain.statusBar().showMessage('%s arena not ready to start. Fish image has not been taken.'%(self.getYokedStr()))
# return False
return True
#---------------------------------------------------
# SLOT CALLBACK METHODS
#---------------------------------------------------
def startstop(self):
self.mutex.acquire()
try:
self.t = time.time()
if self.is_off():
self.begin()
else:
self.quit();
except:
print 'Start Failed:'
traceback.print_exc()
QtCore.pyqtRemoveInputHook()
ipdb.set_trace()
finally:
self.mutex.release()
def getArenaCameraPosition(self):
self.arenaMain.statusBar().showMessage('Click on the corners of the arena on side 1.')
self.currArenaClick = 0
self.arenaCamCorners = []
self.arenaMain.ftDisp.clicked.connect(self.handleArenaClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleArenaClicks(self, x, y):
self.currArenaClick+=1
if self.currArenaClick<5:
self.arenaCamCorners.append((x,y))
if self.currArenaClick==1:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 1.')
elif self.currArenaClick==2:
self.arenaMain.statusBar().showMessage('Now, click on the corners of the arena on side 2.')
elif self.currArenaClick==3:
self.arenaMain.statusBar().showMessage('Click on the other corner of the arena on side 2.')
elif self.currArenaClick==4:
self.arenaMain.ftDisp.clicked.disconnect(self.handleArenaClicks)
self.arenaMain.statusBar().showMessage('')
self.setArenaCamCorners(self.arenaCamCorners)
def setArenaCamCorners(self, corners):
#corners must be tuple of tuples (not list or np.array)
self.currArenaclick=4
self.arenaCamCorners = corners
print corners
[self.arenaMidLine, self.arenaSide1Sign] = self.processArenaCorners(self.arenaCamCorners, .5)
#compute bounding box with vertical and horizontal sides.
self.arenaBB = [[min([p[0] for p in corners]), min([p[1] for p in corners])],
[max([p[0] for p in corners]), max([p[1] for p in corners])]]
#bounding box needs to have even heights and widths in order to save as mp4
if (self.arenaBB[1][0] - self.arenaBB[0][0]) % 2:
self.arenaBB[1][0] += 1
if (self.arenaBB[1][1] - self.arenaBB[0][1]) % 2:
self.arenaBB[1][1] += 1
self.getArenaMask()
self.trackWidget.setBackgroundImage()
def getFishSize(self):
self.arenaMain.statusBar().showMessage('Click on the tip of the fish tail.')
self.currFishClick = 0
self.fishSize = []
self.arenaMain.ftDisp.clicked.connect(self.handleFishClicks)
@QtCore.pyqtSlot(int, int) #not critical but could practice to specify which functions are slots.
def handleFishClicks(self, x, y):
self.currFishClick+=1
if self.currFishClick == 1:
self.fishSize.append((x,y))
self.arenaMain.statusBar().showMessage('Click on the tip of the fish head.')
elif self.currFishClick == 2:
self.arenaMain.ftDisp.clicked.disconnect(self.handleFishClicks)
self.fishSize.append((x,y))
self.fishImg = cv.CloneImage(self.currCvFrame)
self.arenaMain.statusBar().showMessage('')
#---------------------------------------------------
# HELPER METHODS
#---------------------------------------------------
def processArenaCorners(self, arenaCorners, linePosition):
#return the line dividing the center of the arena, and a definition of side 1.
a = 1-linePosition
b = linePosition
ac = np.array(arenaCorners)
#arenaDivideLine = [tuple(np.mean(ac[(0,3),:],axis = 0)),tuple(np.mean(ac[(1,2),:],axis = 0))]
arenaDivideLine = [(a*ac[0,0]+b*ac[3,0], a*ac[0,1]+b*ac[3,1]),(a*ac[1,0]+b*ac[2,0], a*ac[1,1]+b*ac[2,1])]
side1Sign = 1
if not self.isOnSide(arenaCorners[1], arenaDivideLine, side1Sign):
side1Sign = -1
return (arenaDivideLine, side1Sign)
def isOnSide(self, point, line, sideSign):
#return if the fish is on side1 of the arena.
side = (line[1][0] - line[0][0]) * (point[1] - line[0][1]) - (line[1][1] - line[0][1]) * (point[0] - line[0][0])
return cmp(side,0)==sideSign
#convert the arena corners into a color mask image (arena=255, not=0)
def getArenaMask(self):
if self.arenaView:
cvImg = self.currCvFrame
self.arenaCvMask = cv.CreateImage((cvImg.width,cvImg.height), cvImg.depth, cvImg.channels)
cv.SetZero(self.arenaCvMask)
cv.FillConvexPoly(self.arenaCvMask, self.arenaCamCorners, (255,)*cvImg.channels)
self.trackWidget.setTrackMask(self.arenaCvMask)
def initArenaData(self):
#prepare output data structure
self.arenaData = {}
self.arenaData['fishbirthday'] = str(self.infoDOB.date().toPyDate())
self.arenaData['fishage'] = (datetime.date.today() - self.infoDOB.date().toPyDate()).days
self.arenaData['fishstrain'] = str(self.infoType.text())
self.arenaData['fishsize'] = self.fishSize
self.arenaData['parameters'] = { 'Num Sessions':self.paramNumSession.value(),
'Between (s)':self.paramBetween.value(),
'OMRPerSession':self.paramOMRPerSession.value(),
'OMR Duration (s)':self.paramOMRDuration.value(),
'OMR Interval (s)':self.paramOMRInterval.value(),
'OMR Period (mm)': self.paramOMRPeriod.value(),
'OMR Duty Cycle': self.paramOMRDutyCycle.value(),
'OMR Velocity':self.paramOMRVelocity.value(),
'NeutralColor':str(self.paramColorNeutral.currentText()),
'OMRColor':str(self.paramColorOMR.currentText()),
'states':self.states.keys(),
'CodeVersion':None }
self.arenaData['trackingParameters'] = self.trackWidget.getParameterDictionary()
self.arenaData['trackingParameters']['arenaPoly'] = self.arenaCamCorners
self.arenaData['trackingParameters']['arenaDivideLine'] = self.arenaMidLine
self.arenaData['trackingParameters']['arenaSide1Sign'] = self.arenaSide1Sign
self.arenaData['projectorParameters'] = {'position':[self.projX.value(), self.projY.value()],
'size':[self.projLen.value(), self.projWid.value()],
'rotation':self.projRot.value()}
self.arenaData['tankSize_mm'] = [self.tankLength.value(), self.tankWidth.value()]
self.arenaData['trials'] = list() #outcome on each trial
self.arenaData['tracking'] = list() #list of tuples (frametime, posx, posy)
self.arenaData['video'] = list() #list of tuples (frametime, filename)
self.arenaData['stateinfo'] = list() #list of times at switch stimulus flipped.
self.arenaData['shockinfo'] = list()
t = datetime.datetime.now()
#save experiment images
self.bcvImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_BackImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.bcvImgFileName, self.trackWidget.getBackgroundImage())
if self.fishImg:
self.fishImgFileName = str(self.infoDir.text()) + os.sep + self.fnResults + '_FishImg_' + t.strftime('%Y-%m-%d-%H-%M-%S') + '.tiff'
cv.SaveImage(self.fishImgFileName, self.fishImg)
def saveResults(self):
f = open(name=self.jsonFileName, mode='w')
json.dump(self.arenaData,f)
f.close()
def getBrush(self, colorNdx):
if colorNdx == 0:
return QtGui.QBrush(QtCore.Qt.white)
elif colorNdx == 1:
return QtGui.QBrush(QtCore.Qt.red)
elif colorNdx == 2:
return QtGui.QBrush(QtCore.Qt.blue)
elif colorNdx == 3:
return QtGui.QBrush(QtGui.QColor(128,128,128))
elif colorNdx == 4:
return QtGui.QBrush(QtCore.Qt.magenta)
elif colorNdx == 5:
return QtGui.QBrush(QtCore.Qt.black)
else:
return QtGui.QBrush(QtCore.Qt.black)
def useDebugParams(self):
pass
class FishTrackerWidget(QtGui.QGroupBox):
def __init__(self, parent, ftDisp):
#ftDisp is necessary so the FishTrackerWidget can collect mouse clicks
super(FishTrackerWidget, self).__init__(parent)
self.ftDisp = ftDisp
self.currCvFrame = None
self.arenaCvMask = None #track mask
self.bcvImg = None #background image for subtraction
self.currG = None
self.maskG = None
self.backG = None
self.backG32 = None
self.diffG = None
self.thrsG = None
self.thrsMG = None
self.tracEG = None
self.tracDG = None
self.fmask = None
self.bFixBackgroundNow = False
self.setTitle('Tracking Parameters')
self.trackLayout = QtGui.QGridLayout(self)
self.trackLayout.setHorizontalSpacing(3)
self.trackLayout.setVerticalSpacing(3)
self.viewLabel = QtGui.QLabel('Track Display Mode:')
self.viewLabel.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
self.viewLabel.adjustSize()
self.selView = QtGui.QComboBox(self)
self.selView.addItem('Raw')
self.selView.addItem('Background')
self.selView.addItem('Subtracted')
self.selView.addItem('Thresholded')
self.selView.addItem('Masked')
self.selView.addItem('Eroded')
self.selView.addItem('Dilated')
self.selView.addItem('Mask')
self.trackLayout.addWidget(self.viewLabel, 0,0,1,2)
self.trackLayout.addWidget(self.selView, 0,2,1,2)
self.bgButton = QtGui.QPushButton('Get Background')
self.bgButton.clicked.connect(self.setBackgroundImage)
self.trackLayout.addWidget(self.bgButton, 1,0,1,2)
self.fixButton = QtGui.QPushButton('Fix Background')
self.fixButton.clicked.connect(self.fixBackground)
self.trackLayout.addWidget(self.fixButton, 1,2,1,2)
self.trackThreshold = LabeledSpinBox(None,'Threshold',0,255,3,60)
self.trackErode = LabeledSpinBox(None,'Erode',0,20,0,60)
self.trackDilate = LabeledSpinBox(None,'Dilate',0,20,0,60)
self.trackMinArea = LabeledSpinBox(None,'MinArea',0,600000,0,75)
self.trackMaxArea = LabeledSpinBox(None,'MaxArea',0,600000,600000,75)
self.updateCheckbox = QtGui.QCheckBox('Use Updating BG')
self.updateCheckbox.setChecked(True)
self.numFish = LabeledSpinBox(None,'NumFish',0,20,1,60)
self.fishSize = LabeledSpinBox(None,'FG Size',0,1000,30,60)
self.learningRate = LabeledDoubleSpinBox(None,'Learning rate',0,1,0.1,60)
self.trackLayout.addWidget(self.trackThreshold,2,0,1,2)
self.trackLayout.addWidget(self.trackErode, 2,2,1,2)
self.trackLayout.addWidget(self.trackDilate, 3,0,1,2)
self.trackLayout.addWidget(self.trackMinArea, 3,2,1,2)
self.trackLayout.addWidget(self.trackMaxArea, 4,0,1,2)
self.trackLayout.addWidget(self.updateCheckbox,5,0,1,2)
self.trackLayout.addWidget(self.numFish,5,2,1,2)
self.trackLayout.addWidget(self.fishSize,6,0,1,2)
self.trackLayout.addWidget(self.learningRate,6,2,1,2)
self.setLayout(self.trackLayout)
def setBackgroundImage(self):
if self.currCvFrame:
self.bcvImg = cv.CloneImage(self.currCvFrame)
if not self.maskG:
self.backG = cv.CreateImage(cv.GetSize(self.currCvFrame), cv.IPL_DEPTH_8U, 1)
self.backG32 = cv.CreateImage(cv.GetSize(self.currCvFrame), cv.IPL_DEPTH_32F, 1)
cv.CvtColor(self.bcvImg, self.backG, cv.CV_BGR2GRAY)
cv.ConvertScale(self.backG, self.backG32)
else:
self.backG = cv.CreateImage(cv.GetSize(self.maskG), cv.IPL_DEPTH_8U, 1)
self.backG32 = cv.CreateImage(cv.GetSize(self.maskG), cv.IPL_DEPTH_32F, 1)
cv.CvtColor(self.bcvImg[self.maskBoundedR0:self.maskBoundedR0+self.maskBoundedHeight,
self.maskBoundedC0:self.maskBoundedC0+self.maskBoundedWidth], self.backG, cv.CV_BGR2GRAY)
cv.ConvertScale(self.backG, self.backG32)
def getBackgroundImage(self):
return self.bcvImg
def fixBackground(self):
self.ftDisp.clicked.connect(self.handleFixBackgroundClick)
self.bFixBackgroundNow = False
def handleFixBackgroundClick(self,x,y):
self.correctFish = (x - self.maskBoundedC0, y - self.maskBoundedR0)
self.ftDisp.clicked.disconnect(self.handleFixBackgroundClick)
self.bFixBackgroundNow = True
def setTrackMask(self, img):
#Convert mask to gray scale.
self.arenaCvMask = img
self.maskG = cv.CreateImage((self.arenaCvMask.width, self.arenaCvMask.height), cv.IPL_DEPTH_8U, 1)
cv.CvtColor(self.arenaCvMask, self.maskG, cv.CV_BGR2GRAY)
#get bounding rectangle of mask and resize it
maskArray = np.asarray(self.maskG[:,:])
self.maskBoundedC0 = np.min(np.nonzero(np.max(maskArray,0))) #X = width = cols
self.maskBoundedR0 = np.min(np.nonzero(np.max(maskArray,1))) #Y = height = rows
self.maskBoundedWidth = np.max(np.nonzero(np.max(maskArray,0))) - self.maskBoundedC0 + 1
self.maskBoundedHeight = np.max(np.nonzero(np.max(maskArray,1))) - self.maskBoundedR0 + 1
self.maskG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
cv.CvtColor(self.arenaCvMask[self.maskBoundedR0:self.maskBoundedR0+self.maskBoundedHeight,
self.maskBoundedC0:self.maskBoundedC0+self.maskBoundedWidth], self.maskG, cv.CV_BGR2GRAY)
#recrop background image
if self.bcvImg:
self.backG = cv.CreateImage(cv.GetSize(self.maskG), cv.IPL_DEPTH_8U, 1)
self.backG32 = cv.CreateImage(cv.GetSize(self.maskG), cv.IPL_DEPTH_32F, 1)
cv.CvtColor(self.bcvImg[self.maskBoundedR0:self.maskBoundedR0+self.maskBoundedHeight,
self.maskBoundedC0:self.maskBoundedC0+self.maskBoundedWidth], self.backG, cv.CV_BGR2GRAY)
cv.ConvertScale(self.backG, self.backG32)
#resize temp images to bounded rect size
self.currG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.diffG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.thrsG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.thrsMG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.tracEG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.tracDG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.tracG = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
self.fmask = cv.CreateImage((self.maskBoundedWidth, self.maskBoundedHeight), cv.IPL_DEPTH_8U, 1)
def getParameterDictionary(self):
pd = {}
pd['nDiffThreshold'] = self.trackThreshold.value()
pd['nErode'] = self.trackErode.value()
pd['nDilate'] = self.trackDilate.value()
return pd
def getTrackDisplay(self):
if self.currCvFrame == None:
return None
dispImg = cv.CreateImage(cv.GetSize(self.currCvFrame), cv.IPL_DEPTH_8U, 1)
cv.CvtColor(self.currCvFrame, dispImg, cv.CV_BGR2GRAY)
if self.maskG:
cv.SetImageROI(dispImg, (self.maskBoundedC0, self.maskBoundedR0, self.maskBoundedWidth, self.maskBoundedHeight))
dispModeNdx = self.selView.currentIndex()
if dispModeNdx == 1 and self.backG:
cv.Copy(self.backG, dispImg)
elif dispModeNdx == 2 and self.diffG:
cv.Copy(self.diffG, dispImg)
elif dispModeNdx == 3 and self.thrsG:
cv.Copy(self.thrsG, dispImg)
elif dispModeNdx == 4 and self.thrsMG:
cv.Copy(self.thrsMG, dispImg)
elif dispModeNdx == 5 and self.tracEG:
cv.Copy(self.tracEG, dispImg)
elif dispModeNdx == 6 and self.tracDG:
cv.Copy(self.tracDG, dispImg)
elif dispModeNdx == 7 and self.maskG:
cv.Copy(self.maskG, dispImg)
else:
return self.currCvFrame
cv.ResetImageROI(dispImg)
return dispImg
def findFish(self, cvImg):
####NEED TO MODIFY TO USE CV2
self.currCvFrame = cvImg
foundFish = False
fishPos = (0,0)
allFish = []
if not self.bcvImg == None and not self.arenaCvMask == None:
#Background subtract, threshold, mask, erode and dilate
cv.CvtColor(self.currCvFrame[self.maskBoundedR0:self.maskBoundedR0+self.maskBoundedHeight,
self.maskBoundedC0:self.maskBoundedC0+self.maskBoundedWidth], self.currG, cv.CV_BGR2GRAY)
cv.AbsDiff(self.currG, self.backG, self.diffG) #difference
cv.Threshold ( self.diffG , self.thrsG , self.trackThreshold.value() , 255 , cv.CV_THRESH_BINARY ) #threshold
cv.And( self.thrsG, self.maskG, self.thrsMG ) #mask
cv.Erode(self.thrsMG, self.tracEG, erodeKernal, self.trackErode.value()) #erode
cv.Dilate(self.tracEG, self.tracDG, dilateKernal, self.trackDilate.value()) #dilate
cv.Copy(self.tracDG, self.tracG)
#Get List of connected components
seq = cv.FindContours(self.tracG, cv.CreateMemStorage(0), cv.CV_RETR_EXTERNAL, cv.CV_CHAIN_APPROX_NONE)
#Get each connected components area and center of mass (first moments)
areas = []
moments = []
if len(seq)>0:
while seq != None:
#bx, by, bwidth, bheight = cv.BoundingRect(seq, 0)
areas.append(cv.ContourArea(seq))
moments.append(cv.Moments(seq))
seq = seq.h_next()
#get the largest connected component
ndx = areas.index(max(areas))
if (moments[ndx].m00 > self.trackMinArea.value() and moments[ndx].m00 < self.trackMaxArea.value()):
foundFish = True
fishPos = (moments[ndx].m10/moments[ndx].m00 + self.maskBoundedC0, moments[ndx].m01/moments[ndx].m00 + self.maskBoundedR0)
#get all fish sorted by size.
sndx = np.argsort(areas)[::-1]
for bn in sndx:
if (moments[bn].m00 > self.trackMinArea.value() and moments[bn].m00 < self.trackMaxArea.value()):
allFish.append((moments[bn].m10/moments[bn].m00 + self.maskBoundedC0, moments[bn].m01/moments[bn].m00 + self.maskBoundedR0))
del seq
#update background image:
if self.bFixBackgroundNow:
cv.Rectangle(self.fmask,(0,0),cv.GetSize(self.fmask),255,-1)
cv.Circle(self.fmask, tuple([int(x) for x in self.correctFish]), self.fishSize.value(), 0, -1)
cv.Copy(self.currG, self.backG, self.fmask)
cv.ConvertScale(self.backG, self.backG32)
self.bFixBackgroundNow = False
elif self.updateCheckbox.isChecked():
if foundFish:
cv.Rectangle(self.fmask,(0,0),cv.GetSize(self.fmask),255,-1)
for nFish in range(min(self.numFish.value(),len(allFish))):
pos = (int(allFish[nFish][0] - self.maskBoundedC0), int(allFish[nFish][1] - self.maskBoundedR0))
cv.Circle(self.fmask, pos, self.fishSize.value(),0,-1)
cv.RunningAvg(self.currG, self.backG32, self.learningRate.value(), self.fmask)
cv.ConvertScaleAbs(self.backG32,self.backG)
return (foundFish, fishPos, self.getTrackDisplay(), allFish)