app.py

import gradio as gr
import tensorflow as tf
import tensorflow.keras.backend as K
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import os
import shutil
from skimage.metrics import structural_similarity as ssim

beta = 1.0

# Loss for reveal network
def rev_loss(s_true, s_pred):
    # Loss for reveal network is: beta * |S-S'|
    return beta * K.sum(K.square(s_true - s_pred))

# Loss for the full model, used for preparation and hidding networks
def full_loss(y_true, y_pred):
    # Loss for the full model is: |C-C'| + beta * |S-S'|
    s_true, c_true = y_true[...,0:3], y_true[...,3:6]
    s_pred, c_pred = y_pred[...,0:3], y_pred[...,3:6]

    s_loss = rev_loss(s_true, s_pred)
    c_loss = K.sum(K.square(c_true - c_pred))
    return s_loss + c_loss

model = tf.keras.models.load_model("model.h5", custom_objects={'full_loss':
full_loss})


def preprocess_image(img):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)
    img = img.resize((124, 124), Image.ANTIALIAS)
    img = np.array(img)
    img = img / 255.0
    return img

def steganography_image(imageO, imageF):
    # Preprocess images
    img_S = preprocess_image(imageO)
    img_C = preprocess_image(imageF)

    # Add batch dimension
    img_S = np.expand_dims(img_S, axis=0)
    img_C = np.expand_dims(img_C, axis=0)

    # Predict with pre/loaded model
    decoded = model.predict([img_S, img_C])
    decoded_S, decoded_C = decoded[..., 0:3], decoded[..., 3:6]

    # Post-process outputs
    decoded_S = np.squeeze(decoded_S, axis=0)  # Remove batch dimension
    decoded_C = np.squeeze(decoded_C, axis=0)  # Remove batch dimension
    decoded_S = (decoded_S * 255).astype(np.uint8)
    decoded_C = (decoded_C * 255).astype(np.uint8)

    # Calculate absolute differences
    diff_S = np.abs(decoded_S - (img_S.squeeze() * 255)).astype(np.uint8)
    diff_C = np.abs(decoded_C - (img_C.squeeze() * 255)).astype(np.uint8)

    # Create a plot of differences
    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
    ax[0].imshow(diff_S)
    ax[0].set_title('Difference in Secret Image')
    ax[0].axis('off')
    ax[1].imshow(diff_C)
    ax[1].set_title('Difference in Cover Image')
    ax[1].axis('off')
    plt.tight_layout()

    # Return images and plot
    return decoded_S, decoded_C, fig


#Function to clear a folder
def clear_folder(path):
    if os.path.exists(path):
        shutil.rmtree(path)
    os.makedirs(path)

#Function to extract every frame of a video and save them in a folder
def extractImages(pathIn, pathOut):
    clear_folder(pathOut)
    if not os.path.exists(pathOut):
        os.makedirs(pathOut)

    vidcap = cv2.VideoCapture(pathIn)
    success, image = vidcap.read()
    count = 0

    while success:
        frame_path = os.path.join(pathOut, f"frame{count}.jpg")
        success, image = vidcap.read()

        if success:
            resized_image = cv2.resize(image, (124, 124), interpolation=cv2.INTER_AREA)
            cv2.imwrite(frame_path, resized_image)
            print(f'Saved frame {count} to {frame_path}')
        else:
            print(f'Failed to read frame at count {count}')

        count += 1

#Function to create a new video based on a folder of frames
def rebuildVideo(framesPath, outputPath, fps=30):
    frame_files = sorted([f for f in os.listdir(framesPath) if f.endswith('.jpg')],
                         key=lambda x: int(x[5:-4]))

    first_frame_path = os.path.join(framesPath, frame_files[0])
    frame = cv2.imread(first_frame_path)
    height, width, layers = frame.shape

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(outputPath, fourcc, fps, (width, height))

    for file in frame_files:
        frame_path = os.path.join(framesPath, file)
        frame = cv2.imread(frame_path)
        out.write(frame)
    out.release()


def calculate_ssim(img1, img2):
    img1_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    img2_gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    score, _ = ssim(img1_gray, img2_gray, full=True)
    return score

def plot_metrics(metrics):
    fig, ax = plt.subplots()
    ax.plot(metrics, label="SSIM")
    ax.set_xlabel("Frame")
    ax.set_ylabel("SSIM")
    ax.set_title("SSIM over Frames")
    ax.legend()
    ax.grid(True)
    return fig

