feat: add support for ARW, DNG, CR2 raw

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"

rclip/main.py

@@ -9,6 +9,7 @@
 from tqdm import tqdm
 import PIL
 from PIL import Image, ImageFile
+from rawpy import imread
 
 from rclip import db, fs, model
 from rclip.utils.preview import preview
@@ -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(r'^.+\.(jpe?g|png|webp|arw|dng|cr2)$', re.I)
   DB_IMAGES_BEFORE_COMMIT = 50_000
 
   class SearchResult(NamedTuple):
@@ -62,18 +63,38 @@ def __init__(
     excluded_dirs = '|'.join(re.escape(dir) for dir in exclude_dirs or self.EXCLUDE_DIRS_DEFAULT)
     self._exclude_dir_regex = re.compile(f'^.+\\{os.path.sep}({excluded_dirs})(\\{os.path.sep}.+)?$')
 
+  def _read_raw_image_file(self, path: str):
+    image = None
+    try:
+      raw = imread(path)
+      rgb = raw.postprocess()
+      image = Image.fromarray(np.array(rgb))
+    except Exception as ex:
+      print(f'not a valid raw file {path}', ex, file=sys.stderr)
+    return image
+
+  def _read_image_file(self, path: str):
+    image = None
+    try:
+      image = Image.open(path)
+    except PIL.UnidentifiedImageError as ex:
+      print(f'unidentified image error {path}:', ex, file=sys.stderr)
+    except Exception as ex:
+      print(f'error loading image {path}:', ex, file=sys.stderr)
+    return image
+
   def _index_files(self, filepaths: List[str], metas: List[ImageMeta]):
     images: List[Image.Image] = []
     filtered_paths: List[str] = []
     for path in filepaths:
-      try:
-        image = Image.open(path)
+      if os.path.splitext(path)[1].lower() in fs.PRIORITIZED_IMAGE_EXTENSIONS:
+        image = self._read_image_file(path)
+      else:
+        image = self._read_raw_image_file(path)
+
+      if image:
         images.append(image)
         filtered_paths.append(path)
-      except PIL.UnidentifiedImageError as ex:
-        pass
-      except Exception as ex:
-        print(f'error loading image {path}:', ex, file=sys.stderr)
 
     try:
       features = self._model.compute_image_features(images)