[python] PyQt5로 이미지 뷰어 만들기 (간단한 라벨링 툴 만들기 1단계)

키보드 입력으로 이전 다음 이미지를 볼 수 있고 추가로 해당 이미지에 라벨링하는 이미지 라벨링 툴을 만들어보자. 이번 편에서는 이미지 뷰까지만 다루고 다음 편에 각종 라벨링 기능을 추가하는 것을 다루어보겠다.

먼저 PyQt5를 설치해주자.

pip install pyqt5

 

그 다음 ImageViewer 클래스를 만들고 QMainWindow를 상속받는다. 그 다음 __init__ 함수에서 폴더 경로라던가 색상 테이블 등의 변수를 선언해주면 되고 액션 생성, 창 제목 및 크기도 설정해주면 된다. 나머지 함수는 스크롤 기능과 이미지를 윈도우 크기에 최적화 하여 보여주기 위한 함수들이다.

x
from PyQt5.QtCore import QDir, Qt
from PyQt5.QtGui import QImage, QPainter, QPalette, QPixmap, qRgb
from PyQt5.QtWidgets import (QAction, QApplication, QFileDialog, QLabel,
        QMainWindow, QMenu, QMessageBox, QScrollArea, QSizePolicy, QInputDialog)
from PyQt5.QtPrintSupport import QPrintDialog, QPrinter
​
class ImageViewer(QMainWindow):
    def __init__(self):
        super(ImageViewer, self).__init__()
​
        self.gray_color_table = [qRgb(i, i, i) for i in range(256)]
​
        self.base_path = './image_path'
        self.img_list = glob.glob(os.path.join(self.base_path, '*.jpg'))
        self.pos = 0
        self.total = len(self.img_list)
​
        self.printer = QPrinter()
        self.width = 620
        self.height = 620
        
        self.imageLabel = QLabel()
        self.imageLabel.setBackgroundRole(QPalette.Base)
        self.imageLabel.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Ignored)
        self.imageLabel.setScaledContents(True)
​
        self.scrollArea = QScrollArea()
        self.scrollArea.setBackgroundRole(QPalette.Dark)
        self.scrollArea.setWidget(self.imageLabel)
        self.setCentralWidget(self.scrollArea)
​
        self.createActions()
        
        self.setWindowTitle("Image Viewer")
        self.resize(self.width, self.height)
        
        image = cv2.imread(self.img_list[self.pos])
        self.openImage(image=self.toQImage(image))
    
    def normalSize(self):
        self.imageLabel.adjustSize()
​
    def fitToWindow(self):
        fitToWindow = self.fitToWindowAct.isChecked()
        self.scrollArea.setWidgetResizable(fitToWindow)
        if not fitToWindow:
            self.normalSize()
        self.updateActions()
        
    def createActions(self):
        self.normalSizeAct = QAction("&Normal Size", self, shortcut="Ctrl+S",
                enabled=False, triggered=self.normalSize)
​
        self.fitToWindowAct = QAction("&Fit to Window", self, enabled=False,
                checkable=True, shortcut="Ctrl+F", triggered=self.fitToWindow)
        
    def updateActions(self):
        self.normalSizeAct.setEnabled(not self.fitToWindowAct.isChecked())
       
    def adjustScrollBar(self, scrollBar, factor):
        scrollBar.setValue(int(factor * scrollBar.value()
                                + ((factor - 1) * scrollBar.pageStep()/2)))

__init__ 함수 마지막에 imageLabel에 opencv2로 불러온 이미지를 보여주게 된다. 이때 파일 경로를 넘긴 뒤 openImage 함수 내에서 QImage 함수 호출로 이미지를 불러와도 되지만 나중에 이미지에 목적에 따른 여러 가지 처리를 하고 보여주기 위해 이런 식으로 numpy 배열로 넘기는 것이 편리하다. 물론 그냥 numpy 배열로 넘기면 안 되고 toQImage 함수를 통해 QImage로 변환하여 넘겨야 한다.

이제 주요함수만 살펴보자.

 

toQImage 함수

