commented app.py

app.py

@@ -1,3 +1,5 @@
+import time
+import os
 import logging
 import queue
 from pathlib import Path
@@ -9,82 +11,156 @@ import numpy as np
 import streamlit as st
 from streamlit_webrtc import WebRtcMode, webrtc_streamer
 
-from utils.download import download_file 
+from utils.download import download_file
 from utils.turn import get_ice_servers
 
-from mtcnn import MTCNN
-from PIL import Image, ImageDraw
-from transformers import pipeline
+from mtcnn import MTCNN  # Import MTCNN for face detection
+from PIL import Image, ImageDraw  # Import PIL for image processing
+from transformers import pipeline  # Import Hugging Face transformers pipeline
+
+import requests
+from io import BytesIO  # Import for handling byte streams
+
+# CHANGE THIS FUNCTION, USE TO REPLACE WITH YOUR WANTED ANALYSIS.
+#
+#
+# Function to analyze an input frame and generate an analyzed frame
+# This function takes an input video frame, detects faces in it using MTCNN,
+# then for each detected face, it analyzes the sentiment (emotion) using the analyze_sentiment function,
+# draws a rectangle around the face, and overlays the detected emotion on the frame.
+# It also records the time taken to process the frame and stores it in a global container.
+# Constants for text and line size in the output image
+TEXT_SIZE = 1
+LINE_SIZE = 2
+
+
+def analyze_frame(frame):
+    start_time = time.time()  # Start timing the analysis
+    img_container["input"] = frame  # Store the input frame
+    frame = frame.copy()  # Create a copy of the frame to modify
+
+    results = mtcnn.detect_faces(frame)  # Detect faces in the frame
+    for result in results:
+        x, y, w, h = result["box"]  # Get the bounding box of the detected face
+        face = frame[y : y + h, x : x + w]  # Extract the face from the frame
+        sentiment = analyze_sentiment(face)  # Analyze the sentiment of the face
+        # Draw a rectangle around the face
+        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), LINE_SIZE)
+        text_size = cv2.getTextSize(sentiment, cv2.FONT_HERSHEY_SIMPLEX, TEXT_SIZE, 2)[
+            0
+        ]
+        text_x = x
+        text_y = y - 10
+        background_tl = (text_x, text_y - text_size[1])
+        background_br = (text_x + text_size[0], text_y + 5)
+        # Draw a black background for the text
+        cv2.rectangle(frame, background_tl, background_br, (0, 0, 0), cv2.FILLED)
+        # Put the sentiment text on the image
+        cv2.putText(
+            frame,
+            sentiment,
+            (text_x, text_y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            TEXT_SIZE,
+            (255, 255, 255),
+            2,
+        )
+
+    end_time = time.time()  # End timing the analysis
+    execution_time_ms = round(
+        (end_time - start_time) * 1000, 2
+    )  # Calculate execution time in milliseconds
+    img_container["analysis_time"] = execution_time_ms  # Store the execution time
+
+    result_queue.put(results)  # Put the results in the result queue
+    img_container["analyzed"] = frame  # Store the analyzed frame
+
+    return  # End of the function
+
+
+# Function to analyze the sentiment (emotion) of a detected face
+# This function converts the face from BGR to RGB format, then converts it to a PIL image,
+# uses a pre-trained emotion detection model to get emotion predictions,
+# and finally returns the most dominant emotion detected.
+def analyze_sentiment(face):
+    rgb_face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)  # Convert face to RGB format
+    pil_image = Image.fromarray(rgb_face)  # Convert to PIL image
+    results = emotion_pipeline(pil_image)  # Run emotion detection on the image
+    dominant_emotion = max(results, key=lambda x: x["score"])[
+        "label"
+    ]  # Get the dominant emotion
+    return dominant_emotion  # Return the detected emotion
+
+
+# Suppress FFmpeg logs
+os.environ["FFMPEG_LOG_LEVEL"] = "quiet"
+
+# Suppress TensorFlow or PyTorch progress bars
+import tensorflow as tf
+
+tf.get_logger().setLevel("ERROR")
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+# Suppress PyTorch logs
+import torch
+
+logging.getLogger().setLevel(logging.WARNING)
+torch.set_num_threads(1)
+logging.getLogger("torch").setLevel(logging.ERROR)
+
+# Suppress Streamlit logs using the logging module
+logging.getLogger("streamlit").setLevel(logging.ERROR)
+
 
 # Initialize the Hugging Face pipeline for facial emotion detection
 emotion_pipeline = pipeline("image-classification", model="trpakov/vit-face-expression")
 
-img_container = {"webcam": None, "analyzed": None, "uploaded": None}
+# Container to hold image data and analysis results
+img_container = {"input": None, "analyzed": None, "analysis_time": None}
 
 # Initialize MTCNN for face detection
 mtcnn = MTCNN()
 
