Skip to content

Commit

Permalink
fix(video): remove opencv-contrib-python and use pillow to draw text
Browse files Browse the repository at this point in the history 
- fixes #20
- fixes #21
- fixes #24
- fixes #30
  • Loading branch information
@kkoomen
kkoomen committed Nov 5, 2022
1 parent 9d8e9b8 commit aab1f37
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 69 deletions.
8 changes: 2 additions & 6 deletions requirements.txt
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
cffi==1.14.4
future==0.18.2
kociemba==1.2.1
numpy==1.22.0
opencv-python==4.5.1.48
opencv-contrib-python==4.5.1.48
pycparser==2.20
opencv-python==4.6.0.66
pillow==9.2.0
pydocstyle==5.1.1
python-i18n==0.3.9
snowballstemmer==2.1.0
141 changes: 78 additions & 63 deletions src/video.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
from config import config
from helpers import get_next_locale
import i18n
from PIL import ImageFont, ImageDraw, Image
import numpy as np
from constants import (
COLOR_PLACEHOLDER,
LOCALES,
Expand Down Expand Up @@ -54,7 +56,7 @@ def __init__(self):
self.current_color_to_calibrate_index = 0
self.done_calibrating = False

def draw_stickers(self, frame, stickers, offset_x, offset_y):
def draw_stickers(self, stickers, offset_x, offset_y):
"""Draws the given stickers onto the given frame."""
index = -1
for row in range(3):
Expand All @@ -67,7 +69,7 @@ def draw_stickers(self, frame, stickers, offset_x, offset_y):

# shadow
cv2.rectangle(
frame,
self.frame,
(x1, y1),
(x2, y2),
(0, 0, 0),
Expand All @@ -76,21 +78,21 @@ def draw_stickers(self, frame, stickers, offset_x, offset_y):

# foreground color
cv2.rectangle(
frame,
self.frame,
(x1 + 1, y1 + 1),
(x2 - 1, y2 - 1),
color_detector.get_prominent_color(stickers[index]),
-1
)

def draw_preview_stickers(self, frame):
def draw_preview_stickers(self):
"""Draw the current preview state onto the given frame."""
self.draw_stickers(frame, self.preview_state, STICKER_AREA_OFFSET, STICKER_AREA_OFFSET)
self.draw_stickers(self.preview_state, STICKER_AREA_OFFSET, STICKER_AREA_OFFSET)

def draw_snapshot_stickers(self, frame):
def draw_snapshot_stickers(self):
"""Draw the current snapshot state onto the given frame."""
y = STICKER_AREA_TILE_SIZE * 3 + STICKER_AREA_TILE_GAP * 2 + STICKER_AREA_OFFSET * 2
self.draw_stickers(frame, self.snapshot_state, STICKER_AREA_OFFSET, y)
self.draw_stickers(self.snapshot_state, STICKER_AREA_OFFSET, y)

def find_contours(self, dilatedFrame):
"""Find the contours of a 3x3x3 cube."""
Expand Down Expand Up @@ -214,17 +216,17 @@ def scanned_successfully(self):
invalid_colors = [k for k, v in color_count.items() if v != 9]
return len(invalid_colors) == 0

def draw_contours(self, frame, contours):
def draw_contours(self, contours):
"""Draw contours onto the given frame."""
if self.calibrate_mode:
# Only show the center piece contour.
(x, y, w, h) = contours[4]
cv2.rectangle(frame, (x, y), (x + w, y + h), STICKER_CONTOUR_COLOR, 2)
cv2.rectangle(self.frame, (x, y), (x + w, y + h), STICKER_CONTOUR_COLOR, 2)
else:
for index, (x, y, w, h) in enumerate(contours):
cv2.rectangle(frame, (x, y), (x + w, y + h), STICKER_CONTOUR_COLOR, 2)
cv2.rectangle(self.frame, (x, y), (x + w, y + h), STICKER_CONTOUR_COLOR, 2)

def update_preview_state(self, frame, contours):
def update_preview_state(self, contours):
"""
Get the average color value for the contour for every X amount of frames
to prevent flickering and more precise results.
Expand All @@ -244,7 +246,7 @@ def update_preview_state(self, frame, contours):
self.preview_state[index] = eval(most_common_color)
break

roi = frame[y+7:y+h-7, x+14:x+w-14]
roi = self.frame[y+7:y+h-7, x+14:x+w-14]
avg_bgr = color_detector.get_dominant_color(roi)
closest_color = color_detector.get_closest_color(avg_bgr)['color_bgr']
self.preview_state[index] = closest_color
Expand All @@ -253,58 +255,63 @@ def update_preview_state(self, frame, contours):
else:
self.average_sticker_colors[index] = [closest_color]

def update_snapshot_state(self, frame):
def update_snapshot_state(self):
"""Update the snapshot state based on the current preview state."""
self.snapshot_state = list(self.preview_state)
center_color_name = color_detector.get_closest_color(self.snapshot_state[4])['color_name']
self.result_state[center_color_name] = self.snapshot_state
self.draw_snapshot_stickers(frame)
self.draw_snapshot_stickers()

def get_freetype2_font(self):
"""Get the freetype2 font, load it and return it."""
def get_font(self, size=TEXT_SIZE):
"""Load the truetype font with the specified text size."""
font_path = '{}/assets/arial-unicode-ms.ttf'.format(ROOT_DIR)
ft2 = cv2.freetype.createFreeType2()
ft2.loadFontData(font_path, 0)
return ft2
return ImageFont.truetype(font_path, size)

def render_text(self, frame, text, pos, color=(255, 255, 255), size=TEXT_SIZE, bottomLeftOrigin=False):
"""Render text with a shadow."""
ft2 = self.get_freetype2_font()
self.get_text_size(text)
ft2.putText(frame, text, pos, fontHeight=size, color=(0, 0, 0), thickness=2, line_type=cv2.LINE_AA, bottomLeftOrigin=bottomLeftOrigin)
ft2.putText(frame, text, pos, fontHeight=size, color=color, thickness=-1, line_type=cv2.LINE_AA, bottomLeftOrigin=bottomLeftOrigin)
def render_text(self, text, pos, color=(255, 255, 255), size=TEXT_SIZE, anchor='lt'):
"""
Render text with a shadow using the pillow module.
"""
font = self.get_font(size)

# Convert opencv frame (np.array) to PIL Image array.
frame = Image.fromarray(self.frame)

# Draw the text onto the image.
draw = ImageDraw.Draw(frame)
draw.text(pos, text, font=font, fill=color, anchor=anchor,
stroke_width=1, stroke_fill=(0, 0, 0))

# Convert the pillow frame back to a numpy array.
self.frame = np.array(frame)

def get_text_size(self, text, size=TEXT_SIZE):
"""Get text size based on the default freetype2 loaded font."""
ft2 = self.get_freetype2_font()
return ft2.getTextSize(text, size, thickness=-1)
return self.get_font(size).getsize(text)

def draw_scanned_sides(self, frame):
def draw_scanned_sides(self):
"""Display how many sides are scanned by the user."""
text = i18n.t('scannedSides', num=len(self.result_state.keys()))
self.render_text(frame, text, (20, self.height - 20), bottomLeftOrigin=True)
self.render_text(text, (20, self.height - 20), anchor='lb')

def draw_current_color_to_calibrate(self, frame):
def draw_current_color_to_calibrate(self):
"""Display the current side's color that needs to be calibrated."""
y_offset = 20
offset_y = 20
font_size = int(TEXT_SIZE * 1.25)
if self.done_calibrating:
messages = [
i18n.t('calibratedSuccessfully'),
i18n.t('quitCalibrateMode', keyValue=CALIBRATE_MODE_KEY),
]
for index, text in enumerate(messages):
font_size
(textsize_width, textsize_height), _ = self.get_text_size(text, font_size)
y = y_offset + (textsize_height + 10) * index
self.render_text(frame, text, (int(self.width / 2 - textsize_width / 2), y), size=font_size)
_, textsize_height = self.get_text_size(text, font_size)
y = offset_y + (textsize_height + 10) * index
self.render_text(text, (int(self.width / 2), y), size=font_size, anchor='mt')
else:
current_color = self.colors_to_calibrate[self.current_color_to_calibrate_index]
text = i18n.t('currentCalibratingSide.{}'.format(current_color))
(textsize_width, textsize_height), _ = self.get_text_size(text, font_size)
self.render_text(frame, text, (int(self.width / 2 - textsize_width / 2), y_offset), size=font_size)
self.render_text(text, (int(self.width / 2), offset_y), size=font_size, anchor='mt')

def draw_calibrated_colors(self, frame):
def draw_calibrated_colors(self):
"""Display all the colors that are calibrated while in calibrate mode."""
offset_y = 20
for index, (color_name, color_bgr) in enumerate(self.calibrated_colors.items()):
Expand All @@ -315,7 +322,7 @@ def draw_calibrated_colors(self, frame):

# shadow
cv2.rectangle(
frame,
self.frame,
(x1, y1),
(x2, y2),
(0, 0, 0),
Expand All @@ -324,30 +331,29 @@ def draw_calibrated_colors(self, frame):

# foreground
cv2.rectangle(
frame,
self.frame,
(x1 + 1, y1 + 1),
(x2 - 1, y2 - 1),
tuple([int(c) for c in color_bgr]),
-1
)
self.render_text(frame, i18n.t(color_name), (20, y1 + 3))
self.render_text(i18n.t(color_name), (20, y1 + STICKER_AREA_TILE_SIZE / 2 - 3), anchor='lm')

def reset_calibrate_mode(self):
"""Reset calibrate mode variables."""
self.calibrated_colors = {}
self.current_color_to_calibrate_index = 0
self.done_calibrating = False

def draw_current_language(self, frame):
def draw_current_language(self):
text = '{}: {}'.format(
i18n.t('language'),
LOCALES[config.get_setting('locale')]
)
(textsize_width, textsize_height), _ = self.get_text_size(text)
offset = 20
self.render_text(frame, text, (self.width - textsize_width - offset, offset))
self.render_text(text, (self.width - offset, offset), anchor='rt')

def draw_2d_cube_state(self, frame):
def draw_2d_cube_state(self):
"""
Create a 2D cube state visualization and draw the self.result_state.
Expand Down Expand Up @@ -393,8 +399,16 @@ def draw_2d_cube_state(self, frame):
for row in range(3):
for col in range(3):
index += 1
x1 = int((offset_x + MINI_STICKER_AREA_TILE_SIZE * col) + (MINI_STICKER_AREA_TILE_GAP * col) + ((side_size + side_offset) * grid_x))
y1 = int((offset_y + MINI_STICKER_AREA_TILE_SIZE * row) + (MINI_STICKER_AREA_TILE_GAP * row) + ((side_size + side_offset) * grid_y))
x1 = int(
(offset_x + MINI_STICKER_AREA_TILE_SIZE * col) +
(MINI_STICKER_AREA_TILE_GAP * col) +
((side_size + side_offset) * grid_x)
)
y1 = int(
(offset_y + MINI_STICKER_AREA_TILE_SIZE * row) +
(MINI_STICKER_AREA_TILE_GAP * row) +
((side_size + side_offset) * grid_y)
)
x2 = int(x1 + MINI_STICKER_AREA_TILE_SIZE)
y2 = int(y1 + MINI_STICKER_AREA_TILE_SIZE)

Expand All @@ -404,7 +418,7 @@ def draw_2d_cube_state(self, frame):

# shadow
cv2.rectangle(
frame,
self.frame,
(x1, y1),
(x2, y2),
(0, 0, 0),
Expand All @@ -413,7 +427,7 @@ def draw_2d_cube_state(self, frame):

# foreground color
cv2.rectangle(
frame,
self.frame,
(x1 + 1, y1 + 1),
(x2 - 1, y2 - 1),
foreground_color,
Expand Down Expand Up @@ -455,6 +469,7 @@ def run(self):
"""
while True:
_, frame = self.cam.read()
self.frame = frame
key = cv2.waitKey(10) & 0xff

# Quit on escape.
Expand All @@ -464,7 +479,7 @@ def run(self):
if not self.calibrate_mode:
# Update the snapshot when space bar is pressed.
if key == 32:
self.update_snapshot_state(frame)
self.update_snapshot_state()

# Switch to another language.
if key == ord(SWITCH_LANGUAGE_KEY):
Expand All @@ -477,21 +492,21 @@ def run(self):
self.reset_calibrate_mode()
self.calibrate_mode = not self.calibrate_mode

grayFrame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
grayFrame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY)
blurredFrame = cv2.blur(grayFrame, (3, 3))
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)
self.draw_contours(contours)
if not self.calibrate_mode:
self.update_preview_state(frame, contours)
elif key == 32 and self.done_calibrating == False:
self.update_preview_state(contours)
elif key == 32 and self.done_calibrating is False:
current_color = self.colors_to_calibrate[self.current_color_to_calibrate_index]
(x, y, w, h) = contours[4]
roi = frame[y+7:y+h-7, x+14:x+w-14]
roi = self.frame[y+7:y+h-7, x+14:x+w-14]
avg_bgr = color_detector.get_dominant_color(roi)
self.calibrated_colors[current_color] = avg_bgr
self.current_color_to_calibrate_index += 1
Expand All @@ -501,16 +516,16 @@ def run(self):
config.set_setting(CUBE_PALETTE, color_detector.cube_color_palette)

if self.calibrate_mode:
self.draw_current_color_to_calibrate(frame)
self.draw_calibrated_colors(frame)
self.draw_current_color_to_calibrate()
self.draw_calibrated_colors()
else:
self.draw_current_language(frame)
self.draw_preview_stickers(frame)
self.draw_snapshot_stickers(frame)
self.draw_scanned_sides(frame)
self.draw_2d_cube_state(frame)
self.draw_current_language()
self.draw_preview_stickers()
self.draw_snapshot_stickers()
self.draw_scanned_sides()
self.draw_2d_cube_state()

cv2.imshow("Qbr - Rubik's cube solver", frame)
cv2.imshow("Qbr - Rubik's cube solver", self.frame)

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

0 comments on commit aab1f37

Please sign in to comment.