app/models/DEEP_STEGO/Utils/enhance.py

import numpy as np
import argparse
import sys
import math
from PIL import ImageFilter, Image
import cv2

'''
Enhance video by applying denoise and sharpen filters
'''

# Construct argument parser
parser = argparse.ArgumentParser(description='Use block shuffle')
parser.add_argument('--denoise', action='store_true', default=False)
parser.add_argument('--sharpen', action='store_true', default=False)
parser.add_argument("--input_video", required=True, help="path to input video")
args = vars(parser.parse_args())

# Start video enhancement
print('\nEnhancing video ...\n')


# Update progress bar
def update_progress(current_frame, total_frames):
    progress = math.ceil((current_frame / total_frames) * 100)
    sys.stdout.write('\rProgress: [{0}] {1}%'.format('>' * math.ceil(progress / 10), progress))


# Open the input video
vid_cap = cv2.VideoCapture(args['input_video'])

# Total input video frames
num_frames = int(vid_cap.get(cv2.CAP_PROP_FRAME_COUNT))
print("Total frames in input video:", num_frames)

# Initialize the video writer
enhanced_video = cv2.VideoWriter('results/enhanced_secret_300.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 15,
                                 (300, 300))

# Initialize the frame buffer
frames = []

# Load the frames to buffer
while vid_cap.isOpened():
    success, image = vid_cap.read()
    if success:
        frames.append(image)
    else:
        break

# Set up the start frame index
start_frame = 5

# Enhance and save video frame-by-frame
for i in range(start_frame, len(frames) - (start_frame + 1)):
    output = frames[i]
    if args["denoise"]:
        output = cv2.fastNlMeansDenoisingColoredMulti(frames, i, 11)
    if args["sharpen"]:
        output = np.array(Image.fromarray(output).filter(ImageFilter.DETAIL))
    enhanced_video.write(output)
    update_progress(i, num_frames - (start_frame + 1))

# Finish video enhancement
print('\n\nSuccessfully enhanced video !!!\n')

'''
Sample run:-
python enhance.py --input_video results/secret_outvid_300.avi --denoise --sharpen
'''