rawforge 0.1.0__tar.gz → 0.2.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rawforge
-Version: 0.1.0
+Version: 0.2.0
 Summary: A compute backend/CLI application for using machine learning models on raw images.
 Author: Ryan Mueller
 License: MIT License
@@ -40,12 +40,20 @@ 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
+Provides-Extra: cpu
+Requires-Dist: onnxruntime; extra == "cpu"
+Provides-Extra: cuda
+Requires-Dist: onnxruntime-gpu; extra == "cuda"
+Provides-Extra: web
+Requires-Dist: onnxruntime-web; extra == "web"
+Provides-Extra: directml
+Requires-Dist: onnxruntime-directml; extra == "directml"
 Dynamic: license-file
 
 # RawRefinery

rawforge-0.2.0/RawForge/application/ImageSaver.py

@@ -0,0 +1,55 @@
+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':
+                _, 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.0/RawForge/application/InferenceWorker.py

@@ -0,0 +1,79 @@
+import onnxruntime as ort
+import numpy as np
+from blended_tiling_numpy import TilingModule
+from RawForge.application.postprocessing import match_colors_linear
+from tqdm import tqdm
+import rawpy
+
+class InferenceWorker():
+    def __init__(self, model, model_params, conditioning, tile_size=512, tile_overlap=0.25, batch_size=2, disable_tqdm=False):
+        super().__init__()
+        self.model = model
+        self.model_params = model_params
+        self.conditioning = conditioning
+        self.tile_size = tile_size
+        if 'tile_size' in model_params:
+            self.tile_size =  model_params['tile_size']
+        self.tile_overlap = tile_overlap
+        self.batch_size = batch_size
+        if 'batch_size' in model_params:
+            self.batch_size = model_params['batch_size']
+        self._is_cancelled = False
+        self.disable_tqdm = disable_tqdm
+
+    def cancel(self):
+        self._is_cancelled = True
+    
+
+    def _tile_process(self, image_RGB, model_params):
+        # Prepare Data
+        full_size = [image_RGB.shape[2], image_RGB.shape[3]]
+        tile_size = [self.tile_size, self.tile_size]
+        overlap = [self.tile_overlap, self.tile_overlap]
+        # Tiling Setup
+        tiling_module_rgb = TilingModule(tile_size=[s for s in tile_size], tile_overlap=overlap, base_size=[s for s in full_size])
+
+        tiles_rgb = tiling_module_rgb.split_into_tiles(image_RGB)
+        
+        batches_rgb = [tiles_rgb[i : i + self.batch_size] 
+                        for i in range(0, len(tiles_rgb), self.batch_size)]
+        # Conditioning Setup
+        cond_tensor = np.array([self.conditioning]).astype(np.float32)
+        cond_tensor[:, 0] /= 6400.
+        cond_tensor[:, 1] = 0.
+        cond_tensor = cond_tensor[:, 0:1]
+        cond_tensor = cond_tensor.astype(np.float16)
+        if 'cond_scale' in model_params:
+            cond_tensor *= cond_tensor * model_params['cond_scale']
+
+        processed_batches = []
+        
+        # Inference Loop
+        for i, (batch_rgb) in tqdm(enumerate(batches_rgb), disable=self.disable_tqdm):
+            if self._is_cancelled: return None, None
+            
+            B = batch_rgb.shape[0]
+            # Expand conditioning to match batch size
+            curr_cond = np.broadcast_to(cond_tensor, (B, cond_tensor.shape[-1])).astype(np.float16)
+
+            payload = {
+                "input": batch_rgb,
+                "cond": curr_cond
+            }
+            # Filter based on inputs
+            model_inputs = {i.name for i in self.model.get_inputs()}
+            filtered_inputs = {k: v for k, v in payload.items() if k in model_inputs}
+            output = self.model.run(["output"], filtered_inputs)
+            
+            processed_batches.append(output[0])
+                    
+        # Rebuild
+        tiles_out = np.concat(processed_batches, axis=0)
+        stitched = tiling_module_rgb.rebuild_with_masks(tiles_out)
+
+        return image_RGB, stitched
+    
+    def run(self, model_params, image_RGB):
+        img, denoised_img = self._tile_process(image_RGB, model_params)
+        return img, denoised_img
+            

rawforge-0.2.0/RawForge/application/MODEL_REGISTRY.py

