import os import sys if "APP_PATH" in os.environ: app_path = os.path.abspath(os.environ["APP_PATH"]) if os.getcwd() != app_path: # fix sys.path for import os.chdir(app_path) if app_path not in sys.path: sys.path.append(app_path) import gradio as gr import torch import torchaudio import torchvision import matplotlib.pyplot as plt import re import math import random import string import ffmpeg import subprocess import numpy as np import tqdm from audioseal import AudioSeal import videoseal from videoseal.utils.display import save_video_audio_to_mp4 # Load video_model if not already loaded in reload mode if 'video_model' not in globals(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load the VideoSeal model video_model = videoseal.load("videoseal") video_model.eval() video_model.to(device) video_model_nbytes = int(video_model.embedder.msg_processor.nbits / 8) # Load the AudioSeal model # Load audio_generator if not already loaded in reload mode if 'audio_generator' not in globals(): audio_generator = AudioSeal.load_generator("audioseal_wm_16bits") audio_generator = audio_generator.to(device) audio_generator_nbytes = int(audio_generator.msg_processor.nbits / 8) # Load audio_detector if not already loaded in reload mode if 'audio_detector' not in globals(): audio_detector = AudioSeal.load_detector("audioseal_detector_16bits") audio_detector = audio_detector.to(device) def generate_msg_pt_by_format_string(format_string, bytes_count): msg_hex = format_string.replace("-", "") hex_length = bytes_count * 2 binary_list = [] for i in range(0, len(msg_hex), hex_length): chunk = msg_hex[i:i+hex_length] binary = bin(int(chunk, 16))[2:].zfill(bytes_count * 8) binary_list.append([int(b) for b in binary]) # torch.randint(0, 2, (1, 16), dtype=torch.int32) msg_pt = torch.tensor(binary_list, dtype=torch.int32) return msg_pt.to(device) def generate_format_string_by_msg_pt(msg_pt, bytes_count): if msg_pt is None: return '', None hex_length = bytes_count * 2 binary_int = 0 for bit in msg_pt: binary_int = (binary_int << 1) | int(bit.item()) hex_string = format(binary_int, f'0{hex_length}x') split_hex = [hex_string[i:i + 4] for i in range(0, len(hex_string), 4)] format_hex = "-".join(split_hex) return hex_string, format_hex def generate_hex_format_regex(bytes_count): hex_length = bytes_count * 2 hex_string = 'F' * hex_length split_hex = [hex_string[i:i + 4] for i in range(0, len(hex_string), 4)] format_like = "-".join(split_hex) regex_pattern = '^' + '-'.join([r'[0-9A-Fa-f]{4}'] * len(split_hex)) + '$' return format_like, regex_pattern def generate_hex_random_message(bytes_count): hex_length = bytes_count * 2 hex_string = ''.join(random.choice(string.hexdigits) for _ in range(hex_length)) split_hex = [hex_string[i:i + 4] for i in range(0, len(hex_string), 4)] random_str = "-".join(split_hex) return random_str, "".join(split_hex) def embed_video_clip( model, clip: np.ndarray, msgs: torch.Tensor ) -> np.ndarray: clip_tensor = torch.tensor(clip, dtype=torch.float32).to(device).permute(0, 3, 1, 2) / 255.0 outputs = model.embed(clip_tensor, msgs=msgs, is_video=True) processed_clip = outputs["imgs_w"] processed_clip = (processed_clip * 255.0).byte().permute(0, 2, 3, 1).cpu().numpy() return processed_clip def embed_video( model, input_path: str, output_path: str, msgs: torch.Tensor, chunk_size: int, crf: int = 23 ) -> None: # Read video dimensions probe = ffmpeg.probe(input_path) video_info = next(stream for stream in probe['streams'] if stream['codec_type'] == 'video') width = int(video_info['width']) height = int(video_info['height']) fps = float(video_info['r_frame_rate'].split('/')[0]) / float(video_info['r_frame_rate'].split('/')[1]) num_frames = int(video_info['nb_frames']) # Open the input video process1 = ( ffmpeg .input(input_path) .output('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(width, height), r=fps) .run_async(pipe_stdout=True, pipe_stderr=subprocess.PIPE) ) # Open the output video process2 = ( ffmpeg .input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(width, height), r=fps) .output(output_path, vcodec='libx264', pix_fmt='yuv420p', r=fps, crf=crf) .overwrite_output() .run_async(pipe_stdin=True, pipe_stderr=subprocess.PIPE) ) # Process the video frame_size = width * height * 3 chunk = np.zeros((chunk_size, height, width, 3), dtype=np.uint8) frame_count = 0 pbar = tqdm.tqdm(total=num_frames, unit='frame', desc="Watermark video embedding") while True: # TODO block EOF on Windows in_bytes = process1.stdout.read(frame_size) if not in_bytes: break frame = np.frombuffer(in_bytes, np.uint8).reshape([height, width, 3]) chunk[frame_count % chunk_size] = frame frame_count += 1 pbar.update(1) if frame_count % chunk_size == 0: processed_frame = embed_video_clip(model, chunk, msgs) process2.stdin.write(processed_frame.tobytes()) process1.stdout.close() process2.stdin.close() process1.wait() process2.wait() return def get_sample_size(sample_fmt): if sample_fmt == 's16': return 2, np.int16 elif sample_fmt == 's16p': return 2, np.float16 elif sample_fmt == 'flt': return 4, np.int32 elif sample_fmt == 'fltp': return 4, np.float32 elif sample_fmt == 's32': return 4, np.int32 elif sample_fmt == 's32p': return 4, np.float32 elif sample_fmt == 'u8': return 1, np.int8 else: raise ValueError(f"Unsupported sample_fmt: {sample_fmt}") def embed_audio_clip( model, clip: np.ndarray, msgs: torch.Tensor, sample_rate ) -> np.ndarray: clip_tensor = torch.tensor(clip, dtype=torch.float32).to(device) # Resample the audio to 16kHz for watermarking audio_16k = torchaudio.transforms.Resample(sample_rate, 16000).to(device)(clip_tensor) # If the audio has more than one channel, average all channels to 1 channel if audio_16k.shape[0] > 1: audio_16k_mono = torch.mean(audio_16k, dim=0, keepdim=True) else: audio_16k_mono = audio_16k # Add batch dimension to the audio tensor audio_16k_mono_batched = audio_16k_mono.unsqueeze(0) # Get the watermark for the audio with torch.no_grad(): watermark = model.get_watermark( audio_16k_mono_batched, 16000, message=msgs ) # Embed the watermark in the audio audio_16k_w = audio_16k_mono_batched + watermark # Remove batch dimension from the watermarked audio tensor audio_16k_w = audio_16k_w.squeeze(0) # If the original audio had more than one channel, duplicate the watermarked audio to all channels if audio_16k.shape[0] > 1: audio_16k_w = audio_16k_w.repeat(audio_16k.shape[0], 1) # Resample the watermarked audio back to the original sample rate audio_w = torchaudio.transforms.Resample(16000, sample_rate).to(device)(audio_16k_w) processed_clip = audio_w.cpu().numpy() return processed_clip def embed_audio( model, input_path: str, output_path: str, msgs: torch.Tensor, chunk_size: int ) -> None: # Read audio dimensions probe = ffmpeg.probe(input_path) audio_info = next(stream for stream in probe['streams'] if stream['codec_type'] == 'audio') sample_rate = int(audio_info['sample_rate']) sample_fmt = audio_info['sample_fmt'] channels = int(audio_info['channels']) duration = float(audio_info['duration']) # CASE 1 Read audio all at once # audio_data, stderr_output = ( # ffmpeg # .input(input_path, loglevel='debug') # .output('pipe:', format='f32le', acodec='pcm_f32le', ar=sample_rate, ac=channels) # .run(capture_stdout=True, capture_stderr=True) # ) # audio_data = process.stdout.read() # print("audio numpy total size:", len(audio_data)) # process.stdout.close() # process.wait() # stderr_output = process.stderr.read().decode('utf-8') # print(stderr_output) # CASE 2 Read async # NOTE loglevel='debug' not work on Windows # NOTE format='wav' data size(4104768) bigger than format='s16le'(4104688) # process = ( # ffmpeg # .input(input_path, loglevel='debug') # .output('pipe:', format='f32le', acodec='pcm_f32le', ar=sample_rate, ac=channels) # .run_async(pipe_stdout=True, pipe_stderr=subprocess.PIPE) # ) # audio_data = process.stdout.read() # print("audio numpy total size:", len(audio_data)) # process.stdout.close() # process.wait() # stderr_output = process.stderr.read().decode('utf-8') # print(stderr_output) # stderr_output example: # # # AVIOContext @ 0x5d878ea02e80] Statistics: 4104688 bytes written, 0 seeks, 251 writeouts # # [out#0/f32le @ 0x5d878eaf31c0] Output file #0 (pipe:): # # [out#0/f32le @ 0x5d878eaf31c0] Output stream #0:0 (audio): 251 frames encoded (513086 samples); 251 packets muxed (4104688 bytes); # # [out#0/f32le @ 0x5d878eaf31c0] Total: 251 packets (4104688 bytes) muxed # CASE 3 Read by torchaudio # NOTE torchvision read audio format is f32le # _, audio, info = torchvision.io.read_video(input_path, output_format="TCHW") # print("audio numpy total size:", audio.nbytes) # Open the input audio process1 = ( ffmpeg .input(input_path) .output('pipe:', format='f32le', acodec='pcm_f32le', ac=channels, ar=sample_rate) .run_async(pipe_stdout=True, pipe_stderr=subprocess.PIPE) ) # Open the output audio process2 = ( ffmpeg .input('pipe:', format='f32le', ac=channels, ar=sample_rate) .output(output_path, format='wav', acodec='pcm_f32le', ac=channels, ar=sample_rate) # not work # .output(output_path, acodec='libfdk_aac', ac=channels, ar=sample_rate) .overwrite_output() .run_async(pipe_stdin=True, pipe_stderr=subprocess.PIPE) ) # CASE read all and write all # while True: # audio_data = process1.stdout.read() # if not audio_data: # break # try: # process2.stdin.write(audio_data) # except BrokenPipeError: # print("Broken pipe: process2 has closed the input.") # break # Process the audio sample_size, sample_type = get_sample_size(sample_fmt) second_size = sample_size * channels * sample_rate chunk = np.zeros((chunk_size, sample_rate, channels), dtype=sample_type) second_count = 0 pbar = tqdm.tqdm(total=math.ceil(duration), unit='second', desc="Watermark audio embedding") while True: in_bytes = process1.stdout.read(second_size) if not in_bytes: break frame = np.frombuffer(in_bytes, sample_type) frame = frame.reshape((-1, channels)) chunk[second_count % chunk_size, :len(frame)] = frame second_count += 1 pbar.update(1) if second_count % chunk_size == 0: if msgs is None: process2.stdin.write(in_bytes) else: clip = np.concatenate(chunk, axis=0).T processed_frame = embed_audio_clip(model, clip, msgs, sample_rate) process2.stdin.write(processed_frame.T.tobytes()) process1.stdout.close() process2.stdin.close() process1.wait() process2.wait() # CASE print stderr # stderr_output1 = process1.stderr.read().decode('utf-8') # stderr_output2 = process2.stderr.read().decode('utf-8') # print("Process 1 stderr:") # print(stderr_output1) # print("Process 2 stderr:") # print(stderr_output2) return def embed_watermark(input_path, output_path, msg_v, msg_a, video_only, progress): output_path_video = output_path + ".video.mp4" embed_video(video_model, input_path, output_path_video, msg_v, 16) output_path_audio = output_path + ".audio.m4a" if video_only: msg_a = None embed_audio(audio_generator, input_path, output_path_audio, msg_a, 3) # Use FFmpeg to add audio to the video final_command = [ 'ffmpeg', '-i', output_path_video, '-i', output_path_audio, '-c:v', 'copy', '-c:a', 'aac', '-strict', 'experimental', '-y', output_path ] subprocess.run(final_command, check=True) return def detect_video_clip( model, clip: np.ndarray ) -> torch.Tensor: clip_tensor = torch.tensor(clip, dtype=torch.float32).permute(0, 3, 1, 2) / 255.0 outputs = model.detect(clip_tensor, is_video=True) output_bits = outputs["preds"][:, 1:] # exclude the first which may be used for detection return output_bits def detect_video( model, input_path: str, chunk_size: int ) -> None: # Read video dimensions probe = ffmpeg.probe(input_path) video_info = next(stream for stream in probe['streams'] if stream['codec_type'] == 'video') width = int(video_info['width']) height = int(video_info['height']) fps = float(video_info['r_frame_rate'].split('/')[0]) / float(video_info['r_frame_rate'].split('/')[1]) num_frames = int(video_info['nb_frames']) # Open the input video process1 = ( ffmpeg .input(input_path) .output('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(width, height), r=fps) .run_async(pipe_stdout=True, pipe_stderr=subprocess.PIPE) ) # Process the video frame_size = width * height * 3 chunk = np.zeros((chunk_size, height, width, 3), dtype=np.uint8) frame_count = 0 soft_msgs = [] pbar = tqdm.tqdm(total=num_frames, unit='frame', desc="Watermark video detecting") while True: in_bytes = process1.stdout.read(frame_size) if not in_bytes: break frame = np.frombuffer(in_bytes, np.uint8).reshape([height, width, 3]) chunk[frame_count % chunk_size] = frame frame_count += 1 pbar.update(1) if frame_count % chunk_size == 0: soft_msgs.append(detect_video_clip(model, chunk)) process1.stdout.close() process1.wait() soft_msgs = torch.cat(soft_msgs, dim=0) return soft_msgs def detect_audio_clip( model, clip: np.ndarray, sample_rate ) -> torch.Tensor: clip_tensor = torch.tensor(clip, dtype=torch.float32).to(device) # Resample the audio to 16kHz for watermarking audio_16k = torchaudio.transforms.Resample(sample_rate, 16000).to(device)(clip_tensor) # If the audio has more than one channel, average all channels to 1 channel if audio_16k.shape[0] > 1: audio_16k_mono = torch.mean(audio_16k, dim=0, keepdim=True) else: audio_16k_mono = audio_16k # Add batch dimension to the audio tensor audio_16k_mono_batched = audio_16k_mono.unsqueeze(0) # Detect watermarks in the audio with torch.no_grad(): result, message = model.detect_watermark( audio_16k_mono_batched, 16000 ) # pred_prob is a tensor of size batch x 2 x frames, indicating the probability (positive and negative) of watermarking for each frame # A watermarked audio should have pred_prob[:, 1, :] > 0.5 # message_prob is a tensor of size batch x 16, indicating of the probability of each bit to be 1. # message will be a random tensor if the detector detects no watermarking from the audio pred_prob, message_prob = model(audio_16k_mono_batched, sample_rate) # print(f"Detection result for audio: {result}") # _, format_msg = generate_format_string_by_msg_pt(message[0], audio_generator_nbytes) # print(f"Extracted message from audio: {message}: {format_msg}") # print(f"Extracted pred_prob from audio: {pred_prob.shape}") # print(f"Extracted message_prob from audio: {message_prob}") # print(f"Extracted shape from audio 16k: {audio_16k_mono_batched.shape}") # print(f"Extracted shape from audio original: {clip_tensor.shape}") return result, message, pred_prob, message_prob def detect_audio( model, input_path: str, chunk_size: int ) -> None: # Read audio dimensions probe = ffmpeg.probe(input_path) audio_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'audio'] if len(audio_streams) == 0: gr.Warning("No audio stream found in the input file.") return None, None, None, None audio_info = audio_streams[0] sample_rate = int(audio_info['sample_rate']) sample_fmt = audio_info['sample_fmt'] channels = int(audio_info['channels']) duration = float(audio_info['duration']) # Open the input audio process1 = ( ffmpeg .input(input_path) .output('pipe:', format='f32le', acodec='pcm_f32le', ac=channels, ar=sample_rate) .run_async(pipe_stdout=True, pipe_stderr=subprocess.PIPE) ) # Process the audio sample_size, sample_type = get_sample_size(sample_fmt) second_size = sample_size * channels * sample_rate chunk = np.zeros((chunk_size, sample_rate, channels), dtype=sample_type) second_count = 0 soft_result = [] soft_message = [] soft_pred_prob = [] soft_message_prob = [] pbar = tqdm.tqdm(total=math.ceil(duration), unit='second', desc="Watermark audio detecting") while True: in_bytes = process1.stdout.read(second_size) if not in_bytes: break frame = np.frombuffer(in_bytes, sample_type) frame = frame.reshape((-1, channels)) chunk[second_count % chunk_size, :len(frame)] = frame second_count += 1 pbar.update(1) if second_count % chunk_size == 0: clip = np.concatenate(chunk, axis=0).T # print(f"Detection audio second: {second_count-chunk_size}-{second_count}") result, message, pred_prob, message_prob = detect_audio_clip(model, clip, sample_rate) soft_result.append(result) soft_message.append(message) soft_pred_prob.append(pred_prob) soft_message_prob.append(message_prob) process1.stdout.close() process1.wait() soft_message = torch.cat(soft_message, dim=0) soft_pred_prob = torch.cat(soft_pred_prob, dim=0) soft_message_prob = torch.cat(soft_message_prob, dim=0) return (soft_result, soft_message, soft_pred_prob, soft_message_prob) def detect_watermark(input_path, video_only): msgs_v_frame = detect_video(video_model, input_path, 16) msgs_v_avg = msgs_v_frame.mean(dim=0) # Average the predictions across all frames msgs_v_frame = (msgs_v_frame > 0).to(int) msgs_v_avg = (msgs_v_avg > 0).to(int) msgs_v_unique, msgs_v_counts = torch.unique(msgs_v_frame, dim=0, return_counts=True) msgs_v_most = None if len(msgs_v_frame) > len(msgs_v_counts) > 0: msgs_v_most_idx = torch.argmax(msgs_v_counts) msgs_v_most = msgs_v_unique[msgs_v_most_idx] msgs_a_most = msgs_a_res = msgs_a_frame = msgs_a_pred = msgs_a_prob = None if not video_only: msgs_a_res, msgs_a_frame, msgs_a_pred, msgs_a_prob = detect_audio(audio_detector, input_path, 1) if msgs_a_res is not None: msgs_a_res_not_zero = [i for i, x in enumerate(msgs_a_res) if x > 0.5] msgs_a_frame_not_zero = msgs_a_frame[msgs_a_res_not_zero] msgs_a_unique, msgs_a_counts = torch.unique(msgs_a_frame_not_zero, dim=0, return_counts=True) if len(msgs_a_counts) > 0: msgs_a_most_idx = torch.argmax(msgs_a_counts) msgs_a_most = msgs_a_unique[msgs_a_most_idx] return msgs_v_most, msgs_v_avg, msgs_v_frame, msgs_a_most, msgs_a_res, msgs_a_frame, msgs_a_pred, msgs_a_prob with gr.Blocks(title="VideoSeal") as demo: gr.Markdown(""" # VideoSeal Demo For video, each frame will be watermarked and detected. For audio, each 3 seconds will be watermarked, and each second will be detected. **NOTE: The watermarked process will modify both audio and video. The video will be re-encoded to yuv420p using libx264, and the audio will be duplicated from mono 16kHz and resampled back to the original channel sample rate.** Find the project [here](https://github.com/facebookresearch/videoseal.git). """) with gr.Tabs(): with gr.TabItem("Embed Watermark"): with gr.Row(): with gr.Column(): embedding_vid = gr.Video(label="Input Video") with gr.Row(): with gr.Column(): embedding_type = gr.Radio(["random", "input"], value="random", label="Type", info="Type of watermarks") format_like_v, regex_pattern_v = generate_hex_format_regex(video_model_nbytes) msg_v, _ = generate_hex_random_message(video_model_nbytes) embedding_msg_v = gr.Textbox( label=f"Message ({video_model_nbytes} bytes hex string)", info=f"format like {format_like_v}", value=msg_v, interactive=False, show_copy_button=True) with gr.Column(): embedding_only_vid = gr.Checkbox(label="Only Video", value=False) format_like_a, regex_pattern_a = generate_hex_format_regex(audio_generator_nbytes) msg_a, _ = generate_hex_random_message(audio_generator_nbytes) embedding_msg_a = gr.Textbox( label=f"Audio Message ({audio_generator_nbytes} bytes hex string)", info=f"format like {format_like_a}", value=msg_a, interactive=False, show_copy_button=True) embedding_btn = gr.Button("Embed Watermark") with gr.Column(): marked_vid = gr.Video(label="Output Audio", show_download_button=True) def change_embedding_type(video_only): return gr.update(visible=not video_only) embedding_only_vid.change( fn=change_embedding_type, inputs=[embedding_only_vid], outputs=[embedding_msg_a] ) def change_embedding_type(type): if type == "random": msg_v, _ = generate_hex_random_message(video_model_nbytes) msg_a,_ = generate_hex_random_message(audio_generator_nbytes) return [gr.update(interactive=False, value=msg_v),gr.update(interactive=False, value=msg_a)] else: return [gr.update(interactive=True),gr.update(interactive=True)] embedding_type.change( fn=change_embedding_type, inputs=[embedding_type], outputs=[embedding_msg_v, embedding_msg_a] ) def check_embedding_msg(msg_v, msg_a): if not re.match(regex_pattern_v, msg_v): gr.Warning( f"Invalid format. Please use like '{format_like_v}'", duration=0) if not re.match(regex_pattern_a, msg_a): gr.Warning( f"Invalid format. Please use like '{format_like_a}'", duration=0) embedding_msg_v.change( fn=check_embedding_msg, inputs=[embedding_msg_v, embedding_msg_a], outputs=[] ) embedding_msg_a.change( fn=check_embedding_msg, inputs=[embedding_msg_v, embedding_msg_a], outputs=[] ) def run_embed_watermark(input_path, video_only, msg_v, msg_a, progress=gr.Progress(track_tqdm=True)): if input_path is None: raise gr.Error("No file uploaded", duration=5) if not re.match(regex_pattern_v, msg_v): raise gr.Error(f"Invalid format. Please use like '{format_like_v}'", duration=5) if not re.match(regex_pattern_a, msg_a): raise gr.Error(f"Invalid format. Please use like '{format_like_a}'", duration=5) msg_pt_v = generate_msg_pt_by_format_string(msg_v, video_model_nbytes) msg_pt_a = generate_msg_pt_by_format_string(msg_a, audio_generator_nbytes) if video_only: output_path = os.path.join(os.path.dirname(input_path), "__".join([msg_v]) + '.mp4') else: output_path = os.path.join(os.path.dirname(input_path), "__".join([msg_v, msg_a]) + '.mp4') embed_watermark(input_path, output_path, msg_pt_v, msg_pt_a, video_only, progress) return output_path embedding_btn.click( fn=run_embed_watermark, inputs=[embedding_vid, embedding_only_vid, embedding_msg_v, embedding_msg_a], outputs=[marked_vid] ) with gr.TabItem("Detect Watermark"): with gr.Row(): with gr.Column(): detecting_vid = gr.Video(label="Input Video") detecting_only_vid = gr.Checkbox(label="Only Video", value=False) detecting_btn = gr.Button("Detect Watermark") with gr.Column(): predicted_messages = gr.JSON(label="Detected Messages") def run_detect_watermark(file, video_only, progress=gr.Progress(track_tqdm=True)): if file is None: raise gr.Error("No file uploaded", duration=5) msgs_v_most, msgs_v_avg, msgs_v_frame, msgs_a_most, msgs_a_res, msgs_a_frame, msgs_a_pred, msgs_a_prob = detect_watermark(file, video_only) _, format_msg_v_most = generate_format_string_by_msg_pt(msgs_v_most, video_model_nbytes) _, format_msg_v_avg = generate_format_string_by_msg_pt(msgs_v_avg, video_model_nbytes) format_msg_v_frames = {} for idx, msg in enumerate(msgs_v_frame): _, format_msg = generate_format_string_by_msg_pt(msg, video_model_nbytes) format_msg_v_frames[f"{idx}"] = format_msg video_json = { "most": format_msg_v_most, "avg": format_msg_v_avg, "frames": format_msg_v_frames } if msgs_a_res is None: audio_json = None else: _, format_msg_a_most = generate_format_string_by_msg_pt(msgs_a_most, audio_generator_nbytes) format_msg_a_seconds = {} for idx, (result, message, pred_prob, message_prob) in enumerate(zip(msgs_a_res, msgs_a_frame, msgs_a_pred, msgs_a_prob)): _, format_msg = generate_format_string_by_msg_pt(message, audio_generator_nbytes) sum_above_05 = (pred_prob[1, :] > 0.5).sum(dim=0) format_msg_a_seconds[f"{idx}"] = { "socre": result, "message": format_msg, "frames_count_all": pred_prob.shape[1], "frames_count_above_05": sum_above_05.item(), "bits_probability": message_prob.tolist(), "bits_massage": message.tolist() } audio_json = { "most": format_msg_a_most, "seconds": format_msg_a_seconds } # Create message output as JSON message_json = { "video": video_json, "audio:": audio_json } return message_json detecting_btn.click( fn=run_detect_watermark, inputs=[detecting_vid, detecting_only_vid], outputs=[predicted_messages] ) if __name__ == "__main__": demo.launch()