-HERE = Path(__file__).parent
-ROOT = HERE.parent
-
+# Logger for debugging and information
 logger = logging.getLogger(__name__)
 
+
+# Named tuple to store detection results
 class Detection(NamedTuple):
     class_id: int
     label: str
     score: float
     box: np.ndarray
 
-result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
-
-# Function to analyze sentiment
-def analyze_sentiment(face):
-    rgb_face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
-    pil_image = Image.fromarray(rgb_face)
-    results = emotion_pipeline(pil_image)
-    dominant_emotion = max(results, key=lambda x: x['score'])['label']
-    return dominant_emotion
 
-TEXT_SIZE = 1
-LINE_SIZE = 2
+# Queue to store detection results
+result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
 
-# Function to detect faces, analyze sentiment, and draw a red box around them
-def detect_and_draw_faces(frame):
-    results = mtcnn.detect_faces(frame)
-    for result in results:
-        x, y, w, h = result['box']
-        face = frame[y:y+h, x:x+w]
-        sentiment = analyze_sentiment(face)
-        cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 0, 255), LINE_SIZE)
-        text_size = cv2.getTextSize(sentiment, cv2.FONT_HERSHEY_SIMPLEX, TEXT_SIZE, 2)[0]
-        text_x = x
-        text_y = y - 10
-        background_tl = (text_x, text_y - text_size[1])
-        background_br = (text_x + text_size[0], text_y + 5)
-        cv2.rectangle(frame, background_tl, background_br, (0, 0, 0), cv2.FILLED)
-        cv2.putText(frame, sentiment, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, TEXT_SIZE, (255, 255, 255), 2)
-    result_queue.put(results)
-    return frame
 
+# Callback function to process video frames
+# This function is called for each video frame in the WebRTC stream.
+# It converts the frame to a numpy array in RGB format, analyzes the frame,
+# and returns the original frame.
 def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
-    img = frame.to_ndarray(format="bgr24")
-    img_container["webcam"] = img
-    frame_with_boxes = detect_and_draw_faces(img.copy())
-    img_container["analyzed"] = frame_with_boxes
-    return frame
+    img = frame.to_ndarray(format="rgb24")  # Convert frame to numpy array in RGB format
+    analyze_frame(img)  # Analyze the frame
+    return frame  # Return the original frame
 
+
+# Get ICE servers for WebRTC
 ice_servers = get_ice_servers()
 
-# Streamlit UI
+# Streamlit UI configuration
+st.set_page_config(layout="wide")
+
+# Custom CSS for the Streamlit page
 st.markdown(
     """
     <style>
         .main {
-            background-color: #F7F7F7;
             padding: 2rem;
         }
         h1, h2, h3 {
-            color: #333333;
             font-family: 'Arial', sans-serif;
         }
         h1 {
@@ -99,31 +175,24 @@ st.markdown(
             font-weight: 500;
             font-size: 1.5rem;
         }
-        .stButton button {
-            background-color: #E60012;
-            color: white;
-            border-radius: 5px;
-            font-size: 16px;
-            padding: 0.5rem 1rem;
-        }
     </style>
     """,
-    unsafe_allow_html=True
+    unsafe_allow_html=True,
 )
 
-st.title("Computer Vision Test Lab")
+# Streamlit page title and subtitle
+st.title("Computer Vision Playground")
 st.subheader("Facial Sentiment Analysis")
 
-show_labels = st.checkbox("Show the detected labels", value=True)
-
 # Columns for input and output streams
 col1, col2 = st.columns(2)
 
 with col1:
     st.header("Input Stream")
-    st.subheader("Webcam")
+    st.subheader("input")
+    # WebRTC streamer to get video input from the webcam
     webrtc_ctx = webrtc_streamer(
-        key="object-detection",
+        key="input-webcam",
         mode=WebRtcMode.SENDRECV,
         rtc_configuration=ice_servers,
         video_frame_callback=video_frame_callback,
@@ -131,50 +200,131 @@ with col1:
         async_processing=True,
     )
 
+    # File uploader for images
     st.subheader("Upload an Image")
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-with col2:
-    st.header("Analysis")
-    input_subheader_placeholder = st.empty()
-    input_placeholder = st.empty()
-    
-    output_subheader_placeholder = st.empty()
-    output_placeholder = st.empty()
+    # Text input for image URL
+    st.subheader("Or Enter Image URL")
+    image_url = st.text_input("Image URL")
 
-if webrtc_ctx.state.playing:
-    labels_placeholder = st.empty()
-    input_subheader_placeholder.subheader("Input Frame")
-    output_subheader_placeholder.subheader("Output Frame")
+    # File uploader for videos
+    st.subheader("Upload a Video")
+    uploaded_video = st.file_uploader(
+        "Choose a video...", type=["mp4", "avi", "mov", "mkv"]
+    )
 
