Spaces:
Sleeping
Sleeping
| import cv2 | |
| import cvzone | |
| import numpy as np | |
| import os | |
| import gradio as gr | |
| # Load the YuNet model | |
| model_path = 'face_detection_yunet_2023mar.onnx' | |
| face_detector = cv2.FaceDetectorYN.create(model_path, "", (320, 320)) | |
| # Initialize the glass number | |
| num = 1 | |
| overlay = cv2.imread(f'glasses/glass{num}.png', cv2.IMREAD_UNCHANGED) | |
| # Count glasses files | |
| def count_files_in_directory(directory): | |
| file_count = 0 | |
| for root, dirs, files in os.walk(directory): | |
| file_count += len(files) | |
| return file_count | |
| directory_path = 'glasses' | |
| total_glass_num = count_files_in_directory(directory_path) | |
| # Change glasses | |
| def change_glasses(): | |
| global num, overlay | |
| num += 1 | |
| if num > total_glass_num: | |
| num = 1 | |
| overlay = cv2.imread(f'glasses/glass{num}.png', cv2.IMREAD_UNCHANGED) | |
| return overlay | |
| # Process frame for overlay | |
| def process_frame(frame): | |
| global overlay | |
| # Ensure the frame is writable | |
| frame = np.array(frame, copy=True) | |
| height, width = frame.shape[:2] | |
| face_detector.setInputSize((width, height)) | |
| _, faces = face_detector.detect(frame) | |
| if faces is not None: | |
| for face in faces: | |
| x, y, w, h = face[:4].astype(int) | |
| face_landmarks = face[4:14].reshape(5, 2).astype(int) # Facial landmarks | |
| # Get the nose position | |
| nose_x, nose_y = face_landmarks[2].astype(int) | |
| # Left and right eye positions | |
| left_eye_x, left_eye_y = face_landmarks[0].astype(int) | |
| right_eye_x, right_eye_y = face_landmarks[1].astype(int) | |
| # Calculate the midpoint between the eyes | |
| eye_center_x = (left_eye_x + right_eye_x) // 2 | |
| eye_center_y = (left_eye_y + right_eye_y) // 2 | |
| # Calculate the angle of rotation | |
| delta_x = right_eye_x - left_eye_x | |
| delta_y = right_eye_y - left_eye_y | |
| angle = np.degrees(np.arctan2(delta_y, delta_x)) | |
| # Resize the overlay | |
| overlay_resize = cv2.resize(overlay, (int(w * 1.15), int(h * 0.8))) | |
| # Rotate the overlay | |
| overlay_center = (overlay_resize.shape[1] // 2, overlay_resize.shape[0] // 2) | |
| rotation_matrix = cv2.getRotationMatrix2D(overlay_center, angle, 1.0) | |
| overlay_rotated = cv2.warpAffine( | |
| overlay_resize, rotation_matrix, | |
| (overlay_resize.shape[1], overlay_resize.shape[0]), | |
| flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0) | |
| ) | |
| # Calculate the position to center the glasses on the eyes | |
| overlay_x = eye_center_x - overlay_rotated.shape[1] // 2 | |
| overlay_y = eye_center_y - overlay_rotated.shape[0] // 2 | |
| # Overlay the glasses | |
| try: | |
| frame = cvzone.overlayPNG(frame, overlay_rotated, [overlay_x, overlay_y]) | |
| except Exception as e: | |
| print(f"Error overlaying glasses: {e}") | |
| return frame | |
| # Gradio webcam input | |
| def webcam_input(frame): | |
| frame = process_frame(frame) | |
| return frame | |
| # Gradio Interface | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Input", sources="webcam", streaming=True) | |
| next_button = gr.Button("Next Glasses") | |
| input_img.stream(webcam_input, [input_img], [input_img], stream_every=0.1, concurrency_limit=30) | |
| next_button.click(change_glasses, [], []) | |
| if __name__ == "__main__": | |
| demo.launch(share=True) |