feat: switch to objects
This commit is contained in:
parent
d6d53166e8
commit
58cac43a37
125
src/environment.py
Normal file
125
src/environment.py
Normal file
|
@ -0,0 +1,125 @@
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import mediapipe as mp
|
||||||
|
import pygame as pg
|
||||||
|
|
||||||
|
from parts.eye import Eye
|
||||||
|
|
||||||
|
|
||||||
|
class Environment:
|
||||||
|
def __init__(self, camera: cv2.VideoCapture) -> None:
|
||||||
|
# store reference to webcam
|
||||||
|
self.cam = camera
|
||||||
|
|
||||||
|
# mediapipe stuff
|
||||||
|
self.mp_drawing = mp.solutions.drawing_utils # type: ignore
|
||||||
|
self.mp_drawing_styles = mp.solutions.drawing_styles # type: ignore
|
||||||
|
self.mp_face_mesh = mp.solutions.face_mesh # type: ignore
|
||||||
|
|
||||||
|
# init pygame
|
||||||
|
pg.init()
|
||||||
|
|
||||||
|
# get screen size from webcam
|
||||||
|
self.screen_width = camera.get(3)
|
||||||
|
self.screen_height = camera.get(4)
|
||||||
|
|
||||||
|
# create screen
|
||||||
|
self.screen: pg.Surface = pg.display.set_mode(
|
||||||
|
(self.screen_width, self.screen_height),
|
||||||
|
pg.DOUBLEBUF | pg.HWSURFACE,
|
||||||
|
) # type: ignore
|
||||||
|
pg.display.set_caption("Projet APP")
|
||||||
|
|
||||||
|
# create body parts
|
||||||
|
self.body_parts = {
|
||||||
|
"left_eye": Eye(False),
|
||||||
|
"right_eye": Eye(True),
|
||||||
|
}
|
||||||
|
|
||||||
|
def start(self) -> None:
|
||||||
|
while self.cam.isOpened():
|
||||||
|
# read webcam
|
||||||
|
success, self.frame = self.cam.read()
|
||||||
|
if not success:
|
||||||
|
logging.debug("Ignoring empty camera frame.")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# stop if q is pressed (opencv)
|
||||||
|
if cv2.waitKey(5) & 0xFF == ord("q"):
|
||||||
|
break
|
||||||
|
|
||||||
|
# quit the game
|
||||||
|
if any(
|
||||||
|
[event.type == pg.KEYDOWN and event.key == pg.K_q for event in pg.event.get()],
|
||||||
|
):
|
||||||
|
break
|
||||||
|
|
||||||
|
# detect keypoints on frame
|
||||||
|
self.detect_keypoints()
|
||||||
|
|
||||||
|
# draw keypoints on top of fram
|
||||||
|
self.draw_keypoints()
|
||||||
|
|
||||||
|
# draw avatar
|
||||||
|
self.draw_avatar()
|
||||||
|
|
||||||
|
def detect_keypoints(self) -> None:
|
||||||
|
with self.mp_face_mesh.FaceMesh(
|
||||||
|
max_num_faces=1,
|
||||||
|
refine_landmarks=True,
|
||||||
|
min_detection_confidence=0.5,
|
||||||
|
min_tracking_confidence=0.5,
|
||||||
|
) as face_mesh:
|
||||||
|
|
||||||
|
# perf, mark image as not writeable to pass by reference
|
||||||
|
self.frame.flags.writeable = False
|
||||||
|
|
||||||
|
# convert the BGR image to RGB before processing
|
||||||
|
self.frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2RGB)
|
||||||
|
|
||||||
|
# process the image to detect keypoints
|
||||||
|
self.results = face_mesh.process(self.frame)
|
||||||
|
|
||||||
|
# re-enable writeable flag
|
||||||
|
self.frame.flags.writeable = True
|
||||||
|
|
||||||
|
# convert the image back to BGR format
|
||||||
|
self.frame = cv2.cvtColor(self.frame, cv2.COLOR_RGB2BGR)
|
||||||
|
|
||||||
|
def draw_keypoints(self) -> None:
|
||||||
|
# draw the face mesh annotations on the image.
|
||||||
|
if self.results.multi_face_landmarks:
|
||||||
|
for face_landmarks in self.results.multi_face_landmarks:
|
||||||
|
self.mp_drawing.draw_landmarks(
|
||||||
|
image=self.frame,
|
||||||
|
landmark_list=face_landmarks,
|
||||||
|
connections=self.mp_face_mesh.FACEMESH_TESSELATION,
|
||||||
|
landmark_drawing_spec=None,
|
||||||
|
connection_drawing_spec=self.mp_drawing_styles.get_default_face_mesh_tesselation_style(),
|
||||||
|
)
|
||||||
|
self.mp_drawing.draw_landmarks(
|
||||||
|
image=self.frame,
|
||||||
|
landmark_list=face_landmarks,
|
||||||
|
connections=self.mp_face_mesh.FACEMESH_CONTOURS,
|
||||||
|
landmark_drawing_spec=None,
|
||||||
|
connection_drawing_spec=self.mp_drawing_styles.get_default_face_mesh_contours_style(),
|
||||||
|
)
|
||||||
|
self.mp_drawing.draw_landmarks(
|
||||||
|
image=self.frame,
|
||||||
|
landmark_list=face_landmarks,
|
||||||
|
connections=self.mp_face_mesh.FACEMESH_IRISES,
|
||||||
|
landmark_drawing_spec=None,
|
||||||
|
connection_drawing_spec=self.mp_drawing_styles.get_default_face_mesh_iris_connections_style(),
|
||||||
|
)
|
||||||
|
|
||||||
|
# flip the image horizontally for a selfie-view display
|
||||||
|
cv2.imshow("MediaPipe Face Mesh", cv2.flip(self.frame, 1))
|
||||||
|
|
||||||
|
def draw_avatar(self) -> None:
|
||||||
|
self.screen.fill(pg.Color("green"))
|
||||||
|
|
||||||
|
for part in self.body_parts.values():
|
||||||
|
part.draw(self.screen)
|
||||||
|
|
||||||
|
pg.display.flip()
|
59
src/main.py
59
src/main.py
|
@ -1,67 +1,20 @@
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import mediapipe as mp
|
|
||||||
|
from environment import Environment
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
"""Main function."""
|
"""Main function."""
|
||||||
mp_drawing = mp.solutions.drawing_utils # type: ignore
|
|
||||||
mp_drawing_styles = mp.solutions.drawing_styles # type: ignore
|
|
||||||
mp_face_mesh = mp.solutions.face_mesh # type: ignore
|
|
||||||
|
|
||||||
# open webcam
|
# open webcam
|
||||||
cap = cv2.VideoCapture(0)
|
cap = cv2.VideoCapture(0)
|
||||||
|
|
||||||
with mp_face_mesh.FaceMesh(
|
# create env
|
||||||
max_num_faces=2, refine_landmarks=True, min_detection_confidence=0.5, min_tracking_confidence=0.5
|
env = Environment(cap)
|
||||||
) as face_mesh:
|
|
||||||
|
|
||||||
while cap.isOpened():
|
# start env
|
||||||
|
env.start()
|
||||||
# read webcam
|
|
||||||
success, image = cap.read()
|
|
||||||
if not success:
|
|
||||||
print("Ignoring empty camera frame.")
|
|
||||||
# If loading a video, use 'break' instead of 'continue'
|
|
||||||
continue
|
|
||||||
|
|
||||||
# to improve performance, optionally mark the image as not writeable to pass by reference.
|
|
||||||
image.flags.writeable = False
|
|
||||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
||||||
results = face_mesh.process(image)
|
|
||||||
|
|
||||||
# draw the face mesh annotations on the image.
|
|
||||||
image.flags.writeable = True
|
|
||||||
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
||||||
if results.multi_face_landmarks:
|
|
||||||
for face_landmarks in results.multi_face_landmarks:
|
|
||||||
mp_drawing.draw_landmarks(
|
|
||||||
image=image,
|
|
||||||
landmark_list=face_landmarks,
|
|
||||||
connections=mp_face_mesh.FACEMESH_TESSELATION,
|
|
||||||
landmark_drawing_spec=None,
|
|
||||||
connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_tesselation_style(),
|
|
||||||
)
|
|
||||||
mp_drawing.draw_landmarks(
|
|
||||||
image=image,
|
|
||||||
landmark_list=face_landmarks,
|
|
||||||
connections=mp_face_mesh.FACEMESH_CONTOURS,
|
|
||||||
landmark_drawing_spec=None,
|
|
||||||
connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_contours_style(),
|
|
||||||
)
|
|
||||||
mp_drawing.draw_landmarks(
|
|
||||||
image=image,
|
|
||||||
landmark_list=face_landmarks,
|
|
||||||
connections=mp_face_mesh.FACEMESH_IRISES,
|
|
||||||
landmark_drawing_spec=None,
|
|
||||||
connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_iris_connections_style(),
|
|
||||||
)
|
|
||||||
|
|
||||||
# flip the image horizontally for a selfie-view display.
|
|
||||||
cv2.imshow("MediaPipe Face Mesh", cv2.flip(image, 1))
|
|
||||||
if cv2.waitKey(5) & 0xFF == 27:
|
|
||||||
break
|
|
||||||
|
|
||||||
# close webcam
|
# close webcam
|
||||||
cap.release()
|
cap.release()
|
||||||
|
|
Loading…
Reference in a new issue