x
def toQImage(self, im, copy=False):
    if im is None:
        return QImage()
​
    if im.dtype == np.uint8:
        if len(im.shape) == 2:
            qim = QImage(im.data, im.shape[1], im.shape[0], im.strides[0], QImage.Format_Indexed8)
            qim.setColorTable(self.gray_color_table)
            return qim.copy() if copy else qim
​
        elif len(im.shape) == 3:
            if im.shape[2] == 3:
                qim = QImage(im.data, im.shape[1], im.shape[0], im.strides[0], QImage.Format_RGB888)
                return qim.copy() if copy else qim
            elif im.shape[2] == 4:
                qim = QImage(im.data, im.shape[1], im.shape[0], im.strides[0], QImage.Format_ARGB32)
                return qim.copy() if copy else qim

numpy 배열을 넘겨받아서 QImage로 변환해주는 함수이다.

 

openImage 함수

x
def openImage(self, image=None, fileName=None):
        if image == None:
            image = QImage(fileName)
        if image.isNull():
            QMessageBox.information(self, "Image Viewer",
                                    "Cannot load %s." % fileName)
            return
​
        self.imageLabel.setPixmap(QPixmap.fromImage(image))
​
        self.fitToWindowAct.setEnabled(True)
        self.updateActions()
        if not self.fitToWindowAct.isChecked():
            self.imageLabel.adjustSize()

QPixmap.fromImage를 통해 imageLabel에 이미지를 보여줄 수 있도록 변환하고 imageLabel의 setPixmap을 통해 이미지를 보여준다. 그 다음 window를 업데이트해준다.

 

keyPressEvent 함수

x
def keyPressEvent(self, e):
    if e.key() == 65:
        if not self.pos == 0:
            self.pos -= 1
            image = cv2.imread(self.img_list[self.pos])
            """
            이미지 처리
            """
            self.openImage(image=self.toQImage(image))
            print('\r' + self.img_list[self.pos], end="")
                                            
    elif e.key() == 68:
        self.pos += 1
        if self.total == self.pos:
            self.pos -= 1
        image = cv2.imread(self.img_list[self.pos])
        """
        이미지 처리
        """
        self.openImage(image=self.toQImage(image))            
        print('\r' + self.img_list[self.pos], end="")

keyPressEvent 함수는 키 입력에 따른 이벤트를 처리해주는 함수다. 여기에 이전, 다음 키로 자신이 편한 것으로 설정하면 된다. 필자의 경우 A, D 키로 설정했으며 이전과 다음을 누를 때마다 이미지를 불러와 새롭게 처리하는 방식으로 구현했다. 미리 이미지를 다 불러와 보여줘도 되겠지만 굉장히 복잡한 처리가 포함되지 않는 이상 넘기는데 딜레이를 체감하지 못 하기 때문에 이렇게 구현하였다.

 

main 구문

if __name__ == '__main__':
​
    import sys
​
    app = QApplication(sys.argv)
    imageViewer = ImageViewer()
    imageViewer.show()
    sys.exit(app.exec_())

마지막으로 main 구문을 위과 같이 작성하면 이미지 뷰어의 기본 기능만 갖춘 아주 간단한 뷰어가 완성된다.

위 그림처럼 스크롤 기능이 있고 창의 확대, 축소가 자유자재로 가능하며 이전, 다음 이미지를 볼 수 있다. 이미지 라벨링 툴을 만들기 위한 전 단계로 이미지를 보여주는 것에 초점을 두어 설명하였으며, 사실 이미지 뷰어만 완성되어있으면 이미지를 처리하여 라벨링하기 쉽게 보여주는 작업의 난이도는 목적에 달렸다. 마우스 클릭을 하여야 하는 것도 있고 키보드 입력으로 해야 되는 것도 있는데 이제 뷰어가 완성된 이상 pyqt5와는 거의 별개의 문제다.

다음에는 마우스로 그은 선의 절대 각도 측정, 이미지 내 구역 선택, 키보드 입력 등을 구현하여 간단한 라벨링 툴 구현을 설명할 것이다.