From 138344b4112355e71515c5f1fca7ddc256dddc5a Mon Sep 17 00:00:00 2001 From: Bryan Gotti Date: Sat, 1 Nov 2025 01:40:48 +0100 Subject: [PATCH] run black format, allow input subdirectories, fix bugs --- watermark/README.md | 14 +- watermark/helpers/__init__.py | 2 +- watermark/helpers/file_operations.py | 62 ++-- watermark/helpers/image_manipulation.py | 445 +++++++++++++----------- watermark/helpers/others.py | 248 ++++++++----- watermark/watermark.py | 263 ++++++++------ 6 files changed, 614 insertions(+), 420 deletions(-) diff --git a/watermark/README.md b/watermark/README.md index ca9071c..8495112 100644 --- a/watermark/README.md +++ b/watermark/README.md @@ -4,7 +4,17 @@ Par David Resin, ESN EPF Lausanne, 2021-2023 ## Avant l'utilisation -1. Installer Python 3 +1. Installer Python 3.9 2. Exécuter `pip install -r requirements.txt` -## Pour utiliser \ No newline at end of file +## Pour utiliser + +1. Insérer les images à traiter dans le dossier `input/`. +2. Exécuter le script avec la commande suivante : + ```bash + python watermark.py + ``` +3. Les images traitées seront sauvegardées dans le dossier `output/`. Les images invalides seront déplacées dans le dossier `invalid/`. + +> [!NOTE] +> Les images déposées dans le dossier `input/` peuvent être organisées en sous-dossiers. La structure des dossiers sera préservée dans les dossiers `output/` et `invalid/`. \ No newline at end of file diff --git a/watermark/helpers/__init__.py b/watermark/helpers/__init__.py index 30ba21c..e273aae 100644 --- a/watermark/helpers/__init__.py +++ b/watermark/helpers/__init__.py @@ -1,3 +1,3 @@ from .file_operations import * from .image_manipulation import * -from .others import * \ No newline at end of file +from .others import * diff --git a/watermark/helpers/file_operations.py b/watermark/helpers/file_operations.py index f8f64fe..43bb4ae 100644 --- a/watermark/helpers/file_operations.py +++ b/watermark/helpers/file_operations.py @@ -1,4 +1,5 @@ # Default libraries +import os import shutil from pathlib import Path @@ -10,24 +11,24 @@ from helpers.image_manipulation import tilt_img -OTHER_EXTS = ('.jpg', '.png', '.jpeg', '.ico', '.webp') -HEI_EXTS = ('.heic', '.heif') -RAWPY_EXTS = ('.nef',) +OTHER_EXTS = (".jpg", ".png", ".jpeg", ".ico", ".webp") +HEI_EXTS = (".heic", ".heif") +RAWPY_EXTS = (".nef",) IMG_EXTS = OTHER_EXTS + HEI_EXTS + RAWPY_EXTS -IGNORE_EXTS = ('.ds_store') +IGNORE_EXTS = ".ds_store" INVALID_COUNT = 0 # Glob all filenames in a given path with a given pattern, but exclude the patterns in the exclusion list def glob_all_except(path, base_pattern="*", excluded_patterns=[]): - matches = set(path.glob(base_pattern)) + matches = set(path.glob(base_pattern)) - for pattern in excluded_patterns: - matches = matches - set(path.glob(pattern)) + for pattern in excluded_patterns: + matches = matches - set(path.glob(pattern)) - return list(matches) + return list(matches) def extension_match(image_path, extension_list): @@ -36,26 +37,26 @@ def extension_match(image_path, extension_list): # Flush the output directory def flush_output(path_out: Path, exts: tuple[str]) -> None: - for deletion_candidate in path_out.iterdir(): - if extension_match(deletion_candidate, exts): - deletion_candidate.unlink() + for deletion_candidate in path_out.iterdir(): + if extension_match(deletion_candidate, exts): + deletion_candidate.unlink() # Move an invalid picture out def invalidate_path(image_path, path_invalid): - global INVALID_COUNT - shutil.move(image_path, path_invalid) - INVALID_COUNT += 1 + global INVALID_COUNT + shutil.move(image_path, path_invalid) + INVALID_COUNT += 1 # Create a directory if it is missing def create_dir_if_missing(dir_path): - try: - dir_path.mkdir() - except FileExistsError: - return True + try: + dir_path.mkdir() + except FileExistsError: + return True - return False + return False def open_rawpy_image(image_path): @@ -73,11 +74,11 @@ def open_hei_image(image_path): def universal_load_image(image_path): image = None - flag = 'img' + flag = "img" is_hei = False if extension_match(image_path, IGNORE_EXTS): - flag = 'ignore' + flag = "ignore" elif extension_match(image_path, RAWPY_EXTS): image = open_rawpy_image(image_path) elif extension_match(image_path, HEI_EXTS): @@ -86,14 +87,14 @@ def universal_load_image(image_path): elif extension_match(image_path, OTHER_EXTS): image = Image.open(image_path) else: - flag = 'invalid' + flag = "invalid" return image, flag, is_hei def attempt_open_image_attempt_tilt(image): # TODO : Implement no-tilt option in CLI arguments - # For non-HEI file types, try to re-orient the picture if it is allowed and orientation data is available + # For non-HEI file types, try to re-orient the picture if it is allowed and orientation data is available try: # TODO : Potentiellement getexif existe partout alors que _getexif pas image._getexif() @@ -101,19 +102,26 @@ def attempt_open_image_attempt_tilt(image): except AttributeError: # TODO : Setting for what to do when image can't be rotated pass - + return image def attempt_open_image(image_path, path_invalid, attempt_rotate): image, flag, is_hei = universal_load_image(image_path) - if flag == 'ignore': + if flag == "ignore": return None - elif flag == 'invalid': - invalidate_path(image_path, path_inv) + elif flag == "invalid": + invalidate_path(image_path, path_invalid) if not is_hei and attempt_rotate: image = attempt_open_image_attempt_tilt(image) return image + + +def scandir(dirname): + subfolders = [f.path for f in os.scandir(dirname) if f.is_dir()] + for dirname in list(subfolders): + subfolders.extend(scandir(dirname)) + return subfolders diff --git a/watermark/helpers/image_manipulation.py b/watermark/helpers/image_manipulation.py index 44a35ef..b6cd5f4 100644 --- a/watermark/helpers/image_manipulation.py +++ b/watermark/helpers/image_manipulation.py @@ -2,262 +2,305 @@ from PIL import ImageDraw # Custom libraries -from helpers.others import color_mapping_from_setting # Needs to disappear +from helpers.others import color_mapping_from_setting # Needs to disappear TILT_MAP = { - 0: 0, - 1: 180, - 2: 270, - 3: 90, + 0: 0, + 1: 180, + 2: 270, + 3: 90, } EXIF_ORIENTATION_TAG = 274 ESN_CIRCLE_COLOR_MAP = { - "white": "color", - None: "white", + "white": "color", + None: "white", } def get_any_dict_value(dictionary): - return next(iter(dictionary.values())) + return next(iter(dictionary.values())) def get_dict_value_or_none_value(dictionary: dict, key): - return dictionary.get(key, dictionary[None]) + return dictionary.get(key, dictionary[None]) # Auromatically tilt an image based on its EXIF data def tilt_img(image): - try: - exif = image._getexif() - except AttributeError: - exif = image.getexif() + try: + exif = image._getexif() + except AttributeError: + exif = image.getexif() - tilt = exif.get(EXIF_ORIENTATION_TAG) + tilt = exif.get(EXIF_ORIENTATION_TAG) - # Don't tilt if exif orientation value is not between 1 and 8 - if tilt is None or tilt not in range(1, 9): - return image - - tilt_idx = TILT_MAP[(tilt - 1) // 2] - return image.rotate(tilt_idx, expand=True) + # Don't tilt if exif orientation value is not between 1 and 8 + if tilt is None or tilt not in range(1, 9): + return image + + tilt_idx = TILT_MAP[(tilt - 1) // 2] + return image.rotate(tilt_idx, expand=True) def logo_dims_from_image_and_ratio(logo_size, image_size, image_watermark_ratio): - image_w, image_h = image_size - src_logo_w, src_logo_h = logo_size + image_w, image_h = image_size + src_logo_w, src_logo_h = logo_size - # TODO : Make this configurable (h or w) - tgt_logo_h = image_watermark_ratio * min(image_h, image_w) - tgt_logo_w = tgt_logo_h / src_logo_h * src_logo_w + # TODO : Make this configurable (h or w) + tgt_logo_h = image_watermark_ratio * min(image_h, image_w) + tgt_logo_w = tgt_logo_h / src_logo_h * src_logo_w - return tgt_logo_w, tgt_logo_h + return tgt_logo_w, tgt_logo_h def nearest_integer_scale(values, scale_factor): - return [int(scale_factor * value) for value in values] + return [int(scale_factor * value) for value in values] def scale_logos_with_supersampling(logos, target_dims, ss_factor=1): - ss_dims = nearest_integer_scale(target_dims, scale_factor=ss_factor) - return {key: logo.resize(ss_dims) for key, logo in logos.items()} + ss_dims = nearest_integer_scale(target_dims, scale_factor=ss_factor) + return {key: logo.resize(ss_dims) for key, logo in logos.items()} def crop_image_with_supersampling(image, bbox, ss_factor=1): - crop = image.crop(box=bbox) - ss_dims = nearest_integer_scale(crop.size, scale_factor=ss_factor) - return crop.resize(ss_dims) + crop = image.crop(box=bbox) + ss_dims = nearest_integer_scale(crop.size, scale_factor=ss_factor) + return crop.resize(ss_dims) def draw_ellipse_with_supersampling(image, bbox, color, ss_factor=1): - ss_bbox = nearest_integer_scale(bbox, scale_factor=ss_factor) - ImageDraw.Draw(image).ellipse(ss_bbox, fill=color) - return image + ss_bbox = nearest_integer_scale(bbox, scale_factor=ss_factor) + ImageDraw.Draw(image).ellipse(ss_bbox, fill=color) + return image def dims_from_bbox(bbox): - x0, y0, x1, y1 = bbox - return x1 - x0, y1 - y0 + x0, y0, x1, y1 = bbox + return x1 - x0, y1 - y0 def resize_to_bbox_size(image, bbox): - dims = dims_from_bbox(bbox) - return image.resize(dims) + dims = dims_from_bbox(bbox) + return image.resize(dims) def paste_image_on_image_at_bbox(image, pasted_image, bbox, copy_image=False): - new_image = image.copy() if copy_image else image - mask = pasted_image if pasted_image.mode == "RGBA" else None - new_image.paste(pasted_image, bbox[:2], mask) - return new_image + new_image = image.copy() if copy_image else image + mask = pasted_image if pasted_image.mode == "RGBA" else None + new_image.paste(pasted_image, bbox[:2], mask) + return new_image # Watermark an image with a given position and color -def generate_watermarked_image(image, logo_ss, circle_color, ss_factor, positioning_data): - watermark_canvas_ss = crop_image_with_supersampling(image, - bbox=positioning_data["watermark_bbox"], - ss_factor=ss_factor) - - if circle_color is not None: - watermark_canvas_ss = draw_ellipse_with_supersampling( watermark_canvas_ss, - bbox=positioning_data["circle_bbox_in_watermark_bbox"], - color=circle_color, - ss_factor=ss_factor) - - watermark_canvas_ss = paste_image_on_image_at_bbox( watermark_canvas_ss, - pasted_image=logo_ss, - bbox=positioning_data["logo_pos_in_watermark_ss_bbox"]) - - watermark_canvas = resize_to_bbox_size( watermark_canvas_ss, - bbox=positioning_data["watermark_bbox"]) - - watermarked_image = paste_image_on_image_at_bbox( image, - pasted_image=watermark_canvas, - bbox=positioning_data["watermark_bbox"], - copy_image=True) - - return watermarked_image +def generate_watermarked_image( + image, logo_ss, circle_color, ss_factor, positioning_data +): + watermark_canvas_ss = crop_image_with_supersampling( + image, bbox=positioning_data["watermark_bbox"], ss_factor=ss_factor + ) + + if circle_color is not None: + watermark_canvas_ss = draw_ellipse_with_supersampling( + watermark_canvas_ss, + bbox=positioning_data["circle_bbox_in_watermark_bbox"], + color=circle_color, + ss_factor=ss_factor, + ) + + watermark_canvas_ss = paste_image_on_image_at_bbox( + watermark_canvas_ss, + pasted_image=logo_ss, + bbox=positioning_data["logo_pos_in_watermark_ss_bbox"], + ) + + watermark_canvas = resize_to_bbox_size( + watermark_canvas_ss, bbox=positioning_data["watermark_bbox"] + ) + + watermarked_image = paste_image_on_image_at_bbox( + image, + pasted_image=watermark_canvas, + bbox=positioning_data["watermark_bbox"], + copy_image=True, + ) + + return watermarked_image # Watermark an image with a given position and a list of colors def watermark_image_pos(image, path, logos_ss, settings, positioning_data): - color_mapping = color_mapping_from_setting(settings["color_setting"]) - - # Loop through the selected colors - for i, (color_name, color) in enumerate(color_mapping.items()): - - logo_ss = logos_ss[get_dict_value_or_none_value(ESN_CIRCLE_COLOR_MAP, color_name)] - circle_color = color if settings["draw_circle"] else None - - watermarked_image = generate_watermarked_image( image, - logo_ss=logo_ss, - circle_color=circle_color, - ss_factor=settings["ss_factor"], - positioning_data=positioning_data) - - if len(color_mapping) == 1: - suffix = "" - else: - suffix = "_" + str(i) - - path_out = settings["output_path"] / (settings["prefix"] + path.stem + suffix + "." + settings["format"]) - watermarked_image.save(path_out, format="png", compress_level=4) - - -def compute_positioning_data(image_size, logo_ss_size, position_str, positioning_settings, ss_factor): - image_w, image_h = image_size - logo_ss_w, logo_ss_h = logo_ss_size - - # Extract positioning settings for code readability - logo_paddings = positioning_settings["logo_paddings"] - circle_radius = positioning_settings["circle_radius"] - circle_offset_abs = positioning_settings["circle_offset_abs"] - - # Set logo center depending on the chosen position - logo_center_x = logo_paddings[0] if "left" in position_str else image_w - logo_paddings[0] - logo_center_y = logo_paddings[1] if "top" in position_str else image_h - logo_paddings[1] - - # Set direction of offsets depending on the chosen position - circle_offset_x = circle_offset_abs[0] * ((-1) ** ("left" in position_str)) - circle_offset_y = circle_offset_abs[1] * ((-1) ** ("top" in position_str)) - - # Get circle center from logo center and offset - circle_center_x = logo_center_x + circle_offset_x - circle_center_y = logo_center_y + circle_offset_y - - # Get circle bounding box from the center and radius - circle_bbox_xs = circle_center_x - circle_radius, circle_center_x + circle_radius - circle_bbox_ys = circle_center_y - circle_radius, circle_center_y + circle_radius - - # Get watermark bounding box by excluding out-of-bounds parts of the circle - watermark_bbox_xs = max(0, circle_bbox_xs[0]), min(image_size[0], circle_bbox_xs[1]) - watermark_bbox_ys = max(0, circle_bbox_ys[0]), min(image_size[1], circle_bbox_ys[1]) - - # Assemble watermark bounding box - watermark_bbox = tuple([int(elem) for elem in [ - watermark_bbox_xs[0], - watermark_bbox_ys[0], - watermark_bbox_xs[1], - watermark_bbox_ys[1], - ]]) - - # Get logo center relative to the watermark - logo_center_in_watermark_x = logo_center_x - watermark_bbox_xs[0] - logo_center_in_watermark_y = logo_center_y - watermark_bbox_ys[0] - - # Get logo top-left corner relative to supersampled watermark - logo_pos_in_watermark_ss_x = int(ss_factor * logo_center_in_watermark_x - logo_ss_w / 2) - logo_pos_in_watermark_ss_y = int(ss_factor * logo_center_in_watermark_y - logo_ss_h / 2) - - # Get circle sub-bounding box - circle_bbox_in_watermark_xs = [x - watermark_bbox_xs[0] for x in circle_bbox_xs] - circle_bbox_in_watermark_ys = [y - watermark_bbox_ys[0] for y in circle_bbox_ys] - - # Assemble circle sub-bounding box - circle_bbox_in_watermark = tuple([int(elem) for elem in [ - circle_bbox_in_watermark_xs[0], - circle_bbox_in_watermark_ys[0], - circle_bbox_in_watermark_xs[1], - circle_bbox_in_watermark_ys[1], - ]]) - - return { - "watermark_bbox": watermark_bbox, - "circle_bbox_in_watermark_bbox": circle_bbox_in_watermark, - "logo_pos_in_watermark_ss_bbox": ( - logo_pos_in_watermark_ss_x, - logo_pos_in_watermark_ss_y, - ), - } + color_mapping = color_mapping_from_setting(settings["color_setting"]) + + # Loop through the selected colors + for i, (color_name, color) in enumerate(color_mapping.items()): + + logo_ss = logos_ss[ + get_dict_value_or_none_value(ESN_CIRCLE_COLOR_MAP, color_name) + ] + circle_color = color if settings["draw_circle"] else None + + watermarked_image = generate_watermarked_image( + image, + logo_ss=logo_ss, + circle_color=circle_color, + ss_factor=settings["ss_factor"], + positioning_data=positioning_data, + ) + + if len(color_mapping) == 1: + suffix = "" + else: + suffix = "_" + str(i) + + path_out = settings["output_path"] / ( + settings["prefix"] + path.stem + suffix + "." + settings["format"] + ) + watermarked_image.save(path_out, format="png", compress_level=4) + + +def compute_positioning_data( + image_size, logo_ss_size, position_str, positioning_settings, ss_factor +): + image_w, image_h = image_size + logo_ss_w, logo_ss_h = logo_ss_size + + # Extract positioning settings for code readability + logo_paddings = positioning_settings["logo_paddings"] + circle_radius = positioning_settings["circle_radius"] + circle_offset_abs = positioning_settings["circle_offset_abs"] + + # Set logo center depending on the chosen position + logo_center_x = ( + logo_paddings[0] if "left" in position_str else image_w - logo_paddings[0] + ) + logo_center_y = ( + logo_paddings[1] if "top" in position_str else image_h - logo_paddings[1] + ) + + # Set direction of offsets depending on the chosen position + circle_offset_x = circle_offset_abs[0] * ((-1) ** ("left" in position_str)) + circle_offset_y = circle_offset_abs[1] * ((-1) ** ("top" in position_str)) + + # Get circle center from logo center and offset + circle_center_x = logo_center_x + circle_offset_x + circle_center_y = logo_center_y + circle_offset_y + + # Get circle bounding box from the center and radius + circle_bbox_xs = circle_center_x - circle_radius, circle_center_x + circle_radius + circle_bbox_ys = circle_center_y - circle_radius, circle_center_y + circle_radius + + # Get watermark bounding box by excluding out-of-bounds parts of the circle + watermark_bbox_xs = max(0, circle_bbox_xs[0]), min(image_size[0], circle_bbox_xs[1]) + watermark_bbox_ys = max(0, circle_bbox_ys[0]), min(image_size[1], circle_bbox_ys[1]) + + # Assemble watermark bounding box + watermark_bbox = tuple( + [ + int(elem) + for elem in [ + watermark_bbox_xs[0], + watermark_bbox_ys[0], + watermark_bbox_xs[1], + watermark_bbox_ys[1], + ] + ] + ) + + # Get logo center relative to the watermark + logo_center_in_watermark_x = logo_center_x - watermark_bbox_xs[0] + logo_center_in_watermark_y = logo_center_y - watermark_bbox_ys[0] + + # Get logo top-left corner relative to supersampled watermark + logo_pos_in_watermark_ss_x = int( + ss_factor * logo_center_in_watermark_x - logo_ss_w / 2 + ) + logo_pos_in_watermark_ss_y = int( + ss_factor * logo_center_in_watermark_y - logo_ss_h / 2 + ) + + # Get circle sub-bounding box + circle_bbox_in_watermark_xs = [x - watermark_bbox_xs[0] for x in circle_bbox_xs] + circle_bbox_in_watermark_ys = [y - watermark_bbox_ys[0] for y in circle_bbox_ys] + + # Assemble circle sub-bounding box + circle_bbox_in_watermark = tuple( + [ + int(elem) + for elem in [ + circle_bbox_in_watermark_xs[0], + circle_bbox_in_watermark_ys[0], + circle_bbox_in_watermark_xs[1], + circle_bbox_in_watermark_ys[1], + ] + ] + ) + + return { + "watermark_bbox": watermark_bbox, + "circle_bbox_in_watermark_bbox": circle_bbox_in_watermark, + "logo_pos_in_watermark_ss_bbox": ( + logo_pos_in_watermark_ss_x, + logo_pos_in_watermark_ss_y, + ), + } # Watermark an image with a list of positions and a list of colors def watermark_image(image, path, logos, position_list, settings): - # Compute logo dimensions from image dimensions and image-watermark ratio - target_logo_w, target_logo_h = logo_dims_from_image_and_ratio( logo_size=logos["color"].size, - image_size=image.size, - image_watermark_ratio=settings["image_watermark_ratio"]) - - # Get scaled and supersampled logos - logos_ss = scale_logos_with_supersampling( logos=logos, - target_dims=(target_logo_w, target_logo_h), - ss_factor=settings["ss_factor"]) - - # Get logo padding from padding ratio and logo height - # TODO : Make this configurable (h or w) - logo_padding = target_logo_h * settings["logo_padding_ratio"] - - # Get circle radius from logo width and logo-circle ratio - # TODO : Make this configurable (h or w) - circle_radius = target_logo_w * settings["logo_circle_ratio"] / 2 - - # Get logo center - logo_padding_x = logo_padding + target_logo_w / 2 - logo_padding_y = logo_padding + target_logo_h / 2 - - # Get absolute circle offset - circle_offset_abs_x = target_logo_w * (settings["circle_offset_ratio_x"] - .5) - circle_offset_abs_y = target_logo_h * (settings["circle_offset_ratio_y"] - .5) - - positioning_settings = { - "logo_paddings": (logo_padding_x, logo_padding_y), - "circle_offset_abs": (circle_offset_abs_x, circle_offset_abs_y), - "circle_radius": circle_radius, - } - - # Iterate through the given positions - for position_str in position_list: - # Get positioning data - positioning_data = compute_positioning_data(image_size=image.size, - logo_ss_size=get_any_dict_value(logos_ss).size, - position_str=position_str, - positioning_settings=positioning_settings, - ss_factor=settings["ss_factor"]) - - watermark_image_pos(image, - path=path, - logos_ss=logos_ss, - positioning_data=positioning_data, - settings=settings) + # Compute logo dimensions from image dimensions and image-watermark ratio + target_logo_w, target_logo_h = logo_dims_from_image_and_ratio( + logo_size=logos["color"].size, + image_size=image.size, + image_watermark_ratio=settings["image_watermark_ratio"], + ) + + # Get scaled and supersampled logos + logos_ss = scale_logos_with_supersampling( + logos=logos, + target_dims=(target_logo_w, target_logo_h), + ss_factor=settings["ss_factor"], + ) + + # Get logo padding from padding ratio and logo height + # TODO : Make this configurable (h or w) + logo_padding = target_logo_h * settings["logo_padding_ratio"] + + # Get circle radius from logo width and logo-circle ratio + # TODO : Make this configurable (h or w) + circle_radius = target_logo_w * settings["logo_circle_ratio"] / 2 + + # Get logo center + logo_padding_x = logo_padding + target_logo_w / 2 + logo_padding_y = logo_padding + target_logo_h / 2 + + # Get absolute circle offset + circle_offset_abs_x = target_logo_w * (settings["circle_offset_ratio_x"] - 0.5) + circle_offset_abs_y = target_logo_h * (settings["circle_offset_ratio_y"] - 0.5) + + positioning_settings = { + "logo_paddings": (logo_padding_x, logo_padding_y), + "circle_offset_abs": (circle_offset_abs_x, circle_offset_abs_y), + "circle_radius": circle_radius, + } + + # Iterate through the given positions + for position_str in position_list: + # Get positioning data + positioning_data = compute_positioning_data( + image_size=image.size, + logo_ss_size=get_any_dict_value(logos_ss).size, + position_str=position_str, + positioning_settings=positioning_settings, + ss_factor=settings["ss_factor"], + ) + + watermark_image_pos( + image, + path=path, + logos_ss=logos_ss, + positioning_data=positioning_data, + settings=settings, + ) diff --git a/watermark/helpers/others.py b/watermark/helpers/others.py index f3ab48f..8892200 100644 --- a/watermark/helpers/others.py +++ b/watermark/helpers/others.py @@ -9,104 +9,196 @@ COLOR_OPTIONS = { - "white": (255, 255, 255), - "black": ( 0, 0, 0), - "magenta": (236, 0, 140), - "orange": (244, 123, 32), - "green": (122, 193, 67), - "cyan": ( 0, 174, 239), - "purple": ( 46, 49, 146), + "white": (255, 255, 255), + "black": (0, 0, 0), + "magenta": (236, 0, 140), + "orange": (244, 123, 32), + "green": (122, 193, 67), + "cyan": (0, 174, 239), + "purple": (46, 49, 146), } POSITION_OPTIONS = [ - "bottom_right", - "bottom_left", - "top_right", - "top_left", - "random", - "all", + "bottom_right", + "bottom_left", + "top_right", + "top_left", + "random", + "all", ] # Color parsing def color_names_list_from_setting(color_setting): - if color_setting == "random": - return [random.choice(list(COLOR_OPTIONS.keys()))] - elif color_setting == "all": - return list(COLOR_OPTIONS.keys()) - else: - return [color_setting] + if color_setting == "random": + return [random.choice(list(COLOR_OPTIONS.keys()))] + elif color_setting == "all": + return list(COLOR_OPTIONS.keys()) + else: + return [color_setting] def color_mapping_from_setting(color_setting): - ret = dict() + ret = dict() - for color_name in color_names_list_from_setting(color_setting): - color = COLOR_OPTIONS.get(color_name) + for color_name in color_names_list_from_setting(color_setting): + color = COLOR_OPTIONS.get(color_name) - if color is None: - try: - print(color_name, color) - color = ImageColor.getrgb(color_name) - except ValueError: - sys.exit("Wrong color format. Official ESN color or #rrggbb hexadecimal format expected.") + if color is None: + try: + print(color_name, color) + color = ImageColor.getrgb(color_name) + except ValueError: + sys.exit( + "Wrong color format. Official ESN color or #rrggbb hexadecimal format expected." + ) - ret[color_name] = color - - return ret + ret[color_name] = color + + return ret # Generate list of positions based on arguments def position_list_from_setting(position): - if position == "random": - ret = [POSITION_OPTIONS[random.randint(0, 3)]] - elif position == "all": - ret = POSITION_OPTIONS[:4] - else: - ret = [position] + if position == "random": + ret = [POSITION_OPTIONS[random.randint(0, 3)]] + elif position == "all": + ret = POSITION_OPTIONS[:4] + else: + ret = [position] - return ret + return ret # Setup argument parser def setup_argparser(default_vals, color_options, pos_choices): - ap = argparse.ArgumentParser(description="ESN Lausanne Watermark Inserter", formatter_class=argparse.RawTextHelpFormatter) - ap.add_argument("-f", "--flush", action="store_true", help="flush output folder") - ap.add_argument("-np", "--no-prefix", action="store_true", help="do not add a '{}' prefix to outputs".format(default_vals["wm_prefix"])) - ap.add_argument("-nr", "--no-rotate", action="store_true", help="do not rotate images if they are not upright") - ap.add_argument("-nc", "--no-circle", action="store_true", help="do not add a colored circle behind the logo (not recommended)") - ap.add_argument("-cc", "--center-circle", action="store_true", help="center the circle around the logo (not recommended)") - ap.add_argument("-i", "--input-dir", action="store", type=str, default=default_vals["input_dir"], help="set a custom input directory path (default is '{}')".format(default_vals["input_dir"])) - ap.add_argument("-o", "--output-dir", action="store", type=str, default=default_vals["output_dir"], help="set a custom output directory path (default is '{}')".format(default_vals["output_dir"])) - ap.add_argument("-wms", "--watermark-size", action="store", type=float, default=default_vals["wm_size"], help="set the size of the watermark compared to the image's size (default is {})".format(default_vals["wm_size"])) - ap.add_argument("-wmr", "--watermark-ratio", action="store", type=float, default=default_vals["wm_ratio"], help="set the size ratio between the logo's width and the circle's diameter, (default is {})".format(default_vals["wm_ratio"])) - ap.add_argument("-wmp", "--watermark-padding", action="store", type=float, default=default_vals["wm_pad"], help="set the padding between the logo and the edge of the picture, as a ratio of the logo's height (default is {})".format(default_vals["wm_pad"])) - ap.add_argument("-ss", "--supersampling", action="store", type=int, default=default_vals["ss_factor"], metavar="FACTOR", help="set the supersampling factor for smoothing the circle (default is {}, smaller means faster execution but less smoothing)".format(default_vals["ss_factor"])) - ap.add_argument("-c", "--color", action="store", - type=str, - default="random", - help=textwrap.dedent("set the color of the circle, options are the following:\n" - + "> 'random' [Random color for each image, default value]\n" - + "> official ESN colors:\n" - + "\t'white'\n" - + "\t'black'\n" - + "\t'magenta'\n" - + "\t'orange'\n" - + "\t'green'\n" - + "\t'cyan'\n" - + "\t'purple'\n" - + "> 'all' [All versions of each image with suffixes]\n" - + "> '#rrggbb' [Any other HEX RGB color, not recommended]")) - ap.add_argument("-p", "--position", type=str, - metavar="POSITION", - default=pos_choices[0], - choices=pos_choices, - help=textwrap.dedent("set the position of the watermark, options are the following:\n" - + "> 'bottom_right' [Default value]\n" - + "> 'bottom_left'\n" - + "> 'top_right'\n" - + "> 'top_left'\n" - + "> 'random' [Random edge for each image]\n" - + "> 'all' [All versions of each image with suffixes]")) - - return ap + ap = argparse.ArgumentParser( + description="ESN Lausanne Watermark Inserter", + formatter_class=argparse.RawTextHelpFormatter, + ) + ap.add_argument("-f", "--flush", action="store_true", help="flush output folder") + ap.add_argument( + "-np", + "--no-prefix", + action="store_true", + help="do not add a '{}' prefix to outputs".format(default_vals["wm_prefix"]), + ) + ap.add_argument( + "-nr", + "--no-rotate", + action="store_true", + help="do not rotate images if they are not upright", + ) + ap.add_argument( + "-nc", + "--no-circle", + action="store_true", + help="do not add a colored circle behind the logo (not recommended)", + ) + ap.add_argument( + "-cc", + "--center-circle", + action="store_true", + help="center the circle around the logo (not recommended)", + ) + ap.add_argument( + "-i", + "--input-dir", + action="store", + type=str, + default=default_vals["input_dir"], + help="set a custom input directory path (default is '{}')".format( + default_vals["input_dir"] + ), + ) + ap.add_argument( + "-o", + "--output-dir", + action="store", + type=str, + default=default_vals["output_dir"], + help="set a custom output directory path (default is '{}')".format( + default_vals["output_dir"] + ), + ) + ap.add_argument( + "-wms", + "--watermark-size", + action="store", + type=float, + default=default_vals["wm_size"], + help="set the size of the watermark compared to the image's size (default is {})".format( + default_vals["wm_size"] + ), + ) + ap.add_argument( + "-wmr", + "--watermark-ratio", + action="store", + type=float, + default=default_vals["wm_ratio"], + help="set the size ratio between the logo's width and the circle's diameter, (default is {})".format( + default_vals["wm_ratio"] + ), + ) + ap.add_argument( + "-wmp", + "--watermark-padding", + action="store", + type=float, + default=default_vals["wm_pad"], + help="set the padding between the logo and the edge of the picture, as a ratio of the logo's height (default is {})".format( + default_vals["wm_pad"] + ), + ) + ap.add_argument( + "-ss", + "--supersampling", + action="store", + type=int, + default=default_vals["ss_factor"], + metavar="FACTOR", + help="set the supersampling factor for smoothing the circle (default is {}, smaller means faster execution but less smoothing)".format( + default_vals["ss_factor"] + ), + ) + ap.add_argument( + "-c", + "--color", + action="store", + type=str, + default="random", + help=textwrap.dedent( + "set the color of the circle, options are the following:\n" + + "> 'random' [Random color for each image, default value]\n" + + "> official ESN colors:\n" + + "\t'white'\n" + + "\t'black'\n" + + "\t'magenta'\n" + + "\t'orange'\n" + + "\t'green'\n" + + "\t'cyan'\n" + + "\t'purple'\n" + + "> 'all' [All versions of each image with suffixes]\n" + + "> '#rrggbb' [Any other HEX RGB color, not recommended]" + ), + ) + ap.add_argument( + "-p", + "--position", + type=str, + metavar="POSITION", + default=pos_choices[0], + choices=pos_choices, + help=textwrap.dedent( + "set the position of the watermark, options are the following:\n" + + "> 'bottom_right' [Default value]\n" + + "> 'bottom_left'\n" + + "> 'top_right'\n" + + "> 'top_left'\n" + + "> 'random' [Random edge for each image]\n" + + "> 'all' [All versions of each image with suffixes]" + ), + ) + + return ap diff --git a/watermark/watermark.py b/watermark/watermark.py index d823d09..16c3aa6 100644 --- a/watermark/watermark.py +++ b/watermark/watermark.py @@ -13,121 +13,162 @@ glob_all_except, flush_output, attempt_open_image, - IMG_EXTS + IMG_EXTS, + scandir, ) -from helpers.others import ( # Needs to become a * import - setup_argparser, - position_list_from_setting, - color_mapping_from_setting, - COLOR_OPTIONS, - POSITION_OPTIONS +from helpers.others import ( # Needs to become a * import + setup_argparser, + position_list_from_setting, + color_mapping_from_setting, + COLOR_OPTIONS, + POSITION_OPTIONS, ) # Main code if __name__ == "__main__": - # TODO : Configure dry run - - print("Start") - - root_path = Path() - logo_path = root_path / "logos" - - # Define default values - default_values = { - "ss_factor": 2, - "wm_size": 0.07, - "wm_ratio": 1.6, - "wm_pad": 0.15, - "wm_prefix": "wm_", - "input_dir": "input", - "output_dir": "output", - "format": "png", - } - - # Other parameters - str_process = "Processing image {} of {} \t({} of {} variations, \tinvalid: {})" - str_end = "Processed {} images successfully!" - str_invalid = " ({} image(s) failed and moved to 'invalid')" - - # Names of logo files - fn_color = "logo_color.png" - fn_white = "logo_white.png" - - # Open logo files - logos = { - "color": Image.open(logo_path / fn_color), - "white": Image.open(logo_path / fn_white), - } - - # Setup argparser and parse arguments - ap = setup_argparser(default_vals=default_values, color_options=COLOR_OPTIONS, pos_choices=POSITION_OPTIONS) - args = vars(ap.parse_args()) - - prefix = "" if args["no_prefix"] else default_values["wm_prefix"] - path_input = root_path / args["input_dir"] - path_output = root_path / args["output_dir"] - path_invalid = root_path / "invalid" - position_setting = args["position"] - - settings = { - "image_watermark_ratio": args["watermark_size"], - "logo_padding_ratio": args["watermark_padding"], - "logo_circle_ratio": args["watermark_ratio"], - "circle_offset_ratio_x": .5 if args["center_circle"] else 3 / 5, - "circle_offset_ratio_y": .5 if args["center_circle"] else 1, - "ss_factor": args["supersampling"], - "draw_circle": not args["no_circle"], - "output_path": path_output, - "color_setting": args["color"], - "prefix": prefix, # TODO : Offer option to customise the prefix - "format": default_values["format"] # TODO : Offer option to change output format - } - - # Create missing folders if needed - create_dir_if_missing(path_output) - create_dir_if_missing(path_invalid) - - # Process filenames - if path_input.is_dir(): - image_paths = glob_all_except(path_input, excluded_patterns=["*.gitkeep"]) - else: - create_dir_if_missing(default_values["input"]) - sys.exit("Input folder not found. Make sure you arguments are correct or use the default '" + default_values["input_dir"] + "' folder.") - - # Flush all images in the output directory if asked to - if args["flush"]: - flush_output(path_output, EXTS) - - # Apply to all images - invalid_count = 0 - - # TODO : Condition this - # Enable the HEIF/HEIC Pillow plugin - register_heif_opener() - - # Loop through images - for processed_count, image_path in enumerate(image_paths): - image = attempt_open_image( image_path=image_path, - path_invalid=path_invalid, - attempt_rotate=not args["no_rotate"]) - - if image is None: - continue - - # Randomize position if asked - position_list = position_list_from_setting(position_setting) - color_mapping = color_mapping_from_setting(settings["color_setting"]) - - print("Processing image", processed_count + 1, "of", len(image_paths), "| Invalid:", invalid_count, end="\r") - - # Watermark picture - watermark_image(image, - path=image_path, - logos=logos, - position_list=position_list, - settings=settings) - - print() - print("Done") \ No newline at end of file + # TODO : Configure dry run + + print("Start") + + root_path = Path() + logo_path = root_path / "logos" + + # Define default values + default_values = { + "ss_factor": 2, + "wm_size": 0.07, + "wm_ratio": 1.6, + "wm_pad": 0.15, + "wm_prefix": "wm_", + "input_dir": "input", + "output_dir": "output", + "format": "png", + } + + # Other parameters + str_process = "Processing image {} of {} \t({} of {} variations, \tinvalid: {})" + str_end = "Processed {} images successfully!" + str_invalid = " ({} image(s) failed and moved to 'invalid')" + + # Names of logo files + fn_color = "logo_color.png" + fn_white = "logo_white.png" + + # Open logo files + logos = { + "color": Image.open(logo_path / fn_color), + "white": Image.open(logo_path / fn_white), + } + + # Setup argparser and parse arguments + ap = setup_argparser( + default_vals=default_values, + color_options=COLOR_OPTIONS, + pos_choices=POSITION_OPTIONS, + ) + args = vars(ap.parse_args()) + + prefix = "" if args["no_prefix"] else default_values["wm_prefix"] + path_input = root_path / args["input_dir"] + path_output = root_path / args["output_dir"] + path_invalid = root_path / "invalid" + position_setting = args["position"] + + settings = { + "image_watermark_ratio": args["watermark_size"], + "logo_padding_ratio": args["watermark_padding"], + "logo_circle_ratio": args["watermark_ratio"], + "circle_offset_ratio_x": 0.5 if args["center_circle"] else 3 / 5, + "circle_offset_ratio_y": 0.5 if args["center_circle"] else 1, + "ss_factor": args["supersampling"], + "draw_circle": not args["no_circle"], + "output_path": path_output, + "color_setting": args["color"], + "prefix": prefix, # TODO : Offer option to customise the prefix + "format": default_values[ + "format" + ], # TODO : Offer option to change output format + } + + if not path_input.is_dir(): + sys.exit( + "Input folder not found. Make sure you arguments are correct or use the default '" + + default_values["input_dir"] + + "' folder." + ) + + all_input_directories = [str(path_input)] + scandir(path_input) + for current_directory in all_input_directories: + current_input_path = Path(current_directory) + current_output_path = path_output / current_input_path.relative_to(path_input) + current_invalid_path = path_invalid / current_input_path.relative_to(path_input) + + # Create missing folders if needed + create_dir_if_missing(current_output_path) + create_dir_if_missing(current_invalid_path) + + all_paths = glob_all_except(current_input_path, excluded_patterns=["*.gitkeep"]) + image_paths = [p for p in all_paths if not p.is_dir()] + + if len(image_paths) == 0: + print(f'Skipping folder "{current_input_path}" (no images found)') + continue + + print(f'Processing folder "{current_input_path}"') + + # Flush all images in the output directory if asked to + if args["flush"]: + flush_output(current_output_path, IMG_EXTS) + + # Apply to all images + invalid_count = 0 + + # TODO : Condition this + # Enable the HEIF/HEIC Pillow plugin + register_heif_opener() + + # Loop through images + for processed_count, image_path in enumerate(image_paths): + try: + image = attempt_open_image( + image_path=image_path, + path_invalid=current_invalid_path, + attempt_rotate=not args["no_rotate"], + ) + + if image is None: + continue + + # Randomize position if asked + position_list = position_list_from_setting(position_setting) + color_mapping = color_mapping_from_setting(settings["color_setting"]) + + print( + "Processing image", + processed_count + 1, + "of", + len(image_paths), + "| Invalid:", + invalid_count, + end="\r", + ) + + # Watermark picture + watermark_image( + image, + path=image_path, + logos=logos, + position_list=position_list, + settings={ + **settings, + "output_path": current_output_path, + }, + ) + + except Exception as e: + print(f'\nFailed to process image "{image_path}": {e}') + + print() + print("Done")