3dtof_viewer/Viewer/CustomFrame.py

from PyQt5.QtWidgets import (QWidget,QVBoxLayout,QTabWidget,QSplitter, QFileDialog,
                             QLineEdit,QCheckBox,QLabel,QFrame,QPushButton,QMenu,QAction)
from PyQt5.QtGui import QPixmap,QImage,QPainter,QResizeEvent,QCloseEvent,QPaintEvent
from PyQt5.QtCore import QSize,QTimer,QRect,Qt
import numpy as np
import def_pb2 as pb
from concurrent.futures import ThreadPoolExecutor, as_completed, wait
import GrpcClient as rpc
import vtk
from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
#from vtk.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor

import threading
class ControlFrame(QFrame):
    def __init__(self):
        QFrame.__init__(self)
        self.threadPool_ = ThreadPoolExecutor(5)
        self.setGeometry(0, 0, 500, 500)
        self.InitUI()
        self.setFrameShape(self.StyledPanel)
        #self.timer_ = QTimer()
        #self.timer_.timeout.connect(self.Update)
        #self.timer_.start(100)
    def DisplayMeasureInfo(self,info : pb.MeasureInfo):
        self.ErrorInfo.setText(info.error)
    def OnAction(self,dtype,id):
        client=rpc.GrpcStream()
        ip=self.IPEdit.text()
        port=int(self.PortEdit.text())
        if dtype==1:
            self.threadPool_.submit(client.GrabImage,ip,port,id)
        if dtype==2:
            self.threadPool_.submit(client.GrabCloud,ip,port,id)
    def InitUI(self):

        self.begstatic = QLabel(self)
        self.begstatic.setText("   IP:")
        self.begstatic.setGeometry(20, 5, 80, 30)
        self.IPEdit = QLineEdit(self)
        self.IPEdit.setText("192.168.2.45")
        self.IPEdit.setGeometry(90, 5, 150, 30)

        self.endstatic = QLabel(self)
        self.endstatic.setText("Port:")
        self.endstatic.setGeometry(20, 50, 80, 30)
        self.PortEdit = QLineEdit(self)
        self.PortEdit.setText("9876")
        self.PortEdit.setGeometry(90, 50, 150, 30)

        self.btnGrab = QPushButton(self)
        self.btnGrab.setGeometry(20, 140, 100, 40)
        self.btnGrab.setText("抓取图像")
        self.btnGrab.clicked.connect(self.btnGrabImageClick)

        self.btnCancel = QPushButton(self)
        self.btnCancel.setGeometry(150, 140, 100, 40)
        self.btnCancel.setText("抓取点云")
        self.btnCancel.clicked.connect(self.btnGrabCloudClick)


        self.btnImageStreamCheck = QCheckBox("实时图像", self)
        self.btnImageStreamCheck.setGeometry(20, 90, 120, 40)
        self.btnImageStreamCheck.clicked.connect(self.RealTimeImagecb)

        self.btnDataStreamCheck = QCheckBox("实时结果打印", self)
        self.btnDataStreamCheck.setGeometry(150, 90, 120, 40)
        self.btnDataStreamCheck.clicked.connect(self.RealTimeDatacb)

        self.ErrorInfo = QLabel(self)
        self.ErrorInfo.setGeometry(20, 200, 250, 400)
        self.ErrorInfo.setWordWrap(True)
        self.ErrorInfo.setAlignment(Qt.AlignTop)
    def RealTimeDatacb(self):
        if self.btnDataStreamCheck.checkState()==Qt.Checked:
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.OpenDataStream,ip,port)
        else:
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.CloseDataStream,ip,port)
    def RealTimeImagecb(self):
        if self.btnImageStreamCheck.checkState()==Qt.Checked:
            self.btnGrab.setEnabled(False)
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.OpenImageStream,ip,port)
        else:
            self.btnGrab.setEnabled(True)
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.CloseImageStream,ip,port)
    def closeEvent(self, a0: QCloseEvent) -> None:
        if self.btnImageStreamCheck.checkState()==Qt.Checked:
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.CloseImageStream,ip,port)
        if self.btnDataStreamCheck.checkState()==Qt.Checked:
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.CloseDataStream,ip,port)
        self.threadPool_.shutdown(wait=False)
        rpc.GrpcStream().close()
    def btnGrabImageClick(self):
        if self.btnImageStreamCheck.checkState()==Qt.Checked:
            print(" 先关闭实时显示")
        else :
            client=rpc.GrpcStream()
            ip=self.IPEdit.text()
            port=int(self.PortEdit.text())
            self.threadPool_.submit(client.GrabImage,ip,port,0)
    def btnGrabCloudClick(self):
        client=rpc.GrpcStream()
        ip=self.IPEdit.text()
        port=int(self.PortEdit.text())
        for id in range(4):
            self.threadPool_.submit(client.GrabCloud,ip,port,id+1)
    def btnCancelClick(self):
        pass
