Blink Detection Support Added, Predicted DIameter Post Analysis Plots Added

app.py

@@ -5,9 +5,11 @@ import os.path as osp
 from PIL import Image
 from io import BytesIO
 import numpy as np
+import pandas as pd
 import streamlit as st
 from PIL import ImageOps
 from matplotlib import pyplot as plt
+import altair as alt
 
 root_path = osp.abspath(osp.join(__file__, osp.pardir))
 sys.path.append(root_path)
@@ -17,7 +19,7 @@ from app_utils import (
     extract_frames,
     is_image,
     is_video,
-    display_results,
+    convert_diameter,
     overlay_text_on_frame,
     process_frames,
     process_video,
@@ -36,6 +38,18 @@ LABEL_MAP = ["left_pupil", "right_pupil"]
 
 def main():
     st.set_page_config(page_title="Pupil Diameter Estimator", layout="wide")
+    st.markdown(
+        """
+        <style>
+            /* Remove the top margin/padding */
+            .block-container {
+                padding-top: 0rem;
+                padding-bottom: 1rem; /* Adjust this as needed */
+            }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
     st.title("EyeDentify Playground")
     cols = st.columns((1, 1))
     cols[0].header("Input")
@@ -77,6 +91,8 @@ def main():
     )
     tv_model = st.sidebar.selectbox("Classification model", ["ResNet18", "ResNet50"], help="Supported Models")
 
+    blink_detection = st.sidebar.checkbox("Detect Blinks")
+
     if st.sidebar.button("Predict Diameter & Compute CAM"):
         if uploaded_file is None:
             st.sidebar.error("Please upload an image or video")
@@ -90,8 +106,7 @@ def main():
                         tv_model,
                         pupil_selection,
                         cam_method=CAM_METHODS[-1],
-                        output_path=None,
-                        codec=None,
+                        blink_detection=blink_detection,
                     )
                     # for ff in face_frames:
                     #     if ff["has_face"]:
@@ -115,11 +130,63 @@ def main():
 
                 elif is_video(file_extension):
                     output_video_path = f"{root_path}/tmp.webm"
-                    process_video(
-                        cols, video_frames, tv_model, pupil_selection, output_video_path, cam_method=CAM_METHODS[-1]
+                    input_frames, output_frames, predicted_diameters, face_frames = process_video(
+                        cols,
+                        video_frames,
+                        tv_model,
+                        pupil_selection,
+                        output_video_path,
+                        cam_method=CAM_METHODS[-1],
+                        blink_detection=blink_detection,
                     )
                     os.remove(video_path)
 
+                    num_columns = len(predicted_diameters)
+
+                    # Create a layout for the charts
+                    cols = st.columns(num_columns)
+
+                    colors = ["#2ca02c", "#d62728", "#1f77b4", "#ff7f0e"]  # Green, Red, Blue, Orange
+
+                    # Iterate through categories and assign charts to columns
+                    for i, (category, values) in enumerate(predicted_diameters.items()):
+                        with cols[i]:  # Directly use the column index
+                            # st.subheader(category)  # Add a subheader for the category
+
+                            # Convert values to numeric, replacing non-numeric values with None
+                            values = [convert_diameter(value) for value in values]
+
+                            # Create a DataFrame from the values for Altair
+                            df = pd.DataFrame(values, columns=[category])
+                            df["Frame"] = range(1, len(values) + 1)  # Create a frame column starting from 1
+
+                            # Get the min and max values for y-axis limits, ignoring None
+                            min_value = min(filter(lambda x: x is not None, values), default=None)
+                            max_value = max(filter(lambda x: x is not None, values), default=None)
+
+                            # Create an Altair chart with y-axis limits
+                            chart = (
+                                alt.Chart(df)
+                                .mark_line(point=True, color=colors[i])
+                                .encode(
+                                    x=alt.X("Frame:Q", title="Frame Number"),
+                                    y=alt.Y(
+                                        f"{category}:Q",
+                                        title="Diameter",
+                                        scale=alt.Scale(domain=[min_value, max_value]),
+                                    ),
+                                    tooltip=[
+                                        alt.Tooltip("Frame:Q", title="Frame Number"),
+                                        alt.Tooltip(f"{category}:Q", title="Diameter"),
+                                    ],
+                                )
+                                .properties(title=f"{category} - Predicted Diameters")
+                                .configure_axis(grid=True)
+                            )
+
+                            # Display the Altair chart
+                            st.altair_chart(chart, use_container_width=True)
+
 
 if __name__ == "__main__":
     main()

app_utils.py

@@ -110,7 +110,7 @@ def overlay_text_on_frame(frame, text, position=(16, 20)):
     return cv2.putText(frame, text, position, cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1, cv2.LINE_AA)
 
 
-def process_frames(cols, input_imgs, tv_model, pupil_selection, cam_method, output_path, codec):
+def get_configs(blink_detection=False):
     upscale = "-"
     upscale_method_or_model = "-"
     if upscale == "-":
@@ -123,14 +123,21 @@ def process_frames(cols, input_imgs, tv_model, pupil_selection, cam_method, outp
     config_file = {
         "sr_configs": sr_configs,
         "feature_extraction_configs": {
-            "blink_detection": False,
+            "blink_detection": blink_detection,
             "upscale": upscale,
             "extraction_library": "mediapipe",
         },
     }
+
+    return config_file
+
+
+def setup(cols, pupil_selection, tv_model, output_path):
+
     left_pupil_model = None
+    left_pupil_cam_extractor = None
     right_pupil_model = None
-    face_frames = []
+    right_pupil_cam_extractor = None
     output_frames = {}
     input_frames = {}
     predicted_diameters = {}
@@ -163,21 +170,59 @@ def process_frames(cols, input_imgs, tv_model, pupil_selection, cam_method, outp
             input_frames[eye_type] = []
             predicted_diameters[eye_type] = []
 
+    video_input_placeholders = {}
+    video_output_placeholders = {}
+    video_predictions_placeholders = {}
+
     if output_path:
         video_cols = cols[1].columns(len(input_frames.keys()))
 
-        video_input_placeholders = {}
         for i, eye_type in enumerate(list(input_frames.keys())):
             video_input_placeholders[eye_type] = video_cols[i].empty()
 
-        video_output_placeholders = {}
         for i, eye_type in enumerate(list(input_frames.keys())):
             video_output_placeholders[eye_type] = video_cols[i].empty()
 
-        video_predictions_placeholders = {}
         for i, eye_type in enumerate(list(input_frames.keys())):
             video_predictions_placeholders[eye_type] = video_cols[i].empty()
 
+    return (
+        selected_eyes,
+        input_frames,
+        output_frames,
+        predicted_diameters,
+        video_input_placeholders,
+        video_output_placeholders,
+        video_predictions_placeholders,
+        left_pupil_model,
+        left_pupil_cam_extractor,
+        right_pupil_model,
+        right_pupil_cam_extractor,
+    )
+
+
+def process_frames(
+    cols, input_imgs, tv_model, pupil_selection, cam_method, output_path=None, codec=None, blink_detection=False
+):
+
+    config_file = get_configs(blink_detection)
+
+    face_frames = []
+
+    (
+        selected_eyes,
+        input_frames,
+        output_frames,
+        predicted_diameters,
+        video_input_placeholders,
+        video_output_placeholders,
+        video_predictions_placeholders,
+        left_pupil_model,
+        left_pupil_cam_extractor,
+        right_pupil_model,
+        right_pupil_cam_extractor,
+    ) = setup(cols, pupil_selection, tv_model, output_path)
+
     ds_creation = EyeDentityDatasetCreation(
         feature_extraction_configs=config_file["feature_extraction_configs"],
         sr_configs=config_file["sr_configs"],
@@ -212,17 +257,16 @@ def process_frames(cols, input_imgs, tv_model, pupil_selection, cam_method, outp
 
         if ds_results is not None and "eyes" in ds_results.keys():
             blinked = ds_results["eyes"]["blinked"]
-            if not blinked:
-                if "left_eye" in ds_results["eyes"].keys() and ds_results["eyes"]["left_eye"] is not None:
-                    left_eye = ds_results["eyes"]["left_eye"]
-                    left_eye = to_pil_image(left_eye).convert("RGB")
-                    left_eye = preprocess_function(left_eye)
-                    left_eye = left_eye.unsqueeze(0)
-                if "right_eye" in ds_results["eyes"].keys() and ds_results["eyes"]["right_eye"] is not None:
-                    right_eye = ds_results["eyes"]["right_eye"]
-                    right_eye = to_pil_image(right_eye).convert("RGB")
-                    right_eye = preprocess_function(right_eye)
-                    right_eye = right_eye.unsqueeze(0)
+            if "left_eye" in ds_results["eyes"].keys() and ds_results["eyes"]["left_eye"] is not None:
+                left_eye = ds_results["eyes"]["left_eye"]
+                left_eye = to_pil_image(left_eye).convert("RGB")
+                left_eye = preprocess_function(left_eye)
+                left_eye = left_eye.unsqueeze(0)
+            if "right_eye" in ds_results["eyes"].keys() and ds_results["eyes"]["right_eye"] is not None:
+                right_eye = ds_results["eyes"]["right_eye"]
+                right_eye = to_pil_image(right_eye).convert("RGB")
+                right_eye = preprocess_function(right_eye)
+                right_eye = right_eye.unsqueeze(0)
         else:
             input_img = preprocess_function(input_img)
             input_img = input_img.unsqueeze(0)
@@ -235,57 +279,67 @@ def process_frames(cols, input_imgs, tv_model, pupil_selection, cam_method, outp
                 right_eye = input_img
 
         for i, eye_type in enumerate(selected_eyes):
-            if left_eye is not None and eye_type == "left_eye":
-                if left_pupil_cam_extractor is None:
-                    if tv_model == "ResNet18":
-                        target_layer = left_pupil_model.resnet.layer4[-1].conv2
-                    elif tv_model == "ResNet50":
-                        target_layer = left_pupil_model.resnet.layer4[-1].conv3
-                    else:
-                        raise Exception(f"No target layer available for selected model: {tv_model}")
-                    left_pupil_cam_extractor = torchcam_methods.__dict__[cam_method](
-                        left_pupil_model,
-                        target_layer=target_layer,
-                        fc_layer=left_pupil_model.resnet.fc,
-                        input_shape=left_eye.shape,
-                    )
-                output = left_pupil_model(left_eye)
-                predicted_diameter = output[0].item()
-                act_maps = left_pupil_cam_extractor(0, output)
-                activation_map = act_maps[0] if len(act_maps) == 1 else left_pupil_cam_extractor.fuse_cams(act_maps)
-                input_image_pil = to_pil_image(left_eye.squeeze(0))
-            elif right_eye is not None and eye_type == "right_eye":
-                if right_pupil_cam_extractor is None:
-                    if tv_model == "ResNet18":
-                        target_layer = right_pupil_model.resnet.layer4[-1].conv2
-                    elif tv_model == "ResNet50":
-                        target_layer = right_pupil_model.resnet.layer4[-1].conv3
-                    else:
-                        raise Exception(f"No target layer available for selected model: {tv_model}")
-                    right_pupil_cam_extractor = torchcam_methods.__dict__[cam_method](
-                        right_pupil_model,
-                        target_layer=target_layer,
-                        fc_layer=right_pupil_model.resnet.fc,
-                        input_shape=right_eye.shape,
-                    )
-                output = right_pupil_model(right_eye)
-                predicted_diameter = output[0].item()
-                act_maps = right_pupil_cam_extractor(0, output)
-                activation_map = act_maps[0] if len(act_maps) == 1 else right_pupil_cam_extractor.fuse_cams(act_maps)
-                input_image_pil = to_pil_image(right_eye.squeeze(0))
 
             if blinked:
-                zeros_img = to_pil_image(np.zeros((256, 256, 3), dtype=np.uint8))
-                input_image_pil = zeros_img
-                result = zeros_img
-                predicted_diameter = 0
+
+                if left_eye is not None and eye_type == "left_eye":
+                    _, height, width = left_eye.squeeze(0).shape
+                    input_image_pil = to_pil_image(left_eye.squeeze(0))
+                elif right_eye is not None and eye_type == "right_eye":
+                    _, height, width = right_eye.squeeze(0).shape
+                    input_image_pil = to_pil_image(right_eye.squeeze(0))
+
+                input_img_np = np.array(input_image_pil)
+                zeros_img = to_pil_image(np.zeros((height, width, 3), dtype=np.uint8))
+                output_img_np = overlay_text_on_frame(np.array(zeros_img), "blink")
+                predicted_diameter = "blink"
             else:
+                if left_eye is not None and eye_type == "left_eye":
+                    if left_pupil_cam_extractor is None:
+                        if tv_model == "ResNet18":
+                            target_layer = left_pupil_model.resnet.layer4[-1].conv2
+                        elif tv_model == "ResNet50":
+                            target_layer = left_pupil_model.resnet.layer4[-1].conv3
+                        else:
+                            raise Exception(f"No target layer available for selected model: {tv_model}")
+                        left_pupil_cam_extractor = torchcam_methods.__dict__[cam_method](
+                            left_pupil_model,
+                            target_layer=target_layer,
+                            fc_layer=left_pupil_model.resnet.fc,
+                            input_shape=left_eye.shape,
+                        )
+                    output = left_pupil_model(left_eye)
+                    predicted_diameter = output[0].item()
+                    act_maps = left_pupil_cam_extractor(0, output)
+                    activation_map = act_maps[0] if len(act_maps) == 1 else left_pupil_cam_extractor.fuse_cams(act_maps)
+                    input_image_pil = to_pil_image(left_eye.squeeze(0))
+                elif right_eye is not None and eye_type == "right_eye":
+                    if right_pupil_cam_extractor is None:
+                        if tv_model == "ResNet18":
+                            target_layer = right_pupil_model.resnet.layer4[-1].conv2
+                        elif tv_model == "ResNet50":
+                            target_layer = right_pupil_model.resnet.layer4[-1].conv3
+                        else:
+                            raise Exception(f"No target layer available for selected model: {tv_model}")
+                        right_pupil_cam_extractor = torchcam_methods.__dict__[cam_method](
+                            right_pupil_model,
+                            target_layer=target_layer,
+                            fc_layer=right_pupil_model.resnet.fc,
+                            input_shape=right_eye.shape,
+                        )
+                    output = right_pupil_model(right_eye)
+                    predicted_diameter = output[0].item()
+                    act_maps = right_pupil_cam_extractor(0, output)
+                    activation_map = (
+                        act_maps[0] if len(act_maps) == 1 else right_pupil_cam_extractor.fuse_cams(act_maps)
+                    )
+                    input_image_pil = to_pil_image(right_eye.squeeze(0))
+
                 # Create CAM overlay
                 activation_map_pil = to_pil_image(activation_map, mode="F")
                 result = overlay_mask(input_image_pil, activation_map_pil, alpha=0.5)
-
-            input_img_np = np.array(input_image_pil)
-            output_img_np = np.array(result)
+                input_img_np = np.array(input_image_pil)
+                output_img_np = np.array(result)
 
             # Add frame and predicted diameter to lists
             input_frames[eye_type].append(input_img_np)
@@ -295,7 +349,10 @@ def process_frames(cols, input_imgs, tv_model, pupil_selection, cam_method, outp
             if output_path:
                 height, width, _ = output_img_np.shape
                 frame = np.zeros((height, width, 3), dtype=np.uint8)
-                text = f"{predicted_diameter:.2f}"
+                if not isinstance(predicted_diameter, str):
+                    text = f"{predicted_diameter:.2f}"
+                else:
+                    text = predicted_diameter
                 frame = overlay_text_on_frame(frame, text)
 
                 video_input_placeholders[eye_type].image(input_img_np, use_column_width=True)
@@ -385,7 +442,10 @@ def show_pred_text_frames(output_frames, output_path, predicted_diameters, codec
 
         for diameter in predicted_diameters[eye_type]:
             frame = np.zeros((height, width, 3), dtype=np.uint8)
-            text = f"{diameter:.2f}"
+            if not isinstance(diameter, str):
+                text = f"{diameter:.2f}"
+            else:
+                text = diameter
             frame = overlay_text_on_frame(frame, text)
             out.write(frame)
         out.release()
@@ -398,7 +458,7 @@ def show_pred_text_frames(output_frames, output_path, predicted_diameters, codec
         os.remove(output_path)
 
 
-def process_video(cols, video_frames, tv_model, pupil_selection, output_path, cam_method):
+def process_video(cols, video_frames, tv_model, pupil_selection, output_path, cam_method, blink_detection=False):
 
     resized_frames = []
     for i, frame in enumerate(video_frames):
@@ -408,4 +468,16 @@ def process_video(cols, video_frames, tv_model, pupil_selection, output_path, ca
     file_format = output_path.split(".")[-1]
     codec, extension = get_codec_and_extension(file_format)
 
-    process_frames(cols, resized_frames, tv_model, pupil_selection, cam_method, output_path, codec)
+    input_frames, output_frames, predicted_diameters, face_frames = process_frames(
+        cols, resized_frames, tv_model, pupil_selection, cam_method, output_path, codec, blink_detection
+    )
+
+    return input_frames, output_frames, predicted_diameters, face_frames
+
+
+# Function to convert string values to float or None
+def convert_diameter(value):
+    try:
+        return float(value)
+    except (ValueError, TypeError):
+        return None  # Return None if conversion fails

image.py

@@ -1,32 +0,0 @@
-import cv2
-import numpy as np
-
-# Load the original face image
-face_image = cv2.imread("path_to_face_image.jpg")
-
-# Suppose CAM_left and CAM_right are the CAM results for the eyes (each 32x64)
-CAM_left = cv2.imread("path_to_CAM_left.jpg")  # or generated by your model
-CAM_right = cv2.imread("path_to_CAM_right.jpg")  # or generated by your model
-
-# Example bounding boxes for the left and right eye
-left_eye_bbox = (x_left, y_left, width_left, height_left)
-right_eye_bbox = (x_right, y_right, width_right, height_right)
-
-# Resize CAM images if needed (they should be 32x64, but resize to match bbox size)
-CAM_left_resized = cv2.resize(CAM_left, (width_left, height_left))
-CAM_right_resized = cv2.resize(CAM_right, (width_right, height_right))
-
-# Create a copy of the face image to overlay the CAM results
-face_with_CAM = face_image.copy()
-
-# Overlay left eye CAM
-face_with_CAM[y_left : y_left + height_left, x_left : x_left + width_left] = CAM_left_resized
-
-# Overlay right eye CAM
-face_with_CAM[y_right : y_right + height_right, x_right : x_right + width_right] = CAM_right_resized
-
-# Save or display the result
-cv2.imwrite("face_with_CAM_overlay.jpg", face_with_CAM)
-cv2.imshow("Face with CAM Overlay", face_with_CAM)
-cv2.waitKey(0)
-cv2.destroyAllWindows()

registrations/models.py

@@ -11,38 +11,6 @@ sys.path.append(root_path)
 # ============================= ResNets =============================
 
 
-# @MODEL_REGISTRY.register()
-# class ResNet18(nn.Module):
-#     def __init__(self, model_args):
-#         super(ResNet18, self).__init__()
-#         self.num_classes = model_args.get("num_classes", 1)
-#         self.resnet = models.resnet18(weights=None, num_classes=self.num_classes)
-
-#     def forward(self, x, masks=None):
-#         return self.resnet(x)
-
-
-# @MODEL_REGISTRY.register()
-# class ResNet18(nn.Module):
-#     def __init__(self, model_args):
-#         super(ResNet18, self).__init__()
-#         self.num_classes = model_args.get("num_classes", 1)
-#         self.resnet = models.resnet18(weights=None, num_classes=self.num_classes)
-
-#     def forward(self, x, masks=None):
-#         # Calculate the padding dynamically based on the input size
-#         height, width = x.shape[2], x.shape[3]
-#         pad_height = max(0, (224 - height) // 2)
-#         pad_width = max(0, (224 - width) // 2)
-
-#         # Apply padding
-#         x = F.pad(
-#             x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
-#         )
-#         x = self.resnet(x)
-#         return x
-
-
 @MODEL_REGISTRY.register()
 class ResNet18(nn.Module):
     def __init__(self, model_args):
@@ -58,46 +26,12 @@ class ResNet18(nn.Module):
         pad_width = max(0, (224 - width) // 2)
 
         # Apply padding
-        x = F.pad(
-            x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
-        )
+        x = F.pad(x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0)
         x = self.resnet(x)
         x = self.regression_head(x)
         return x
 
 
-# @MODEL_REGISTRY.register()
-# class ResNet50(nn.Module):
-#     def __init__(self, model_args):
-#         super(ResNet50, self).__init__()
-#         self.num_classes = model_args.get("num_classes", 1)
-#         self.resnet = models.resnet50(weights=None, num_classes=self.num_classes)
-
-#     def forward(self, x, masks=None):
-#         return self.resnet(x)
-
-
-# @MODEL_REGISTRY.register()
-# class ResNet50(nn.Module):
-#     def __init__(self, model_args):
-#         super(ResNet50, self).__init__()
-#         self.num_classes = model_args.get("num_classes", 1)
-#         self.resnet = models.resnet50(weights=None, num_classes=self.num_classes)
-
-#     def forward(self, x, masks=None):
-#         # Calculate the padding dynamically based on the input size
-#         height, width = x.shape[2], x.shape[3]
-#         pad_height = max(0, (224 - height) // 2)
-#         pad_width = max(0, (224 - width) // 2)
-
-#         # Apply padding
-#         x = F.pad(
-#             x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
-#         )
-#         x = self.resnet(x)
-#         return x
-
-
 @MODEL_REGISTRY.register()
 class ResNet50(nn.Module):
     def __init__(self, model_args):
@@ -113,9 +47,7 @@ class ResNet50(nn.Module):
         pad_width = max(0, (224 - width) // 2)
 
         # Apply padding
-        x = F.pad(
-            x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
-        )
+        x = F.pad(x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0)
         x = self.resnet(x)
         x = self.regression_head(x)
         return x

video.py

@@ -1,48 +0,0 @@
-import cv2
-import torch
-
-# Load the video
-video_path = "path_to_video.mp4"
-cap = cv2.VideoCapture(video_path)
-
-# Video properties
-frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-fps = cap.get(cv2.CAP_PROP_FPS)
-
-# Create a VideoWriter object for the output video
-out = cv2.VideoWriter("output_with_CAM.mp4", cv2.VideoWriter_fourcc(*"mp4v"), fps, (frame_width, frame_height))
-
-# Process each frame
-while True:
-    ret, frame = cap.read()
-    if not ret:
-        break  # End of the video
-
-    # Detect landmarks for left and right eye bounding boxes (example)
-    left_eye_bbox = (x_left, y_left, width_left, height_left)
-    right_eye_bbox = (x_right, y_right, width_right, height_right)
-
-    # Crop the eyes
-    left_eye = frame[y_left : y_left + height_left, x_left : x_left + width_left]
-    right_eye = frame[y_right : y_right + height_right, x_right : x_right + width_right]
-
-    # Generate CAMs for left and right eyes
-    CAM_left = generate_CAM(left_eye)  # Use your model here
-    CAM_right = generate_CAM(right_eye)  # Use your model here
-
-    # Resize CAMs if necessary
-    CAM_left_resized = cv2.resize(CAM_left, (width_left, height_left))
-    CAM_right_resized = cv2.resize(CAM_right, (width_right, height_right))
-
-    # Overlay the CAMs onto the original frame
-    frame[y_left : y_left + height_left, x_left : x_left + width_left] = CAM_left_resized
-    frame[y_right : y_right + height_right, x_right : x_right + width_right] = CAM_right_resized
-
-    # Write the processed frame to the output video
-    out.write(frame)
-
-# Release resources
-cap.release()
-out.release()
-cv2.destroyAllWindows()