PyCS/pycs/util/file_ops.py

import os
import typing as T

import cv2

from PIL import Image

from pycs.database.File import File

# pylint: disable=too-many-arguments, too-many-locals
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 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
    """

    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=80)

    return os.path.split(target_path)


def resize_file(file: File,
                project_root: str,
                max_width: int,
                max_height: int) -> T.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 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) -> T.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=80)
    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()