Update app.py

app.py

@@ -4,128 +4,133 @@ import numpy as np
 from PIL import Image
 import time
 import cvlib as cv
-from cvlib.object_detection import draw_bbox
 
-# Session state handling
-if "camera_image" not in st.session_state:
-    st.session_state.camera_image = None
+# Set page config
+st.set_page_config(page_title="Face Mask Detection", layout="wide")
+
+# Initialize session state
+if "image_captured" not in st.session_state:
+    st.session_state.image_captured = None
 if "camera_key" not in st.session_state:
     st.session_state.camera_key = str(time.time())
 
-# App setup
-st.set_page_config(page_title="Face Mask Detection", layout="wide")
+# Header
 st.markdown("<h1 style='text-align: center;'>😷 Face Mask Detection</h1>", unsafe_allow_html=True)
 input_method = st.selectbox("Choose Input Method", ["Camera Capture", "Upload Image"])
 
-# Face detection with dummy classifier
+# Dummy mask classifier (for now)
+def classify_face_dummy(face_img):
+    mean = np.mean(face_img)
+    if mean % 2 < 1:
+        return "Mask", 0.90
+    else:
+        return "No Mask", 0.60
+
 def detect_and_classify_faces(img_np):
     faces, confidences = cv.detect_face(img_np)
     results = []
-
-    for face in faces:
+    for i, face in enumerate(faces):
         (startX, startY) = face[0], face[1]
         (endX, endY) = face[2], face[3]
+        face_crop = img_np[startY:endY, startX:endX]
 
-        # Simulated mask classification (replace with real model)
-        face_img = img_np[startY:endY, startX:endX]
-        mean_pixel = np.mean(face_img)
-        if mean_pixel % 2 < 1:
-            label = "Mask"
-            confidence = 0.90
-        else:
-            label = "No Mask"
-            confidence = 0.65
+        if face_crop.size == 0:
+            continue
 
+        label, conf = classify_face_dummy(face_crop)
         results.append({
             "box": [startX, startY, endX - startX, endY - startY],
             "label": label,
-            "confidence": confidence
+            "confidence": conf
         })
-
     return results
 
-# Columns
+# Layout
 col1, col2 = st.columns(2)
 
-# --------------- CAMERA INPUT MODE ---------------
+# ------------------- CAMERA CAPTURE MODE -------------------
 if input_method == "Camera Capture":
     with col1:
         st.markdown("### 📸 Capturing Image")
-
         camera_img = st.camera_input("Take a photo", key=st.session_state.camera_key)
 
         if camera_img:
-            st.session_state.camera_image = camera_img
-            img = Image.open(st.session_state.camera_image)
-            st.image(img, caption="Captured Image")
+            st.session_state.image_captured = camera_img
 
+            # Clear Button
             if st.button("❌ Clear photo"):
-                st.session_state.camera_image = None
+                st.session_state.image_captured = None
                 st.session_state.camera_key = str(time.time())
                 st.experimental_rerun()
 
-    if st.session_state.camera_image:
+            image = Image.open(camera_img)
+            st.image(image, caption="Captured Image", use_container_width=True)
+
+    if st.session_state.image_captured:
         with col2:
             st.markdown("### 🧠 Detection Result")
+            image = Image.open(st.session_state.image_captured)
+            img_np = np.array(image)
 
-            img = Image.open(st.session_state.camera_image)
-            img_np = np.array(img)
-            faces = detect_and_classify_faces(img_np)
+            results = detect_and_classify_faces(img_np)
 
-            for face in faces:
-                x, y, w, h = face["box"]
-                label = face["label"]
-                conf = face["confidence"]
+            for res in results:
+                x, y, w, h = res["box"]
+                label = res["label"]
+                conf = res["confidence"]
                 color = (0, 255, 0) if label == "Mask" else (0, 0, 255)
 
                 cv2.rectangle(img_np, (x, y), (x + w, y + h), color, 2)
                 cv2.putText(img_np, f"{label} ({conf*100:.2f}%)", (x, y - 10),
                             cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
 
-            st.image(img_np, caption="Result Image", use_container_width=True)
+            st.image(img_np, caption="Result Image", channels="BGR", use_container_width=True)
 
-            # Summary
-            for i, face in enumerate(faces, 1):
+            for i, face in enumerate(results, 1):
                 label_color = "green" if face["label"] == "Mask" else "red"
-                st.markdown(f"**Face {i}:** <span style='color:{label_color}'>{face['label']}</span> ({face['confidence']*100:.2f}%)", unsafe_allow_html=True)
+                st.markdown(
+                    f"**Face {i}:** <span style='color:{label_color}'>{face['label']}</span> ({face['confidence']*100:.2f}%)",
+                    unsafe_allow_html=True)
 
-            if any(face["label"] == "No Mask" for face in faces):
+            if any(face["label"] == "No Mask" for face in results):
                 st.error("⚠️ One or more people not wearing a mask!")
             else:
                 st.success("✅ All faces have masks.")
 
-# --------------- UPLOAD IMAGE MODE ---------------
+# ------------------- UPLOAD IMAGE MODE -------------------
 elif input_method == "Upload Image":
     with col1:
-        st.markdown("### 📷 Captured Image")
-        uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-        if uploaded_file:
-            img = Image.open(uploaded_file)
-            st.image(img, caption="Uploaded Image")
+        st.markdown("### 📁 Upload Image")
+        uploaded_img = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+        if uploaded_img:
+            image = Image.open(uploaded_img)
+            st.image(image, caption="Uploaded Image", use_container_width=True)
 
-    if uploaded_file:
+    if uploaded_img:
         with col2:
             st.markdown("### 🧠 Detection Result")
-            img_np = np.array(Image.open(uploaded_file))
-            faces = detect_and_classify_faces(img_np)
+            img_np = np.array(Image.open(uploaded_img))
+            results = detect_and_classify_faces(img_np)
 
-            for face in faces:
-                x, y, w, h = face["box"]
-                label = face["label"]
-                conf = face["confidence"]
+            for res in results:
+                x, y, w, h = res["box"]
+                label = res["label"]
+                conf = res["confidence"]
                 color = (0, 255, 0) if label == "Mask" else (0, 0, 255)
 
                 cv2.rectangle(img_np, (x, y), (x + w, y + h), color, 2)
                 cv2.putText(img_np, f"{label} ({conf*100:.2f}%)", (x, y - 10),
                             cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
 
-            st.image(img_np, caption="Result Image", use_container_width=True)
+            st.image(img_np, caption="Result Image", channels="BGR", use_container_width=True)
 
-            for i, face in enumerate(faces, 1):
+            for i, face in enumerate(results, 1):
                 label_color = "green" if face["label"] == "Mask" else "red"
-                st.markdown(f"**Face {i}:** <span style='color:{label_color}'>{face['label']}</span> ({face['confidence']*100:.2f}%)", unsafe_allow_html=True)
+                st.markdown(
+                    f"**Face {i}:** <span style='color:{label_color}'>{face['label']}</span> ({face['confidence']*100:.2f}%)",
+                    unsafe_allow_html=True)
 
-            if any(face["label"] == "No Mask" for face in faces):
+            if any(face["label"] == "No Mask" for face in results):
                 st.error("⚠️ One or more people not wearing a mask!")
             else:
                 st.success("✅ All faces have masks.")