utils/utils.py

import ffmpeg
import cv2
import os

ALLOWED_EXTENSIONS = set(['txt', 'pdf', 'png', 'jpg', 'jpeg', 'gif'])


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def check_rotation(path_video_file):
    meta_dict = ffmpeg.probe(path_video_file)
    rotateCode = None
    print(int(meta_dict['streams'][0]['tags']['rotate']))
    if int(meta_dict['streams'][0]['tags']['rotate']) == 90:
        rotateCode = cv2.ROTATE_180
    elif int(meta_dict['streams'][0]['tags']['rotate']) == 180:
        rotateCode = cv2.ROTATE_90_COUNTERCLOCKWISE
    elif int(meta_dict['streams'][0]['tags']['rotate']) == 270:
        rotateCode = cv2.ROTATE_90_CLOCKWISE
    return rotateCode


def correct_rotation(frame, rotateCode):
    return cv2.rotate(frame, rotateCode)


def create_folder(dir_path):
    if not os.path.isdir(dir_path):
        os.mkdir(dir_path)
    subfolder = ['frames', 'ss', 'sts', 'ir']
    subfolder = [os.path.join(dir_path, subf) for subf in subfolder]
    for subf in subfolder:
        if not os.path.isdir(subf):
            os.mkdir(subf)
    return subfolder