File size: 5,568 Bytes
df80275 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
import math
from typing import List
try:
from typing import Literal, cast
except ImportError:
from typing_extensions import Literal # type: ignore
import av
import cv2
import numpy as np
import streamlit as st
from streamlit_webrtc import (
WebRtcMode,
create_mix_track,
create_process_track,
webrtc_streamer,
)
st.markdown(
"""
Mix multiple inputs with different video filters into one stream.
"""
)
VideoFilterType = Literal["noop", "cartoon", "edges", "rotate"]
def make_video_frame_callback(_type: VideoFilterType):
def callback(frame: av.VideoFrame) -> av.VideoFrame:
img = frame.to_ndarray(format="bgr24")
if _type == "noop":
pass
elif _type == "cartoon":
# prepare color
img_color = cv2.pyrDown(cv2.pyrDown(img))
for _ in range(6):
img_color = cv2.bilateralFilter(img_color, 9, 9, 7)
img_color = cv2.pyrUp(cv2.pyrUp(img_color))
# prepare edges
img_edges = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
img_edges = cv2.adaptiveThreshold(
cv2.medianBlur(img_edges, 7),
255,
cv2.ADAPTIVE_THRESH_MEAN_C,
cv2.THRESH_BINARY,
9,
2,
)
img_edges = cv2.cvtColor(img_edges, cv2.COLOR_GRAY2RGB)
# combine color and edges
img = cv2.bitwise_and(img_color, img_edges)
elif _type == "edges":
# perform edge detection
img = cv2.cvtColor(cv2.Canny(img, 100, 200), cv2.COLOR_GRAY2BGR)
elif _type == "rotate":
# rotate image
rows, cols, _ = img.shape
M = cv2.getRotationMatrix2D((cols / 2, rows / 2), frame.time * 45, 1)
img = cv2.warpAffine(img, M, (cols, rows))
return av.VideoFrame.from_ndarray(img, format="bgr24")
return callback
def mixer_callback(frames: List[av.VideoFrame]) -> av.VideoFrame:
buf_w = 640
buf_h = 480
buffer = np.zeros((buf_h, buf_w, 3), dtype=np.uint8)
n_inputs = len(frames)
n_cols = math.ceil(math.sqrt(n_inputs))
n_rows = math.ceil(n_inputs / n_cols)
grid_w = buf_w // n_cols
grid_h = buf_h // n_rows
for i in range(n_inputs):
frame = frames[i]
if frame is None:
continue
grid_x = (i % n_cols) * grid_w
grid_y = (i // n_cols) * grid_h
img = frame.to_ndarray(format="bgr24")
src_h, src_w = img.shape[0:2]
aspect_ratio = src_w / src_h
window_w = min(grid_w, int(grid_h * aspect_ratio))
window_h = min(grid_h, int(window_w / aspect_ratio))
window_offset_x = (grid_w - window_w) // 2
window_offset_y = (grid_h - window_h) // 2
window_x0 = grid_x + window_offset_x
window_y0 = grid_y + window_offset_y
window_x1 = window_x0 + window_w
window_y1 = window_y0 + window_h
buffer[window_y0:window_y1, window_x0:window_x1, :] = cv2.resize(
img, (window_w, window_h)
)
new_frame = av.VideoFrame.from_ndarray(buffer, format="bgr24")
return new_frame
COMMON_RTC_CONFIG = {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
st.header("Input 1")
input1_ctx = webrtc_streamer(
key="input1_ctx",
mode=WebRtcMode.SENDRECV,
rtc_configuration=COMMON_RTC_CONFIG,
media_stream_constraints={"video": True, "audio": False},
)
filter1_type = st.radio(
"Select transform type",
("noop", "cartoon", "edges", "rotate"),
key="mix-filter1-type",
)
callback = make_video_frame_callback(cast(VideoFilterType, filter1_type))
input1_video_process_track = None
if input1_ctx.output_video_track:
input1_video_process_track = create_process_track(
input_track=input1_ctx.output_video_track,
frame_callback=callback,
)
st.header("Input 2")
input2_ctx = webrtc_streamer(
key="input2_ctx",
mode=WebRtcMode.SENDRECV,
rtc_configuration=COMMON_RTC_CONFIG,
media_stream_constraints={"video": True, "audio": False},
)
filter2_type = st.radio(
"Select transform type",
("noop", "cartoon", "edges", "rotate"),
key="mix-filter2-type",
)
callback = make_video_frame_callback(cast(VideoFilterType, filter2_type))
input2_video_process_track = None
if input2_ctx.output_video_track:
input2_video_process_track = create_process_track(
input_track=input2_ctx.output_video_track, frame_callback=callback
)
st.header("Input 3 (no filter)")
input3_ctx = webrtc_streamer(
key="input3_ctx",
mode=WebRtcMode.SENDRECV,
rtc_configuration=COMMON_RTC_CONFIG,
media_stream_constraints={"video": True, "audio": False},
)
st.header("Mixed output")
mix_track = create_mix_track(kind="video", mixer_callback=mixer_callback, key="mix")
mix_ctx = webrtc_streamer(
key="mix",
mode=WebRtcMode.RECVONLY,
rtc_configuration=COMMON_RTC_CONFIG,
source_video_track=mix_track,
desired_playing_state=input1_ctx.state.playing
or input2_ctx.state.playing
or input3_ctx.state.playing,
)
if mix_ctx.source_video_track and input1_video_process_track:
mix_ctx.source_video_track.add_input_track(input1_video_process_track)
if mix_ctx.source_video_track and input2_video_process_track:
mix_ctx.source_video_track.add_input_track(input2_video_process_track)
if mix_ctx.source_video_track and input3_ctx.output_video_track:
# Input3 is sourced without any filter.
mix_ctx.source_video_track.add_input_track(input3_ctx.output_video_track)
|