def process_frame(imageO, imageF):
    img_S = preprocess_image(imageO)
    img_C = preprocess_image(imageF)
    img_S = np.expand_dims(img_S, axis=0)
    img_C = np.expand_dims(img_C, axis=0)
    decoded = model.predict([img_S, img_C])
    decoded_S, decoded_C = decoded[..., 0:3], decoded[..., 3:6]
    decoded_S = np.squeeze(decoded_S, axis=0)
    decoded_C = np.squeeze(decoded_C, axis=0)
    decoded_S = (decoded_S * 255).astype(np.uint8)
    decoded_C = (decoded_C * 255).astype(np.uint8)
    return decoded_S, decoded_C

def steganography_video(video_path1, video_path2):
    input_frames_path = "Frames1"
    input_frames_path2 = "Frames2"
    output_frames_path = "Frames3"
    output_frames_path2 = "Frames4"
    output_video_path = "output_video.mp4"
    output_video_path2 = "output_video2.mp4"
    extractImages(video_path1, input_frames_path)
    extractImages(video_path2, input_frames_path2)

    input_frame_files = sorted([f for f in os.listdir(input_frames_path) if f.endswith('.jpg')],
                               key=lambda x: int(x[5:-4]))

    clear_folder(output_frames_path)
    clear_folder(output_frames_path2)
    i = 0
    ssim_scores = []
    ssim_scores2 = []

    for file in input_frame_files:
        frame_path = os.path.join(input_frames_path, file)
        frame_path2 = os.path.join(input_frames_path2, f"frame{i}.jpg")
        frame = cv2.imread(frame_path)
        try:
          frame2 = cv2.imread(frame_path2)
        except:
          print("Second video is too short, will be cut up to the length of the first one")
          break
        if frame2 is None:
          break
        decoded_S, decoded_C = process_frame(frame, frame2)
        decoded_S_path = os.path.join(output_frames_path, file)
        cv2.imwrite(decoded_S_path, decoded_S)
        decoded_C_path = os.path.join(output_frames_path2, file)
        cv2.imwrite(decoded_C_path, decoded_C)
        print(frame.shape)
        print(decoded_S.shape)
        print(frame2.shape)
        print(decoded_C.shape)
        ssim_scores.append(calculate_ssim(frame, decoded_S))
        ssim_scores2.append(calculate_ssim(frame2, decoded_C))
        i+=1

    rebuildVideo(output_frames_path, output_video_path, fps=20)
    rebuildVideo(output_frames_path2, output_video_path2, fps=20)

    return output_video_path, output_video_path2, ssim_scores, ssim_scores2

example_secret_image = "Examples/secret.jpg"
example_cover_image = "Examples/cover.jpg"
example_cover_video = "Examples/cover.mp4"
example_secret_video = "Examples/secret.mp4"

with gr.Blocks() as demo:
    with gr.Tab("Image Processing"):
        image_input1 = gr.Image(label="Cover Image")
        image_input2 = gr.Image(label="Secret Image")
        image_output1 = gr.Image(label="Decoded Cover Image")
        image_output2 = gr.Image(label="Decoded Secret Image")
        plot = gr.Plot(label = "Noise behind each image")
        btn_image = gr.Button("Process Images")

        btn_image.click(
            fn=steganography_image,
            inputs=[image_input1, image_input2],
            outputs=[image_output1, image_output2, plot]
        )

    with gr.Tab("Video Processing"):
        video_input = gr.Video(label="Input Cover Video")
        video_input2 = gr.Video(label="Input Secret Video")
        video_output = gr.Video(label="Output Cover Video")
        video_output2 = gr.Video(label="Output Secret Video")
        plot_output = gr.Plot(label="SSIM over Frames for Cover")
        plot_output2 = gr.Plot(label="SSIM over Frames for Secret")
        btn_video = gr.Button("Process Video")

        def process_video_and_plot(video_path, video_path2):
            video_path, video_path2, ssim_scores, ssim_scores2 = steganography_video(video_path, video_path2)
            plot = plot_metrics(ssim_scores)
            plot2 = plot_metrics(ssim_scores2)
            plot.show()
            return video_path, video_path2, plot, plot2

        btn_video.click(
            fn=process_video_and_plot,
            inputs=[video_input, video_input2],
            outputs=[video_output, video_output2, plot_output, plot_output2]
        )

demo.launch(debug=True, share = True)