rawforge 0.0.1__py3-none-any.whl

RawForge/application/InferenceWorker.py

@@ -0,0 +1,99 @@
+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/application/MODEL_REGISTRY.py

@@ -0,0 +1,30 @@
+MODEL_REGISTRY = {
+    "TreeNetDenoise": {
+        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1.pt",
+        "filename": "ShadowWeightedL1.pt"
+    },
+    "TreeNetDenoiseLight": {
+        "url": "    https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1_light.pt",
+        "filename": "ShadowWeightedL1_light.pt"
+    },
+    "TreeNetDenoiseSuperLight": {
+        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1_super_light.pt",
+        "filename": "ShadowWeightedL1_super_light.pt"
+    },
+
+    "TreeNetDenoiseHeavy": {
+        "url": "https://github.com/rymuelle/RawRefinery/releases/download/v1.2.1-alpha/ShadowWeightedL1_24_deep_500.pt",
+        "filename": "ShadowWeightedL1_24_deep_500.pt"
+    },
+
+    "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,
+    },
+}

RawForge/application/ModelHandler.py

@@ -0,0 +1,188 @@
+import torch
+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 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
+
+key_string = '''-----BEGIN PUBLIC KEY-----
+MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEA8iRGMPqFIFVF0TM/AbMI
+DJUqdjY1S7dGn6rYjLixnhohHLKIo2ZhFUfPaeYrDoqJblP9MxbBLm6a782/Us0A
+vblTQOsdVFHOlVEiDUkG9CJrzh7arqJF+v2LLP9qPIcL5QdIHM+BCKlbNPBU/TJB
+49b6a+1FfKCEeY1z9F8H6GCHGeRB43lz5/1yMoBnq//Rc7NrvinwlNcFYHHM1oj6
+Hk6KPkgitya11QgTTva+XimR7cbw7h9/vJKbrS7tValApio3Ypmx7AKf6/k16S9K
+BCFDN3cyWmjItQNzEWbO2nuM9d3PX2O4FcZVfsA/GU0qSuKFUrrN0KcxKGglLdu4
+3Nt3JmOh+VebVWPSTeMzn2R1LDs2CsDpGG+KnHso80HBBq6RuHTugTiUZ2EwjiXN
+lRS7olKFQOPwT0tm1EVkH8IxQgV4KJbCb6hAScvWfsDdsP+bu4R+QI9hfU6HCWG3
+a8w1AY+5GT7zp1pzKifmnXgMXF3VnAPTGRhpIvPQfum2+tppLZueXlalobK0MDzi
+n36TNhRELao1W7Tvc18fxyZn37BBgKs89JO85/cjD72yhVowW7Hy9lL7RnB+etaN
+ehXoYFsJReNmD5KNgRtmXbsCUJ+D8v7BVYNGl1UgebmQnMdMWyiU/3l1Uuy8HS3L
+1QJYp42f5QqONttCqVzgzrECAwEAAQ==
+-----END PUBLIC KEY-----'''
+
+
+class ModelHandler():
+    """
+    Manages the LifeCycle of the Model, the RawHandler, and the Worker Thread.
+    """
+    def __init__(self):
+        super().__init__()
+
+        self.model = None
+        self.rh = None
+        self.iso = 100
+        self.colorspace = 'lin_rec2020'
+
+        # 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.pub = serialization.load_pem_public_key(key_string.encode('utf-8'))
+
+    def load_rh(self, path):
+        """Loads the raw file handler"""
+        self.rh = RawHandler(path, colorspace=self.colorspace)
+        if 'EXIF ISOSpeedRatings' in self.rh.full_metadata:
+            self.iso = int(self.rh.full_metadata['EXIF ISOSpeedRatings'].values[0])
+        else:
+            self.iso = 100
+        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"]
+
+        # Handle Download
+        if not model_path.is_file():
+            if conf["url"]:
+                print(f"Downloading {model_key}...")
+                if not self._download_file(conf["url"], model_path):
+                    print("Failed to download model.")
+                    return
+            else:
+                 print(f"Model file not found at {model_path}")
+                 return
+
+        try:
+            print(f"Loading model: {model_path}")
+            # 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)
+        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
+
+        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()
+            signature = Path(sig_path).read_bytes()
+            self.pub.verify(
+                signature,
+                data,
+                padding.PSS(
+                    mgf=padding.MGF1(hashes.SHA256()),
+                    salt_length=padding.PSS.MAX_LENGTH
+                ),
+                hashes.SHA256(),
+            )
+            print(f"Model {dest_path} verified!")
+            return True
+        except Exception as e:
+            print(e)
+            if dest_path.exists():
+                dest_path.unlink()
+            if sig_path.exists():
+                sig_path.unlink()
+            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:
+            r = requests.get(url, stream=True)
+            r.raise_for_status()
+            with open(dest_path, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            # Download model signature
+            r = requests.get(url + '.sig', stream=True)
+            r.raise_for_status()
+            sig_path = dest_path.with_suffix(f'{dest_path.suffix}.sig')
+            with open(sig_path, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print("test verification")
+            return self._verify_model(dest_path, sig_path)
+        
+        except Exception as e:
+            print(e)
+            return False
+        
+
+    def handle_full_image(self, denoised, filename, save_cfa):
+        
+        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/application/dng_utils.py

@@ -0,0 +1,201 @@
+import numpy as np
+from pidng.core import RAW2DNG, DNGTags, Tag
+from pidng.defs import *
+
+def get_ratios(string, rh):
+    return [x.as_integer_ratio() for x in rh.full_metadata[string].values]
+
+
+def get_as_shot_neutral(rh, denominator=10000):
+
+    cam_mul = rh.core_metadata.camera_white_balance
+    
+    if cam_mul[0] == 0 or cam_mul[2] == 0:
+        return [[denominator, denominator], [denominator, denominator], [denominator, denominator]]
+
+    r_neutral = cam_mul[1] / cam_mul[0]
+    g_neutral = 1.0 
+    b_neutral = cam_mul[1] / cam_mul[2]
+
+    return [
+        [int(r_neutral * denominator), denominator],
+        [int(g_neutral * denominator), denominator],
+        [int(b_neutral * denominator), denominator],
+    ]
+def convert_ccm_to_rational(matrix_3x3, denominator=10000):
+
+    numerator_matrix = np.round(matrix_3x3 * denominator).astype(int)
+    numerators_flat = numerator_matrix.flatten()
+    ccm_rational = [[num, denominator] for num in numerators_flat]
+    
+    return ccm_rational
+
+
+   
+def simulate_CFA(image, pattern="RGGB", cfa_type="bayer"):
+    """
+    Simulate a CFA image from an RGB image.
+
+    Args:
+        image: numpy array (H, W, 3), RGB image.
+        pattern: CFA pattern string, one of {"RGGB","BGGR","GRBG","GBRG"} for Bayer,
+                 or ignored if cfa_type="xtrans".
+        cfa_type: "bayer" or "xtrans".
+
+    Returns:
+        cfa: numpy array (H, W) CFA image.
+        sparse_mask:  numpy array (H, W, r), mask of pixels.
+    """
+    width = image.shape[1]
+    height = image.shape[0]
+    cfa = np.zeros((height, width, 3), dtype=image.dtype)
+    sparse_mask = np.zeros((height, width, 3), dtype=image.dtype)
+    if cfa_type == "bayer":
+        # 2×2 Bayer masks
+        masks = {
+            "RGGB": np.array([["R", "G"], ["G", "B"]]),
+            "BGGR": np.array([["B", "G"], ["G", "R"]]),
+            "GRBG": np.array([["G", "R"], ["B", "G"]]),
+            "GBRG": np.array([["G", "B"], ["R", "G"]]),
+        }
+        if pattern not in masks:
+            raise ValueError(f"Unknown Bayer pattern: {pattern}")
+
+        mask = masks[pattern]
+        cmap = {"R": 0, "G": 1, "B": 2}
+         
+        for i in range(2):
+            for j in range(2):
+                ch = cmap[mask[i, j]]
+                cfa[i::2, j::2, ch] = image[i::2, j::2, ch]
+                sparse_mask[i::2, j::2, ch] = 1
+    elif cfa_type == "xtrans":
+        # Fuji X-Trans 6×6 repeating pattern
+        xtrans_pattern = np.array([
+            ["G","B","R","G","R","B"],
+            ["R","G","G","B","G","G"],
+            ["B","G","G","R","G","G"],
+            ["G","R","B","G","B","R"],
+            ["B","G","G","R","G","G"],
+            ["R","G","G","B","G","G"],
+        ])
+        cmap = {"R":0, "G":1, "B":2}
+
+        for i in range(6):
+            for j in range(6):
+                ch = cmap[xtrans_pattern[i, j]]
+                cfa[i::6, j::6, ch] = image[i::6, j::6, ch]
+                sparse_mask[i::2, j::2, ch] = 1
+    else:
+        raise ValueError(f"Unknown CFA type: {cfa_type}")
+
+    return cfa.sum(axis=2), sparse_mask
+
+def to_dng(uint_img, rh, filepath, ccm1, save_cfa=True, convert_to_cfa=True, use_orig_wb_points=False):
+    width = uint_img.shape[1]
+    height = uint_img.shape[0]
+    bpp = 16 
+
+    t = DNGTags()
+
+    if save_cfa:
+      if convert_to_cfa:
+        cfa, _ = simulate_CFA(uint_img, pattern="RGGB", cfa_type="bayer")
+        uint_img = cfa.astype(np.uint16)
+      t.set(Tag.BitsPerSample, bpp)
+      t.set(Tag.SamplesPerPixel, 1) 
+      t.set(Tag.PhotometricInterpretation, PhotometricInterpretation.Color_Filter_Array)
+      t.set(Tag.CFARepeatPatternDim, [2,2])
+      t.set(Tag.CFAPattern, CFAPattern.RGGB)
+      t.set(Tag.BlackLevelRepeatDim, [2,2])
+      # This should not be used except to save testing patches
+      if use_orig_wb_points:
+        bl = rh.core_metadata.black_level_per_channel
+        t.set(Tag.BlackLevel, bl)
+        t.set(Tag.WhiteLevel, rh.core_metadata.white_level)
+      else:
+        t.set(Tag.BlackLevel, [0, 0, 0, 0])
+        t.set(Tag.WhiteLevel, 65535)
+    else:
+      t.set(Tag.BitsPerSample, [bpp, bpp, bpp]) # 3 channels for RGB
+      t.set(Tag.SamplesPerPixel, 3) # 3 for RGB
+      t.set(Tag.PhotometricInterpretation, PhotometricInterpretation.Linear_Raw)
+      t.set(Tag.BlackLevel,[0,0,0])
+      t.set(Tag.WhiteLevel, [65535, 65535, 65535])
+
+    t.set(Tag.ImageWidth, width)
+    t.set(Tag.ImageLength, height)
+    t.set(Tag.PlanarConfiguration, 1) # 1 for chunky (interleaved RGB)
+
+    t.set(Tag.TileWidth, width)
+    t.set(Tag.TileLength, height)
+
+    t.set(Tag.ColorMatrix1, ccm1)
+    t.set(Tag.CalibrationIlluminant1, CalibrationIlluminant.D65)
+    wb = get_as_shot_neutral(rh)
+    t.set(Tag.AsShotNeutral, wb)
+    t.set(Tag.BaselineExposure, [[0,100]])
+
+
+    try:
+      t.set(Tag.Make, rh.full_metadata['Image Make'].values)
+      t.set(Tag.Model, rh.full_metadata['Image Model'].values)
+      t.set(Tag.Orientation, rh.full_metadata['Image Orientation'].values[0])
+      exposures = get_ratios('EXIF ExposureTime', rh)
+      fnumber = get_ratios('EXIF FNumber', rh)
+      ExposureBiasValue = get_ratios('EXIF ExposureBiasValue', rh) 
+      FocalLength = get_ratios('EXIF FocalLength', rh) 
+      t.set(Tag.FocalLength, FocalLength)
+      t.set(Tag.EXIFPhotoLensModel, rh.full_metadata['EXIF LensModel'].values)
+      t.set(Tag.ExposureBiasValue, ExposureBiasValue)
+      t.set(Tag.ExposureTime, exposures)
+      t.set(Tag.FNumber, fnumber)
+      t.set(Tag.PhotographicSensitivity, rh.full_metadata['EXIF ISOSpeedRatings'].values)
+    except:
+      print("Could not save EXIF")
+    t.set(Tag.DNGVersion, DNGVersion.V1_4)
+    t.set(Tag.DNGBackwardVersion, DNGVersion.V1_2)
+    t.set(Tag.PreviewColorSpace, PreviewColorSpace.Adobe_RGB)
+
+    r = RAW2DNG()
+
+    r.options(t, path="", compress=False)
+
+    r.convert(uint_img, filename=filepath)
+
+
+
+def convert_color_matrix(matrix):
+  """
+  Converts a 3x3 NumPy matrix of floats into a list of integer pairs.
+
+  Each float value in the matrix is converted to a fractional representation
+  with a denominator of 10000. The numerator is calculated by scaling the
+  float value by 10000 and rounding to the nearest integer.
+
+  Args:
+    matrix: A 3x3 NumPy array with floating-point numbers.
+
+  Returns:
+    A list of 9 lists, where each inner list contains two integers
+    representing the numerator and denominator.
+  """
+  # Ensure the input is a NumPy array
+  if not isinstance(matrix, np.ndarray):
+    raise TypeError("Input must be a NumPy array.")
+
+  # Flatten the 3x3 matrix into a 1D array of 9 elements
+  flattened_matrix = matrix.flatten()
+
+  # Initialize the list for the converted matrix
+  converted_list = []
+  denominator = 10000
+
+  # Iterate over each element in the flattened matrix
+  for element in flattened_matrix:
+    # Scale the element, round it to the nearest integer, and cast to int
+    numerator = int(round(element * denominator))
+    # Append the [numerator, denominator] pair to the result list
+    converted_list.append([numerator, denominator])
+
+  return converted_list

RawForge/application/postprocessing.py

@@ -0,0 +1,58 @@
+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

RawForge/application/utils.py

@@ -0,0 +1,10 @@
+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/main.py

@@ -0,0 +1,41 @@
+import sys
+import platform
+from  RawForge.application.ModelHandler import ModelHandler 
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(description='A command line utility for processing raw images.')
+    parser.add_argument('model', type=str, help='The name of the model to use.')
+    parser.add_argument('in_file', type=str, help='The name of the file to open.')
+    parser.add_argument('out_file', type=str, help='The name of the file to save.')
+    parser.add_argument('--conditioning', type=str, help='Conditioning array to feed model.')
+    parser.add_argument('--dims', type=int, nargs=4, metavar=("x0", "x1", "y0", "y1"), help='Optional crop dimensions.')
+
+    parser.add_argument('--cfa', action='store_true', help='Save the image as a CFA image (default: False).')
+    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)
+
+    args = parser.parse_args()
+
+    handler = ModelHandler()
+
+    handler.load_model(args.model)
+
+    iso = handler.load_rh(args.in_file)
+
+    if not args.conditioning:
+        conditioning  = [iso, 0]
+
+    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)
+
+    handler.handle_full_image(denoised_image, args.out_file, args.cfa)
+
+
+if __name__ == '__main__':
+    main()

