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("웹캠 스트림을 시작하면 브라우저에서 담배 탐지 시 알림 소리가 재생됩니다.")