Files
WallE/control/camera.py

58 lines
1.7 KiB
Python

from io import BytesIO
import os.path
from typing import Optional
import configparser
import cv2
from PIL import Image
FACE_DATA = next(filter(os.path.exists, [
'/usr/share/opencv/haarcascades/haarcascade_frontalface_default.xml',
'/usr/share/opencv4/haarcascades/haarcascade_frontalface_default.xml',
]))
face_cascade = cv2.CascadeClassifier(FACE_DATA)
def get_rotate_value(input_value):
if input_value == '0':
return None
if input_value == '90':
return cv2.ROTATE_90_CLOCKWISE
if input_value == '-90' or input_value == '270':
return cv2.ROTATE_90_COUNTERCLOCKWISE
if input_value == '180':
return cv2.ROTATE_180
raise ValueError(f"Unsupported rotation value {input_value!r}")
class Camera:
capture: cv2.VideoCapture
def __init__(self):
self.capture = cv2.VideoCapture(0)
parser = configparser.ConfigParser()
parser.read('config.ini')
self.config = parser['camera']
self.rotate = get_rotate_value(self.config['rotate'])
def get_image(self) -> Optional[bytes]:
return_code, image = self.capture.read()
if not return_code:
return None
if self.rotate is not None:
image = cv2.rotate(image, self.rotate)
image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
faces = face_cascade.detectMultiScale(image_gray, 1.3, 5)
for (x, y, w, h) in faces:
image = cv2.rectangle(image, (x, y), (x+w, y+h), (255, 0, 0), 2)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
jpg = Image.fromarray(image_rgb)
byte_stream = BytesIO()
jpg.save(byte_stream, 'JPEG')
return bytes(byte_stream.getbuffer())