rawforge-0.0.1.dist-info/METADATA

@@ -0,0 +1,106 @@
+Metadata-Version: 2.4
+Name: rawforge
+Version: 0.0.1
+Summary: A compute backend/CLI application for using machine learning models on raw images.
+Author: Ryan Mueller
+License: MIT License
+        
+        Copyright (c) 2025 rymuelle
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/rymuelle/RawForge
+Project-URL: Bug Tracker, https://github.com/rymuelle/RawForge/issues
+Keywords: python,Machine Learning,Camera Raw
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pidng~=4.0.9
+Requires-Dist: numpy>=2.3
+Requires-Dist: RawHandler~=0.0.2
+Requires-Dist: colour_demosaicing~=0.2.6
+Requires-Dist: blended_tiling>=0.0.1.dev7
+Requires-Dist: requests~=2.32
+Requires-Dist: platformdirs~=4.5
+Requires-Dist: tqdm~=4.67
+Requires-Dist: cryptography~=46.0
+Dynamic: license-file
+
+# RawRefinery
+
+<!-- [![PyPI version](https://img.shields.io/pypi/v/rawrefinery.svg)](https://pypi.org/project/rawrefinery/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python version](https://img.shields.io/pypi/pyversions/rawrefinery.svg)](https://pypi.org/project/rawrefinery/) -->
+
+**RawForge** is an open-source command line application for **raw image quality refinement and denoising**.
+
+Currently in **alpha release**.
+
+Sorry for the lack of instructions!
+
+---
+Example command line syntax:
+
+```bash
+rawforge TreeNetDenoiseHeavy test.CR2 test_heavy.dng --cfa 
+```
+
+
+```bash 
+usage: rawforge [-h] [--conditioning CONDITIONING] [--dims x0 x1 y0 y1] [--cfa] [--device DEVICE] [--disable_tqdm] [--tile_size TILE_SIZE] model in_file out_file
+
+A command line utility for processing raw images.
+
+positional arguments:
+  model                 The name of the model to use.
+  in_file               The name of the file to open.
+  out_file              The name of the file to save.
+
+options:
+  -h, --help            show this help message and exit
+  --conditioning CONDITIONING
+                        Conditioning array to feed model.
+  --dims x0 x1 y0 y1    Optional crop dimensions.
+  --cfa                 Save the image as a CFA image (default: False).
+  --device DEVICE       Set device backend (cuda, cpu, mps).
+  --disable_tqdm        Disable the progress bar.
+  --tile_size TILE_SIZE
+                        Set tile size. (default: 256)
+```
+
+----
+
+## Acknowledgments
+
+With thanks to:
+
+
+> Brummer, Benoit; De Vleeschouwer, Christophe. (2025).
+> *Raw Natural Image Noise Dataset.*
+> [https://doi.org/10.14428/DVN/DEQCIM](https://doi.org/10.14428/DVN/DEQCIM), Open Data @ UCLouvain, V1.
+
+> Chen, Liangyu; Chu, Xiaojie; Zhang, Xiangyu; Chen, Jianhao. (2022).
+> *NAFNet: Simple Baselines for Image Restoration.*
+> [https://doi.org/10.48550/arXiv.2208.04677](https://doi.org/10.48550/arXiv.2208.04677), arXiv, V1.

rawforge-0.0.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+RawForge/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+RawForge/main.py,sha256=vhA0Q8AT9-5DoZiAEbDXr9pSXDfaGIqvs4Bg12Dvt4o,1696
+RawForge/application/InferenceWorker.py,sha256=7a_zlukd39x9n0Hq1hnIJJq3shLbHvYPONHW5ywLMJo,4216
+RawForge/application/MODEL_REGISTRY.py,sha256=DnVUACayTTXoLQvdCH1DUekGEGQXLgWIlFzvQXs2pmI,1237
+RawForge/application/ModelHandler.py,sha256=j2mlTi_WsCHHW7U808sB9ChQjbeRgn0jbGFLV5Xjfl8,6815
+RawForge/application/dng_utils.py,sha256=F_TFWQ7A19QIBhK3aXAEkUuXpHO8vOA5LixDQNGxGnM,7053
+RawForge/application/postprocessing.py,sha256=_neffnnVbIDp5Ut2LisFk-xX5QIId72qKEBPTSc0c9Q,1493
+RawForge/application/utils.py,sha256=lpGLpGkmhhlGBrav1r0ygBCAbBT0A6kC9fje46_eIu8,205
+rawforge-0.0.1.dist-info/licenses/LICENSE,sha256=C6iC-8U95WxetY9TsV-RHN-CbI9wPI3nFMHMpu_Ve4E,1065
+rawforge-0.0.1.dist-info/METADATA,sha256=WcnBA2HYI68sPOZmNBzKvJVfA_rrH6mXLYsC_ysHADI,4266
+rawforge-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rawforge-0.0.1.dist-info/entry_points.txt,sha256=7yVkkrWfZQqHO9012zxV1igRLOpkNf42KemeslwzhL0,48
+rawforge-0.0.1.dist-info/top_level.txt,sha256=tCW6Kg9PQlghoo0V4yDVw5j49b2LbGZOslrsaRXE1iw,9
+rawforge-0.0.1.dist-info/RECORD,,

rawforge-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

rawforge-0.0.1.dist-info/entry_points.txt

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

rawforge-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 rymuelle
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rawforge-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+RawForge