working webcam

app.py

@@ -1,8 +1,3 @@
-"""Object detection demo with MobileNet SSD.
-This model and code are based on
-https://github.com/robmarkcole/object-detection-app
-"""
-
 import logging
 import queue
 from pathlib import Path
@@ -17,41 +12,24 @@ from streamlit_webrtc import WebRtcMode, webrtc_streamer
 from utils.download import download_file 
 from utils.turn import get_ice_servers
 
-HERE = Path(__file__).parent
-ROOT = HERE.parent
+from mtcnn import MTCNN
+from PIL import Image, ImageDraw
+from transformers import pipeline
 
-logger = logging.getLogger(__name__)
 
+# Initialize the Hugging Face pipeline for facial emotion detection
+emotion_pipeline = pipeline("image-classification", model="trpakov/vit-face-expression")
 
-MODEL_URL = "https://github.com/robmarkcole/object-detection-app/raw/master/model/MobileNetSSD_deploy.caffemodel"  # noqa: E501
-MODEL_LOCAL_PATH = ROOT / "./models/MobileNetSSD_deploy.caffemodel"
-PROTOTXT_URL = "https://github.com/robmarkcole/object-detection-app/raw/master/model/MobileNetSSD_deploy.prototxt.txt"  # noqa: E501
-PROTOTXT_LOCAL_PATH = ROOT / "./models/MobileNetSSD_deploy.prototxt.txt"
-
-CLASSES = [
-    "background",
-    "aeroplane",
-    "bicycle",
-    "bird",
-    "boat",
-    "bottle",
-    "bus",
-    "car",
-    "cat",
-    "chair",
-    "cow",
-    "diningtable",
-    "dog",
-    "horse",
-    "motorbike",
-    "person",
-    "pottedplant",
-    "sheep",
-    "sofa",
-    "train",
-    "tvmonitor",
-]
+img_container = {"webcam": None,
+                 "analyzed": None}
+
+# Initialize MTCNN for face detection
+mtcnn = MTCNN()
+
+HERE = Path(__file__).parent
+ROOT = HERE.parent
 
+logger = logging.getLogger(__name__)
 
 class Detection(NamedTuple):
     class_id: int
@@ -59,94 +37,128 @@ class Detection(NamedTuple):
     score: float
     box: np.ndarray
 
-
-@st.cache_resource  # type: ignore
-def generate_label_colors():
-    return np.random.uniform(0, 255, size=(len(CLASSES), 3))
-
-
-COLORS = generate_label_colors()
-
-download_file(MODEL_URL, MODEL_LOCAL_PATH, expected_size=23147564)
-download_file(PROTOTXT_URL, PROTOTXT_LOCAL_PATH, expected_size=29353)
-
-
-# Session-specific caching
-cache_key = "object_detection_dnn"
-if cache_key in st.session_state:
-    net = st.session_state[cache_key]
-else:
-    net = cv2.dnn.readNetFromCaffe(str(PROTOTXT_LOCAL_PATH), str(MODEL_LOCAL_PATH))
-    st.session_state[cache_key] = net
-
-score_threshold = st.slider("Score threshold", 0.0, 1.0, 0.5, 0.05)
-
 # NOTE: The callback will be called in another thread,
 #       so use a queue here for thread-safety to pass the data
 #       from inside to outside the callback.
 # TODO: A general-purpose shared state object may be more useful.
 result_queue: "queue.Queue[List[Detection]]" = queue.Queue()
 
+# Function to analyze sentiment
+def analyze_sentiment(face):
+    # Convert face to RGB
+    rgb_face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
+    # Convert the face to a PIL image
+    pil_image = Image.fromarray(rgb_face)
+    # Analyze sentiment using the Hugging Face pipeline
+    results = emotion_pipeline(pil_image)
+    # Get the dominant emotion
+    dominant_emotion = max(results, key=lambda x: x['score'])['label']
+    return dominant_emotion
+
+TEXT_SIZE = 1
+LINE_SIZE = 2
+
+# Function to detect faces, analyze sentiment, and draw a red box around them
+def detect_and_draw_faces(frame):
+    # Detect faces using MTCNN
+    results = mtcnn.detect_faces(frame)
+    
+    # Draw on the 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)  # Thicker red box
+        
+        # Calculate position for the text background and the text itself
+        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 black rectangle as background
+        cv2.rectangle(frame, background_tl, background_br, (0, 0, 0), cv2.FILLED)
+        # Draw white text on top
+        cv2.putText(frame, sentiment, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, TEXT_SIZE, (255, 255, 255), 2)
+    
+    result_queue.put(results)
+    return frame
 
 def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
-    image = frame.to_ndarray(format="bgr24")
+    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
 