class ImageViewer(QLabel):
    def __init__(self,parent,OnAction,id):
        super(ImageViewer,self).__init__(parent=parent)
        self.OnAction=OnAction
        self.dtype=1
        self.id=id
        self.image=None
        self.show = False
        self.setStyleSheet("border-width:1px;border-style:solid;border-color:rgb(150,150,150)")

    def ShowImg(self,pixmap):
        self.show=True
        self.image=pixmap
        self.update()

    def updataData(self):
        self.OnAction(self.dtype,self.id)
        print("%d,%d"%(self.dtype,self.id))

    def contextMenuEvent(self, a0):
        menu=QMenu(self)
        updata_act=QAction("更新")
        updata_act.triggered.connect(self.updataData)
        menu.addAction(updata_act)
        ret=menu.exec_(a0.globalPos())

    def paintEvent(self, a0: QPaintEvent) -> None:
        if not self.image == None:
            w, h = self.size().width(),self.size().height()
            iw, ih = self.image.width(), self.image.height()
            painter=QPainter(self)
            painter.drawPixmap(QRect(0,0,w,h),self.image,QRect(0,0,iw,ih))


class VtkPointCloud:
    def __init__(self, zMin=-10.0, zMax=10.0, maxNumPoints=1e7):
        self.maxNumPoints = maxNumPoints
        self.vtkPolyData = vtk.vtkPolyData()
        self.clearPoints()
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(self.vtkPolyData)
        mapper.SetColorModeToDefault()
        mapper.SetScalarRange(zMin, zMax)
        mapper.SetScalarVisibility(1)
        self.vtkActor = vtk.vtkActor()


        self.vtkActor.SetMapper(mapper)
        self.vtkActor.GetProperty().SetPointSize(1)
        self.vtkActor.GetProperty().SetColor(1, 1, 1)

    def addPoint(self, point):
        if self.vtkPoints.GetNumberOfPoints() < self.maxNumPoints:
            pointId = self.vtkPoints.InsertNextPoint(point[0:-1])
            self.vtkDepth.InsertNextValue(point[2])
            self.vtkCells.InsertNextCell(1)
            self.vtkCells.InsertCellPoint(pointId)
        self.vtkCells.Modified()
        self.vtkPoints.Modified()
        self.vtkDepth.Modified()

    def clearPoints(self):
        self.vtkPoints = vtk.vtkPoints()
        self.vtkCells = vtk.vtkCellArray()
        self.vtkDepth = vtk.vtkDoubleArray()
        self.vtkDepth.SetName('DepthArray')
        self.vtkPolyData.SetPoints(self.vtkPoints)
        self.vtkPolyData.SetVerts(self.vtkCells)
        #self.vtkPolyData.GetPointData().SetScalars(self.vtkDepth)
        self.vtkPolyData.GetPointData().SetActiveScalars('DepthArray')

class pclViewer(QVTKRenderWindowInteractor):
    def __init__(self,parent=None):
        super(pclViewer,self).__init__(parent=parent)
    def mousePressEvent(self, ev):
        btn = ev.button()
        if btn == Qt.RightButton:
            print(" r btn")
        else:
            QVTKRenderWindowInteractor.mousePressEvent(self,ev)


