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)