Update app.py

app.py

@@ -4,8 +4,14 @@ import traceback
 import gradio as gr
 import cv2 as cv
 import numpy as np
+import time
 import mediapipe as mp
-from utils import blinkRatio
+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)
@@ -16,18 +22,16 @@ 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, alpha, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS):
-    rgb_frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
-    rgba_frame = cv.cvtColor(frame, cv.COLOR_BGR2RGBA)
-    height, width = rgba_frame.shape[:2]
-    results = face_mesh.process(rgb_frame)
-    if results.multi_face_landmarks:
+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.multi_face_landmarks[0].landmark])
+                [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(rgb_frame, mesh_points, RIGHT_EYE, LEFT_EYE)
+        _, 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))
@@ -60,116 +64,150 @@ def process_frame(frame, overlay, alpha, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_I
     return rgba_frame
 
 def process_image(input_image, overlay_file, alpha=0.3):
-    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)
-    w,h,_ = frame.shape
-    new_h = 500
-    new_w = int((w/h)*new_h)
-    frame = cv.resize(frame, (new_h,new_w), interpolation=cv.INTER_NEAREST)
-    processed_frame = process_frame(frame, overlay, alpha, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS)
-    return cv.cvtColor(processed_frame, cv.COLOR_BGR2RGB)
+    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):
-    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}')
-    # Define the codec and create a VideoWriter object to save the processed video
-    if (not isinstance(overlay,type(None))) & (not isinstance(cap,type(None))):
-        # Get the dimensions of the frame, fps
-        fps=int(output_frame_rate)
-        if fps==0:
-            fps = cap.get(5)
-        ret, frame = cap.read()
-        h,w,_ = frame.shape
-        new_h = 500
-        new_w = int((w/h)*new_h)
-        fourcc = cv.VideoWriter_fourcc(*'mp4v' if output_format == 'mp4' else 'MJPG')
-        out = cv.VideoWriter(output_path, fourcc, fps, (new_w,new_h))  
-        while(cap.isOpened()):
-            ret, frame = cap.read()
-            if ret == True:
-                frame = cv.resize(frame, (new_w,new_h), interpolation=cv.INTER_NEAREST)
-                processed_frame = process_frame(frame,overlay,alpha,LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS)  # Assuming process_frame is a function that processes a single frame
+    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)
-            else:
-                break
-        cap.release()
-        out.release()
-    return output_path
+                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
 
-# Initialize face mesh once and reuse it
-mp_face_mesh = mp.solutions.face_mesh
-face_mesh = mp_face_mesh.FaceMesh(
-        max_num_faces=1,
-        refine_landmarks=True,
-        min_detection_confidence=0.5,
-        min_tracking_confidence=0.5
-    )
 
-LEFT_EYE = [362, 382, 381, 380, 374, 373, 390, 249, 263, 466, 388, 387, 386, 385, 384, 398]
-RIGHT_EYE = [33, 7, 163, 144, 145, 153, 154, 155, 133, 173, 157, 158, 159, 160, 161, 246]
-LEFT_IRIS = [474, 475, 476, 477]
-RIGHT_IRIS = [469, 470, 471, 472]
 
 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]
-# with gr.Blocks() as demo:
-#     with gr.Tab("Image"):
-#         with gr.Row():
-#             overlay_file = gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")
-#             # min_detection_confidence = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Min Detection Confidence")
-#             # min_tracking_confidence = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Min Tracking Confidence")
-#             # alpha = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, label="Overlay Transparency")
-#         with gr.Row():
-#             input_image = gr.Image(height=500,width=400,label="Upload Image")
-#             output_image = gr.Image(label="Processed Image")
-#         process_image_btn = gr.Button("Process Image")
-#         process_image_btn.click(process_image, 
-#                                 inputs=[input_image, overlay_file,], 
-#                                 outputs=output_image)
-    
-#     with gr.Tab("Video"):
-#         with gr.Row():
-#             overlay_file = gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")
-#             # min_detection_confidence = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Min Detection Confidence")
-#             # min_tracking_confidence = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Min Tracking Confidence")
-#             # alpha = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, label="Overlay Transparency")
-#         with gr.Row():
-#             input_video = gr.Video(height=500,width=400,label="Upload Video")
-#             output_video = gr.Video(height=500,label="Processed Video")
-#         process_video_btn = gr.Button("Process Video")
-#         process_video_btn.click(process_video, 
-#                                 inputs=[input_video, overlay_file,], 
-#                                 outputs=output_video)
-
-#     with gr.Tab("Webcam"):
-#         with gr.Row():
-#             overlay_file = gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")
-#             # min_detection_confidence = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Min Detection Confidence")
-#             # min_tracking_confidence = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Min Tracking Confidence")
-#             # alpha = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, label="Overlay Transparency")
-#         with gr.Row():
-#             # input_webcam = gr.Video(sources="webcam", label="Webcam")
-#             webcam = gr.Image(sources="webcam",label="Processed Webcam",streaming=True)
-#         process_webcam_btn = gr.Button("Process Webcam")
-#         process_webcam_btn.click(process_webcam, 
-#                                 inputs=[webcam, overlay_file,], 
-#                                 outputs=webcam)
-
-# demo.launch()
+
+
 overlay_file = gr.Dropdown(choices=overlay_choices, value='Blue', label="Select a color")
-demo = gr.Interface(
+image_interface = gr.Interface(
     process_image, 
-    [gr.Image(sources=["webcam"], streaming=True),overlay_file], 
+    [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.launch()
+
+demo = gr.TabbedInterface([image_interface,video_interface,webcam_interface],['Image','Video','Webcam'])
+
+
+demo.launch()