Tracking2d

# [DocSampleImports]
import sys
import os
import PyIPSDK
import PyIPSDK.IPSDKIPLFeatureDetection as fd
import PyIPSDK.IPSDKIPLRegistration as registration
from PyQt5.QtWidgets import QWidget, QApplication, QLabel
from PyQt5.QtGui import QPainter, QColor, QFont, QPen, QBrush, QImage, QPixmap, QPolygon
from PyQt5.QtCore import Qt, QCoreApplication, QPoint
# [DocSampleImports]

class TrackingViewer(QWidget): 
    
    def __init__(self):
        super(QWidget, self).__init__()
        
# [DocSampleProcessSequence]
        # open sequence image
        imagesSamplePath = PyIPSDK.getIPSDKDirectory(PyIPSDK.eInternalDirectory.eID_Images)
        inputImgPath = os.path.join(imagesSamplePath,  "rotation_seq_2d_1.tif")
        self.seqImg = PyIPSDK.loadTiffImageFile(inputImgPath, PyIPSDK.eTiffDirectoryMode.eTDM_Temporal)
        firstImgPlan = PyIPSDK.extractPlan(0, 0, 0, self.seqImg)        
        
        # extract harris features
        nbSamples = 400
        minDist = 5
        pixels2d = fd.harrisCorner2d(firstImgPlan, nbSamples, minDist)
        
        # process registration
        self.initGrid = PyIPSDK.createRegistrationTracking2dGrid(self.seqImg.getSizeX() * 0.25,
                                                                 self.seqImg.getSizeY() * 0.25,
                                                                 self.seqImg.getSizeX() * 0.75,
                                                                 self.seqImg.getSizeY() * 0.75)
        self.initPts = PyIPSDK.toCoords2dColl(pixels2d)
        self.piMotionTransform2d = registration.intensityBasedTracker2d(self.seqImg, self.initPts,
                                                                        PyIPSDK.eRegistrationMotionModel2d.eRMM2d_Rigid,
                                                                        self.initGrid)
# [DocSampleProcessSequence]

        # setup interface
        self.setGeometry(300, 300, self.seqImg.getSizeX(), self.seqImg.getSizeY())
        self.updateImg(0)
        self.show()
        
    def paintEvent(self, event):
        paint = QPainter()
        paint.begin(self)
        paint.drawImage(QPoint(0,0), self.qImg)
        for curCoord, flag in zip(self.curPts.coll, self.flags):
            if flag:
                paint.setPen(QPen(Qt.green))
                paint.setBrush(QBrush(Qt.green))
            else:
                paint.setPen(QPen(Qt.red))
                paint.setBrush(QBrush(Qt.red))
            paint.drawEllipse(curCoord.x, curCoord.y, 5, 5)
        paint.setPen(QPen(Qt.green, 3))
        grid = QPolygon()
        grid.append(QPoint(self.trackingGrid.pt0.x, self.trackingGrid.pt0.y))
        grid.append(QPoint(self.trackingGrid.pt1.x, self.trackingGrid.pt1.y))
        grid.append(QPoint(self.trackingGrid.pt2.x, self.trackingGrid.pt2.y))
        grid.append(QPoint(self.trackingGrid.pt3.x, self.trackingGrid.pt3.y))
        grid.append(QPoint(self.trackingGrid.pt0.x, self.trackingGrid.pt0.y))
        paint.drawPolyline(grid)
        paint.end()
        
    def applyTransform(self, curMotion):
        # Create a rigid transform from the curMotion parameters
        rtparams = curMotion.params
        rt = PyIPSDK.RigidTransform2d(rtparams[0], rtparams[1], rtparams[2])
        
        # Convert the Coord2dColl to a list
        lPoints = []
        for i in range(0, len(self.curPts.coll)):
            lPoints.append(PyIPSDK.Real64Point2d(self.curPts.coll[i].x, self.curPts.coll[i].y))
        
        # Convert the RegistrationTracking2dGrid to a list
        lGrid = [PyIPSDK.Real64Point2d(self.trackingGrid.pt0.x, self.trackingGrid.pt0.y),
                 PyIPSDK.Real64Point2d(self.trackingGrid.pt1.x, self.trackingGrid.pt1.y),
                 PyIPSDK.Real64Point2d(self.trackingGrid.pt2.x, self.trackingGrid.pt2.y),
                 PyIPSDK.Real64Point2d(self.trackingGrid.pt3.x, self.trackingGrid.pt3.y)]

        # Apply the transform to the point clouds and the grid
        pointClouds = PyIPSDK.apply(rt, PyIPSDK.Real64PointCloud2d(lPoints))
        grid = PyIPSDK.apply(rt, PyIPSDK.Real64PointCloud2d(lGrid))

        # Convert the point cloud to the Coord2dColl attribute
        for i in range(0, len(self.curPts.coll)):
            self.curPts.coll[i].x = pointClouds.getCollection()[i].x
            self.curPts.coll[i].y = pointClouds.getCollection()[i].y

        # Convert the grid list to a RegistrationTracking2dGrid
        self.trackingGrid.pt0.x = grid.getCollection()[0].x
        self.trackingGrid.pt0.y = grid.getCollection()[0].y
        self.trackingGrid.pt1.x = grid.getCollection()[1].x
        self.trackingGrid.pt1.y = grid.getCollection()[1].y
        self.trackingGrid.pt2.x = grid.getCollection()[2].x
        self.trackingGrid.pt2.y = grid.getCollection()[2].y
        self.trackingGrid.pt3.x = grid.getCollection()[3].x
        self.trackingGrid.pt3.y = grid.getCollection()[3].y
        
    def updateImg(self, newT):
        self.curT = newT
        if (self.curT >= self.seqImg.getSizeT()):
            self.curT = self.seqImg.getSizeT() - 1
        if (self.curT < 0):
            self.curT = 0
        curImgPlan = self.seqImg.array[self.curT, :, :]
        self.qImg = QImage(curImgPlan.data, curImgPlan.shape[1], curImgPlan.shape[0], curImgPlan.shape[1], QImage.Format_Grayscale8)
        self.setWindowTitle('Tracking sequence [' + str(self.curT) + " / " + str(self.seqImg.getSizeT()) + "]")
        curMotion = self.piMotionTransform2d.getValue(0, 0, self.curT)
        self.curPts = self.initPts.duplicate()
        self.trackingGrid = self.initGrid.duplicate()
        
        self.applyTransform(curMotion)
        
        self.flags = []
        for curCoord in self.curPts.coll:
            self.flags.append(PyIPSDK.isInside(self.trackingGrid, curCoord))
        self.update()
        
    def keyPressEvent(self, e):

        if e.key() == Qt.Key_Right:
            self.updateImg(self.curT+1)
        if e.key() == Qt.Key_Left:
            self.updateImg(self.curT-1)
        if e.key() == Qt.Key_Up:
            self.updateImg(self.seqImg.getSizeT())
        if e.key() == Qt.Key_Down:
            self.updateImg(0)

if __name__ == "__main__": 
    app = QCoreApplication.instance()
    if app is None:
        app = QApplication([])
    viewer = TrackingViewer()
    sys.exit(app.exec_())