class VtkPointCloudCanvas(QWidget):
    def __init__(self, *args, **kwargs):
        super(VtkPointCloudCanvas, self).__init__(*args, **kwargs)
        self.lock_=threading.RLock()

        self._layout = QVBoxLayout()
        self.setLayout(self._layout)
        self._vtk_widget = pclViewer(self)
        self._vtk_widget.SetInteractorStyle(vtk.vtkInteractorStyleTrackballCamera())  # 设置交互方式==常用的方式 移动摄像机

        self._layout.addWidget(self._vtk_widget)

        self._render = vtk.vtkRenderer()
        self._vtk_widget.GetRenderWindow().AddRenderer(self._render)
        self._iren = self._vtk_widget.GetRenderWindow().GetInteractor()
        self._point_cloud = VtkPointCloud()
        self._render.AddActor(self._point_cloud.vtkActor)

        transform = vtk.vtkTransform()
        transform.Translate(0, 0.0, 0.0)
        axes = vtk.vtkAxesActor()
        axes.SetUserTransform(transform)
        colors = vtk.vtkNamedColors()
        axes.GetXAxisCaptionActor2D().GetCaptionTextProperty().SetColor(colors.GetColor3d("Red"))
        axes.GetYAxisCaptionActor2D().GetCaptionTextProperty().SetColor(colors.GetColor3d("Green"))
        axes.GetZAxisCaptionActor2D().GetCaptionTextProperty().SetColor(colors.GetColor3d("Blue"))
        self._render.AddActor(axes)

        self.camera=self._render.GetActiveCamera()
        self.show()
        self._iren.Initialize()

    def close(self) -> bool:

        print("  tk close")
        self._point_cloud.clearPoints()

    def displayPCL(self,cloud:np.array):
        colors = vtk.vtkUnsignedCharArray()
        colors.SetNumberOfComponents(3)
        colors.SetName("Colors")

        self._point_cloud.clearPoints()
        for point in cloud:
            self._point_cloud.addPoint(point)
            prob=point[-1]
            colors.InsertNextTypedTuple((int(255*prob),int(255-prob*255),255-int(prob*255)))
        self._point_cloud.vtkPolyData.GetPointData().SetScalars(colors)
        self._vtk_widget.update()
    def resetViewer(self):
        self.update()


class PointCLViwer(QSplitter):
    def __init__(self,OnAction,id):
        super(PointCLViwer, self).__init__()
        self.OnAction=OnAction
        self.dtype=2
        self.id=id
        self.pointCloud=[]
        self.InitUI()
    def InitUI(self):
        self.pclViewer=VtkPointCloudCanvas(self)
        self.addWidget(self.pclViewer)
        self.setStretchFactor(0,4)
    def save(self):
        fileName2, ok2 = QFileDialog.getSaveFileName(self, "文件保存", "C:/","All Files (*);;Text Files (*.txt)")
        print(fileName2 ,ok2)
        with open(fileName2,"w+") as f:
            for point in self.pointCloud:
                prob=point[3]
                f.write("%f %f %f %d %d %d\n"%(point[0],point[1],point[2],int(255*prob),int(255-prob*255),255-int(prob*255)))

        print("Save Cloud to %s"%(fileName2))

    def displayCloud(self,points:np.array):
        self.pointCloud=points
        self.pclViewer.displayPCL(points)
    def resetViewer(self):
        self.pclViewer.resetViewer()
    def updataData(self):
        self.OnAction(self.dtype,self.id)
    def contextMenuEvent(self, a0):
        menu=QMenu(self)
        updata_act=QAction("更新点云")
        updata_act.triggered.connect(self.updataData)
        menu.addAction(updata_act)

        #act=QAction("还原视野")
        #act.triggered.connect(self.resetViewer)
        #menu.addAction(act)

        actSave=QAction("保存")
        actSave.triggered.connect(self.save)
        menu.addAction(actSave)
        ret=menu.exec_(a0.globalPos())


