diff --git a/README.md b/README.md index 5737655..fc617ea 100644 --- a/README.md +++ b/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. - -\ \ No newline at end of file +> [https://doi.org/10.14428/DVN/DEQCIM](https://doi.org/10.14428/DVN/DEQCIM), Open Data @ UCLouvain, V1. \ No newline at end of file diff --git a/examples/dataset_for_noisy_images.ipynb b/examples/dataset_for_noisy_images.ipynb index 4c7b569..17d7a37 100644 --- a/examples/dataset_for_noisy_images.ipynb +++ b/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", diff --git a/examples/simple_demosaicing.ipynb b/examples/simple_demosaicing.ipynb index 30964e4..e4772f8 100644 --- a/examples/simple_demosaicing.ipynb +++ b/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,7 +150,14 @@ "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", @@ -158,7 +165,13 @@ "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": { diff --git a/pyproject.toml b/pyproject.toml index 6feab43..995d6fe 100644 --- a/pyproject.toml +++ b/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"}, diff --git a/src/RawHandler/RawDataset.py b/src/RawHandler/RawDataset.py index e0fbd5d..ad8fe48 100644 --- a/src/RawHandler/RawDataset.py +++ b/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 diff --git a/src/RawHandler/RawHandler.py b/src/RawHandler/RawHandler.py index 949926d..8d0f143 100644 --- a/src/RawHandler/RawHandler.py +++ b/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, )