import streamlit as st
from streamlit_webrtc import webrtc_streamer
from streamlit_webrtc import VideoProcessorBase
import cv2
import av #strealing video library
from loguru import logger
import numpy as np

st.title('Streamlit App Test')
st.write('Gray Scale -> Color')

class VideoProcessor(VideoProcessorBase):
    def __init__(self) -> None:
        self.is_gray_state = None
        self.snapshot = False

        # image1の白色のピクセルを透明に変換する
        self.right_overray = cv2.imread('example_overray.png', cv2.IMREAD_UNCHANGED)
        white_pixels = np.all(self.right_overray[:, :, :3] == [255, 255, 255], axis=-1)
        self.right_overray[white_pixels, 3] = 0  # アルファチャンネルを0に設定

    def _add_overray(self, img):
        # 画像をoverrayする関数
        img_alpha = np.ones(img.shape[:2], dtype=img.dtype) * 255
        img = cv2.merge([img, img_alpha])
        height, width = img.shape[:2]

        rate = img.shape[:2][0] / self.right_overray.shape[:2][0]
        new_width = int(rate * self.right_overray.shape[1])
        right_overray = cv2.resize(self.right_overray, (new_width, height))

        # right_overrayの幅がimgの幅より大きい場合、imgの幅に合わせる
        if right_overray.shape[1] > width:
            right_overray = right_overray[:, :width]
        # right_overrayの幅がimgの幅より小さい場合、残りの部分を埋める
        elif right_overray.shape[1] < width:
            padding_left_width = (width - right_overray.shape[1]) // 2
            padding_right_width = width - right_overray.shape[1] - padding_left_width
            padding_left = np.zeros((height, padding_left_width, 4), dtype=img.dtype)
            padding_right = np.zeros((height, padding_right_width, 4), dtype=img.dtype)
            right_overray = np.hstack((padding_left, right_overray, padding_right))
        combined_image = np.where(right_overray[:, :, 3:4] != 0, right_overray, img)

        # アルファチャンネルを削除
        combined_image_rgb = combined_image[:, :, :3]
        vframe = av.VideoFrame.from_ndarray(combined_image_rgb, format='bgr24')

        return vframe

    def recv(self, frame: av.VideoFrame) -> av.VideoFrame:
        img = frame.to_ndarray(format = 'bgr24')
        
        if self.is_gray_state == True:
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
            vframe = av.VideoFrame.from_ndarray(img, format='gray')
        else:
            vframe = av.VideoFrame.from_ndarray(img, format='bgr24')

        # スナップショットボタンが押されたら画像を保存
        if self.snapshot:
            try:
                cv2.imwrite('output_image.png', img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
            except Exception as exc:
                logger.exception(exc)
            self.snapshot = False

        try:
            vframe = self._add_overray(img)
        except Exception as exc:
            logger.exception(exc)

        return vframe

ctx = webrtc_streamer(key='example', video_processor_factory=VideoProcessor) # type: ignore

if ctx.video_processor:
    ctx.video_processor.is_gray_state = st.checkbox('Gray Scale -> Color ')
    if st.button("Snapshot"):
        ctx.video_processor.snapshot = True