Can download data as csv file and some UI changes

app.py

@@ -1,251 +1,60 @@
-import os
 import sys
-import tempfile
 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)
 
 from registry_utils import import_registered_modules
 from app_utils import (
-    extract_frames,
     is_image,
     is_video,
-    convert_diameter,
-    overlay_text_on_frame,
-    process_frames,
-    process_video,
-    resize_frame,
+    process_image_and_vizualize_data,
+    process_video_and_visualize_data,
+    set_frames_processed_count_placeholder,
+    set_input_image_on_ui,
+    set_input_video_on_ui,
+    set_page_info,
+    set_sidebar_info,
 )
 
 import_registered_modules()
 
-CAM_METHODS = ["CAM"]
-TV_MODELS = ["ResNet18", "ResNet50"]
-SR_METHODS = ["GFPGAN", "CodeFormer", "RealESRGAN", "SRResNet", "HAT"]
-UPSCALE = [2, 4]
-UPSCALE_METHODS = ["BILINEAR", "BICUBIC"]
-LABEL_MAP = ["left_pupil", "right_pupil"]
-
 
 def main():
-    st.set_page_config(page_title="Pupil Diameter Estimator", layout="wide")
-    st.title("EyeDentify Playground")
-    cols = st.columns((1, 1))
-    cols[0].header("Input")
-    cols[-1].header("Prediction")
-
-    st.sidebar.title("Upload Face or Eye")
-    uploaded_file = st.sidebar.file_uploader(
-        "Upload Image or Video", type=["png", "jpeg", "jpg", "mp4", "avi", "mov", "mkv", "webm"]
-    )
+    cols = set_page_info()
+    uploaded_file, pupil_selection, tv_model, blink_detection = set_sidebar_info()
 
     if uploaded_file is not None:
         file_extension = uploaded_file.name.split(".")[-1]
+        st.session_state["file_extension"] = file_extension
 
         if is_image(file_extension):
-            input_img = Image.open(BytesIO(uploaded_file.read())).convert("RGB")
-            # NOTE: images taken with phone camera has an EXIF data field which often rotates images taken with the phone in a tilted position. PIL has a utility function that removes this data and ‘uprights’ the image.
-            input_img = ImageOps.exif_transpose(input_img)
-            input_img = resize_frame(input_img, max_width=640, max_height=480)
-            input_img = resize_frame(input_img, max_width=640, max_height=480)
-            cols[0].image(input_img, use_column_width=True)
-            st.session_state.total_frames = 1
-
+            input_img = set_input_image_on_ui(uploaded_file, cols)
+            st.session_state["input_img"] = input_img
         elif is_video(file_extension):
-            tfile = tempfile.NamedTemporaryFile(delete=False)
-            tfile.write(uploaded_file.read())
-            video_path = tfile.name
-            video_frames = extract_frames(video_path)
-            cols[0].video(video_path)
-            st.session_state.total_frames = len(video_frames)
-
-        st.session_state.current_frame = 0
-        st.session_state.frame_placeholder = cols[0].empty()
-        txt = f"<p style='font-size:20px;'> Number of Frames Processed: <strong>{st.session_state.current_frame} / {st.session_state.total_frames}</strong> </p>"
-        st.session_state.frame_placeholder.markdown(txt, unsafe_allow_html=True)
-
-    st.sidebar.title("Setup")
-    pupil_selection = st.sidebar.selectbox(
-        "Pupil Selection", ["both"] + LABEL_MAP, help="Select left or right pupil OR both for diameter estimation"
-    )
-    tv_model = st.sidebar.selectbox("Classification model", ["ResNet18", "ResNet50"], help="Supported Models")
+            video_frames, video_path = set_input_video_on_ui(uploaded_file, cols)
+            st.session_state["video_frames"] = video_frames
+            st.session_state["video_path"] = video_path
 