-    while True:
+    # Text input for video URL
+    st.subheader("Or Enter Video URL")
+    video_url = st.text_input("Video URL")
+
+
+# Function to initialize the analysis UI
+# This function sets up the placeholders and UI elements in the analysis section.
+# It creates placeholders for input and output frames, analysis time, and detected labels.
+def analysis_init():
+    global analysis_time, show_labels, labels_placeholder, input_placeholder, output_placeholder
+
+    with col2:
+        st.header("Analysis")
+        st.subheader("Input Frame")
+        input_placeholder = st.empty()  # Placeholder for input frame
+
+        st.subheader("Output Frame")
+        output_placeholder = st.empty()  # Placeholder for output frame
+        analysis_time = st.empty()  # Placeholder for analysis time
+        show_labels = st.checkbox(
+            "Show the detected labels", value=True
+        )  # Checkbox to show/hide labels
+        labels_placeholder = st.empty()  # Placeholder for labels
+
+
+# Function to publish frames and results to the Streamlit UI
+# This function retrieves the latest frames and results from the global container and result queue,
+# and updates the placeholders in the Streamlit UI with the current input frame, analyzed frame, analysis time, and detected labels.
+def publish_frame():
+    if not result_queue.empty():
         result = result_queue.get()
         if show_labels:
-            labels_placeholder.table(result)
+            labels_placeholder.table(
+                result
+            )  # Display labels if the checkbox is checked
+
+    img = img_container["input"]
+    if img is None:
+        return
+    input_placeholder.image(img, channels="RGB")  # Display the input frame
+
+    analyzed = img_container["analyzed"]
+    if analyzed is None:
+        return
+    output_placeholder.image(analyzed, channels="RGB")  # Display the analyzed frame
+
+    time = img_container["analysis_time"]
+    if time is None:
+        return
+    analysis_time.text(f"Analysis Time: {time} ms")  # Display the analysis time
+
+
+# If the WebRTC streamer is playing, initialize and publish frames
+if webrtc_ctx.state.playing:
+    analysis_init()  # Initialize the analysis UI
+    while True:
+        publish_frame()  # Publish the frames and results
+        time.sleep(0.1)  # Delay to control frame rate
+
+
+# If an image is uploaded or a URL is provided, process the image
+if uploaded_file is not None or image_url:
+    analysis_init()  # Initialize the analysis UI
+
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)  # Open the uploaded image
+        img = np.array(image.convert("RGB"))  # Convert the image to RGB format
+    else:
+        response = requests.get(image_url)  # Download the image from the URL
+        image = Image.open(BytesIO(response.content))  # Open the downloaded image
+        img = np.array(image.convert("RGB"))  # Convert the image to RGB format
+
+    analyze_frame(img)  # Analyze the image
+    publish_frame()  # Publish the results
+
+
+# Function to process video files
+# This function reads frames from a video file, analyzes each frame for face detection and sentiment analysis,
+# and updates the Streamlit UI with the current input frame, analyzed frame, and detected labels.
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)  # Open the video file
+    while cap.isOpened():
+        ret, frame = cap.read()  # Read a frame from the video
+        if not ret:
+            break  # Exit the loop if no more frames are available
 
-        img = img_container["webcam"]
-        frame_with_boxes = img_container["analyzed"]
+        input_placeholder.image(frame)  # Display the current frame as the input frame
+        analyze_frame(
+            frame
+        )  # Analyze the frame for face detection and sentiment analysis
+        publish_frame()  # Publish the results
 
-        if img is None:
-            continue
+        if not result_queue.empty():
+            result = result_queue.get()
+            if show_labels:
+                labels_placeholder.table(
+                    result
+                )  # Display labels if the checkbox is checked
 
-        input_placeholder.image(img, channels="BGR")
-        output_placeholder.image(frame_with_boxes, channels="BGR")
+    cap.release()  # Release the video capture object
 
-if uploaded_file is not None:
-    input_subheader_placeholder.subheader("Input Frame")
-    output_subheader_placeholder.subheader("Output Frame")
 
-    image = Image.open(uploaded_file)
-    img = np.array(image.convert("RGB"))  # Ensure image is in RGB format
-    img_container["uploaded"] = img
-    analyzed_img = detect_and_draw_faces(img.copy())
-    input_placeholder.image(img)
-    output_placeholder.image(analyzed_img)
+# If a video is uploaded or a URL is provided, process the video
+if uploaded_video is not None or video_url:
+    analysis_init()  # Initialize the analysis UI
 
-    result = result_queue.get()
-    if show_labels:
-        labels_placeholder = st.empty()
-        labels_placeholder.table(result)
+    if uploaded_video is not None:
+        video_path = uploaded_video.name  # Get the name of the uploaded video
+        with open(video_path, "wb") as f:
+            f.write(uploaded_video.getbuffer())  # Save the uploaded video to a file
+    else:
+        video_path = download_file(video_url)  # Download the video from the URL
 
-        
+    process_video(video_path)  # Process the video