app.py

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