@@ -0,0 +1,135 @@
+MODEL_REGISTRY = {
+"TreeNetDenoise": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/ShadowWeightedL1.onnx",
+        "filename": "ShadowWeightedL1.onnx",
+        "max_iso": 65535,
+        "demosaicing": "Malvar2004",
+    },
+    "TreeNetDenoiseLight": {
+        "url": "    https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/ShadowWeightedL1_light.onnx",
+        "filename": "ShadowWeightedL1_light.onnx",
+        "max_iso": 65535,
+        "demosaicing": "Malvar2004",
+    },
+    "TreeNetDenoiseSuperLight": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/ShadowWeightedL1_super_light.onnx",
+        "filename": "ShadowWeightedL1_super_light.onnx",
+        "max_iso": 65535,
+        "demosaicing": "Malvar2004",
+    },
+    "TreeNetDenoiseHeavy": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/ShadowWeightedL1_24_deep_500.onnx",
+        "filename": "ShadowWeightedL1_24_deep_500.onnx",
+        "max_iso": 65535,
+        "demosaicing": "Malvar2004",
+    },
+    "Deblur": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/realblur_gamma_140.onnx",
+        "filename": "realblur_gamma_140.onnx",
+        "demosaicing": "Malvar2004",
+    },
+    "DeepSharpen": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/Deblur_deep_24.onnx",
+        "filename": "Deblur_deep_24.onnx",
+        "demosaicing": "Malvar2004",
+    },
+    "TreeNetDenoiseXTrans": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/xtrans_fixed_exposure_no_conditioning_400.onnx",
+        "filename": "xtrans_fixed_exposure_no_conditioning_400.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+        "cond_scale": 0,
+    },
+    "RestormerXTrans": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/restormer.onnx",
+        "filename": "restormer.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+    }, 
+    "XFormerXTrans2": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/xformer_fp16.onnx",
+        "filename": "xformer_fp16.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+    }, 
+
+    "XFormerDenoise": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/xformer_static.onnx",
+        "filename": "xformer_static.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+    "NAFDenoise": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/NAF_static.onnx",
+        "filename": "NAF_static.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+
+    "NAFDenoise_dynamic": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/NAF_dynamic.onnx",
+        "filename": "NAF_dynamic.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+    "NAFDenoise_dynamic_fp16": {
+        "url": "https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/NAF_dynamic_fp16.onnx",
+        "filename": "NAF_dynamic_fp16.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+    "NAF_static_fp16": {
+        "url": "    https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/NAF_static_fp16.onnx",
+        "filename": "NAF_static_fp16.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+    "NAF_trace_test": {
+        "url": "    https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/xtrans_fixed_exposure_no_conditioning_400.onnx",
+        "filename": "xtrans_fixed_exposure_no_conditioning_400.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+    "non_trace_test": {
+        "url": "    https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/non_trace_test.onnx",
+        "filename": "non_trace_test.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    }, 
+    "trace_test": {
+        "url": "    https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/trace_test.onnx",
+        "filename": "trace_test.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    },   
+    "non_trace_test_opset_20_static_simp": {
+        "url": "    https://github.com/rymuelle/RawForge/releases/download/onnx_v1.0.0/non_trace_test_opset_20_static_simp.onnx",
+        "filename": "non_trace_test_opset_20_static_simp.onnx",
+        "demosaicing": "rawpy",
+        "conditioning": "false",
+        "batch_size": 1, 
+        "crop_size": 256,
+    },   
+    
+    
+}

{rawforge-0.1.0 → rawforge-0.2.0}/RawForge/application/ModelHandler.py

@@ -1,21 +1,15 @@
-import torch
+import onnxruntime as ort
+import os
 import numpy as np
 from pathlib import Path
 from platformdirs import user_data_dir
-from time import perf_counter
 import requests
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import padding
 
-from RawHandler.RawHandler import RawHandler
-from RawHandler.RawHandlerRawpy import RawHandlerRawpy
-
-from RawForge.application.dng_utils import convert_color_matrix, to_dng
-from RawForge.application.utils import can_use_gpu
-
 from RawForge.application.MODEL_REGISTRY import MODEL_REGISTRY
 from RawForge.application.InferenceWorker import InferenceWorker
