Esempio n. 1
0
    def run(self):
        """
        Open up the webcam and present the user with the Qbr user interface.

        Returns a string of the scanned state in rubik's cube notation.
        """
        while True:
            _, frame = self.cam.read()
            key = cv2.waitKey(10) & 0xff

            # Quit on escape.
            if key == 27:
                break

            # Update the snapshot when space bar is pressed.
            if key == 32 and not self.calibrate_mode:
                self.update_snapshot_state(frame)

            # Toggle calibrate mode.
            if key == ord(CALIBRATE_MODE_KEY):
                self.reset_calibrate_mode()
                self.calibrate_mode = not self.calibrate_mode

            grayFrame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            blurredFrame = cv2.blur(grayFrame, (5, 5))
            cannyFrame = cv2.Canny(blurredFrame, 30, 60, 3)
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
            dilatedFrame = cv2.dilate(cannyFrame, kernel)

            contours = self.find_contours(dilatedFrame)
            if len(contours) == 9:
                self.draw_contours(frame, contours)
                if not self.calibrate_mode:
                    self.update_preview_state(frame, contours)
                elif key == 32 and self.done_calibrating == False:
                    current_color = self.cube_sides[
                        self.current_color_to_calibrate_index]
                    (x, y, w, h) = contours[4]
                    roi = frame[y + 7:y + h - 7, x + 14:x + w - 14]
                    avg_bgr = ColorDetector.get_dominant_color(roi)
                    self.calibrated_colors[current_color] = avg_bgr
                    self.current_color_to_calibrate_index += 1
                    self.done_calibrating = self.current_color_to_calibrate_index == len(
                        self.cube_sides)
                    if self.done_calibrating:
                        ColorDetector.set_cube_color_pallete(
                            self.calibrated_colors)

            if self.calibrate_mode:
                self.display_current_color_to_calibrate(frame)
                self.display_calibrated_colors(frame)
            else:
                self.draw_preview_stickers(frame)
                self.draw_snapshot_stickers(frame)
                self.display_scanned_sides(frame)

            cv2.imshow('default', frame)

        self.cam.release()
        cv2.destroyAllWindows()

        if len(self.sides.keys()) != 6:
            return False

        if not self.scanned_successfully():
            return False

        # Convert all the sides and their BGR colors to cube notation.
        notation = dict(self.sides)
        for side, preview in notation.items():
            for sticker_index, bgr in enumerate(preview):
                notation[side][
                    sticker_index] = ColorDetector.convert_bgr_to_notation(bgr)

        # Join all the sides together into one single string.
        # Order must be URFDLB (white, red, green, yellow, orange, blue)
        combined = ''
        for side in ['white', 'red', 'green', 'yellow', 'orange', 'blue']:
            combined += ''.join(notation[side])
        return combined