PyCS/pycs/util/FileOperations.py

import os
import typing as T

from collections import namedtuple
from pathlib import Path

import cv2

from PIL import Image
from tqdm import tqdm

from pycs import app
from pycs.database.File import File

DEFAULT_JPEG_QUALITY = 80


BoundingBox = namedtuple("BoundingBox", "x y w h")
Size = namedtuple("Size", "max_width max_height")


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 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 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')
    resized = resize_image(abs_file_path, abs_target_path, max_width, max_height)

    # return path
    if resized:
        return os.path.split(abs_target_path)

    return os.path.split(abs_file_path)


def crop_file(file: File, project_root: str, box: BoundingBox,
              max_width: int, max_height: int) -> T.Tuple[str, 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 max_width: maximum image or thumbnail width
    :param max_height: maximum image or thumbnail height
    :param box: BoundingBox with relative x, y coordinates and
        the relative height and width of the bounding box

    :return: directory and file name of the cropped patch
    """
    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

    bounding_box_suffix = f"{box.x}_{box.y}_{box.w}_{box.h}"
    # crop image file
    abs_target_path = os.path.join(os.getcwd(),
                                   project_root,
                                   'temp',
                                   f'{file.uuid}_{bounding_box_suffix}.jpg')
    cropped = crop_image(abs_file_path, abs_target_path, box)

    if cropped:
        abs_file_path = abs_target_path

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

    if resized:
        abs_file_path = abs_target_path

    # return image
    return os.path.split(abs_file_path)


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()


def resize_image(file_path: str, target_path: str, max_width: int, max_height: int) -> bool:
    """
    Resize an image so width < `max_width` and height < `max_height` applies.
    If the image is already smaller than the given dimensions no new file is stored.

    :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: `True` if a resize operation was performed or the target file already exists
    """
    # 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:
        image.close()
        return False

    # 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)

    # close opened files.
    image.close()
    return True


def crop_image(file_path: str, target_path: str, box: BoundingBox) -> bool:
    """
    Crop an image with the given coordinates, width and height.
    If however no crop is applied no new file is stored.

    :param file_path: path to source file
    :param target_path: path to target file
    :param box: BoundingBox with relative x, y coordinates and
        the relative height and width of the bounding box
    :return: `True` if a crop operation was performed or the target file already exists
    """
    # return if file exists
    if os.path.exists(target_path):
        return True

    # abort if no crop would be applied
    if box.x <= 0 and box.y <= 0 and box.w >= 1 and box.h >= 1:
        return False

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

    # calculate absolute crop position
    crop_x1 = int(img_width * box.x)
    crop_y1 = int(img_height * box.y)
    crop_x2 = min(int(img_width * box.w) + crop_x1, img_width)
    crop_y2 = min(int(img_height * box.h) + crop_y1, img_height)

    # crop image
    cropped_image = image.crop((crop_x1, crop_y1, crop_x2, crop_y2))

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


def find_images(folder,
                suffixes: T.Optional[T.List[str]] = None) -> T.List[Path]:
    """ walk recursively the folder and find images """

    suffixes = suffixes if suffixes is not None else [".jpg", ".jpeg", ".png"]
    images: T.List[Path] = []
    for root, _, files in os.walk(folder):
        for file in files:
            fpath = Path(root, file)
            if fpath.suffix.lower() not in suffixes:
                continue

            images.append(fpath)

    return images


def generate_thumbnails(project: "Project", sizes = None):
    """ generates thumbnails for all image files in the given  """

    if sizes is None:
        sizes = [Size(200, 200), Size(2000, 1200)]

    app.logger.info(f"Generating thumbnails for project \"{project.name}\"")

    files = list(project.files)
    for file in tqdm(files):
        for size in sizes:
            resize_file(file,
                project.root_folder,
                size.max_width,
                size.max_height)