Cleanup

app.py

@@ -21,7 +21,36 @@ 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.
+
+# Named tuple to store detection results
+class Detection(NamedTuple):
+    class_id: int
+    label: str
+    score: float
+    box: np.ndarray
+
+
+# Queue to store detection results
+result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
+
+# CHANGE CODE BELOW HERE, USE TO REPLACE WITH YOUR WANTED ANALYSIS.
+# Update below string to set display title of analysis
+
+# Appropriate imports needed for analysis
+
+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
+
+# Initialize the Hugging Face pipeline for facial emotion detection
+emotion_pipeline = pipeline("image-classification", model="trpakov/vit-face-expression")
+
+
+# Default title - "Facial Sentiment Analysis"
+
+ANALYSIS_TITLE = "Facial Sentiment Analysis"
+
+# CHANGE THE CONTENTS OF THIS FUNCTION, USE TO REPLACE WITH YOUR WANTED ANALYSIS.
 #
 #
 # Function to analyze an input frame and generate an analyzed frame
@@ -34,7 +63,12 @@ TEXT_SIZE = 1
 LINE_SIZE = 2
 
 
-def analyze_frame(frame):
+# Set analysis results in img_container and result queue for display
+# img_container["input"] - holds the input frame contents - of type np.ndarray
+# img_container["analyzed"] - holds the analyzed frame with any added annotations - of type np.ndarray
+# img_container["analysis_time"] - holds how long the analysis has taken in miliseconds
+# result_queue - holds the analysis metadata results - of type queue.Queue[List[Detection]]
+def analyze_frame(frame: np.ndarray):
     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
@@ -44,6 +78,7 @@ def analyze_frame(frame):
         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
+        result["label"] = sentiment
         # 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)[
@@ -92,6 +127,12 @@ def analyze_sentiment(face):
     return dominant_emotion  # Return the detected emotion
 
 
+#
+#
+# DO NOT TOUCH THE BELOW CODE (NOT NEEDED)
+#
+#
+
 # Suppress FFmpeg logs
 os.environ["FFMPEG_LOG_LEVEL"] = "quiet"
 
@@ -111,10 +152,6 @@ 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")
-
 # Container to hold image data and analysis results
 img_container = {"input": None, "analyzed": None, "analysis_time": None}
 
@@ -125,18 +162,6 @@ mtcnn = MTCNN()
 logger = logging.getLogger(__name__)
 
 
-# Named tuple to store detection results
-class Detection(NamedTuple):
-    class_id: int
-    label: str
-    score: float
-    box: np.ndarray
-
-
-# Queue to store detection results
-result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
-
-
 # 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,
@@ -182,7 +207,7 @@ st.markdown(
 
 # Streamlit page title and subtitle
 st.title("Computer Vision Playground")
-st.subheader("Facial Sentiment Analysis")
+st.subheader(ANALYSIS_TITLE)
 
 # Columns for input and output streams
 col1, col2 = st.columns(2)
@@ -215,7 +240,7 @@ with col1:
     )
 
     # Text input for video URL
-    st.subheader("Or Enter Video URL")
+    st.subheader("Or Enter Video Download URL")
     video_url = st.text_input("Video URL")