# app_streamlit_final.py
import streamlit as st
import cv2
from ultralytics import YOLO
import numpy as np

# Page configuration
st.set_page_config(page_title="Live Mask Detector", layout="wide")
st.title("🎭 Real-Time Face Mask Detection")

# 1. Model load karein
model = YOLO("best.pt")

# 2. UI Layout
col1, col2 = st.columns([2, 1])
with col1:
    # Frame placeholder jahan video nazar aayegi
    frame_placeholder = st.empty()

with col2:
    st.info("System Status: Running")
    stop_button = st.button("Stop Camera")
    conf_slider = st.slider("Confidence Threshold", 0.1, 1.0, 0.4)

# 3. Camera Loop
cap = cv2.VideoCapture(0)

while cap.isOpened() and not stop_button:
    ret, frame = cap.read()
    if not ret:
        st.error("Camera access nahi mil raha. Please check permissions.")
        break

    # Prediction (Wahi logic jo mask_test.py mein tha)
    results = model.predict(frame, conf=conf_slider)
    
    # Annotated frame
    annotated_frame = results[0].plot()

    # BGR to RGB (Streamlit requirements)
    frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)

    # Browser mein display karein
    frame_placeholder.image(frame_rgb, channels="RGB", use_container_width=True)

    if stop_button:
        break

cap.release()
st.warning("Camera Stopped.")