Update steganography.py

steganography.py

@@ -7,6 +7,9 @@ import matplotlib.pyplot as plt
 import numpy as np
 import soundfile as sf
 from PIL import Image, ImageDraw, ImageFont
+import os
+import cv2
+from moviepy.editor import VideoFileClip, AudioFileClip
 
 DEFAULT_FONT_PATH = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
 DEFAULT_SAMPLE_RATE = 22050
@@ -116,6 +119,117 @@ def gradio_image_to_audio_fn(upload_image):
 def gradio_decode_fn(upload_audio):
     return display_audio_spectrogram(upload_audio)
 
+
+def display_progress(percent, message, progress=gr.Progress()):
+    progress(percent, desc=message)
+
+def extract_audio(video_path, progress):
+    display_progress(0.1, "Extracting audio from video", progress)
+    try:
+        video = VideoFileClip(video_path)
+        if video.audio is None:
+            raise ValueError("No audio found in the video")
+        audio_path = "extracted_audio.wav"
+        video.audio.write_audiofile(audio_path)
+        display_progress(0.2, "Audio extracted", progress)
+        return audio_path
+    except Exception as e:
+        display_progress(0.2, f"Failed to extract audio: {e}", progress)
+        return None
+
+def extract_frames(video_path, progress):
+    display_progress(0.3, "Extracting frames from video", progress)
+    try:
+        video = cv2.VideoCapture(video_path)
+        frames = []
+        success, frame = video.read()
+        while success:
+            frames.append(frame)
+            success, frame = video.read()
+        video.release()
+        display_progress(0.4, "Frames extracted", progress)
+        return frames
+    except Exception as e:
+        display_progress(0.4, f"Failed to extract frames: {e}", progress)
+        return None
+
+def frame_to_spectrogram(frame, sr=DEFAULT_SAMPLE_RATE):
+    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    S = np.flipud(gray_frame.astype(np.float32) / 255.0 * 100.0)
+    y = librosa.griffinlim(S)
+    return y
+
+def save_audio(y, sr=DEFAULT_SAMPLE_RATE):
+    audio_path = 'output_frame_audio.wav'
+    sf.write(audio_path, y, sr)
+    return audio_path
+
+def save_spectrogram_image(S, frame_number, temp_dir):
+    plt.figure(figsize=(10, 4))
+    librosa.display.specshow(S)
+    plt.tight_layout()
+    image_path = os.path.join(temp_dir, f'spectrogram_frame_{frame_number}.png')
+    plt.savefig(image_path)
+    plt.close()
+    return image_path
+
+def process_video_frames(frames, sr=DEFAULT_SAMPLE_RATE, temp_dir=None, progress=gr.Progress()):
+    processed_frames = []
+    total_frames = len(frames)
+    for i, frame in enumerate(frames):
+        y = frame_to_spectrogram(frame, sr)
+        S = librosa.feature.melspectrogram(y=y, sr=sr)
+        image_path = save_spectrogram_image(S, i, temp_dir)
+        processed_frame = cv2.imread(image_path)
+        processed_frames.append(processed_frame)
+        display_progress(0.5 + int((i + 1) / total_frames * 0.7), f"Frame processing {i + 1}/{total_frames}", progress)
+    display_progress(0.8, "All frames processed", progress)
+    return processed_frames
+
+def save_video_from_frames(frames, output_path, fps=30):
+    height, width, layers = frames[0].shape
+    video = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    for frame in frames:
+        video.write(frame)
+    video.release()
+
+def add_audio_to_video(video_path, audio_path, output_path, progress):
+    display_progress(0.9, "Adding audio back to video", progress)
+    try:
+        video = VideoFileClip(video_path)
+        audio = AudioFileClip(audio_path)
+        final_video = video.set_audio(audio)
+        final_video.write_videofile(output_path, codec='libx264', audio_codec='aac')
+        display_progress(1, "Video's ready", progress)
+    except Exception as e:
+        display_progress(1, f"Failed to add audio to video: {e}", progress)
+
+def process_video(video_path, progress=gr.Progress()):
+    try:
+        video = VideoFileClip(video_path)
+        if video.duration > 10:
+            video = video.subclip(0, 10)
+            temp_trimmed_video_path = "trimmed_video.mp4"
+            video.write_videofile(temp_trimmed_video_path, codec='libx264')
+            video_path = temp_trimmed_video_path
+    except Exception as e:
+        return f"Failed to load video: {e}"
+
+    audio_path = extract_audio(video_path, progress)
+    if audio_path is None:
+        return "Failed to extract audio from video."
+    frames = extract_frames(video_path, progress)
+    if frames is None:
+        return "Failed to extract frames from video."
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        processed_frames = process_video_frames(frames, temp_dir=temp_dir, progress=progress)
+        temp_video_path = os.path.join(temp_dir, 'processed_video.mp4')
+        save_video_from_frames(processed_frames, temp_video_path)
+        output_video_path = 'output_video_with_audio.mp4'
+        add_audio_to_video(temp_video_path, audio_path, output_video_path, progress)
+        return output_video_path
+
 def create_gradio_interface():
     with gr.Blocks(title="Audio Steganography", css="footer{display:none !important}", theme=gr.themes.Soft(primary_hue="green", secondary_hue="green", spacing_size="sm", radius_size="lg")) as txt2spec:
         with gr.Tab("Text to Spectrogram"):
@@ -147,7 +261,7 @@ def create_gradio_interface():
 
             convert_button.click(gradio_image_to_audio_fn, inputs=[upload_image], outputs=[output_audio_from_image])
 
-        with gr.Tab("Audio Spectrogram"):
+        with gr.Tab("Audio to Spectrogram"):
             with gr.Group():
                 with gr.Column():
                     upload_audio = gr.Audio(type="filepath", label="Upload audio", scale=3)
@@ -158,8 +272,18 @@ def create_gradio_interface():
 
             decode_button.click(gradio_decode_fn, inputs=[upload_audio], outputs=[decoded_image])
 
+        with gr.Tab("Video to Spectrogram"):
+            with gr.Group():
+                video_input = gr.Video(label="Upload video")
+                generate_button = gr.Button("Generate", variant="primary", size="lg")
+
+            with gr.Column(variant="panel"):
+                video_output = gr.Video(label="Video Spectrogram")
+
+            generate_button.click(process_video, inputs=[video_input], outputs=[video_output])
+
     return txt2spec
 
 if __name__ == "__main__":
     txt2spec = create_gradio_interface()
-    txt2spec.launch(share=True)
+    txt2spec.launch(share=True)