class ViewerFrame(QFrame):
    def __init__(self,OnAction):
        super(ViewerFrame, self).__init__()
        self.OnAction=OnAction
        self.InitUI()
    def InitUI(self):

        self.table1=QTabWidget(self)
        self.table2=QTabWidget(self)
        self.table3=QTabWidget(self)
        self.table4=QTabWidget(self)
        self.table1.setTabPosition(QTabWidget.TabPosition.North)
        self.table2.setTabPosition(QTabWidget.TabPosition.North)
        self.table3.setTabPosition(QTabWidget.TabPosition.North)
        self.table4.setTabPosition(QTabWidget.TabPosition.North)

        self.panel1=ImageViewer(self,self.OnAction,1)
        self.panel2=ImageViewer(self,self.OnAction,2)
        self.panel3=ImageViewer(self,self.OnAction,3)
        self.panel4=ImageViewer(self,self.OnAction,4)

        self.pcViewer1=PointCLViwer(self.OnAction,1)
        self.pcViewer2=PointCLViwer(self.OnAction,2)
        self.pcViewer3=PointCLViwer(self.OnAction,3)
        self.pcViewer4=PointCLViwer(self.OnAction,4)

        self.table1.addTab(self.panel1,"图像")
        self.table1.addTab(self.pcViewer1,"点云")
        self.table2.addTab(self.panel2,"图像")
        self.table2.addTab(self.pcViewer2,"点云")
        self.table3.addTab(self.panel3,"图像")
        self.table3.addTab(self.pcViewer3,"点云")
        self.table4.addTab(self.panel4,"图像")
        self.table4.addTab(self.pcViewer4,"点云")

    def closeEvent(self, a0: QCloseEvent) -> None:
        pass
    def resizeEvent(self, a0: QResizeEvent) -> None:
        w, h = self.size().width(), self.size().height()
        w=w-15
        h=h-15
        self.table1.setGeometry(5, 5, w/2, h/2)
        self.table2.setGeometry(10+w/2, 5, w/2, h/2)
        self.table3.setGeometry(5, 10+h/2, w/2, h/2)
        self.table4.setGeometry(10+w/2, 10+h/2, w/2, h/2)

    def PbImg2QPix(self,image:pb.Image):
        img1=image
        w=img1.width
        h=img1.height
        c=img1.channel
        btarry=bytearray(img1.data)
        npdata = np.frombuffer(btarry, dtype=np.uint8)
        if c==3:
            npdata = npdata.reshape([h,w,c])
            qimg=QImage(npdata.data,w,h,QImage.Format_RGB888)
        if c==1:
            npdata = npdata.reshape([h,w])
            qimg=QImage(npdata.data,w,h,QImage.Format_Grayscale8)
        pix=QPixmap(qimg)
        return pix
    def PbCloud2Pts(self,cloud:pb.PointCloud):

        size=cloud.size
        btarry=bytearray(cloud.data)
        npdata = np.frombuffer(btarry, dtype=np.short)
        points= npdata.reshape([size,4])/32750.0 *5.0
        return points
    def DisplayFrame(self,frame : pb.ResFrame):
        if frame.HasField("images"):
            self.DisplayImage(frame.images)
        if frame.HasField("clouds"):
            self.DisplayCloud(frame.clouds)
    def DisplayImage(self,images:pb.ResImage):
        if images.HasField("img1"):
            self.panel1.ShowImg(self.PbImg2QPix(images.img1))
        if images.HasField("img2"):
            self.panel2.ShowImg(self.PbImg2QPix(images.img2))
        if images.HasField("img3"):
            self.panel3.ShowImg(self.PbImg2QPix(images.img3))
        if images.HasField("img4"):
            self.panel4.ShowImg(self.PbImg2QPix(images.img4))
    def DisplayCloud(self,clouds:pb.ResCloud):
        if clouds.HasField("cloud1"):
            self.pcViewer1.displayCloud(self.PbCloud2Pts(clouds.cloud1))
        if clouds.HasField("cloud2"):
            self.pcViewer2.displayCloud(self.PbCloud2Pts(clouds.cloud2))
        if clouds.HasField("cloud3"):
            self.pcViewer3.displayCloud(self.PbCloud2Pts(clouds.cloud3))
        if clouds.HasField("cloud4"):
            self.pcViewer4.displayCloud(self.PbCloud2Pts(clouds.cloud4))