Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import traceback | |
| import gradio as gr | |
| import cv2 as cv | |
| import numpy as np | |
| import time | |
| import mediapipe as mp | |
| from mediapipe.tasks import python | |
| from mediapipe.tasks.python import vision | |
| from utils import blinkRatio,LEFT_EYE,RIGHT_EYE,LEFT_IRIS,RIGHT_IRIS | |
| def custom_excepthook(type, value, tb): | |
| traceback.print_exception(type, value, tb) | |
| sys.__excepthook__(type, value, tb) | |
| sys.excepthook = custom_excepthook | |
| def list_overlay_images(directory): | |
| return [f for f in os.listdir(directory) if f.endswith('.png')] | |
| def process_frame(frame, overlay, results, frame_timestamp_ms=None, task='image', alpha=None): | |
| if results.face_landmarks: | |
| rgba_frame = cv.cvtColor(frame, cv.COLOR_BGR2RGBA) | |
| height, width = rgba_frame.shape[:2] | |
| zero_overlay = np.zeros_like(rgba_frame) | |
| mesh_points = np.array([np.multiply([p.x, p.y], | |
| [width, height]).astype(int) for p in results.face_landmarks[0]]) | |
| iris_mask_left = np.zeros(rgba_frame.shape, dtype=np.uint8) | |
| iris_mask_right = np.zeros(rgba_frame.shape, dtype=np.uint8) | |
| _, re_ratio, le_ratio = blinkRatio(frame, mesh_points, RIGHT_EYE, LEFT_EYE) | |
| (l_cx, l_cy), l_radius = cv.minEnclosingCircle(mesh_points[LEFT_IRIS]) | |
| (r_cx, r_cy), r_radius = cv.minEnclosingCircle(mesh_points[RIGHT_IRIS]) | |
| center_left = (int(l_cx), int(l_cy)) | |
| center_right = (int(r_cx), int(r_cy)) | |
| cv.circle(iris_mask_left, center_left, int(l_radius), (255, 0, 0, 255), -1, cv.LINE_AA) | |
| cv.circle(iris_mask_right, center_right, int(r_radius), (255, 0, 0, 255), -1, cv.LINE_AA) | |
| bbx_size_l = int((l_radius * 2) / 2) | |
| bbx_size_r = int((r_radius * 2) / 2) | |
| resized_overlay_l = cv.resize(overlay, (bbx_size_l * 2, bbx_size_l * 2), interpolation=cv.INTER_CUBIC) | |
| resized_overlay_r = cv.resize(overlay, (bbx_size_r * 2, bbx_size_r * 2), interpolation=cv.INTER_CUBIC) | |
| y1_r = center_right[1] - bbx_size_r | |
| y2_r = center_right[1] + bbx_size_r | |
| x1_r = center_right[0] - bbx_size_r | |
| x2_r = center_right[0] + bbx_size_r | |
| y1_l = center_left[1] - bbx_size_l | |
| y2_l = center_left[1] + bbx_size_l | |
| x1_l = center_left[0] - bbx_size_l | |
| x2_l = center_left[0] + bbx_size_l | |
| if (resized_overlay_l.shape == zero_overlay[y1_l:y2_l, x1_l:x2_l].shape) & (le_ratio < 5.0) & (le_ratio > 2.0): | |
| zero_overlay[y1_l:y2_l, x1_l:x2_l] = resized_overlay_l | |
| if (resized_overlay_r.shape == zero_overlay[y1_r:y2_r, x1_r:x2_r].shape) & (re_ratio < 5.0) & (re_ratio > 2.0): | |
| zero_overlay[y1_r:y2_r, x1_r:x2_r] = resized_overlay_r | |
| eye_mask_left = np.zeros(rgba_frame.shape, dtype=np.uint8) | |
| eye_mask_right = np.zeros(rgba_frame.shape, dtype=np.uint8) | |
| cv.fillPoly(eye_mask_left, [mesh_points[LEFT_EYE]], (255, 0, 0, 255)) | |
| cv.fillPoly(eye_mask_right, [mesh_points[RIGHT_EYE]], (255, 0, 0, 255)) | |
| zero_overlay[np.where((iris_mask_left[:, :, 3] > 0) & (eye_mask_left[:, :, 3] == 0))] = 0 | |
| zero_overlay[np.where((iris_mask_right[:, :, 3] > 0) & (eye_mask_right[:, :, 3] == 0))] = 0 | |
| rgba_frame = cv.addWeighted(rgba_frame, 1, zero_overlay, alpha, 0) | |
| return rgba_frame | |
| def process_image(input_image, overlay_file, alpha=0.3): | |
| model_path = os.path.join(os.getcwd(),'face_landmarker.task') | |
| BaseOptions = mp.tasks.BaseOptions | |
| FaceLandmarker = mp.tasks.vision.FaceLandmarker | |
| FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions | |
| VisionRunningMode = mp.tasks.vision.RunningMode | |
| options = FaceLandmarkerOptions( | |
| base_options=BaseOptions(model_asset_path=model_path), | |
| running_mode=VisionRunningMode.IMAGE) | |
| with FaceLandmarker.create_from_options(options) as landmarker: | |
| overlay_file = overlay_file + '.png' | |
| overlay_path = os.path.join(os.getcwd(),'overlays', overlay_file) | |
| overlay = cv.imread(overlay_path, cv.IMREAD_UNCHANGED) | |
| frame = np.array(input_image) | |
| if frame.dtype != np.uint8: | |
| frame = (frame * 255).astype(np.uint8) | |
| rgb_frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB) | |
| mp_frame = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame) | |
| results = landmarker.detect(mp_frame) | |
| processed_frame = process_frame(frame=frame, overlay=overlay, results=results, alpha=alpha) | |
| return cv.cvtColor(processed_frame, cv.COLOR_BGR2RGB) | |
| def process_video(input_video, overlay_file, alpha=0.3, output_format='mp4', output_frame_rate=30): | |
| model_path = os.path.join(os.getcwd(), 'face_landmarker.task') | |
| BaseOptions = mp.tasks.BaseOptions | |
| FaceLandmarker = mp.tasks.vision.FaceLandmarker | |
| FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions | |
| VisionRunningMode = mp.tasks.vision.RunningMode | |
| options = FaceLandmarkerOptions( | |
| base_options=BaseOptions(model_asset_path=model_path), | |
| running_mode=VisionRunningMode.VIDEO) | |
| with FaceLandmarker.create_from_options(options) as landmarker: | |
| overlay_file = overlay_file + '.png' | |
| overlay_path = os.path.join(os.getcwd(), 'overlays', overlay_file) | |
| overlay = cv.imread(overlay_path, cv.IMREAD_UNCHANGED) | |
| cap = cv.VideoCapture(input_video) | |
| output_path = os.path.join(os.getcwd(), f'video_processed.{output_format}') | |
| if overlay is not None and cap.isOpened(): | |
| fps = int(output_frame_rate) if output_frame_rate > 0 else cap.get(cv.CAP_PROP_FPS) | |
| h, w = None, None | |
| new_h, new_w = None, None | |
| frame_idx = 0 | |
| fourcc = cv.VideoWriter_fourcc(*'mp4v' if output_format == 'mp4' else 'MJPG') | |
| out = cv.VideoWriter(output_path, fourcc, fps, (new_w, new_h)) | |
| start_time = time.time() | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| if h is None or w is None: | |
| h, w, _ = frame.shape | |
| new_h = 800 | |
| new_w = int((w / h) * new_h) | |
| out = cv.VideoWriter(output_path, fourcc, fps, (new_w, new_h)) # Initialize output writer with correct size | |
| frame = cv.resize(frame, (new_w, new_h), interpolation=cv.INTER_NEAREST) | |
| if frame.dtype != np.uint8: | |
| frame = (frame * 255).astype(np.uint8) | |
| rgb_frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB) | |
| mp_frame = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame) | |
| timestamp = int(frame_idx * 1000 / fps) # Convert frame index to milliseconds | |
| results = landmarker.detect_for_video(mp_frame, timestamp) | |
| processed_frame = process_frame(frame=frame, overlay=overlay, results=results, alpha=alpha) | |
| processed_frame = cv.cvtColor(processed_frame, cv.COLOR_RGBA2BGR) | |
| out.write(processed_frame) | |
| frame_idx += 1 | |
| cap.release() | |
| out.release() | |
| end_time = time.time() | |
| execution_time = end_time - start_time | |
| print(f"Execution time: {execution_time} seconds") | |
| return output_path | |
| def process_webcam(frame, overlay_file, alpha=0.3, min_detection_confidence=0.5, min_tracking_confidence=0.5): | |
| BaseOptions = mp.tasks.BaseOptions | |
| FaceLandmarker = mp.tasks.vision.FaceLandmarker | |
| FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions | |
| FaceLandmarkerResult = mp.tasks.vision.FaceLandmarkerResult | |
| VisionRunningMode = mp.tasks.vision.RunningMode | |
| model_path = os.path.join(os.getcwd(), 'face_landmarker.task') | |
| overlay_file = overlay_file + '.png' | |
| overlay_path = os.path.join(os.getcwd(), overlay_file) | |
| overlay = cv.imread(overlay_path, cv.IMREAD_UNCHANGED) | |
| global latest_results | |
| latest_results = None | |
| def return_result(result: FaceLandmarkerResult, output_image: mp.Image, timestamp_ms: int): | |
| global latest_results | |
| latest_results = result | |
| options = FaceLandmarkerOptions( | |
| base_options=BaseOptions(model_asset_path=model_path), | |
| running_mode=VisionRunningMode.LIVE_STREAM, | |
| result_callback=return_result) | |
| with FaceLandmarker.create_from_options(options) as landmarker: | |
| timestamp_ms = int(time.time() * 1000) # Current time in milliseconds | |
| mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame) | |
| landmarker.detect_async(mp_image, timestamp_ms) | |
| while latest_results is None: | |
| time.sleep(0.01) # Wait for the result to be available | |
| processed_frame = process_frame(frame, overlay, latest_results, alpha) | |
| return processed_frame | |
| overlay_dir = os.path.join(os.getcwd(),'overlays') | |
| overlay_files = list_overlay_images(overlay_dir) | |
| overlay_choices = [x.split('.png')[0] for x in overlay_files] | |
| overlay_file = gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color") | |
| image_interface = gr.Interface( | |
| process_image, | |
| [gr.Image(height=500,label="Upload Image"), | |
| gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")], | |
| gr.Image(height=500), | |
| ) | |
| video_interface = gr.Interface( | |
| process_video, | |
| [gr.Video(height=500,label="Upload Video"), | |
| gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")], | |
| gr.Video(height=500,label="Processed Video"), | |
| ) | |
| webcam_interface = gr.Interface( | |
| process_webcam, | |
| [gr.Image(sources=["webcam"], streaming=True), | |
| gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")], | |
| "image", | |
| live=True | |
| ) | |
| demo = gr.TabbedInterface([image_interface,video_interface],['Image','Video']) | |
| demo.launch() | |