PyCS/pycs/util/FileOperations.py

import os
from typing import Optional
from typing import Tuple

import cv2
from PIL import Image

from pycs.database.File import File

DEFAULT_JPEG_QUALITY = 80


def file_info(data_folder: str, file_name: str, file_ext: str):
    """
    Receive file type, frame count and frames per second.
    The last two are always None for images.

    :param data_folder: path to data folder
    :param file_name: file name
    :param file_ext: file extension
    :return: file type, frame count, frames per second
    """
    # determine file type
    if file_ext.lower() in ['.jpg', '.png']:
        file_type = 'image'

    elif file_ext.lower() in ['.mp4']:
        file_type = 'video'

    else:
        raise ValueError(f"Unsupported file extension: {file_ext}!")

    # determine frames and fps for video files
    if file_type == 'image':
        frames = None
        fps = None
    else:
        file_path = os.path.join(data_folder, file_name + file_ext)
        video = cv2.VideoCapture(file_path)

        frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = video.get(cv2.CAP_PROP_FPS)

        video.release()

    # return values
    return file_type, frames, fps


def crop_file(file: File, project_root: str, x: float, y: float, w: float, h: float) -> str:
    """
    gets a file for the given file_id, crops the according image to the
    bounding box and saves the crops in the temp folder of the project.

    :param file: file object
    :param project_root: project root folder path
    :param x: relative x-coordinate of the top left corner
    :param y: relative y-coordinate of the top left corner
    :param w: relative width of the bounding box
    :param h: relative height of the bounding box

    :return: directory and file name of the cropped patch
    """

    # TODO in browser cache
    # TODO not possible for video files

    image = Image.open(file.absolute_path)
    width, height = image.size

    crop_width = int(width * w)
    crop_height = int(height * h)

    x0 = int(width * x)
    y0 = int(height * y)
    x1 = x0 + crop_width
    y1 = y0 + crop_height

    target_path = os.path.join(
        os.getcwd(),
        project_root,
        'temp',
        f'{file.uuid}_{x0}-{y0}_{x1}-{y1}.jpg',
    )

    if not os.path.exists(target_path):
        crop = image.crop((x0, y0, x0 + crop_width, y0 + crop_height))
        crop.save(target_path, quality=DEFAULT_JPEG_QUALITY)

    return os.path.split(target_path)


def resize_file(file: File, project_root: str, max_width: int, max_height: int) -> Tuple[str, str]:
    """
    If file type equals video this function extracts a thumbnail first. It calls resize_image
    to resize and returns the resized files directory and name.

    :param file: file object
    :param project_root: project root folder path
    :param max_width: maximum image or thumbnail width
    :param max_height: maximum image or thumbnail height
    :return: resized file directory, resized file name
    """

    abs_file_path = file.absolute_path

    # extract video thumbnail
    if file.type == 'video':
        abs_target_path = os.path.join(os.getcwd(), project_root, 'temp', f'{file.uuid}.jpg')
        create_thumbnail(abs_file_path, abs_target_path)

        abs_file_path = abs_target_path

    # resize image file
    abs_target_path = os.path.join(os.getcwd(), project_root,
                                   'temp', f'{file.uuid}_{max_width}_{max_height}.jpg')
    result = resize_image(abs_file_path, abs_target_path, max_width, max_height)

    # return path
    if result is not None:
        return os.path.split(abs_target_path)

    return os.path.split(abs_file_path)


def resize_image(file_path: str, target_path: str,
                 max_width: int, max_height: int) -> Optional[bool]:
    """
    resize an image so width < max_width and height < max_height

    :param file_path: path to source file
    :param target_path: path to target file
    :param max_width: maximum image width
    :param max_height: maximum image height
    :return:
    """
    # return if file exists
    if os.path.exists(target_path):
        return True

    # load full size image
    image = Image.open(file_path)
    img_width, img_height = image.size

    # abort if file is smaller than desired
    if img_width < max_width and img_height < max_height:
        return None

    # calculate target size
    target_width = int(max_width)
    target_height = int(max_width * img_height / img_width)

    if target_height > max_height:
        target_height = int(max_height)
        target_width = int(max_height * img_width / img_height)

    # resize image
    resized_image = image.resize((target_width, target_height))

    # save to file
    resized_image.save(target_path, quality=DEFAULT_JPEG_QUALITY)
    return True


def create_thumbnail(file_path: str, target_path: str) -> None:
    """
    extract a thumbnail from a video

    :param file_path: path to source file
    :param target_path: path to target file
    :return:
    """
    # return if file exists
    if os.path.exists(target_path):
        return

    # load video
    video = cv2.VideoCapture(file_path)

    # create thumbnail
    _, image = video.read()
    cv2.imwrite(target_path, image)

    # close video file
    video.release()