import os import typing as T from collections import namedtuple import cv2 from PIL import Image from pycs.database.File import File DEFAULT_JPEG_QUALITY = 80 BoundingBox = namedtuple("BoundingBox", "x y w h") 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