-    # Run inference
-    blob = cv2.dnn.blobFromImage(
-        cv2.resize(image, (300, 300)), 0.007843, (300, 300), 127.5
-    )
-    net.setInput(blob)
-    output = net.forward()
-
-    h, w = image.shape[:2]
-
-    # Convert the output array into a structured form.
-    output = output.squeeze()  # (1, 1, N, 7) -> (N, 7)
-    output = output[output[:, 2] >= score_threshold]
-    detections = [
-        Detection(
-            class_id=int(detection[1]),
-            label=CLASSES[int(detection[1])],
-            score=float(detection[2]),
-            box=(detection[3:7] * np.array([w, h, w, h])),
-        )
-        for detection in output
-    ]
-
-    # Render bounding boxes and captions
-    for detection in detections:
-        caption = f"{detection.label}: {round(detection.score * 100, 2)}%"
-        color = COLORS[detection.class_id]
-        xmin, ymin, xmax, ymax = detection.box.astype("int")
-
-        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
-        cv2.putText(
-            image,
-            caption,
-            (xmin, ymin - 15 if ymin - 15 > 15 else ymin + 15),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.5,
-            color,
-            2,
-        )
-
-    result_queue.put(detections)
-
-    return av.VideoFrame.from_ndarray(image, format="bgr24")
+    return frame
+#    return av.VideoFrame.from_ndarray(frame_with_boxes, format="bgr24")
 
 ice_servers = get_ice_servers()
 
-webrtc_ctx = webrtc_streamer(
-    key="object-detection",
-    mode=WebRtcMode.SENDRECV,
-    rtc_configuration=ice_servers,
-    video_frame_callback=video_frame_callback,
-    media_stream_constraints={"video": True, "audio": False},
-    async_processing=True,
+# Streamlit UI
+st.markdown(
+    """
+    <style>
+        .main {
+            background-color: #F7F7F7;
+            padding: 2rem;
+        }
+        h1, h2, h3 {
+            color: #333333;
+            font-family: 'Arial', sans-serif;
+        }
+        h1 {
+            font-weight: 700;
+            font-size: 2.5rem;
+        }
+        h2 {
+            font-weight: 600;
+            font-size: 2rem;
+        }
+        h3 {
+            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
 )
 
-if st.checkbox("Show the detected labels", value=True):
-    if webrtc_ctx.state.playing:
+st.title("Computer Vision Test Lab")
+st.subheader("Facial Sentiment Analysis")
+
+# Columns for input and output streams
+col1, col2 = st.columns(2)
+
+with col1:
+    st.header("Input Stream")
+    st.subheader("Webcam")
+    webrtc_ctx = webrtc_streamer(
+        key="object-detection",
+        mode=WebRtcMode.SENDRECV,
+        rtc_configuration=ice_servers,
+        video_frame_callback=video_frame_callback,
+        media_stream_constraints={"video": True, "audio": False},
+        async_processing=True,
+    )
+
+with col2:
+    st.header("Analysis")
+    st.subheader("Input Frame")
+    input_placeholder = st.empty()
+    st.subheader("Output Frame")
+    output_placeholder = st.empty()
+
+if webrtc_ctx.state.playing:
+    if st.checkbox("Show the detected labels", value=True):
         labels_placeholder = st.empty()
         # NOTE: The video transformation with object detection and
         # this loop displaying the result labels are running
@@ -157,8 +169,12 @@ if st.checkbox("Show the detected labels", value=True):
             result = result_queue.get()
             labels_placeholder.table(result)
 
-st.markdown(
-    "This demo uses a model and code from "
-    "https://github.com/robmarkcole/object-detection-app. "
-    "Many thanks to the project."
-)
+            img = img_container["webcam"]
+            frame_with_boxes = img_container["analyzed"]
+
+            if img is None:
+                continue
+
+            input_placeholder.image(img, channels="BGR")
+            output_placeholder.image(frame_with_boxes, channels="BGR")
+

object_detection.py

@@ -134,11 +134,12 @@ def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
 
     return av.VideoFrame.from_ndarray(image, format="bgr24")
 
+ice_servers = get_ice_servers()
 
 webrtc_ctx = webrtc_streamer(
     key="object-detection",
     mode=WebRtcMode.SENDRECV,
-    rtc_configuration={"iceServers": get_ice_servers()},
+    rtc_configuration=ice_servers,
     video_frame_callback=video_frame_callback,
     media_stream_constraints={"video": True, "audio": False},
     async_processing=True,

sentiment.py

@@ -14,7 +14,7 @@ logging.getLogger("transformers").setLevel(logging.ERROR)
 
 lock = threading.Lock()
 img_container = {"webcam": None,
-                 "analzyed": None}
+                 "analyzed": None}
 
 # Initialize the Hugging Face pipeline for facial emotion detection
 emotion_pipeline = pipeline("image-classification", model="trpakov/vit-face-expression")