-    blink_detection = st.sidebar.checkbox("Detect Blinks")
-
-    st.markdown("<style>#vg-tooltip-element{z-index: 1000051}</style>", unsafe_allow_html=True)
+        set_frames_processed_count_placeholder(cols)
 
     if st.sidebar.button("Predict Diameter & Compute CAM"):
         if uploaded_file is None:
             st.sidebar.error("Please upload an image or video")
         else:
             with st.spinner("Analyzing..."):
-
-                if is_image(file_extension):
-                    input_frames, output_frames, predicted_diameters, face_frames, eyes_ratios = process_frames(
-                        cols,
-                        [input_img],
-                        tv_model,
-                        pupil_selection,
-                        cam_method=CAM_METHODS[-1],
-                        blink_detection=blink_detection,
+                if is_image(st.session_state.get("file_extension")):
+                    input_img = st.session_state.get("input_img")
+                    process_image_and_vizualize_data(cols, input_img, tv_model, pupil_selection, blink_detection)
+                elif is_video(st.session_state.get("file_extension")):
+                    video_frames = st.session_state.get("video_frames")
+                    video_path = st.session_state.get("video_path")
+                    process_video_and_visualize_data(
+                        cols, video_frames, tv_model, pupil_selection, blink_detection, video_path
                     )
-                    # for ff in face_frames:
-                    #     if ff["has_face"]:
-                    #         cols[1].image(face_frames[0]["img"], use_column_width=True)
-
-                    input_frames_keys = input_frames.keys()
-                    video_cols = cols[1].columns(len(input_frames_keys))
-                    for i, eye_type in enumerate(input_frames_keys):
-                        video_cols[i].image(input_frames[eye_type][-1], use_column_width=True)
-
-                    output_frames_keys = output_frames.keys()
-                    fig, axs = plt.subplots(1, len(output_frames_keys), figsize=(10, 5))
-                    for i, eye_type in enumerate(output_frames_keys):
-                        height, width, c = output_frames[eye_type][0].shape
-                        video_cols[i].image(output_frames[eye_type][-1], use_column_width=True)
-
-                        frame = np.zeros((height, width, c), dtype=np.uint8)
-                        text = f"{predicted_diameters[eye_type][0]:.2f}"
-                        frame = overlay_text_on_frame(frame, text)
-                        video_cols[i].image(frame, use_column_width=True)
-
-                elif is_video(file_extension):
-                    output_video_path = f"{root_path}/tmp.webm"
-                    input_frames, output_frames, predicted_diameters, face_frames, eyes_ratios = 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
-                    colors = ["#1f77b4", "#ff7f0e", "#636363"]  # Blue, Orange, Gray
-
-                    # 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
-                            line_chart = (
-                                alt.Chart(df)
-                                .mark_line(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=[
-                                        "Frame",
-                                        alt.Tooltip(f"{category}:Q", title="Diameter"),
-                                    ],
-                                )
-                                # .properties(title=f"{category} - Predicted Diameters")
-                                # .configure_axis(grid=True)
-                            )
-                            points_chart = line_chart.mark_point(color=colors[i], filled=True)
-
-                            final_chart = (
-                                line_chart.properties(title=f"{category} - Predicted Diameters") + points_chart
-                            ).interactive()
-
-                            final_chart = final_chart.configure_axis(grid=True)
-
-                            # Display the Altair chart
-                            st.altair_chart(final_chart, use_container_width=True)
-
-                    if eyes_ratios is not None and len(eyes_ratios) > 0:
-                        df = pd.DataFrame(eyes_ratios, columns=["EAR"])
-                        df["Frame"] = range(1, len(eyes_ratios) + 1)  # Create a frame column starting from 1
-
-                        # Create an Altair chart for eyes_ratios
-                        line_chart = (
-                            alt.Chart(df)
-                            .mark_line(color=colors[-1])  # Set color of the line
-                            .encode(
-                                x=alt.X("Frame:Q", title="Frame Number"),
-                                y=alt.Y("EAR:Q", title="Eyes Aspect Ratio"),
-                                tooltip=["Frame", "EAR"],
-                            )
-                            # .properties(title="Eyes Aspect Ratios (EARs)")
-                            # .configure_axis(grid=True)
-                        )
-                        points_chart = line_chart.mark_point(color=colors[-1], filled=True)
-
-                        # Create a horizontal rule at y=0.22
-                        line1 = alt.Chart(pd.DataFrame({"y": [0.22]})).mark_rule(color="red").encode(y="y:Q")
-
-                        line2 = alt.Chart(pd.DataFrame({"y": [0.25]})).mark_rule(color="green").encode(y="y:Q")
-
-                        # Add text annotations for the lines
-                        text1 = (
-                            alt.Chart(pd.DataFrame({"y": [0.22], "label": ["Definite Blinks (<=0.22)"]}))
-                            .mark_text(align="left", dx=100, dy=9, color="red", size=16)
-                            .encode(y="y:Q", text="label:N")
-                        )
-
-                        text2 = (
-                            alt.Chart(pd.DataFrame({"y": [0.25], "label": ["No Blinks (>=0.25)"]}))
-                            .mark_text(align="left", dx=-150, dy=-9, color="green", size=16)
-                            .encode(y="y:Q", text="label:N")
-                        )
-
-                        # Add gray area text for the region between red and green lines
-                        gray_area_text = (
-                            alt.Chart(pd.DataFrame({"y": [0.235], "label": ["Gray Area"]}))
-                            .mark_text(align="left", dx=0, dy=0, color="gray", size=16)
-                            .encode(y="y:Q", text="label:N")
-                        )
-
-                        # Combine all elements: line chart, points, rules, and text annotations
-                        final_chart = (
-                            line_chart.properties(title="Eyes Aspect Ratios (EARs)")
-                            + points_chart
-                            + line1
-                            + line2
-                            + text1
-                            + text2
-                            + gray_area_text
-                        ).interactive()
-
-                        # Configure axis properties at the chart level
-                        final_chart = final_chart.configure_axis(grid=True)
-
-                        # Display the Altair chart
-                        # st.subheader("Eyes Aspect Ratios (EARs)")
-                        st.altair_chart(final_chart, use_container_width=True)
 
 
 if __name__ == "__main__":
     main()
+# run: streamlit run app.py --server.enableXsrfProtection false

app_utils.py

@@ -1,16 +1,21 @@
 import base64
 from io import BytesIO
+import io
 import os
 import sys
 import cv2
 from matplotlib import pyplot as plt
 import numpy as np
+import pandas as pd
 import streamlit as st
 import torch
 import tempfile
 from PIL import Image
 from torchvision.transforms.functional import to_pil_image
 from torchvision import transforms
+from PIL import ImageOps
+import altair as alt
+
 
 from torchcam.methods import CAM
 from torchcam import methods as torchcam_methods
@@ -23,6 +28,10 @@ sys.path.append(root_path)
 from preprocessing.dataset_creation import EyeDentityDatasetCreation
 from utils import get_model
 
+CAM_METHODS = ["CAM"]
+# colors = ["#2ca02c", "#d62728", "#1f77b4", "#ff7f0e"]  # Green, Red, Blue, Orange
+colors = ["#1f77b4", "#ff7f0e", "#636363"]  # Blue, Orange, Gray
+
 
 @torch.no_grad()
 def load_model(model_configs, device="cpu"):
@@ -234,12 +243,12 @@ def process_frames(
     )
 
     preprocess_steps = [
-        transforms.ToTensor(),
         transforms.Resize(
             [32, 64],
             interpolation=transforms.InterpolationMode.BICUBIC,
             antialias=True,
         ),
+        transforms.ToTensor(),
     ]
     preprocess_function = transforms.Compose(preprocess_steps)
 
@@ -368,7 +377,11 @@ def process_frames(
 
                 combined_frame = np.vstack((input_img_np, output_img_np, frame))
 
-                video_placeholders[eye_type].image(combined_frame, use_column_width=True)
+                img_base64 = pil_image_to_base64(Image.fromarray(combined_frame))
+                image_html = f'<div style="width: {str(50*len(selected_eyes))}%;"><img src="data:image/png;base64,{img_base64}" style="width: 100%;"></div>'
+                video_placeholders[eye_type].markdown(image_html, unsafe_allow_html=True)
+
+                # video_placeholders[eye_type].image(combined_frame, use_column_width=True)
 
         st.session_state.current_frame = idx + 1
         txt = f"<p style='font-size:20px;'> Number of Frames Processed: <strong>{st.session_state.current_frame} / {st.session_state.total_frames}</strong> </p>"
@@ -383,9 +396,9 @@ def process_frames(
 
 
 # Function to display video with autoplay and loop
-def display_video_with_autoplay(video_col, video_path):
+def display_video_with_autoplay(video_col, video_path, width):
     video_html = f"""
-        <video width="100%" height="auto" autoplay loop muted>
+        <video width="{str(width)}%" height="auto" autoplay loop muted>
             <source src="data:video/mp4;base64,{video_path}" type="video/mp4">
         </video>
     """
@@ -458,7 +471,359 @@ def combine_and_show_frames(input_frames, cam_frames, pred_diameters_frames, out
             video_base64 = base64.b64encode(video_bytes).decode("utf-8")
 
         # Display the combined video
-        display_video_with_autoplay(video_cols[eye_type], video_base64)
+        display_video_with_autoplay(video_cols[eye_type], video_base64, width=len(video_cols) * 50)
 
         # Clean up
         os.remove(output_path)
+
+
+def set_input_image_on_ui(uploaded_file, cols):
+    input_img = Image.open(BytesIO(uploaded_file.read())).convert("RGB")
+    # NOTE: images taken with phone camera has an EXIF data field which often rotates images taken with the phone in a tilted position. PIL has a utility function that removes this data and ‘uprights’ the image.
+    input_img = ImageOps.exif_transpose(input_img)
+    input_img = resize_frame(input_img, max_width=640, max_height=480)
+    input_img = resize_frame(input_img, max_width=640, max_height=480)
+    cols[0].image(input_img, use_column_width=True)
+    st.session_state.total_frames = 1
+    return input_img
+
+
+def set_input_video_on_ui(uploaded_file, cols):
+    tfile = tempfile.NamedTemporaryFile(delete=False)
+    tfile.write(uploaded_file.read())
+    video_path = tfile.name
+    video_frames = extract_frames(video_path)
+    cols[0].video(video_path)
+    st.session_state.total_frames = len(video_frames)
+    return video_frames, video_path
+
+
+def set_frames_processed_count_placeholder(cols):
+    st.session_state.current_frame = 0
+    st.session_state.frame_placeholder = cols[0].empty()
+    txt = f"<p style='font-size:20px;'> Number of Frames Processed: <strong>{st.session_state.current_frame} / {st.session_state.total_frames}</strong> </p>"
+    st.session_state.frame_placeholder.markdown(txt, unsafe_allow_html=True)
+
+
+def set_page_info():
+    st.set_page_config(page_title="Pupil Diameter Estimator", layout="wide")
+    st.title("EyeDentify Playground")
+    cols = st.columns((1, 1))
+    cols[0].header("Input")
+    cols[-1].header("Prediction")
+    return cols
+
+
+def set_sidebar_info():
+    LABEL_MAP = ["left_pupil", "right_pupil"]
+    TV_MODELS = ["ResNet18", "ResNet50"]
+
+    st.sidebar.title("Upload Face or Eye")
+    uploaded_file = st.sidebar.file_uploader(
+        "Upload Image or Video", type=["png", "jpeg", "jpg", "mp4", "avi", "mov", "mkv", "webm"]
+    )
+    st.sidebar.title("Setup")
+    pupil_selection = st.sidebar.selectbox(
+        "Pupil Selection", ["both"] + LABEL_MAP, help="Select left or right pupil OR both for diameter estimation"
+    )
+    tv_model = st.sidebar.selectbox("Classification model", TV_MODELS, help="Supported Models")
+
+    blink_detection = st.sidebar.checkbox("Detect Blinks")
+
+    st.markdown("<style>#vg-tooltip-element{z-index: 1000051}</style>", unsafe_allow_html=True)
+
+    return (uploaded_file, pupil_selection, tv_model, blink_detection)
+
+
+def pil_image_to_base64(img):
+    """Convert a PIL Image to a base64 encoded string."""
+    buffered = io.BytesIO()
+    img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return img_str
+
+
+def process_image_and_vizualize_data(cols, input_img, tv_model, pupil_selection, blink_detection):
+    input_frames, output_frames, predicted_diameters, face_frames, eyes_ratios = process_frames(
+        cols,
+        [input_img],
+        tv_model,
+        pupil_selection,
+        cam_method=CAM_METHODS[-1],
+        blink_detection=blink_detection,
+    )
+    # for ff in face_frames:
+    #     if ff["has_face"]:
+    #         cols[1].image(face_frames[0]["img"], use_column_width=True)
+
+    input_frames_keys = input_frames.keys()
+    video_cols = cols[1].columns(len(input_frames_keys))
+
+    for i, eye_type in enumerate(input_frames_keys):
+        # Check the pupil_selection and set the width accordingly
+        if pupil_selection == "both":
+            video_cols[i].image(input_frames[eye_type][-1], use_column_width=True)
+        else:
+            img_base64 = pil_image_to_base64(Image.fromarray(input_frames[eye_type][-1]))
+            image_html = f'<div style="width: 50%; margin-bottom: 1.2%;"><img src="data:image/png;base64,{img_base64}" style="width: 100%;"></div>'
+            video_cols[i].markdown(image_html, unsafe_allow_html=True)
+
+    output_frames_keys = output_frames.keys()
+    fig, axs = plt.subplots(1, len(output_frames_keys), figsize=(10, 5))
+    for i, eye_type in enumerate(output_frames_keys):
+        height, width, c = output_frames[eye_type][0].shape
+        frame = np.zeros((height, width, c), dtype=np.uint8)
+        text = f"{predicted_diameters[eye_type][0]:.2f}"
+        frame = overlay_text_on_frame(frame, text)
+
+        if pupil_selection == "both":
+            video_cols[i].image(output_frames[eye_type][-1], use_column_width=True)
+            video_cols[i].image(frame, use_column_width=True)
+        else:
+            img_base64 = pil_image_to_base64(Image.fromarray(output_frames[eye_type][-1]))
+            image_html = f'<div style="width: 50%; margin-top: 1.2%; margin-bottom: 1.2%"><img src="data:image/png;base64,{img_base64}" style="width: 100%;"></div>'
+            video_cols[i].markdown(image_html, unsafe_allow_html=True)
+            img_base64 = pil_image_to_base64(Image.fromarray(frame))
+            image_html = f'<div style="width: 50%; margin-top: 1.2%"><img src="data:image/png;base64,{img_base64}" style="width: 100%;"></div>'
+            video_cols[i].markdown(image_html, unsafe_allow_html=True)
+
+    return None
+
+
+def plot_ears(eyes_ratios, eyes_df):
+    eyes_df["EAR"] = eyes_ratios
+    df = pd.DataFrame(eyes_ratios, columns=["EAR"])
+    df["Frame"] = range(1, len(eyes_ratios) + 1)  # Create a frame column starting from 1
+
+    # Create an Altair chart for eyes_ratios
+    line_chart = (
+        alt.Chart(df)
+        .mark_line(color=colors[-1])  # Set color of the line
+        .encode(
+            x=alt.X("Frame:Q", title="Frame Number"),
+            y=alt.Y("EAR:Q", title="Eyes Aspect Ratio"),
+            tooltip=["Frame", "EAR"],
+        )
+        # .properties(title="Eyes Aspect Ratios (EARs)")
+        # .configure_axis(grid=True)
+    )
+    points_chart = line_chart.mark_point(color=colors[-1], filled=True)
+
+    # Create a horizontal rule at y=0.22
+    line1 = alt.Chart(pd.DataFrame({"y": [0.22]})).mark_rule(color="red").encode(y="y:Q")
+
+    line2 = alt.Chart(pd.DataFrame({"y": [0.25]})).mark_rule(color="green").encode(y="y:Q")
+
+    # Add text annotations for the lines
+    text1 = (
+        alt.Chart(pd.DataFrame({"y": [0.22], "label": ["Definite Blinks (<=0.22)"]}))
+        .mark_text(align="left", dx=100, dy=9, color="red", size=16)
+        .encode(y="y:Q", text="label:N")
+    )
+
+    text2 = (
+        alt.Chart(pd.DataFrame({"y": [0.25], "label": ["No Blinks (>=0.25)"]}))
+        .mark_text(align="left", dx=-150, dy=-9, color="green", size=16)
+        .encode(y="y:Q", text="label:N")
+    )
+
+    # Add gray area text for the region between red and green lines
+    gray_area_text = (
+        alt.Chart(pd.DataFrame({"y": [0.235], "label": ["Gray Area"]}))
+        .mark_text(align="left", dx=0, dy=0, color="gray", size=16)
+        .encode(y="y:Q", text="label:N")
+    )
+
+    # Combine all elements: line chart, points, rules, and text annotations
+    final_chart = (
+        line_chart.properties(title="Eyes Aspect Ratios (EARs)")
+        + points_chart
+        + line1
+        + line2
+        + text1
+        + text2
+        + gray_area_text
+    ).interactive()
+
+    # Configure axis properties at the chart level
+    final_chart = final_chart.configure_axis(grid=True)
+
+    # Display the Altair chart
+    # st.subheader("Eyes Aspect Ratios (EARs)")
+    st.altair_chart(final_chart, use_container_width=True)
+    return eyes_df
+
+
+def plot_individual_charts(predicted_diameters, cols):
+    # 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
+            if "left" in category:
+                selected_color = colors[0]
+            elif "right" in category:
+                selected_color = colors[1]
+            else:
+                selected_color = colors[i]
+
+            # Convert values to numeric, replacing non-numeric values with None
+            values = [convert_diameter(value) for value in values]
+
+            if "left" in category:
+                category_name = "Left Pupil Diameter"
+            else:
+                category_name = "Right Pupil Diameter"
+
+            # Create a DataFrame from the values for Altair
+            df = pd.DataFrame(
+                {
+                    "Frame": range(1, len(values) + 1),
+                    category_name: values,
+                }
+            )
+
+            # 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
+            line_chart = (
+                alt.Chart(df)
+                .mark_line(color=selected_color)
+                .encode(
+                    x=alt.X("Frame:Q", title="Frame Number"),
+                    y=alt.Y(
+                        f"{category_name}:Q",
+                        title="Diameter",
+                        scale=alt.Scale(domain=[min_value, max_value]),
+                    ),
+                    tooltip=[
+                        "Frame",
+                        alt.Tooltip(f"{category_name}:Q", title="Diameter"),
+                    ],
+                )
+                # .properties(title=f"{category} - Predicted Diameters")
+                # .configure_axis(grid=True)
+            )
+            points_chart = line_chart.mark_point(color=selected_color, filled=True)
+
+            final_chart = (
+                line_chart.properties(
+                    title=f"{'Left Pupil' if 'left' in category else 'Right Pupil'} - Predicted Diameters"
+                )
+                + points_chart
+            ).interactive()
+
+            final_chart = final_chart.configure_axis(grid=True)
+
+            # Display the Altair chart
+            st.altair_chart(final_chart, use_container_width=True)
+    return df
+
+
+def plot_combined_charts(predicted_diameters):
+    all_min_values = []
+    all_max_values = []
+
+    # Create an empty DataFrame to store combined data for plotting
+    combined_df = pd.DataFrame()
+
+    # Iterate through categories and collect data
+    for category, values in predicted_diameters.items():
+        # Convert values to numeric, replacing non-numeric values with None
+        values = [convert_diameter(value) for value in values]
+
+        # 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)
+
+        all_min_values.append(min_value)
+        all_max_values.append(max_value)
+
+        category = "left_pupil" if "left" in category else "right_pupil"
+
+        # Create a DataFrame from the values
+        df = pd.DataFrame(
+            {
+                "Diameter": values,
+                "Frame": range(1, len(values) + 1),  # Create a frame column starting from 1
+                "Category": category,  # Add a column to specify the category
+            }
+        )
+
+        # Append to combined DataFrame
+        combined_df = pd.concat([combined_df, df], ignore_index=True)
+
+    combined_chart = (
+        alt.Chart(combined_df)
+        .mark_line()
+        .encode(
+            x=alt.X("Frame:Q", title="Frame Number"),
+            y=alt.Y(
+                "Diameter:Q",
+                title="Diameter",
+                scale=alt.Scale(domain=[min(all_min_values), max(all_max_values)]),
+            ),
+            color=alt.Color("Category:N", scale=alt.Scale(range=colors), title="Pupil Type"),
+            tooltip=["Frame", "Diameter:Q", "Category:N"],
+        )
+    )
+    points_chart = combined_chart.mark_point(filled=True)
+
+    final_chart = (combined_chart.properties(title="Predicted Diameters") + points_chart).interactive()
+
+    final_chart = final_chart.configure_axis(grid=True)
+
+    # Display the combined chart
+    st.altair_chart(final_chart, use_container_width=True)
+
+    # --------------------------------------------
+    # Convert to a DataFrame
+    left_pupil_values = [convert_diameter(value) for value in predicted_diameters["left_eye"]]
+    right_pupil_values = [convert_diameter(value) for value in predicted_diameters["right_eye"]]
+
+    df = pd.DataFrame(
+        {
+            "Frame": range(1, len(left_pupil_values) + 1),
+            "Left Pupil Diameter": left_pupil_values,
+            "Right Pupil Diameter": right_pupil_values,
+        }
+    )
+
+    # Calculate the difference between left and right pupil diameters
+    df["Difference Value"] = df["Left Pupil Diameter"] - df["Right Pupil Diameter"]
+
+    # Determine the status of the difference
+    df["Difference Status"] = df.apply(
+        lambda row: "L>R" if row["Left Pupil Diameter"] > row["Right Pupil Diameter"] else "L<R",
+        axis=1,
+    )
+
+    return df
+
+
+def process_video_and_visualize_data(cols, video_frames, tv_model, pupil_selection, blink_detection, video_path):
+    output_video_path = f"{root_path}/tmp.webm"
+    input_frames, output_frames, predicted_diameters, face_frames, eyes_ratios = 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)
+    cols = st.columns(num_columns)
+
+    if num_columns == 2:
+        df = plot_combined_charts(predicted_diameters)
+    else:
+        df = plot_individual_charts(predicted_diameters, cols)
+
+    if eyes_ratios is not None and len(eyes_ratios) > 0:
+        df = plot_ears(eyes_ratios, df)
+
+    st.dataframe(df, hide_index=True, use_container_width=True)

requirements.txt

@@ -20,7 +20,7 @@ dlib
 einops
 transformers
 gfpgan
-# streamlit
+streamlit==1.38.0
 mediapipe
 imutils
 scipy