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
from pycs.database.Project import Project

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:
        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)
    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] = list()
    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 = [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)