app.py

import streamlit as st
import av # streamlit-webrtc가 비디오 프레임을 다루기 위해 사용
import cv2
import numpy as np
from ultralytics import YOLO

from streamlit_webrtc import webrtc_streamer, VideoTransformerBase, RTCConfiguration, WebRtcMode

# --- 설정 ---
MODEL_PATH = 'trained_model.pt' # 학습한 YOLO 모델 파일 경로
CONFIDENCE_THRESHOLD = 0.4 # 객체 탐지 최소 신뢰도 (모델에 따라 조정)
SEND_ALERT_INTERVAL = 30 # 담배 탐지 시 데이터 채널 메시지를 몇 프레임마다 보낼지 (너무 짧으면 클라이언트 부담)

# --- YOLO 모델 로드 ---
@st.cache_resource # Streamlit의 캐싱 기능 사용: 앱 실행 중 모델을 한 번만 로드
def load_yolo_model(model_path):
    try:
        model = YOLO(model_path)
        if hasattr(model, 'model') and model.model is not None:
             st.success(f"YOLO 모델 로드 성공: {model_path}")
             return model
        else:
             st.error(f"YOLO 모델 로드 실패 또는 객체 초기화 문제: {model_path}")
             st.stop()
    except FileNotFoundError:
        st.error(f"오류: 모델 파일이 없습니다. '{model_path}' 경로를 확인하세요.")
        st.stop() # 모델 파일 없으면 앱 중지
    except Exception as e:
        st.error(f"YOLO 모델 로드 중 오류 발생: {e}")
        st.stop() # 모델 로드 실패 시 앱 중지

model = load_yolo_model(MODEL_PATH)

# --- Streamlit-WebRTC를 위한 비디오 변환 클래스 ---
# 이 클래스의 인스턴스가 비디오 프레임마다 호출됩니다.
class YOLOVideoTransformer(VideoTransformerBase):
    # __init__에 data_channel 인자를 추가하여 클라이언트와 통신
    def __init__(self, model, confidence_thresh, send_interval, data_channel):
        self.model = model
        self.confidence_thresh = confidence_thresh
        self.send_interval = send_interval
        self._data_channel = data_channel # 클라이언트와 통신할 데이터 채널 객체

        self.detected_in_prev_frame = False # 이전 프레임에서 탐지되었는지 여부
        self.frame_counter = 0 # 프레임 카운터

    # 각 비디오 프레임을 처리하는 메서드 (비동기 함수로 정의)
    async def recv(self, frame: av.VideoFrame) -> av.VideoFrame:
        self.frame_counter += 1

        # AV 프레임을 OpenCV(numpy) 이미지로 변환
        img = frame.to_ndarray(format="bgr24")

        # YOLOv8 모델로 객체 탐지
        # verbose=False: 콘솔에 탐지 결과 출력 안 함
        results = self.model(img, conf=self.confidence_thresh, verbose=False)

        cigarette_detected_in_current_frame = False
        # 결과에서 'cigarette' 객체가 탐지되었는지 확인
        if results and len(results) > 0:
             for box in results[0].boxes:
                 class_id = int(box.cls[0])
                 confidence = float(box.conf[0])
                 class_name = self.model.names[class_id]

                 if class_name == 'cigarette' and confidence >= self.confidence_thresh:
                     cigarette_detected_in_current_frame = True
                     break # 하나라도 탐지되면 더 이상 확인할 필요 없음

        # --- 클라이언트 소리 알림 로직 (데이터 채널 사용) ---
        # 담배가 현재 프레임에서 탐지되었고, 데이터 채널이 열려 있으며,
        # 메시지 전송 간격에 도달했을 때 메시지 전송
        if cigarette_detected_in_current_frame and self._data_channel and self._data_channel.readyState == "open":
             if not self.detected_in_prev_frame or self.frame_counter % self.send_interval == 0:
                 # print("Sending DETECT_CIGARETTE message to client...") # 디버그 출력
                 await self._data_channel.send("DETECT_CIGARETTE") # 클라이언트로 메시지 전송


        # 다음 프레임을 위해 현재 탐지 상태 저장
        self.detected_in_prev_frame = cigarette_detected_in_current_frame

        # 탐지된 결과를 이미지에 표시 (바운딩 박스, 라벨 등)
        annotated_img = results[0].plot()

        # 처리된 이미지(numpy)를 다시 AV 프레임으로 변환하여 반환
        return av.VideoFrame.from_ndarray(annotated_img, format="bgr24")