-from RawForge.application.InferenceWorkerRawpy import InferenceWorkerRawpy
+from RawForge.application.helpers.utils import get_best_providers
 
 key_string = '''-----BEGIN PUBLIC KEY-----
 MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEA8iRGMPqFIFVF0TM/AbMI
@@ -40,46 +34,23 @@ class ModelHandler():
     def __init__(self):
         super().__init__()
 
-        self.model = None
-        self.rh = None
-        self.iso = 100
-        self.colorspace = 'lin_rec2020'
-
+        self.model = None        
+        app_name = "RawForge"
+        self.data_dir = Path(user_data_dir(app_name))
+        self.cache_dir = Path(os.path.expanduser(f"{self.data_dir}/model_cache")).resolve()
         # Manage devices
-        devices = {
-                   "cuda": can_use_gpu(),
-                   "mps": torch.backends.mps.is_available(),
-                   "cpu": lambda : True
-        }
-        self.devices = [d for d, is_available in devices.items() if is_available]
-        self.set_device(self.devices[0])
-        
-        self.filename = None
-        self.start_time = None
-        self.model_params = {}
-
+        self.providers = get_best_providers(cache_dir=self.cache_dir)
         self.pub = serialization.load_pem_public_key(key_string.encode('utf-8'))
 
-    def load_rh(self, path):
-        """Loads the raw file handler"""
-        if 'backend' in self.model_params and self.model_params['backend'] == 'rawpy':
-                    self.rh = RawHandlerRawpy(path)
-        else:
-                    self.rh = RawHandler(path)
-        self.iso = self.rh.full_metadata.get_ISO()
-        return self.iso
 
     def load_model(self, model_key):
         """Loads a model by key from the registry"""
         if model_key not in MODEL_REGISTRY:
             print(f"Model {model_key} not found in registry.")
             return
-
         conf = MODEL_REGISTRY[model_key]
         self.model_params = conf
-        app_name = "RawForge"
-        data_dir = Path(user_data_dir(app_name))
-        model_path = data_dir / conf["filename"]
+        model_path = self.data_dir / conf["filename"]
 
         # Handle Download
         if not model_path.is_file():
@@ -97,41 +68,15 @@ class ModelHandler():
             # Verify model before load
             self._verify_model(model_path, model_path.with_suffix(f'{model_path.suffix}.sig'))
             
-            loaded = torch.jit.load(model_path, map_location='cpu')
-            self.model = loaded.eval().to(self.device)
+            session = ort.InferenceSession(
+               model_path, 
+               providers=self.providers, 
+            )
+            print("Loaded!")
+            self.model = session
         except Exception as e:
             print(f"Failed to load model: {e}")
 
-    def set_device(self, device):
-        self.device = torch.device(device)
-        if self.model:
-            self.model.to(self.device)
-        print(f"Using Device {self.device} from {device}")
-
-    def run_inference(self, conditioning, dims=None, inference_kwargs={}):
-        """Starts the worker thread"""
-        if not self.model or not self.rh:
-            print("Model or Image not loaded.")
-            return
-        
-        # Some older models were trained with a lower max iso
-        if "max_iso" in self.model_params:
-            conditioning[0] = min(conditioning[0], self.model_params["max_iso"])
-
-        if 'backend' in self.model_params and self.model_params['backend'] == 'rawpy':
-            worker = InferenceWorkerRawpy(self.model, self.model_params, self.device, self.rh, conditioning, dims, **inference_kwargs)
-        else:
-            worker = InferenceWorker(self.model, self.model_params, self.device, self.rh, conditioning, dims, **inference_kwargs)
-        img, final_denoised =  worker.run()
-
-        return img, final_denoised
-
-
-    def generate_thumbnail(self, size=400):
-        if not self.rh: return None
-        thumb = self.rh.generate_thumbnail(min_preview_size=size, clip=True)
-        return thumb
-
     def _verify_model(self, dest_path, sig_path):
         try:
             data = Path(dest_path).read_bytes()
@@ -156,7 +101,6 @@ class ModelHandler():
             print(f"Model {dest_path} not verified! Deleting.")
             return False
 
-
     def _download_file(self, url, dest_path):
         dest_path.parent.mkdir(parents=True, exist_ok=True)
         try:
@@ -179,31 +123,4 @@ class ModelHandler():
         except Exception as e:
             print(e)
             return False
-        
-
-    def handle_full_image(self, denoised, filename, save_cfa):
-        # Compute CFA
-        if 'backend' in self.model_params and self.model_params['backend'] == 'rawpy':
-            _, mask = self.rh.compute_mask_and_sparse(dims=(0, 99999, 0, 99999))
-            denoised = denoised.transpose(2, 0, 1)
-            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.colorspace)
-                    )
-
-            CCM = self.rh.rgb_colorspace_transform(colorspace='XYZ')
-            CCM = np.linalg.inv(CCM)
-
-            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.0/RawForge/application/helpers/censored_fit.py

@@ -0,0 +1,100 @@
+import numpy as np
+import math
+
+def norm_cdf(z):
+    erf_vec = np.vectorize(math.erf)
+    return 0.5 * (1 + erf_vec(z / np.sqrt(2)))
+
+def norm_pdf(z):
+    return (1 / np.sqrt(2 * np.pi)) * np.exp(-0.5 * z**2)
+
+
+def censored_linear_fit_twosided(x, y, clip_low=0, clip_high=1,
+                                 max_iter=200, tol=1e-6, include_offset=True):
+    """
+    Fit y ≈ a + b*x + ε, ε ~ N(0, σ²) under two-sided censoring:
+    clip_low ≤ y_true ≤ clip_high
+    Observed y are clipped to [clip_low, clip_high].
+    Returns (a, b, sigma) estimated via EM.
+
+    Parameters
+    ----------
+    x, y : array_like
+        Input data.
+    clip_low, clip_high : float or None
+        Lower/upper clip levels. Can be None for one-sided clipping.
+    max_iter : int
+        Maximum EM iterations.
+    tol : float
+        Relative tolerance for convergence.
+    include_offset : bool
+        If False, forces intercept a=0 (fit y ≈ b*x).
+
+    Returns
+    -------
+    a, b, sigma : floats
+        Estimated regression parameters.
+    """
+    x = np.asarray(x).ravel()
+    y = np.asarray(y).ravel()
+    mask = np.isfinite(x) & np.isfinite(y)
+    x, y = x[mask], y[mask]
+
+    if len(x) < 3:
+        raise ValueError("Not enough data points.")
+
+    # Initial guess (ordinary least squares) 
+    if include_offset:
+        A = np.vstack([np.ones_like(x), x]).T
+        a, b = np.linalg.lstsq(A, y, rcond=None)[0]
+    else:
+        b = np.dot(x, y) / np.dot(x, x)
+        a = 0.0
+    sigma = np.std(y - (a + b*x))
+
+    for _ in range(max_iter):
+        mu = a + b*x
+        y_exp = y.copy()
+
+        # Handle right-censoring (high clip)
+        if clip_high is not None:
+            high_mask = y >= clip_high - 1e-12
+            if np.any(high_mask):
+                z = (clip_high - mu[high_mask]) / sigma
+                Phi = norm_cdf(z)
+                phi = norm_pdf(z)
+                one_minus_Phi = 1.0 - Phi
+                lambda_ = np.zeros_like(z)
+                valid = one_minus_Phi > 1e-15
+                lambda_[valid] = phi[valid] / one_minus_Phi[valid]
+                y_exp[high_mask] = mu[high_mask] + sigma * lambda_
+
+        # Handle left-censoring (low clip)
+        if clip_low is not None:
+            low_mask = y <= clip_low + 1e-12
+            if np.any(low_mask):
+                z = (clip_low - mu[low_mask]) / sigma
+                Phi = norm_cdf(z)
+                phi = norm_pdf(z)
+                lambda_ = np.zeros_like(z)
+                valid = Phi > 1e-15
+                lambda_[valid] = -phi[valid] / Phi[valid]
+                y_exp[low_mask] = mu[low_mask] + sigma * lambda_
+
+        # M-step: re-fit with imputed expectations
+        if include_offset:
+            A = np.vstack([np.ones_like(x), x]).T
+            a_new, b_new = np.linalg.lstsq(A, y_exp, rcond=None)[0]
+        else:
+            b_new = np.dot(x, y_exp) / np.dot(x, x)
+            a_new = 0.0
+
+        sigma_new = np.std(y_exp - (a_new + b_new*x))
+
+        if np.allclose([a, b, sigma], [a_new, b_new, sigma_new],
+                       rtol=tol, atol=tol):
+            break
+
+        a, b, sigma = a_new, b_new, sigma_new
+
+    return a, b, sigma

rawforge-0.2.0/RawForge/application/helpers/get_image.py

@@ -0,0 +1,41 @@
+import numpy as np
+import rawpy
+from RawHandler.RawHandler import RawHandler
+from RawHandler.RawHandlerRawpy import RawHandlerRawpy
+
+def load_rh(path, demosaicing):
+    """Loads the raw file handler"""
+    if demosaicing == 'rawpy':
+                rh = RawHandlerRawpy(path)
+    elif demosaicing == 'Malvar2004':
+                rh = RawHandler(path)
+    iso = rh.full_metadata.get_ISO()
+    return rh, iso
+
+def get_image(path, model_params, dims=None):
+        demosaicing = model_params['demosaicing']
+        rh, iso = load_rh(path, demosaicing)
+        
+        if demosaicing=='Malvar2004':
+            from colour_demosaicing import demosaicing_CFA_Bayer_Malvar2004
+            image_RGB = rh.as_rgb(dims=dims, 
+                                  demosaicing_func=demosaicing_CFA_Bayer_Malvar2004, 
+                                  colorspace='lin_rec2020', clip=True)
+        elif demosaicing=='rawpy':
+            image_RGB = rh.rawpy_object.postprocess(
+                        user_wb=[1, 1, 1, 1],
+                        output_color=rawpy.ColorSpace.raw,
+                        demosaic_algorithm= rawpy.DemosaicAlgorithm(3),
+                        no_auto_bright=True,
+                        use_camera_wb=False,
+                        use_auto_wb=False,
+                        gamma=(1, 1),
+                        user_flip=0,
+                        output_bps=16,
+                        no_auto_scale=True,
+                    ) / rh.rawpy_object.white_level
+            if dims is not None:
+                image_RGB = image_RGB[dims[0]:dims[1], dims[2]:dims[3]]
+            image_RGB = image_RGB.transpose(2, 0, 1)
+        image_RGB = np.expand_dims(image_RGB, axis=0).astype(np.float16)
+        return rh, image_RGB, iso

rawforge-0.2.0/RawForge/application/helpers/utils.py

@@ -0,0 +1,17 @@
+import onnxruntime as ort
+import os
+
+def get_best_providers(cache_dir):
+    os.makedirs(cache_dir, exist_ok=True)
+    available = ort.get_available_providers()
+    return available
+    priority = [
+        "CUDAExecutionProvider",      # NVIDIA
+        "ROCMExecutionProvider",      # AMD (Direct)
+        "MIGraphXExecutionProvider",  # AMD (Optimized)
+        "CoreMLExecutionProvider",    # Apple Silicon
+        "DmlExecutionProvider",       # Windows (AMD/Intel/Generic GPU)
+        "CPUExecutionProvider"        # The fallback
+    ]
+    
+    return [p for p in priority if p in available]

rawforge-0.2.0/RawForge/application/postprocessing.py

@@ -0,0 +1,92 @@
+import numpy as np
+from RawForge.application.helpers.censored_fit import censored_linear_fit_twosided
+
+def match_colors_linear(
+    src: np.ndarray, 
+    tgt: np.ndarray, 
+    sample_fraction: float = 0.05,
+    censored_fit: bool = True
+):
+    """
+    Fit per-channel affine color transforms using NumPy.
+    src/tgt: [B, C, H, W]
+    """
+    B, C, H, W = src.shape
+
+    # Flatten spatial dims: [B, C, H*W]
+    src_flat = src.reshape(B, C, -1)
+    tgt_flat = tgt.reshape(B, C, -1)
+
+    # Sample subset of pixels
+    N = src_flat.shape[-1]
+    k = max(64, int(N * sample_fraction))
+    
+    # Generate random indices
+    idx = np.random.randint(0, N, size=(k,))
+
+    src_s = src_flat[..., idx]  # [B, C, k]
+    tgt_s = tgt_flat[..., idx]
+
+    # Compute scale and bias using least squares
+    src_mean = src_s.mean(axis=-1, keepdims=True)
+    tgt_mean = tgt_s.mean(axis=-1, keepdims=True)
+
+    src_centered = src_s - src_mean
+    tgt_centered = tgt_s - tgt_mean
+
+    var_src = (src_centered ** 2).mean(axis=-1)
+    cov = (src_centered * tgt_centered).mean(axis=-1)
+
+    scale = cov / (var_src + 1e-8)  # [B, C]
+    # Squeeze the mean to [B, C] for bias calculation
+    bias = tgt_mean.squeeze(-1) - scale * src_mean.squeeze(-1)
+
+    # Apply correction: reshape for broadcasting to [B, C, H, W]
+    scale_ = scale[:, :, np.newaxis, np.newaxis]
+    bias_ = bias[:, :, np.newaxis, np.newaxis]
+    transformed = src * scale_ + bias_
+
+    if censored_fit:
+        a, b, sigma = censored_linear_fit_twosided(tgt_s.astype(np.float32()), src_s.astype(np.float32()))
+        transformed =  a + b * transformed
+
+    return transformed, scale, bias
+
+
+def scaled_dot_product(x1, x2, eps=1e-6):
+    # Sum along axis 2 (C if input is B, H, C)
+    dot = (x1 * x2).sum(axis=2, keepdims=True)
+    x1_mag = np.sqrt((x1 * x1).sum(axis=2, keepdims=True))
+    x2_mag = np.sqrt((x2 * x2).sum(axis=2, keepdims=True))
+    return dot / (x1_mag + x2_mag + eps)
+
+
+def postprocess(img, denoised, lumi_blend=0, chroma_blend=0, eps=1e-6,
+                clip_highlights=False, affine=False):
+    
+    if affine:
+        _denoised = np.expand_dims(denoised.transpose(2, 0, 1), axis=0)
+        _img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
+        _denoised, _, _ = match_colors_linear(_denoised, _img)
+        denoised = _denoised[0].transpose(1, 2, 0)
+
+    dot = (img * denoised).sum(axis=2, keepdims=True)
+    denoised_mag = (denoised * denoised).sum(axis=2, keepdims=True) ** .5
+    
+    # Project denoised along original image vector
+    lumi = dot / (denoised_mag ** 2 + eps) * denoised
+    chroma = img - lumi 
+    output = (1 - lumi_blend) * denoised + lumi * (lumi_blend) + chroma_blend * chroma
+
+    if clip_highlights:
+        output = clip_highlights_func(img, output)
+
+
+
+    return output
+
+def clip_highlights_func(img, denoised):
+   out = np.copy(denoised)
+   mask = (img == 1)
+   out[mask] = 1
+   return out

{rawforge-0.1.0 → rawforge-0.2.0}/RawForge/main.py

@@ -1,8 +1,15 @@
 import sys
 import platform
+import argparse
+from pathlib import Path
 from  RawForge.application.ModelHandler import ModelHandler 
 from RawForge.application.postprocessing import postprocess
-import argparse
+from  RawForge.application.MODEL_REGISTRY import MODEL_REGISTRY 
+from RawForge.application.InferenceWorker import InferenceWorker
+from RawForge.application.ImageSaver import ImageSaver
+from RawForge.application.helpers.get_image import get_image
+import numpy as np
+
 # import glob
 
 def main():
@@ -17,13 +24,12 @@ def main():
     parser.add_argument('--device', type=str, help='Set device backend (cuda, cpu, mps).')
     parser.add_argument('--disable_tqdm', action='store_true', help='Disable the progress bar.')
     parser.add_argument('--tile_size', type=int, help='Set tile size. (default: 256)', default=256)
+    parser.add_argument('--tile_overlap', type=float, help='Set tile overlap. (default: 0.25)', default=0.25)
 
     parser.add_argument('--lumi', type=float, help='Lumi noise (0-1).', default=0)
     parser.add_argument('--chroma', type=float, help='Chroma noise (0-1).', default=0)
     parser.add_argument('--clip_highlights', action='store_true', help='Do not run model on clipped highlights.')
-
-    args = parser.parse_args()
-
+    parser.add_argument('--affine', action='store_true', help='Affine fit the model to the input.')
     # # Glob handeling
     # in_files = sorted(glob.glob(args.in_file))
     # if not in_files:
@@ -33,30 +39,39 @@ def main():
     #         raise ValueError(
     #             "When using glob input, out_file must be a directory."
     #         )
-    
-    handler = ModelHandler()
-
-    handler.load_model(args.model)
-
-    iso = handler.load_rh(args.in_file)
+    args = parser.parse_args()
 
+    # Initialize
+    models = args.model.split(',')
+    model_params = MODEL_REGISTRY[models[0]]
+    rh, image_RGB, iso = get_image(args.in_file, model_params, dims=args.dims)
     if not args.conditioning:
         conditioning  = [iso, 0]
     else: 
         conditioning = [int(x) for x in args.conditioning.split(',')]
-
-    if args.device:
-        handler.set_device(args.device)
-
-
     inference_kwargs = {"disable_tqdm": args.disable_tqdm,
-                        "tile_size": args.tile_size}
-    img, denoised_image = handler.run_inference(conditioning=conditioning, dims=args.dims, inference_kwargs=inference_kwargs)
-
-    output = postprocess(img, denoised_image, lumi_blend=args.lumi, chroma_blend=args.chroma, eps=1e-6,
-                         clip_highlights=args.clip_highlights)
-    handler.handle_full_image(output, args.out_file, args.cfa)
+                        "tile_size": args.tile_size,
+                        "tile_overlap": args.tile_overlap}
+    
+    # Run processing
+    output_img = image_RGB
+    for model in models:
+        handler = ModelHandler()
+        handler.load_model(model)
+        worker = InferenceWorker(handler.model, handler.model_params, conditioning, **inference_kwargs)
+        _, output_img = worker._tile_process(output_img, handler.model_params)
 
+    # Postprocess
+    output = postprocess(image_RGB, output_img, lumi_blend=args.lumi, chroma_blend=args.chroma, eps=1e-6,
+                         clip_highlights=args.clip_highlights, affine=args.affine)
+    
+    # Save image
+    saver = ImageSaver(model_params, rh, dims=args.dims)
+    apply_ccm = model_params['demosaicing'] == 'rawpy'
+    if Path(args.out_file).suffix=='.tiff':
+        saver.to_tiff(output, args.out_file, apply_ccm=apply_ccm)
+    else:
+        saver.to_raw(output, args.out_file, args.cfa)
 
 if __name__ == '__main__':
     main()

{rawforge-0.1.0 → rawforge-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rawforge" 
-version = "0.1.0"
+version = "0.2.0"
 description = "A compute backend/CLI application for using machine learning models on raw images."
 authors = [
     { name = "Ryan Mueller"},
@@ -21,7 +21,7 @@ dependencies = [
     "numpy >=2.2", 
     "RawHandler ~=0.2.0",
     "colour_demosaicing ~=0.2.6",
-    "blended_tiling >=0.0.1.dev7",
+    "blended-tiling-numpy",
     "requests ~=2.32",
     "platformdirs ~=4.5",
     "tqdm ~=4.67",
@@ -30,6 +30,13 @@ dependencies = [
 ]
 
 
+[project.optional-dependencies]
+cpu = ["onnxruntime"]
+cuda = ["onnxruntime-gpu"]
+web = ["onnxruntime-web"]
+directml = ["onnxruntime-directml"]
+
+
 [project.urls]
 "Homepage" = "https://github.com/rymuelle/RawForge"
 "Bug Tracker" = "https://github.com/rymuelle/RawForge/issues"
@@ -43,4 +50,4 @@ where = ["."]
 include = ["RawForge*"]
 
 [project.scripts]
-rawforge = "RawForge.main:main"
+rawforgeonnx = "RawForge.main:main"

{rawforge-0.1.0 → rawforge-0.2.0}/rawforge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rawforge
-Version: 0.1.0
+Version: 0.2.0
 Summary: A compute backend/CLI application for using machine learning models on raw images.
 Author: Ryan Mueller
 License: MIT License
@@ -40,12 +40,20 @@ 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
+Provides-Extra: cpu
+Requires-Dist: onnxruntime; extra == "cpu"
+Provides-Extra: cuda
+Requires-Dist: onnxruntime-gpu; extra == "cuda"
+Provides-Extra: web
+Requires-Dist: onnxruntime-web; extra == "web"
+Provides-Extra: directml
+Requires-Dist: onnxruntime-directml; extra == "directml"
 Dynamic: license-file
 
 # RawRefinery

{rawforge-0.1.0 → rawforge-0.2.0}/rawforge.egg-info/SOURCES.txt

@@ -3,13 +3,15 @@ README.md
 pyproject.toml
 RawForge/__init__.py
 RawForge/main.py
+RawForge/application/ImageSaver.py
 RawForge/application/InferenceWorker.py
-RawForge/application/InferenceWorkerRawpy.py
 RawForge/application/MODEL_REGISTRY.py
 RawForge/application/ModelHandler.py
 RawForge/application/dng_utils.py
 RawForge/application/postprocessing.py
-RawForge/application/utils.py
+RawForge/application/helpers/censored_fit.py
+RawForge/application/helpers/get_image.py
+RawForge/application/helpers/utils.py
 RawForge/scripts/generate_keys.py
 RawForge/scripts/sign_models.py
 rawforge.egg-info/PKG-INFO

rawforge-0.2.0/rawforge.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+rawforgeonnx = RawForge.main:main

{rawforge-0.1.0 → rawforge-0.2.0}/rawforge.egg-info/requires.txt

@@ -1,9 +1,21 @@
 numpy>=2.2
 RawHandler~=0.2.0
 colour_demosaicing~=0.2.6
-blended_tiling>=0.0.1.dev7
+blended-tiling-numpy
 requests~=2.32
 platformdirs~=4.5
 tqdm~=4.67
 cryptography>=46.0
 tifffile
+
+[cpu]
+onnxruntime
+
+[cuda]
+onnxruntime-gpu
+
+[directml]
+onnxruntime-directml
+
+[web]
+onnxruntime-web

rawforge-0.1.0/RawForge/application/InferenceWorker.py

@@ -1,98 +0,0 @@
-import torch
-from blended_tiling import TilingModule
-from colour_demosaicing import demosaicing_CFA_Bayer_Malvar2004
-from RawForge.application.postprocessing import match_colors_linear
-from tqdm import tqdm
-
-class InferenceWorker():
-    def __init__(self, model, model_params, device, rh, conditioning, dims, tile_size=256, tile_overlap=0.25, batch_size=2, disable_tqdm=False):
-        super().__init__()
-        self.model = model
-        self.model_params = model_params
-        self.device = device
-        self.rh = rh
-        self.conditioning = conditioning
-        self.dims = dims
-        # Quick and dirty hack to force to be even
-        self.tile_size = tile_size
-        self.tile_overlap = tile_overlap
-        self.batch_size = batch_size
-        self._is_cancelled = False
-        self.disable_tqdm = disable_tqdm
-
-    def cancel(self):
-        self._is_cancelled = True
-
-    def _tile_process(self):
-        # Prepare Data
-        image_RGGB = self.rh.as_rggb(dims=self.dims, colorspace='lin_rec2020')
-        image_RGB = self.rh.as_rgb(dims=self.dims, demosaicing_func=demosaicing_CFA_Bayer_Malvar2004, colorspace='lin_rec2020', clip=True)
-        tensor_image = torch.from_numpy(image_RGGB).unsqueeze(0).contiguous()
-        tensor_RGB = torch.from_numpy(image_RGB).unsqueeze(0).contiguous()
-
-        full_size = [image_RGGB.shape[1], image_RGGB.shape[2]]
-        tile_size = [self.tile_size // 2, self.tile_size // 2]
-        overlap = [self.tile_overlap, self.tile_overlap]
-
-        # Tiling Setup
-        tiling_module = TilingModule(tile_size=tile_size, tile_overlap=overlap, base_size=full_size)
-        tiling_module_rgb = TilingModule(tile_size=[s*2 for s in tile_size], tile_overlap=overlap, base_size=[s*2 for s in full_size])
-        tiling_module_rebuild = TilingModule(tile_size=[s*2 for s in tile_size], tile_overlap=overlap, base_size=[s*2 for s in full_size])
-
-        tiles = tiling_module.split_into_tiles(tensor_image).float().to(self.device)
-        tiles_rgb = tiling_module_rgb.split_into_tiles(tensor_RGB).float().to(self.device)
-        
-        batches = torch.split(tiles, self.batch_size)
-        batches_rgb = torch.split(tiles_rgb, self.batch_size)
-
-        # Conditioning Setup
-        cond_tensor = torch.as_tensor(self.conditioning, device=self.device).float().unsqueeze(0)
-        cond_tensor[:, 0] /= 6400
-        cond_tensor[:, 1] = 0
-        cond_tensor = cond_tensor[:, 0:1]
-
-        processed_batches = []
-        
-        # Determine Dtype
-        dtype_map = {'mps': torch.float16, 'cuda': torch.float16, 'cpu': torch.bfloat16}
-        autocast_dtype = dtype_map.get(self.device.type, torch.float32)
-
-        total_batches = len(batches_rgb)
-        
-        # Inference Loop
-        with torch.no_grad():
-            with torch.autocast(device_type=self.device.type, dtype=autocast_dtype):
-                for i, (batch, batch_rgb) in tqdm(enumerate(zip(batches, batches_rgb)), disable=self.disable_tqdm):
-                    if self._is_cancelled: return None, None
-                    
-                    B = batch.shape[0]
-                    # Expand conditioning to match batch size
-                    curr_cond = cond_tensor.expand(B, -1)
-                    
-                    output = self.model(batch_rgb, curr_cond)
-
-                    # Output processing
-                    if "affine" in self.model_params:
-                        output, _, _ = match_colors_linear(output, batch_rgb)
-                    processed_batches.append(output.cpu())
-                    
-        # Rebuild
-        tiles_out = torch.cat(processed_batches, dim=0)
-        stitched = tiling_module_rebuild.rebuild_with_masks(tiles_out).detach().cpu().numpy()[0]
-
-        torch.cuda.empty_cache()
-
-        return image_RGB.transpose(1, 2, 0), stitched.transpose(1, 2, 0)
-    
-    def run(self):
-        try:
-            img, denoised_img = self._tile_process()
-
-            # Post-process blending
-            blend_alpha = self.conditioning[1] / 100
-            final_denoised = (denoised_img * (1 - blend_alpha)) + (img * blend_alpha)
-            
-            return img, final_denoised
-            
-        except Exception as e:
-            print(str(e))

rawforge-0.1.0/RawForge/application/InferenceWorkerRawpy.py

@@ -1,108 +0,0 @@
-import torch
-from blended_tiling import TilingModule
-from RawForge.application.postprocessing import match_colors_linear
-from tqdm import tqdm
-import rawpy
-
-class InferenceWorkerRawpy():
-    def __init__(self, model, model_params, device, rh, conditioning, dims, tile_size=512, tile_overlap=0.25, batch_size=2, disable_tqdm=False):
-        super().__init__()
-        self.model = model
-        self.model_params = model_params
-        self.device = device
-        self.rh = rh
-        self.conditioning = conditioning
-        self.dims = dims
-        self.tile_size = tile_size
-        if 'tile_size' in model_params:
-            self.tile_size =  model_params['tile_size']
-        self.tile_overlap = tile_overlap
-        self.batch_size = batch_size
-        if 'batch_size' in model_params:
-            self.batch_size = model_params['batch_size']
-        self._is_cancelled = False
-        self.disable_tqdm = disable_tqdm
-
-    def cancel(self):
-        self._is_cancelled = True
-
-    def _tile_process(self):
-        # Prepare Data
-        image_RGB = self.rh.rawpy_object.postprocess(
-                    user_wb=[1, 1, 1, 1],
-                    output_color=rawpy.ColorSpace.raw,
-                    demosaic_algorithm= rawpy.DemosaicAlgorithm(3),
-                    no_auto_bright=True,
-                    use_camera_wb=False,
-                    use_auto_wb=False,
-                    gamma=(1, 1),
-                    user_flip=0,
-                    output_bps=16,
-                    no_auto_scale=True,
-                ) / self.rh.rawpy_object.white_level
-
-        image_RGB = image_RGB.transpose(2, 0, 1)
-
-        tensor_RGB = torch.from_numpy(image_RGB).unsqueeze(0).contiguous()
-
-        full_size = [image_RGB.shape[1], image_RGB.shape[2]]
-        tile_size = [self.tile_size, self.tile_size]
-        overlap = [self.tile_overlap, self.tile_overlap]
-
-        # Tiling Setup
-        tiling_module_rgb = TilingModule(tile_size=[s for s in tile_size], tile_overlap=overlap, base_size=[s for s in full_size])
-
-        tiles_rgb = tiling_module_rgb.split_into_tiles(tensor_RGB).float().to(self.device)
-        
-        batches_rgb = torch.split(tiles_rgb, self.batch_size)
-
-        # Conditioning Setup
-        cond_tensor = torch.as_tensor(self.conditioning, device=self.device).float().unsqueeze(0)
-        cond_tensor[:, 0] /= 6400
-        cond_tensor[:, 1] = 0
-        cond_tensor = cond_tensor[:, 0:1]
-
-        processed_batches = []
-        
-        # Determine Dtype
-        dtype_map = {'mps': torch.float16, 'cuda': torch.float16, 'cpu': torch.bfloat16}
-        autocast_dtype = dtype_map.get(self.device.type, torch.float32)
-
-        total_batches = len(batches_rgb)
-        
-        # Inference Loop
-        with torch.no_grad():
-            with torch.autocast(device_type=self.device.type, dtype=autocast_dtype):
-                for i, (batch_rgb) in tqdm(enumerate(batches_rgb), disable=self.disable_tqdm):
-                    if self._is_cancelled: return None, None
-                    
-                    B = batch_rgb.shape[0]
-                    # Expand conditioning to match batch size
-                    curr_cond = cond_tensor.expand(B, -1)
-                    output = self.model(batch_rgb, curr_cond*0)
-                    processed_batches.append(output.cpu())
-                    
-        # Rebuild
-        tiles_out = torch.cat(processed_batches, dim=0)
-        stitched = tiling_module_rgb.rebuild_with_masks(tiles_out).detach().cpu()
-
-        if "affine" in self.model_params:
-            stitched, _, _ = match_colors_linear(stitched, tensor_RGB)
-
-        stitched = stitched.numpy()[0]
-        torch.cuda.empty_cache()
-
-        return image_RGB.transpose(1, 2, 0), stitched.transpose(1, 2, 0)
-    
-    def run(self):
-        try:
-            img, denoised_img = self._tile_process()
-
-            # Post-process blending
-            blend_alpha = self.conditioning[1] / 100
-            final_denoised = (denoised_img * (1 - blend_alpha)) + (img * blend_alpha)
-            
-            return img, final_denoised
-            
-        except Exception as e:
-            print(str(e))

rawforge-0.1.0/RawForge/application/MODEL_REGISTRY.py

@@ -1,59 +0,0 @@
-MODEL_REGISTRY = {
-    "TreeNetDenoise": {
-        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1.pt",
-        "filename": "ShadowWeightedL1.pt",
-        "max_iso": 65535,
-    },
-    "TreeNetDenoiseLight": {
-        "url": "    https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1_light.pt",
-        "filename": "ShadowWeightedL1_light.pt",
-        "max_iso": 65535,
-    },
-    "TreeNetDenoiseSuperLight": {
-        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1_super_light.pt",
-        "filename": "ShadowWeightedL1_super_light.pt",
-        "max_iso": 65535,
-    },
-    "TreeNetDenoiseHeavy": {
-        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1_24_deep_500.pt",
-        "filename": "ShadowWeightedL1_24_deep_500.pt",
-        "max_iso": 65535,
-    },
-    "Deblur": {
-        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/realblur_gamma_140.pt",
-        "filename": "realblur_gamma_140.pt",
-        "affine": True,
-    },
-    "DeepSharpen": {
-        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/Deblur_deep_24.pt",
-        "filename": "Deblur_deep_24.pt",
-        "affine": True,
-    },
-    "TreeNetDenoiseXTrans": {
-        "url": "https://github.com/rymuelle/RawForge/releases/download/xtrans_v1.0.0/xtrans_fixed_exposure_no_conditioning_400.pt",
-        "filename": "xtrans_fixed_exposure_no_conditioning_400.pt",
-        "backend": "rawpy",
-        "conditioning": "false",
-    },
-    "XFormerXTrans": {
-        "url": "https://github.com/rymuelle/RawForge/releases/download/xtrans_v1.0.0/xformer.pt",
-        "filename": "xformer.pt",
-        "backend": "rawpy",
-        "conditioning": "false",
-        "batch_size": 1, 
-    }, 
-    "XFormerXTrans352": {
-        "url": "https://github.com/rymuelle/RawForge/releases/download/xtrans_v1.0.0/xformer_352.pt",
-        "filename": "xformer_352.pt",
-        "backend": "rawpy",
-        "conditioning": "false",
-        "batch_size": 1, 
-    }, 
-    "RestormerXTrans": {
-        "url": "https://github.com/rymuelle/RawForge/releases/download/xtrans_v1.0.0/restormer.pt",
-        "filename": "restormer.pt",
-        "backend": "rawpy",
-        "conditioning": "false",
-        "batch_size": 1, 
-    }, 
-}

rawforge-0.1.0/RawForge/application/postprocessing.py

@@ -1,87 +0,0 @@
-import torch
-
-def match_colors_linear(
-    src: torch.Tensor, 
-    tgt: torch.Tensor, 
-    sample_fraction: float = 0.05
-):
-    """
-    Fit per-channel affine color transforms:
-        tgt ≈ scale * src + bias
-
-    Args:
-        src: [B, C, H, W] source tensor
-        tgt: [B, C, H, W] target tensor
-        sample_fraction: fraction of pixels to use for fitting
-
-    Returns:
-        transformed_src: source after color matching
-        scale: [B, C]
-        bias:  [B, C]
-    """
-
-    B, C, H, W = src.shape
-    device = src.device
-
-    # Flatten spatial dims
-    src_flat = src.view(B, C, -1)
-    tgt_flat = tgt.view(B, C, -1)
-
-    # Sample subset of pixels
-    N = src_flat.shape[-1]
-    k = max(64, int(N * sample_fraction))
-
-    idx = torch.randint(0, N, (k,), device=device)
-
-    src_s = src_flat[..., idx]  # [B, C, k]
-    tgt_s = tgt_flat[..., idx]
-
-    # Compute scale and bias using least squares
-    # scale = cov(src, tgt) / var(src)
-    src_mean = src_s.mean(-1, keepdim=True)
-    tgt_mean = tgt_s.mean(-1, keepdim=True)
-
-    src_centered = src_s - src_mean
-    tgt_centered = tgt_s - tgt_mean
-
-    var_src = (src_centered ** 2).mean(-1)
-    cov = (src_centered * tgt_centered).mean(-1)
-
-    scale = cov / (var_src + 1e-8)            # [B, C]
-    bias = tgt_mean.squeeze(-1) - scale * src_mean.squeeze(-1)
-
-    # Apply correction
-    scale_ = scale.view(B, C, 1, 1)
-    bias_ = bias.view(B, C, 1, 1)
-    transformed = src * scale_ + bias_
-
-    return transformed, scale, bias
-
-
-def scaled_dot_product(x1, x2, eps=1e-6):
-    dot = (x1 * x2).sum(axis=2, keepdims=True)
-    x1_mag = (x1 * x1).sum(axis=2, keepdims=True) ** .5
-    x2_mag = (x2 * x2).sum(axis=2, keepdims=True) ** .5
-    return dot/(x1_mag+x2_mag+eps)
-
-
-def postprocess(img, denoised, lumi_blend=0, chroma_blend=0, eps=1e-6,
-                clip_highlights=False):
-    # Suggested by Jakob Andrén
-    dot = (img * denoised).sum(axis=2, keepdims=True)
-    img_mag = (img * img).sum(axis=2, keepdims=True) ** .5
-    denoised_mag = (denoised * denoised).sum(axis=2, keepdims=True) ** .5
-    # Project denoised along original image vector
-    lumi = dot / (denoised_mag ** 2 + eps) * denoised
-    chroma = img - lumi 
-    output = (1-lumi_blend) * denoised + lumi * (lumi_blend) + chroma_blend * chroma
-
-    if clip_highlights:
-        output = clip_highlights_func(img, output)
-
-    return output
-
-def clip_highlights_func(img, denoised):
-   mask = img == 1
-   denoised[mask] = 1
-   return denoised

rawforge-0.1.0/RawForge/application/utils.py

@@ -1,10 +0,0 @@
-import torch
-
-def can_use_gpu():
-    if not torch.cuda.is_available():
-        return False
-    try:
-        x = torch.zeros(1, device="cuda")
-        return True
-    except Exception:
-        return False

rawforge-0.1.0/rawforge.egg-info/entry_points.txt

@@ -1,2 +0,0 @@
-[console_scripts]
-rawforge = RawForge.main:main