sksurgerycharucotest_demo.py 4.47 KB
Newer Older
1 2 3 4 5 6 7
# coding=utf-8

""" Demo app, to demo charuco markers."""
import sys
import six
import cv2
from cv2 import aruco
8
from PySide2 import QtCore, QtWidgets
9
from PySide2.QtCore import Slot
10
import sksurgeryimage.calibration.charuco as ar
11
import sksurgeryutils.utils.image_utils as iu
12 13 14 15 16 17

# pylint: disable=too-many-instance-attributes


class CharucoDemoGui(QtWidgets.QWidget):
    """ Demo GUI, with 2 QLabel side by side."""
18
    def __init__(self, camera, width, height, rows, columns, dictionary):
19 20
        super().__init__()

21
        if not isinstance(camera, str) and camera < 0:
22 23
            raise ValueError('The camera number should be >= 0')

24 25 26 27
        if not isinstance(dictionary, int):
            raise TypeError('dictionary should be an int >= 0')

        if dictionary < 0:
28
            raise ValueError('dictionary should be an ArUco enum of dictionary')
29

30
        self.cap = cv2.VideoCapture(camera)  # pylint: disable=no-member
31 32
        self.cap.set(3, width)
        self.cap.set(4, height)
33 34 35 36 37 38 39 40

        if not self.cap.isOpened():
            raise ValueError("Unable to open camera:" + str(camera))

        grabbed, self.frame = self.cap.read()
        if not grabbed:
            raise RuntimeError("Failed to grab first frame.")

41 42 43 44 45 46 47 48
        if self.frame.shape[0] != height:
            raise ValueError("Grabbed image has wrong number of rows:"
                             + str(self.frame.shape[0]))

        if self.frame.shape[1] != width:
            raise ValueError("Grabbed image has wrong number of columns:"
                             + str(self.frame.shape[1]))

49
        # pylint: disable=no-member
50
        self.dictionary = aruco.Dictionary_get(dictionary)
51 52
        self.board = aruco.CharucoBoard_create(rows,
                                               columns,
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
                                               2,
                                               1,
                                               self.dictionary)

        self.layout = QtWidgets.QHBoxLayout()

        self.image_label = QtWidgets.QLabel("Image")
        self.image_label.setAlignment(QtCore.Qt.AlignCenter)
        self.layout.addWidget(self.image_label)

        self.setLayout(self.layout)

        self.grab = QtCore.QTimer()
        self.grab.setInterval(50)

        # pylint: disable=maybe-no-member
        self.grab.timeout.connect(self.update_image)
        self.grab.start()

    @Slot()
    def update_image(self):

        """ Updates the image. """

        # Then grab image
        grabbed = self.cap.grab()
        if not grabbed:
            raise RuntimeError("Failed to grab frame number:"
                               + str(self.number_frames))
        # Then decode image
        self.cap.retrieve(self.frame)

        # OpenCV does BGR, Qt wants RGB
        # pylint: disable=no-member
        rgb_frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2RGB)
        greyscale_frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY)

        # Run aruco to detect markers
91 92
        # pylint: disable=unbalanced-tuple-unpacking
        _, _, \
93 94
            chessboard_corners, chessboard_ids \
            = ar.detect_charuco_points(self.dictionary,
95 96
                                       self.board,
                                       greyscale_frame
97 98 99 100 101
                                       )

        # Draw detected corners back on RGB image.
        if chessboard_corners is not None \
            and chessboard_ids is not None \
102 103
                and chessboard_corners.shape[0] \
                and chessboard_ids.shape[0]:
104 105 106 107 108 109 110 111 112 113 114 115
            annotated_frame = ar.draw_charuco_corners(rgb_frame,
                                                      chessboard_corners,
                                                      chessboard_ids
                                                      )
        else:
            annotated_frame = rgb_frame

        # Display annotated RGB image.
        pixmap = iu.image_to_pixmap(annotated_frame)
        self.image_label.setPixmap(pixmap)


116
def run_demo(camera, width, height, rows, columns, dictionary):
117 118 119 120 121

    """ Prints command line args, and launches main screen."""
    six.print_("Camera:" + str(camera))
    six.print_("  Width:" + str(width))
    six.print_("  Height:" + str(height))
122 123 124 125
    six.print_("Board:")
    six.print_("  Rows:" + str(rows))
    six.print_("  Columns:" + str(columns))
    six.print_("  Dictionary:" + str(dictionary))
126 127 128

    app = QtWidgets.QApplication([])

129
    widget = CharucoDemoGui(camera, width, height, rows, columns, dictionary)
130 131 132
    widget.show()

    return sys.exit(app.exec_())