Added functionality to store the original white balance and rotations and flips

README.md

@@ -58,7 +58,7 @@ pip install .
 
 ---
 
-## 📘 Example
+## Example
 
 A simple demo notebook is available:
 
@@ -82,6 +82,4 @@ Special thanks to the authors of **RawNIND**:
 
 > 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.
-
-\
+> [https://doi.org/10.14428/DVN/DEQCIM](https://doi.org/10.14428/DVN/DEQCIM), Open Data @ UCLouvain, V1.

examples/dataset_for_noisy_images.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "405d14a3",
    "metadata": {},
    "outputs": [],
@@ -15,7 +15,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "2bee5068",
    "metadata": {},
    "outputs": [],
@@ -37,10 +37,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "49b48e2c",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "AttributeError",
+     "evalue": "'BaseRawHandler' object has no attribute 'input_handler'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 8\u001b[39m\n\u001b[32m      5\u001b[39m high_iso_rh = RawHandler(\u001b[33m\"\u001b[39m\u001b[33mgtBark_12800.arw\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m      6\u001b[39m dims = (\u001b[32m1500\u001b[39m, \u001b[32m1500\u001b[39m + \u001b[32m200\u001b[39m, \u001b[32m4500\u001b[39m, \u001b[32m4500\u001b[39m + \u001b[32m200\u001b[39m)\n\u001b[32m----> \u001b[39m\u001b[32m8\u001b[39m offset = \u001b[43malign_images\u001b[49m\u001b[43m(\u001b[49m\u001b[43mhigh_iso_rh\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlow_iso_rh\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdims\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moffset\u001b[49m\u001b[43m=\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m      9\u001b[39m offset\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Develop/RawHandler/src/RawHandler/utils.py:40\u001b[39m, in \u001b[36malign_images\u001b[39m\u001b[34m(rh1, rh2, dims, offset, max_iters, step_sizes)\u001b[39m\n\u001b[32m     36\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34malign_images\u001b[39m(\n\u001b[32m     37\u001b[39m     rh1, rh2, dims, offset=(\u001b[32m0\u001b[39m, \u001b[32m0\u001b[39m, \u001b[32m0\u001b[39m, \u001b[32m0\u001b[39m), max_iters=\u001b[32m100\u001b[39m, step_sizes=[\u001b[32m16\u001b[39m, \u001b[32m8\u001b[39m, \u001b[32m4\u001b[39m, \u001b[32m2\u001b[39m]\n\u001b[32m     38\u001b[39m ):\n\u001b[32m     39\u001b[39m     offset = np.array(offset)\n\u001b[32m---> \u001b[39m\u001b[32m40\u001b[39m     bayer1 = \u001b[43mrh1\u001b[49m\u001b[43m.\u001b[49m\u001b[43minput_handler\u001b[49m(dims=dims)\n\u001b[32m     41\u001b[39m     img_shape = rh1.raw.shape[-\u001b[32m2\u001b[39m:]\n\u001b[32m     43\u001b[39m     loss = get_loss(bayer1, rh2.input_handler(dims=dims + offset))\n",
+      "\u001b[31mAttributeError\u001b[39m: 'BaseRawHandler' object has no attribute 'input_handler'"
+     ]
+    }
+   ],
    "source": [
     "# A simple code to align image is also provided for use in supervised learning.\n",
     "# The first raw handler is the target to be used to align to, the second is aligned.\n",

examples/simple_demosaicing.ipynb

@@ -51,7 +51,7 @@
    "source": [
     "# We can acccess the underlying bayer data\n",
     "# dims = (h1, h2, w1, w2)\n",
-    "dims = (1500, 1500 + 200, 4500, 4500 + 200)\n",
+    "dims = (1800, 1800 + 200, 2700, 2700 + 200)\n",
     "bayer = rh.raw\n",
     "plt.imshow(bayer[dims[0] : dims[1], dims[2] : dims[3]])"
    ]
@@ -150,15 +150,28 @@
    "id": "c2309479",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "## We can flip the image (axis=0 for vertical, axis=1 for horizontal)\n",
+    "\n",
+    "rh.flip(axis=0)\n",
+    "rh.flip(axis=1)\n",
+    "img = rh.generate_thumbnail(clip=True)\n",
+    "plt.imshow(linear_to_srgb(img))"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "6aaf94c4",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "## We can rotate the image in 90 degree increments. k determines the number of increments.\n",
+    "rh.rotate(k=1)\n",
+    "rgb = rh.as_rgb(dims=dims)\n",
+    "plt.imshow(linear_to_srgb(rgb).transpose(1, 2, 0))\n",
+    "rgb.min(), rgb.max()"
+   ]
   }
  ],
  "metadata": {

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "RawHandler" 
-version = "0.0.2"
+version = "0.0.3"
 description = "A basic library to handle camera raw files for use in machine learning. Built on rawpy and cv2."
 authors = [
     { name = "Ryan Mueller"},

src/RawHandler/RawDataset.py

@@ -1,8 +1,10 @@
 import random
 from torch.utils.data import Dataset
+import torch
 
 from RawHandler.RawHandler import RawHandler
-from RawHandler.utils import align_images
+
+import re
 
 
 class RawDataset(Dataset):
@@ -25,8 +27,8 @@ def __getitem__(self, idx):
         # Crop and align
         H, W = noisy_rh.raw.shape[-2:]
         half_crop = self.crop_size // 2
-        H_center = random.randint(0 + half_crop, H - half_crop)
-        W_center = random.randint(0 + half_crop, W - half_crop)
+        H_center = random.randint(0 + half_crop * 2, H - half_crop * 2)
+        W_center = random.randint(0 + half_crop * 2, W - half_crop * 2)
         crop = (
             H_center - half_crop,
             H_center + half_crop,
@@ -36,16 +38,33 @@ def __getitem__(self, idx):
         if self.offsets is None:
             offset = (0, 0, 0, 0)
         else:
-            offset = self.offsets[idx]
-        offset = align_images(noisy_rh, gt_rh, crop, offset=(0, 0, 0, 0))
+            offset = self.offsets[idx][0][0]
 
         # Adjust exposure
-        noisy_rggb = noisy_rh.as_rggb(dims=crop)
-        noisy_rgb = noisy_rh.as_rgb(dims=crop)
-        clean_rgb = gt_rh.as_rgb(dims=crop)
+        gain = (
+            noisy_rh.adjust_bayer_bw_levels(dims=crop).mean()
+            / gt_rh.adjust_bayer_bw_levels(dims=crop).mean()
+        )
+        gt_rh.gain = gain
+
+        # offset = align_images(noisy_rh, gt_rh, crop, offset=offset, step_sizes=[2])
+
+        noisy_rggb = noisy_rh.as_rggb_colorspace(dims=crop, colorspace="AdobeRGB")
+        noisy_rgb = noisy_rh.as_rgb_colorspace(dims=crop, colorspace="AdobeRGB")
+        clean_rgb = gt_rh.as_rgb_colorspace(dims=crop + offset, colorspace="AdobeRGB")
+
+        iso = re.findall("_ISO([0-9]+)_", noisy_file)
+        if len(iso) == 1:
+            iso = int(iso[0])
+        else:
+            iso = -100
+
+        iso_conditioning = iso / 65535
 
         if self.transform:
-            noisy_rggb = self.transform(noisy_rggb)
-            noisy_rgb = self.transform(noisy_rgb)
-            clean_rgb = self.transform(clean_rgb)
-        return noisy_rggb, noisy_rgb, clean_rgb, offset
+            noisy_rggb = self.transform(noisy_rggb.transpose(1, 2, 0))
+            noisy_rgb = self.transform(noisy_rgb.transpose(1, 2, 0))
+            clean_rgb = self.transform(clean_rgb.transpose(1, 2, 0))
+            iso_conditioning = torch.tensor([iso_conditioning])
+
+        return noisy_rggb, noisy_rgb, clean_rgb, offset, iso_conditioning

src/RawHandler/RawHandler.py

@@ -20,6 +20,7 @@ class CoreRawMetadata(NamedTuple):
     white_level: int
     rgb_xyz_matrix: np.ndarray
     raw_pattern: np.ndarray
+    camera_white_balance: np.ndarray
     iheight: int
     iwidth: int
 
@@ -57,6 +58,22 @@ def _remove_masked_pixels(self, img: np.ndarray) -> np.ndarray:
         """Removes masked pixels from the image based on core_metadata.iheight and core_metadata.iwidth."""
         return img[:, 0 : self.core_metadata.iheight, 0 : self.core_metadata.iwidth]
 
+    def flip(self, axis=1):
+        raw = np.flip(self.raw, axis=axis)
+        if axis == 1:
+            self.raw = safe_crop(raw, dx=1, dy=0)
+        else:
+            self.raw = safe_crop(raw, dx=0, dy=1)
+
+    def rotate(self, k=1):
+        raw = np.rot90(self.raw, k=k)
+        if k == 1:
+            self.raw = safe_crop(raw, dx=0, dy=1)
+        if k == 2:
+            self.raw = safe_crop(raw, dx=1, dy=1)
+        if k == 3:
+            self.raw = safe_crop(raw, dx=1, dy=0)
+
     def _input_handler(self, dims=None) -> np.ndarray:
         """
         Crops bayer array.
@@ -75,7 +92,7 @@ def _input_handler(self, dims=None) -> np.ndarray:
         else:
             return img
 
-    def _adjust_bayer_bw_levels(self, dims=None) -> np.ndarray:
+    def _adjust_bayer_bw_levels(self, dims=None, clip=False) -> np.ndarray:
         """
         Adjusts black and white levels of Bayer data.
         """
@@ -90,7 +107,8 @@ def _adjust_bayer_bw_levels(self, dims=None) -> np.ndarray:
                 self.core_metadata.white_level
                 - self.core_metadata.black_level_per_channel[channel]
             )
-        img = np.clip(img, 0, 1)
+        if clip:
+            img = np.clip(img, 0, 1)
         return img
 
     def _make_bayer_map(self, bayer: np.ndarray) -> np.ndarray:
@@ -240,6 +258,7 @@ def __new__(cls, path: str, **kwargs):
             white_level=rawpy_object.white_level,
             rgb_xyz_matrix=rawpy_object.rgb_xyz_matrix,
             raw_pattern=rawpy_object.raw_pattern,
+            camera_white_balance=np.array(rawpy_object.camera_whitebalance),
             iheight=rawpy_object.sizes.iheight,
             iwidth=rawpy_object.sizes.iwidth,
         )