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()