NVIDIA Jetson PyQt5에서 추론도구 사용하기 -Python

추론 결과를 GUI에 보여주기 위해 PyQt5를 사용합니다.

# PyQt5가 설치되어 있는 경우
sudo apt upgrade python3-pyqt5*

# PyQt5를 새로 설치해야 하는 경우
sudo apt install python3-pyqt5*

실행 결과

스크립트

뷰어로는 QGraphicsView를 사용합니다. 비디오를 표시할 때 라벨로 표시하는 방법을 많이 사용하지만 이 때는 위젯의 repaint를 너무 많이 발생하게 만들 수 있다는 이야기를 들어서 그래픽스 위젯을 사용했습니다. 비디오 스트림 작업은 QThread를 사용했습니다.

샘플 동영상은 앞서 동영상 객체인식 예제를 만들 때 사용했던 인텔의 샘플 비디오를 사용했습니다. 앞의 예제에서 설명했던 내용들은 제외하고 주석을 첨부하였습니다.

"""PyQt에서 Jetson-Inference 사용
"""
import glob
import os

import jetson.inference
import jetson.utils
from PyQt5.QtCore import QRectF, QThread, pyqtSignal, Qt
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import QGraphicsPixmapItem, QGraphicsScene, QGraphicsView


class VideoStream(QThread):
    sigVideoClose = pyqtSignal()  # 비디오 종료 시그널
    sigGetFrame = pyqtSignal(QPixmap)  # 프레임 업데이트 시그널
    __index = 0  # 비디오 파일 인덱스
    __net = jetson.inference.detectNet("ssd-mobilenet-v2", threshold=0.4)
    __path = [os.path.abspath(x) for x in glob.glob("sample-videos/*.mp4")]

    def run(self):
        """인덱스에 해당하는 비디오파일을 읽어서 뷰어로 프레임을 보냅니다"""
        index = self.__index % len(self.__path)
        path = self.__path[index]
        self.source = jetson.utils.videoSource(path)
        while True:
            try:
                # 프레임레이트마다 조금 쉽니다.
                self.msleep(int(1000 / self.source.GetFrameRate()))
                cuda_img = self.source.Capture()
            except:
                break
            self.__net.Detect(cuda_img)
            pix = self._cuda_to_qimg(cuda_img)
            self.sigGetFrame.emit(pix)
        self.source.Close()
        self.__index += 1
        self.sigVideoClose.emit()
        return

    def _cuda_to_qimg(self, cuda_img: jetson.utils.cudaImage):
        """캡쳐한 cuda image를 QPixmap으로 변환"""
        cv_img = jetson.utils.cudaToNumpy(cuda_img)
        h, w, c = cv_img.shape
        qimg = QImage(cv_img.data, w, h, c * w, QImage.Format_RGB888)
        qpix = QPixmap.fromImage(qimg)
        return qpix


class VideoViewer(QGraphicsView):
    def __init__(self):
        super().__init__()
        self.pixmap = QGraphicsPixmapItem()
        self.stream = VideoStream(self)

        self.setScene(QGraphicsScene())
        self.scene().addItem(self.pixmap)
        self.setFixedSize(640, 480)
        # 뷰어의 스크롤바를 제거합니다
        self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)

        # 이벤트 연결
        self.stream.sigGetFrame.connect(self.frame_update)
        self.stream.sigVideoClose.connect(self.stream_end)

        # 비디오 재생 스레드 실행
        self.stream.start()
        self.show()

    def frame_update(self, pix: QPixmap):
        """프레임 업데이트 시그널을 위한 슬롯"""
        # 뷰어에 이미지를 맞춤
        self.fitInView(QRectF(pix.rect()), mode=Qt.KeepAspectRatio)
        # Pixmap 아이템 업데이트
        self.pixmap.setPixmap(pix)
        return

    def stream_end(self) -> None:
        """비디오 종료 시그널을 위한 슬롯"""
        self.stream.start()  # 다음 비디오 재생
        return


if __name__ == "__main__":
    import sys
    from PyQt5.QtWidgets import QApplication

    app = QApplication(sys.argv)
    win = VideoViewer()
    app.exec_()