rawforge 0.1.0__tar.gz → 0.2.1__tar.gz

{rawforge-0.1.0 → rawforge-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rawforge
-Version: 0.1.0
+Version: 0.2.1
 Summary: A compute backend/CLI application for using machine learning models on raw images.
 Author: Ryan Mueller
 License: MIT License
@@ -40,12 +40,27 @@ License-File: LICENSE
 Requires-Dist: numpy>=2.2
 Requires-Dist: RawHandler~=0.2.0
 Requires-Dist: colour_demosaicing~=0.2.6
-Requires-Dist: blended_tiling>=0.0.1.dev7
+Requires-Dist: blended-tiling-numpy
 Requires-Dist: requests~=2.32
 Requires-Dist: platformdirs~=4.5
 Requires-Dist: tqdm~=4.67
 Requires-Dist: cryptography>=46.0
 Requires-Dist: tifffile
+Requires-Dist: pyexiftool
+Provides-Extra: cpu
+Requires-Dist: onnxruntime; extra == "cpu"
+Provides-Extra: cuda
+Requires-Dist: onnxruntime-gpu; extra == "cuda"
+Provides-Extra: web
+Requires-Dist: onnxruntime-webgpu; extra == "web"
+Provides-Extra: directml
+Requires-Dist: onnxruntime-directml; extra == "directml"
+Provides-Extra: torch
+Requires-Dist: torch; extra == "torch"
+Provides-Extra: fastapi
+Requires-Dist: fastapi; extra == "fastapi"
+Requires-Dist: pydantic; extra == "fastapi"
+Requires-Dist: uvicorn; extra == "fastapi"
 Dynamic: license-file
 
 # RawRefinery
@@ -65,6 +80,41 @@ Currently in **alpha release**.
 pip install rawforge
 ```
 
+
+## New: ONNX backend!
+
+### NVIDIA: For cuda (tested, must have CUDA installed):
+
+```bash
+pip install rawforge[cuda]
+```
+
+### AMD and NVIDIA on Windows:
+
+```bash
+pip install rawforge[directml]
+```
+
+
+### All others, including M1+ Macs (CoreML)
+
+```bash
+pip install rawforge[cpu]
+```
+
+### For web:
+
+```bash
+pip install rawforge[web]
+```
+
+Useage is the same, but use the CLI "rawforgeonnx".
+
+
+```bash
+rawforgeonnx TreeNetDenoiseHeavy test.CR2 test_heavy.dng --cfa 
+```
+
 ---
 ## Example command line syntax:
 
@@ -97,6 +147,9 @@ options:
   --chroma CHROMA       Chroma noise (0-1).
 ```
 
+
+
+
 ----
 
 ## Acknowledgments

{rawforge-0.1.0 → rawforge-0.2.1}/README.md

@@ -15,6 +15,41 @@ Currently in **alpha release**.
 pip install rawforge
 ```
 
+
+## New: ONNX backend!
+
+### NVIDIA: For cuda (tested, must have CUDA installed):
+
+```bash
+pip install rawforge[cuda]
+```
+
+### AMD and NVIDIA on Windows:
+
+```bash
+pip install rawforge[directml]
+```
+
+
+### All others, including M1+ Macs (CoreML)
+
+```bash
+pip install rawforge[cpu]
+```
+
+### For web:
+
+```bash
+pip install rawforge[web]
+```
+
+Useage is the same, but use the CLI "rawforgeonnx".
+
+
+```bash
+rawforgeonnx TreeNetDenoiseHeavy test.CR2 test_heavy.dng --cfa 
+```
+
 ---
 ## Example command line syntax:
 
@@ -47,6 +82,9 @@ options:
   --chroma CHROMA       Chroma noise (0-1).
 ```
 
+
+
+
 ----
 
 ## Acknowledgments

rawforge-0.2.1/RawForge/application/ImageSaver.py

@@ -0,0 +1,70 @@
+import numpy as np
+from RawForge.application.dng_utils import convert_color_matrix, to_dng
+import tifffile
+
+
+class ImageSaver:
+    def __init__(self, model_params, rh, dims=None):
+        self.rh = rh
+        self.model_params = model_params
+        self.dims = dims
+
+    def to_tiff(self, image, filename, apply_ccm=True):
+        image = image
+        if apply_ccm:
+            transform_matrix = self.rh.rgb_colorspace_transform(
+                colorspace="lin_rec2020"
+            )
+            image = image[0].transpose(1, 2, 0)
+            transformed = image @ transform_matrix.T
+        else:
+            transformed = image[0].transpose(1, 2, 0)
+
+        transformed = np.clip(transformed, 0, 1)
+        transformed = transformed ** (1 / 2.2)
+        transformed = transformed * (2**8 - 1)
+
+        uint_img = transformed.astype(np.uint8)
+
+        tifffile.imwrite(
+            filename,
+            uint_img,
+            photometric="rgb",  # Explicitly define the color space
+            compression="deflate",  # Optional: Lossless compression supported by darktable
+        )
+
+    def to_raw(self, denoised, filename, save_cfa):
+        # Compute CFA
+        if self.model_params["demosaicing"] == "rawpy" or self.model_params["demosaicing"] == "sixchan":
+            if self.dims is None:
+                self.dims = [0, 9999999, 0, 9999999]
+            _, mask = self.rh.compute_mask_and_sparse(dims=self.dims)
+            denoised = denoised[0]
+            denoised = denoised.clip(0, 1)
+
+            denoised = np.where(mask, denoised, 0)
+            denoised = denoised.sum(axis=0)
+            denoised = (
+                denoised * (self.rh.core_metadata.white_level)
+                + self.rh.core_metadata.black_level_per_channel[0]
+            )
+            self.rh.to_dng(filename, uint_img=denoised)
+        else:
+            transform_matrix = np.linalg.inv(
+                self.rh.rgb_colorspace_transform(colorspace=self.rh.colorspace)
+            )
+
+            CCM = self.rh.rgb_colorspace_transform(colorspace="XYZ")
+            CCM = np.linalg.inv(CCM)
+            denoised = denoised[0].transpose(1, 2, 0)
+            transformed = denoised @ transform_matrix.T
+            uint_img = np.clip(transformed * 2**16 - 1, 0, 2**16 - 1).astype(np.uint16)
+            ccm1 = convert_color_matrix(CCM)
+            to_dng(
+                uint_img,
+                self.rh,
+                filename,
+                ccm1,
+                save_cfa=save_cfa,
+                convert_to_cfa=True,
+            )

rawforge-0.2.1/RawForge/application/InferenceWorker.py

@@ -0,0 +1,84 @@
+import numpy as np
+from blended_tiling_numpy import TilingModule
+from tqdm import tqdm
+
+
+class InferenceWorker:
+    def __init__(
+        self,
+        backend,
+        model_params,
+        conditioning,
+        disable_tqdm=False,
+    ):
+        self.backend = backend
+        self.model_params = model_params
+        self.conditioning = conditioning
+
+        self.tile_size = model_params.get("tile_size", 256)
+        self.tile_overlap = model_params.get("tile_overlap", 0.25)
+        self.batch_size = model_params.get("batch_size", 2)
+
+        self.disable_tqdm = disable_tqdm
+        self._is_cancelled = False
+
+    def cancel(self):
+        self._is_cancelled = True
+
+    def _build_conditioning(self):
+        cond = np.array([self.conditioning], dtype=np.float32)
+
+        if "max_iso" in self.model_params:
+            cond[:, 0] = min(cond[:, 0], self.model_params["max_iso"])
+        cond[:, 0] /= 6400.0
+        cond[:, 1] = 0.0
+        cond = cond[:, :1].astype(np.float16)
+
+        if "cond_scale" in self.model_params:
+            cond *= cond * self.model_params["cond_scale"]
+
+        return self.backend.prepare_conditioning(cond)
+
+    def run(self, image_RGB, progress_callback=None):
+
+        full_size = image_RGB.shape[2:]
+
+        tiler = TilingModule(
+            tile_size=[self.tile_size] * 2,
+            tile_overlap=[self.tile_overlap] * 2,
+            base_size=list(full_size),
+        )
+
+        tiles = tiler.split_into_tiles(image_RGB)
+
+        batches = [
+            tiles[i:i+self.batch_size]
+            for i in range(0, len(tiles), self.batch_size)
+        ]
+
+        cond = self._build_conditioning()
+
+        outputs = []
+
+        for i, batch in tqdm(
+            enumerate(batches),
+            total=len(batches),
+            disable=self.disable_tqdm,
+        ):
+
+            if self._is_cancelled:
+                return None, None
+
+            outputs.append(
+                self.backend.infer(batch, cond)
+            )
+
+            if progress_callback:
+                progress_callback((i + 1) / len(batches))
+
+        stitched = tiler.rebuild_with_masks(
+            np.concatenate(outputs, axis=0)
+        )
+
+        return image_RGB, stitched
+