Update app.py

app.py

@@ -1,174 +1,181 @@
-import os
-import sys
-import traceback
-import gradio as gr
-import cv2 as cv
-import numpy as np
-import mediapipe as mp
-from utils import blinkRatio
-
-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, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
-                  min_detection_confidence, min_tracking_confidence, alpha):
-    try:
-        mp_face_mesh = mp.solutions.face_mesh
-        with mp_face_mesh.FaceMesh(
-            max_num_faces=1,
-            refine_landmarks=True,
-            min_detection_confidence=min_detection_confidence,
-            min_tracking_confidence=min_tracking_confidence
-        ) as face_mesh:
-            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:
-                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])
-                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)
-                (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
-    except Exception as e:
-        print(f"Error in process_frame: {e}")
-        traceback.print_exc()
-
-def process_image(input_image, overlay_file, alpha=0.3, min_detection_confidence=0.5, min_tracking_confidence=0.5):
-    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)
-    processed_frame = process_frame(frame, overlay, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
-                                    min_detection_confidence, min_tracking_confidence, 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,
-                                    min_detection_confidence=0.5, min_tracking_confidence=0.5):
-    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()
-        height, width, _ = frame.shape
-        fourcc = cv.VideoWriter_fourcc(*'mp4v' if output_format == 'mp4' else 'MJPG')
-        out = cv.VideoWriter(output_path, fourcc, fps, (width, height))  
-        while(cap.isOpened()):
-            ret, frame = cap.read()
-            if ret == True:
-                processed_frame = process_frame(frame,overlay,LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
-                                                float(min_detection_confidence), 
-                                                float(min_tracking_confidence), float(alpha))  # Assuming process_frame is a function that processes a single frame
-                processed_frame = cv.cvtColor(processed_frame, cv.COLOR_RGBA2BGR)
-                out.write(processed_frame)
-            else:
-                break
-        cap.release()
-        out.release()
-    return output_path
-
-
-def process_webcam(frame, overlay_file, alpha=0.3, min_detection_confidence=0.5, min_tracking_confidence=0.5):
-    overlay_file = overlay_file + '.png'
-    overlay_path = os.path.join(os.getcwd(), overlay_file)
-    overlay = cv.imread(overlay_path, cv.IMREAD_UNCHANGED)
-    processed_frame = process_frame(frame, overlay, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
-                                    min_detection_confidence, min_tracking_confidence, alpha)
-    return processed_frame
-
-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(height=500,width=400,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,width=400,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")
-            output_webcam = gr.Image(label="Processed Webcam")
-        process_webcam_btn = gr.Button("Process Webcam")
-        process_webcam_btn.click(process_webcam, 
-                                inputs=[input_webcam, overlay_file,], 
-                                outputs=output_webcam)
-
-demo.launch()
+import os
+import sys
+import traceback
+import gradio as gr
+import cv2 as cv
+import numpy as np
+import mediapipe as mp
+from utils import blinkRatio
+
+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, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
+                  min_detection_confidence, min_tracking_confidence, alpha):
+    try:
+        mp_face_mesh = mp.solutions.face_mesh
+        with mp_face_mesh.FaceMesh(
+            max_num_faces=1,
+            refine_landmarks=True,
+            min_detection_confidence=min_detection_confidence,
+            min_tracking_confidence=min_tracking_confidence
+        ) as face_mesh:
+            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:
+                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])
+                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)
+                (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
+    except Exception as e:
+        print(f"Error in process_frame: {e}")
+        traceback.print_exc()
+
+def process_image(input_image, overlay_file, alpha=0.3, min_detection_confidence=0.5, min_tracking_confidence=0.5):
+    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_CUBIC)
+    processed_frame = process_frame(frame, overlay, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
+                                    min_detection_confidence, min_tracking_confidence, 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,
+                                    min_detection_confidence=0.5, min_tracking_confidence=0.5):
+    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,LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
+                                                float(min_detection_confidence), 
+                                                float(min_tracking_confidence), float(alpha))  # Assuming process_frame is a function that processes a single frame
+                processed_frame = cv.cvtColor(processed_frame, cv.COLOR_RGBA2BGR)
+                out.write(processed_frame)
+            else:
+                break
+        cap.release()
+        out.release()
+    return output_path
+
+
+def process_webcam(frame, overlay_file, alpha=0.3, min_detection_confidence=0.5, min_tracking_confidence=0.5):
+    overlay_file = overlay_file + '.png'
+    overlay_path = os.path.join(os.getcwd(), overlay_file)
+    overlay = cv.imread(overlay_path, cv.IMREAD_UNCHANGED)
+    processed_frame = process_frame(frame, overlay, LEFT_EYE, RIGHT_EYE, LEFT_IRIS, RIGHT_IRIS, 
+                                    min_detection_confidence, min_tracking_confidence, alpha)
+    yield processed_frame
+
+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()