# --- 클라이언트 측 JavaScript 코드 (데이터 채널 메시지 수신 및 소리 재생) ---
# webrtc_streamer의 on_data_channel 인자에 전달될 JavaScript 함수 정의
# 이 함수는 data channel 객체를 인자로 받습니다.
# 메시지를 받으면 웹 오디오 API로 사인파를 생성하여 재생합니다.
JS_CLIENT_SOUND_SCRIPT = """
(channel) => {
    // 오디오 컨텍스트 생성 (클릭 등 사용자 상호작용 후에 생성해야 할 수 있음)
    // webrtc_streamer 시작 버튼이 이미 상호작용 역할을 합니다.
    const audioContext = new (window.AudioContext || window.webkitAudioContext)();
    let lastPlayTime = 0; // 마지막 소리 재생 시간 (ms)
    const playCooldown = 200; // 소리 재생 최소 간격 (ms)

    // 사인파 소리를 재생하는 함수
    const playSineWaveAlert = () => {
        const now = audioContext.currentTime * 1000; // 현재 시간을 밀리초로 변환
        if (now - lastPlayTime < playCooldown) {
            // console.log("Cooldown active. Skipping sound."); // 디버그 출력
            return; // 쿨다운 중이면 재생하지 않음
        }
        lastPlayTime = now; // 마지막 재생 시간 업데이트

        try {
            const oscillator = audioContext.createOscillator();
            const gainNode = audioContext.createGain();

            oscillator.type = 'sine'; // 사인파
            oscillator.frequency.setValueAtTime(600, audioContext.currentTime); // 주파수 (예: 600 Hz)
            gainNode.gain.setValueAtTime(0.3, audioContext.currentTime); // 볼륨 (0.0 ~ 1.0)

            oscillator.connect(gainNode);
            gainNode.connect(audioContext.destination);

            oscillator.start();
            oscillator.stop(audioContext.currentTime + 0.2); // 0.2초 재생

            // console.log("Playing sine wave sound."); // 디버그 출력

        } catch (e) {
            console.error("Error playing sine wave:", e);
        }
    };

    // 데이터 채널로부터 메시지를 수신했을 때 실행될 콜백 함수
    channel.onmessage = (event) => {
        // console.log("Received message:", event.data); // 수신 메시지 확인

        if (event.data === "DETECT_CIGARETTE") {
            // 서버에서 담배 탐지 메시지를 받으면 소리 재생
            playSineWaveAlert();
        }
    };

    // 데이터 채널이 열렸을 때
    channel.onopen = () => {
        console.log("Data channel opened!");
    };

    // 데이터 채널이 닫혔을 때
    channel.onclose = () => {
        console.log("Data channel closed.");
    };

    // 데이터 채널 에러 발생 시
    channel.onerror = (error) => {
        console.error("Data channel error:", error);
    };
}
"""


# --- Streamlit 앱 레이아웃 구성 ---
st.title("🚬 실시간 담배 탐지 웹 애플리케이션 (클라이언트 소리)")

st.write("""
웹캠 피드를 통해 담배 객체를 실시간으로 탐지하고 영상에 표시합니다.
담배가 탐지되면 **사용자의 브라우저**에서 알림 소리(사인파)가 재생됩니다.

**주의:**
*   이 앱은 사용자의 브라우저 웹캠 및 오디오 재생 권한이 필요합니다. 브라우저 요청 시 허용해주세요.
*   네트워크 상태 및 컴퓨터 성능에 따라 영상 처리에 지연이 발생할 수 있습니다.
*   `trained_model.pt` 파일이 스크립트 파일과 같은 디렉토리에 있는지 확인하세요.
""")

st.write("---")

st.subheader("웹캠 스트림 및 담배 탐지 결과")

# RTC 설정 (NAT 통과를 위해 필요, Google STUN 서버 사용)
# 대부분의 경우 기본 설정으로 충분하나, 명시적으로 설정할 수 있습니다.
rtc_configuration = RTCConfiguration({"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]})

# Streamlit-WebRTC 컴포넌트 추가
webrtc_ctx = webrtc_streamer(
    key="yolo-detection-client-sound", # 고유 키
    mode=WebRtcMode.SENDRECV, # 비디오를 보내고 (SEND) 서버에서 처리된 비디오를 다시 받음 (RECV)
    video_processor_factory=lambda: YOLOVideoTransformer( # 비디오 변환 클래스 팩토리
        model=model,
        confidence_thresh=CONFIDENCE_THRESHOLD,
        send_interval=SEND_ALERT_INTERVAL,
        # NOTE: video_processor_factory가 lambda 함수로 사용될 때,
        # webrtc_streamer 내부적으로 생성된 data_channel 객체가 VideoTransformer 인스턴스에 전달됩니다.
        # 명시적으로 lambda 인자로 channel을 받지 않아도 됩니다.
        # (webrtc_streamer의 구현 방식에 따라 다를 수 있으므로 문서 확인 필요)
        # 최신 버전에서는 __init__에 data_channel=data_channel 형태로 전달됨
        data_channel=None # 초기값은 None, webrtc_streamer가 인스턴스 생성 시 실제 객체 주입
                          # -> 아님, lambda 팩토리가 인자를 받도록 변경해야 함
                          # lambda channel: YOLOVideoTransformer(..., data_channel=channel) 이 더 정확함
    ),
    rtc_configuration=rtc_configuration,
    media_stream_constraints={"video": True, "audio": False}, # 웹캠 비디오만 사용
    async_processing=True, # 비디오 처리를 비동기로 실행
    on_data_channel=JS_CLIENT_SOUND_SCRIPT # 데이터 채널 관련 클라이언트 JS 코드
)

# 위 video_processor_factory lambda 부분을 다음과 같이 명시적으로 data_channel을 받도록 수정합니다.
# lambda channel: YOLOVideoTransformer(
#     model=model,
#     confidence_thresh=CONFIDENCE_THRESHOLD,
#     send_interval=SEND_ALERT_INTERVAL,
#     data_channel=channel # 데이터 채널 객체 전달
# ),


st.write("---")

st.info("웹캠 스트림을 시작하면 브라우저에서 담배 탐지 시 알림 소리가 재생됩니다.")