python

"""
This script add haze to a video or a single image or an image folder.
"""

import cv2
import os
import numpy as np
import random
from sys import argv


IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp'  # include image suffixes

def hazeImg(src, t=0.6, isVideo=False):
    '''
    @: Add haze to current image.
    para src:    input image data.
    para t:      throw rate (0~1)
    para isVideo:boolean value, if True, we use the same A, to get smoother video.
    
    '''

    A = [random.uniform(0.8,1.0), 0.8][isVideo]
    
    image = src * t + A * 255*(1-t)

    return image.astype(np.uint8)


def process_single_image(imgPath):
    '''
    @: Add haze to a single image.
    para imgPath: Full path to an image.
    '''
    print('*'*10 + 'Start process a single image' + '*'*10)
    img = cv2.imread(imgPath)

    result = hazeImg(img, 0.4)

    savePath = imgPath.rsplit('.',1)[0] + "_haze.png"

    cv2.imwrite(savePath, result)
    
    print("save to :", savePath)
    print('*'*10 + 'End process a single image' + '*'*10)
    return 



def process_image_folder(imgDir):
    '''
    @： Add haze to images in a directory.
    para imgDir: Directory of images
    '''
    print('*'*10 + 'Start process a directory' + '*'*10)
    allFiles = os.listdir(imgDir)
    for f in allFiles:
        imgPath = os.path.join(imgDir, f)
        img = cv2.imread(imgPath)
        t = random.uniform(0.1,1.0)
        result = hazeImg(img, t)

        savePath = imgPath.rsplit('.',1)[0] + "_haze.png"
        cv2.imwrite(savePath, result)

    print('*'*10 + 'End process a directory' + '*'*10)
    return 


def process_video(vPath):
    '''
    @: Add haze to a video.
    para vPath: Full path to a video
    '''
    print('*'*10 + 'Start process a video' + '*'*10)
    #savePath = vPath.split(".")[0] + ".avi"
    savePath = vPath.rsplit('.',1)[0] + ".avi"
    

    cap = cv2.VideoCapture(vPath)

    if not cap.isOpened():
        print("Unable to read frames!")

    w,h = int(cap.get(3)), int(cap.get(4))

    fps = 30
    sz = (w,h)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    vout = cv2.VideoWriter()
    vout.open(savePath, fourcc, fps, sz, True)

    while 1:
        ret, frame = cap.read()

        if ret:
            result = hazeImg(frame,isVideo=True)
            vout.write(result)
        else:
            break

    vout.release()
    cap.release()
    print("save to :", savePath)
    print('*'*10 + 'End process a video' + '*'*10)

    return 


def main(ipath):
    '''
    @: The main function of the script, which runs from here
    para ipath: input image or video path.
    
    '''
    if os.path.isdir(ipath):
        process_image_folder(imgDir=ipath)
    else:
        just_run = [process_video,process_single_image][ipath.split('.')[-1].lower() in IMG_FORMATS]

        just_run(ipath)


    return 

if __name__ == '__main__':
    from sys import argv


    main(argv[1])