diff --git a/poetry.lock b/poetry.lock index 34926f23..d1542ed9 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. [[package]] name = "certifi" @@ -811,6 +811,38 @@ files = [ {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, ] +[[package]] +name = "rawpy" +version = "0.23.2" +description = "RAW image processing for Python, a wrapper for libraw" +optional = false +python-versions = "*" +files = [ + {file = "rawpy-0.23.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5687a050af9e026c88bff62741214c32c7a4491244fa642c6a1b9807dc18e29c"}, + {file = "rawpy-0.23.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:14faecb9b69aa4eaea9711e06eb817b716132bb77b47318efd11273d1115f438"}, + {file = "rawpy-0.23.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:80251d3cdfcf04d3d56f3209697a231a730c54e48f0d69acc086a8a47875aefe"}, + {file = "rawpy-0.23.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc52ce40f68ac07cd44c4239dd0b6effff303ba731587e82d938dc6cb58e5df8"}, + {file = "rawpy-0.23.2-cp310-cp310-win_amd64.whl", hash = "sha256:d00d97f722f42e9e3173eb4446465a77ea2c374b8b3043f2cf6e6e0b893bcc40"}, + {file = "rawpy-0.23.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:dd77b7ad389b47d374b522cb93ff2ed07f54230c92be8c13bfd7a620f7b9c566"}, + {file = "rawpy-0.23.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d1e00dea7636031244633606a5ef95f7087c6807f3cecc9829cfe09809c6b99a"}, + {file = "rawpy-0.23.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9e58c8d33381e9638f22909554eb442f6043fdc0b475701c424b8ca05afe5d36"}, + {file = "rawpy-0.23.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:59c0701667f13d1773f3b6bab60b9a7a6e22ff184d25f4c1d1bfb444fef98914"}, + {file = "rawpy-0.23.2-cp311-cp311-win_amd64.whl", hash = "sha256:cf4790b20a0b27232b9511669e8777d839863efc9acc38e82db87dfec0384b4a"}, + {file = "rawpy-0.23.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b4fbda4ba01ee3bde82d3940e025f226c715fc963bb876532d640490de133c0e"}, + {file = "rawpy-0.23.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5998a4f7cd13236bbde72811dff3a9be7b723ef25078263ecef8ec826328483e"}, + {file = "rawpy-0.23.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d138c6ebc2796b44d57bbae306ec75076acd0b59091de8b886d1e9e6c30b966b"}, + {file = "rawpy-0.23.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b0ce160d06b353769abc5ac6298fbf812d218c49943f1fe9fa97229e8846ca6d"}, + {file = "rawpy-0.23.2-cp312-cp312-win_amd64.whl", hash = "sha256:75fd34b63bed8ed9f585fb8d783b6260c32b45a25c202e812af642dc05837bd7"}, + {file = "rawpy-0.23.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:aa2301297d8c95a6b49e41798240c14619f4753fe35f80d565fb303ea95412d1"}, + {file = "rawpy-0.23.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:490df6aed38a63e01bec71be6c3dcaa8541b23911459f0160cb56d9b745ce74c"}, + {file = "rawpy-0.23.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b17f4f7d53b7cf8fc33879a7e36439b6cbac152ab923238245cc1ee010d86cc"}, + {file = "rawpy-0.23.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5851361e31669f7e08493dd0ed70817ea82db18a74437abf97b44cf02881fd3c"}, + {file = "rawpy-0.23.2-cp39-cp39-win_amd64.whl", hash = "sha256:540eac17a50f6d9f84500ba1ec8d4e53472f2ee82293d9af99a7c76fcdf819a2"}, +] + +[package.dependencies] +numpy = ">=1.26.0" + [[package]] name = "regex" version = "2024.5.15" @@ -1353,4 +1385,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = ">=3.9 <3.13" -content-hash = "ffbdbfd76077322c443179339fb5037cb8a8437286546d753d2a241dfcf285d2" +content-hash = "92df2052244a7720dbf326c77ef18d44868149b3a1eaaa51530db2b6b96822a8" diff --git a/pyproject.toml b/pyproject.toml index 9a70d664..63c2f444 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ torchvision = [ { version = "==0.17.2+cpu", source = "pytorch-cpu", markers = "sys_platform == 'linux' and platform_machine != 'aarch64'" } ] tqdm = "^4.65.0" +rawpy = "^0.23.2" [tool.poetry.group.dev.dependencies] pycodestyle = ">=2.7,<3.0" diff --git a/rclip/const.py b/rclip/const.py index 7740e18d..b98114b2 100644 --- a/rclip/const.py +++ b/rclip/const.py @@ -4,3 +4,8 @@ IS_MACOS = sys.platform == 'darwin' IS_LINUX = sys.platform.startswith('linux') IS_WINDOWS = sys.platform == 'win32' or sys.platform == 'cygwin' + +# these images are always processed +IMAGE_EXT = ["jpg", "jpeg", "png", "webp"] +# RAW images are processed only if there is no processed image alongside it +IMAGE_RAW_EXT = ["arw", "cr2"] diff --git a/rclip/main.py b/rclip/main.py index 53a9a328..918e652c 100644 --- a/rclip/main.py +++ b/rclip/main.py @@ -11,6 +11,7 @@ from PIL import Image, ImageFile from rclip import db, fs, model +from rclip.const import IMAGE_EXT, IMAGE_RAW_EXT from rclip.utils.preview import preview from rclip.utils.snap import check_snap_permissions, is_snap from rclip.utils import helpers @@ -41,7 +42,7 @@ def is_image_meta_equal(image: db.Image, meta: ImageMeta) -> bool: class RClip: EXCLUDE_DIRS_DEFAULT = ['@eaDir', 'node_modules', '.git'] - IMAGE_REGEX = re.compile(r'^.+\.(jpe?g|png|webp)$', re.I) + IMAGE_REGEX = re.compile(f'^.+\\.({"|".join([*IMAGE_EXT, *IMAGE_RAW_EXT])})$', re.I) DB_IMAGES_BEFORE_COMMIT = 50_000 class SearchResult(NamedTuple): @@ -67,7 +68,7 @@ def _index_files(self, filepaths: List[str], metas: List[ImageMeta]): filtered_paths: List[str] = [] for path in filepaths: try: - image = Image.open(path) + image = helpers.read_image(path) images.append(image) filtered_paths.append(path) except PIL.UnidentifiedImageError as ex: @@ -88,6 +89,18 @@ def _index_files(self, filepaths: List[str], metas: List[ImageMeta]): vector=vector.tobytes() ), commit=False) + def _does_processed_image_exist_for_raw(self, raw_path: str) -> bool: + """Check if there is a processed image alongside the raw one; doesn't support mixed-case extensions, + e.g. it won't detect the .JpG image, but will detect .jpg or .JPG""" + + image_path = os.path.splitext(raw_path)[0] + for ext in IMAGE_EXT: + if os.path.isfile(image_path + "." + ext): + return True + if os.path.isfile(image_path + "." + ext.upper()): + return True + return False + def ensure_index(self, directory: str): print( 'checking images in the current directory for changes;' @@ -113,7 +126,13 @@ def update_total_images(count: int): metas: List[ImageMeta] = [] for entry in fs.walk(directory, self._exclude_dir_regex, self.IMAGE_REGEX): filepath = entry.path - image = self._db.get_image(filepath=filepath) + + file_ext = helpers.get_file_extension(filepath) + if file_ext in IMAGE_RAW_EXT and self._does_processed_image_exist_for_raw(filepath): + images_processed += 1 + pbar.update() + continue + try: meta = get_image_meta(entry) except Exception as ex: @@ -125,6 +144,7 @@ def update_total_images(count: int): images_processed += 1 pbar.update() + image = self._db.get_image(filepath=filepath) if image and is_image_meta_equal(image, meta): self._db.remove_indexing_flag(filepath, commit=False) continue diff --git a/rclip/utils/helpers.py b/rclip/utils/helpers.py index 161cf108..6f37b38a 100644 --- a/rclip/utils/helpers.py +++ b/rclip/utils/helpers.py @@ -4,11 +4,13 @@ import textwrap from PIL import Image, UnidentifiedImageError import re +import numpy as np +import rawpy import requests import sys from importlib.metadata import version -from rclip.const import IS_LINUX, IS_MACOS, IS_WINDOWS +from rclip.const import IMAGE_RAW_EXT, IS_LINUX, IS_MACOS, IS_WINDOWS MAX_DOWNLOAD_SIZE_BYTES = 50_000_000 @@ -186,15 +188,29 @@ def download_image(url: str) -> Image.Image: return img +def get_file_extension(path: str) -> str: + return os.path.splitext(path)[1].lower()[1:] + + +def read_raw_image_file(path: str): + raw = rawpy.imread(path) + rgb = raw.postprocess() + return Image.fromarray(np.array(rgb)) + + def read_image(query: str) -> Image.Image: path = remove_prefix(query, 'file://') try: - img = Image.open(path) + file_ext = get_file_extension(path) + if file_ext in IMAGE_RAW_EXT: + image = read_raw_image_file(path) + else: + image = Image.open(path) except UnidentifiedImageError as e: # by default the filename on the UnidentifiedImageError is None e.filename = path raise e - return img + return image def is_http_url(path: str) -> bool: diff --git a/rclip/utils/preview.py b/rclip/utils/preview.py index 574028a3..f8f26c44 100644 --- a/rclip/utils/preview.py +++ b/rclip/utils/preview.py @@ -3,6 +3,8 @@ import os from PIL import Image +from rclip.utils.helpers import read_image + def _get_start_sequence(): term_env_var = os.getenv('TERM') @@ -19,7 +21,7 @@ def _get_end_sequence(): def preview(filepath: str, img_height_px: int): - with Image.open(filepath) as img: + with read_image(filepath) as img: if img_height_px >= img.height: width_px, height_px = img.width, img.height else: