polyucolor

.gitattributes

@@ -38,3 +38,4 @@ SCBC/Input/IMG_20240215_213619.png filter=lfs diff=lfs merge=lfs -text
 SCBC/Input/IMG_20240215_214449.png filter=lfs diff=lfs merge=lfs -text
 SCBC/Output/IMG_20240215_213330.png filter=lfs diff=lfs merge=lfs -text
 SCBC/Output/IMG_20240215_214449.png filter=lfs diff=lfs merge=lfs -text
+PolyuColor/resources/average_shading.png filter=lfs diff=lfs merge=lfs -text

PolyuColor/.gitignore

@@ -0,0 +1,6 @@
+*.png
+*.jpg
+*.json
+__pycache__
+*.cube
+*.ckpt

PolyuColor/Dockerfile

@@ -0,0 +1,40 @@
+FROM nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y \
+    libpng-dev libjpeg-dev \
+    libopencv-dev ffmpeg \
+    libgl1-mesa-glx && \
+    apt clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN apt update && \
+    apt install -y \
+        wget build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev \
+        libreadline-dev libffi-dev libsqlite3-dev libbz2-dev liblzma-dev && \
+    apt clean && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /temp
+
+RUN wget https://www.python.org/ftp/python/3.9.10/Python-3.9.10.tgz && \
+    tar -xvf Python-3.9.10.tgz
+
+RUN cd Python-3.9.10 && \
+    ./configure --enable-optimizations && \
+    make && \
+    make install
+
+WORKDIR /workspace
+
+RUN rm -r /temp && \
+    ln -s /usr/local/bin/python3 /usr/local/bin/python && \
+    ln -s /usr/local/bin/pip3 /usr/local/bin/pip
+
+COPY requirements.txt .
+RUN python -m pip install --no-cache -r requirements.txt
+RUN pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118 --no-cache
+
+WORKDIR ..
+COPY . .
+CMD ["./run.sh"]

PolyuColor/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Color Reproduction and Synthesis
+
+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.

PolyuColor/README.md

@@ -0,0 +1,78 @@
+# Team: **play1play**
+
+for 'Night Photography Rendering Challenge'
+
+This repo contains the source code of [Night Photography Rendering Challenge 2024](https://nightimaging.org/).
+
+## How to run
+### Run without Docker
+- Install python >= 3.9
+- Install the required packages: `pip install -r requirements.txt`
+- Put the test images in the `data` folder or specify the input image path with `-p` option.
+- Run the script: 
+```
+python run.py -p <input_image_path> -o <output_image_path>
+```
+- The output images will be saved in the `output` folder or specifiy the output path with `-o` option.
+
+### Run with Docker
+- Build the docker image from beginning (optional): 
+```
+docker build -t play1play .
+```
+- Run the docker container with gpu on linux:
+```
+docker run -it --rm --gpus=all -v $(pwd)/data:/data play1play ./run.sh
+```
+`Do not forget to to --gpus flag, our model requires GPU to run.`
+
+## Update
+*2024.3.21:*
+
+Final version v3.0 is released for 3rd validation!
+
+Key features:
+- Utilize the patch-based and calibration-based white-balance algorithm to improve the image quality.
+- Modify the resizing strategy to improve the image quality.
+
+*2024.3.16:*
+
+Release v2.0 for 3rd validation!
+
+Key features:
+- Increase the overall saturation and brightness of the image.
+- Add more contrast to the image.
+- Add more dynamic range to the image.
+
+Algorithm changes:
+- Add luma shading correction (LSC) module
+- Add auto-contrast module (dynamic gamma)
+- LSC, LTM, auto-contrast module can dynamiclly adjust the parameters based on the camera gain from the metadata.
+- Add another white-balance process at the end of the pipeline to further improve the image quality.
+
+Key algorithm parameter explanation:
+- k_th: defines the threshold for the noise level, higher means tolerant to noise, lower means more sensitive to noise. For this sensor, default is 2.5e-3. `Note that this paramerts are shared by all the modules, and hence defined as the member varibale in RawProcessingPipelineDemo class.`
+- s of local_tone_mapping: defines how to apply the gain map to different image channels. Higher means more saturated, lower means less saturated. Default is 0.7. `Currently, s is automatically adjusted based on the camera gain from the metadata.`
+
+*2024.3.7:*
+
+release v1.0 for 3rd validation!
+
+----
+
+## Version-1
+
+- TMO-ratio50
+
+-----
+
+## Version-2
+
+- TMO: ratio50
+
+- Gamma: 1.5
+
+- Contrast:[low=2,high=0.2]
+
+- Post-AWB: GI
+

PolyuColor/raw_prc_pipeline/__init__.py

@@ -0,0 +1,3 @@
+expected_img_ext = '.jpg'
+expected_landscape_img_height = 768
+expected_landscape_img_width = 1024

PolyuColor/raw_prc_pipeline/contrast_enhancement.py

@@ -0,0 +1,159 @@
+import numpy as np
+from utils import *
+
+
+def _global_mean_contrast(img: np.ndarray,
+                          beta: float = 1.2,
+                          copy=True,
+                          channel_wise=False,
+                          protect_ratio=0.95) -> np.ndarray:
+    """
+    Global mean contrast enhancement.
+    
+    :param img: input image
+    :param beta: contrast enhancement factor
+    :return: enhanced image
+    """
+    if copy:
+        img = img.copy()
+    if channel_wise:
+        remain_ratio = (1 - protect_ratio) / 2
+        maxi_value = 1 - remain_ratio
+        mini_value = remain_ratio
+        for i in range(img.shape[-1]):
+            channel = img[:, :, i]
+            mean = np.mean(channel)
+            if protect_ratio > 0:
+                beta = min(beta, (1 - mean) / (maxi_value - mean))
+                beta = min(beta, mean / (mean - mini_value))
+                beta = max(1, beta)
+            gap = channel - mean
+            channel = gap * beta + mean
+            img[:, :, i] = channel
+    else:
+        y = compute_y(img)
+        mean = np.mean(y)
+        beta = min(beta, (1 - mean) / (1 - protect_ratio - mean))
+        new_y = (y - mean) * beta + mean
+        y[y == 0] = 1
+        gain_map = new_y / y
+        img = img * gain_map[:, :, np.newaxis]
+
+    img = np.clip(img, 0, 1)
+    return img
+
+
+def _s_curve_correction(img: np.ndarray,
+                        alpha: float = 0.15,
+                        gamma: float = 1 / 1.3,
+                        copy=True,
+                        channel_wise=False) -> np.ndarray:
+    """
+    S-curve correction.
+    
+    :param img: input image
+    :param alpha: contrast enhancement factor
+    :param gamma: gamma correction factor
+    :return: enhanced image
+    """
+    if copy:
+        img = img.copy()
+    if channel_wise:
+        for i in range(img.shape[-1]):
+            channel = img[:, :, i]
+            mask = channel > alpha
+            channel[mask] = alpha + (1 - alpha) * (((channel[mask] - alpha) /
+                                                    (1 - alpha))**gamma)
+            channel[~mask] = alpha - alpha * (
+                (1 - channel[~mask] / alpha)**gamma)
+            img[:, :, i] = channel
+    else:
+        y = compute_y(img)
+        new_y = y.copy()
+        mask = new_y > alpha
+        new_y[mask] = alpha + (1 - alpha) * ((new_y[mask] - alpha) /
+                                             (1 - alpha))**gamma
+        new_y[~mask] = alpha - alpha * ((1 - new_y[~mask] / alpha)**gamma)
+        y[y == 0] = 1
+        gain_map = new_y / y
+        img = img * gain_map[:, :, np.newaxis]
+    img = np.clip(img, 0, 1)
+    return img
+
+
+def _hist_stretching(img: np.ndarray,
+                     copy=True,
+                     channel_wise=False) -> np.ndarray:
+    """
+    Histogram stretching.
+    
+    :param img: input image
+    :return: enhanced image
+    """
+    if copy:
+        img = img.copy()
+    if channel_wise:
+        for i in range(img.shape[-1]):
+            channel = img[:, :, i]
+            channel = (channel - channel.min()) / (channel.max() -
+                                                   channel.min())
+            img[:, :, i] = channel
+    else:
+        y = compute_y(img)
+        y_new = (y - y.min()) / (y.max() - y.min())
+        y[y == 0] = 1
+        gain_map = y_new / y
+        img = img * gain_map[:, :, np.newaxis]
+    img = np.clip(img, 0, 1)
+
+    return img
+
+
+def _conditioanl_contrast_correction(img: np.ndarray, k: float,
+                                     k_th: float) -> np.ndarray:
+    y = compute_y(img)
+    """
+    Conditional contrast correction based on the cameara gain value k
+    
+    Parameters:
+    k: camera gain value
+    k_th: basic threshold of camera gain value
+
+    Returns:
+    enhanced image
+    """
+    mean_y = y.mean()
+    first_k = k_th / 2
+    second_k = k_th
+    third_k = k_th * 2
+    forth_k = k_th * 3
+    if k <= first_k:
+        target_illum = 0.35
+        gamma = np.log(target_illum) / np.log(mean_y)
+        gamma = np.clip(gamma, 1 / 1.6, None)
+    elif first_k < k_th <= second_k:
+        target_illum = 0.32
+        gamma = np.log(target_illum) / np.log(mean_y)
+        gamma = np.clip(gamma, 1 / 1.5, None)
+    elif second_k < k_th <= third_k:
+        target_illum = 0.27
+        gamma = np.log(target_illum) / np.log(mean_y)
+        gamma = np.clip(gamma, 1 / 1.4, None)
+    elif third_k < k_th <= forth_k:
+        target_illum = 0.23
+        gamma = np.log(target_illum) / np.log(mean_y)
+        gamma = np.clip(gamma, 1 / 1.2, None)
+    elif k > forth_k:
+        target_illum = 0.15
+        gamma = np.log(target_illum) / np.log(mean_y)
+        gamma = np.clip(gamma, 1 / 1.1, None)
+    else:
+        gamma = 1.
+    img = img**gamma
+    img = np.clip(img, 0, 1)
+    return img
+
+
+def contrast_enhancement(img: np.ndarray, k: float, k_th: float) -> np.ndarray:
+    img = _conditioanl_contrast_correction(img, k, k_th)
+    return img

PolyuColor/raw_prc_pipeline/exif_data_formats.py

@@ -0,0 +1,22 @@
+class ExifFormat:
+    def __init__(self, id, name, size, short_name):
+        self.id = id
+        self.name = name
+        self.size = size
+        self.short_name = short_name  # used with struct.unpack()
+
+
+exif_formats = {
+    1: ExifFormat(1, 'unsigned byte', 1, 'B'),
+    2: ExifFormat(2, 'ascii string', 1, 's'),
+    3: ExifFormat(3, 'unsigned short', 2, 'H'),
+    4: ExifFormat(4, 'unsigned long', 4, 'L'),
+    5: ExifFormat(5, 'unsigned rational', 8, ''),
+    6: ExifFormat(6, 'signed byte', 1, 'b'),
+    7: ExifFormat(7, 'undefined', 1, 'B'),  # consider `undefined` as `unsigned byte`
+    8: ExifFormat(8, 'signed short', 2, 'h'),
+    9: ExifFormat(9, 'signed long', 4, 'l'),
+    10: ExifFormat(10, 'signed rational', 8, ''),
+    11: ExifFormat(11, 'single float', 4, 'f'),
+    12: ExifFormat(12, 'double float', 8, 'd'),
+}

PolyuColor/raw_prc_pipeline/exif_utils.py

@@ -0,0 +1,208 @@
+"""
+Manual parsing of image file directories (IFDs).
+"""
+
+
+import struct
+from fractions import Fraction
+from raw_prc_pipeline.exif_data_formats import exif_formats
+
+class Ifd:
+    def __init__(self):
+        self.offset = -1
+        self.tags = {}  # <key, tag> dict; tag number will be key.
+
+
+class Tag:
+    def __init__(self):
+        self.offset = -1
+        self.tag_num = -1
+        self.data_format = -1
+        self.num_values = -1
+        self.values = []
+
+
+def parse_exif(image_path, verbose=True):
+    """
+    Parse EXIF tags from a binary file and return IFDs.
+    Returned IFDs include EXIF SubIFDs, if any.
+    """
+
+    def print_(str_):
+        if verbose:
+            print(str_)
+
+    ifds = {}  # dict of <offset, Ifd> pairs; using offset to IFD as key.
+
+    with open(image_path, 'rb') as fid:
+        fid.seek(0)
+        b0 = fid.read(1)
+        _ = fid.read(1)
+        # byte storage direction (endian):
+        # +1: b'M' (big-endian/Motorola)
+        # -1: b'I' (little-endian/Intel)
+        endian = 1 if b0 == b'M' else -1
+        print_("Endian = {}".format(b0))
+        endian_sign = "<" if endian == -1 else ">"  # used in struct.unpack
+        print_("Endian sign = {}".format(endian_sign))
+        _ = fid.read(2)  # 0x002A
+        b4_7 = fid.read(4)  # offset to first IFD
+        offset_ = struct.unpack(endian_sign + "I", b4_7)[0]
+        i = 0
+        ifd_offsets = [offset_]
+        while len(ifd_offsets) > 0:
+            offset_ = ifd_offsets.pop(0)
+            # check if IFD at this offset was already parsed before
+            if offset_ in ifds:
+                continue
+            print_("=========== Parsing IFD # {} ===========".format(i))
+            ifd_ = parse_exif_ifd(fid, offset_, endian_sign, verbose)
+            ifds.update({ifd_.offset: ifd_})
+            print_("=========== Finished parsing IFD # {} ===========".format(i))
+            i += 1
+            # check SubIFDs; zero or more offsets at tag 0x014a
+            sub_idfs_tag_num = int('0x014a', 16)
+            if sub_idfs_tag_num in ifd_.tags:
+                ifd_offsets.extend(ifd_.tags[sub_idfs_tag_num].values)
+            # check Exif SUbIDF; usually one offset at tag 0x8769
+            exif_sub_idf_tag_num = int('0x8769', 16)
+            if exif_sub_idf_tag_num in ifd_.tags:
+                ifd_offsets.extend(ifd_.tags[exif_sub_idf_tag_num].values)
+    return ifds
+
+
+def parse_exif_ifd(binary_file, offset_, endian_sign, verbose=True):
+    """
+    Parse an EXIF IFD.
+    """
+
+    def print_(str_):
+        if verbose:
+            print(str_)
+
+    ifd = Ifd()
+    ifd.offset = offset_
+    print_("IFD offset = {}".format(ifd.offset))
+    binary_file.seek(offset_)
+    num_entries = struct.unpack(endian_sign + "H", binary_file.read(2))[0]  # format H = unsigned short
+    print_("Number of entries = {}".format(num_entries))
+    for t in range(num_entries):
+        print_("---------- Tag {} / {} ----------".format(t + 1, num_entries))
+        if t == 22:
+            ttt = 1
+        tag_ = parse_exif_tag(binary_file, endian_sign, verbose)
+        ifd.tags.update({tag_.tag_num: tag_})  # supposedly, EXIF tag numbers won't repeat in the same IFD
+    # TODO: check for subsequent IFDs by parsing the next 4 bytes immediately after the IFD
+    return ifd
+
+
+def parse_exif_tag(binary_file, endian_sign, verbose=True):
+    """
+    Parse EXIF tag from a binary file starting from the current file pointer and returns the tag values.
+    """
+
+    def print_(str_):
+        if verbose:
+            print(str_)
+
+    tag = Tag()
+
+    # tag offset
+    tag.offset = binary_file.tell()
+    print_("Tag offset = {}".format(tag.offset))
+
+    # tag number
+    bytes_ = binary_file.read(2)
+    tag.tag_num = struct.unpack(endian_sign + "H", bytes_)[0]  # H: unsigned 2-byte short
+    print_("Tag number = {} = 0x{:04x}".format(tag.tag_num, tag.tag_num))
+
+    # data format (some value between [1, 12])
+    tag.data_format = struct.unpack(endian_sign + "H", binary_file.read(2))[0]  # H: unsigned 2-byte short
+    exif_format = exif_formats[tag.data_format]
+    print_("Data format = {} = {}".format(tag.data_format, exif_format.name))
+
+    # number of components/values
+    tag.num_values = struct.unpack(endian_sign + "I", binary_file.read(4))[0]  # I: unsigned 4-byte integer
+    print_("Number of values = {}".format(tag.num_values))
+
+    # total number of data bytes
+    total_bytes = tag.num_values * exif_format.size
+    print_("Total bytes = {}".format(total_bytes))
+
+    # seek to data offset (if needed)
+    data_is_offset = False
+    current_offset = binary_file.tell()
+    if total_bytes > 4:
+        print_("Total bytes > 4; The next 4 bytes are an offset.")
+        data_is_offset = True
+        data_offset = struct.unpack(endian_sign + "I", binary_file.read(4))[0]
+        current_offset = binary_file.tell()
+        print_("Current offset = {}".format(current_offset))
+        print_("Seeking to data offset = {}".format(data_offset))
+        binary_file.seek(data_offset)
+
+    # read values
+    # TODO: need to distinguish between numeric and text values?
+    if tag.num_values == 1 and total_bytes < 4:
+        # special case: data is a single value that is less than 4 bytes inside 4 bytes, take care of endian
+        val_bytes = binary_file.read(4)
+        # if endian_sign == ">":
+        # val_bytes = val_bytes[4 - total_bytes:]
+        # else:
+        # val_bytes = val_bytes[:total_bytes][::-1]
+        val_bytes = val_bytes[:total_bytes]
+        tag.values.append(struct.unpack(endian_sign + exif_format.short_name, val_bytes)[0])
+    else:
+        # read data values one by one
+        for k in range(tag.num_values):
+            val_bytes = binary_file.read(exif_format.size)
+            if exif_format.name == 'unsigned rational':
+                tag.values.append(eight_bytes_to_fraction(val_bytes, endian_sign, signed=False))
+            elif exif_format.name == 'signed rational':
+                tag.values.append(eight_bytes_to_fraction(val_bytes, endian_sign, signed=True))
+            else:
+                tag.values.append(struct.unpack(endian_sign + exif_format.short_name, val_bytes)[0])
+        if total_bytes < 4:
+            # special case: multiple values less than 4 bytes in total, inside the 4 bytes; skip the extra bytes
+            binary_file.seek(4 - total_bytes, 1)
+
+    if verbose:
+        if len(tag.values) > 100:
+            print_("Got more than 100 values; printing first 100 only:")
+            print_("Tag values = {}".format(tag.values[:100]))
+        else:
+            print_("Tag values = {}".format(tag.values))
+    if tag.data_format == 2:
+        print_("Tag values (string) = {}".format(b''.join(tag.values).decode()))
+
+    if data_is_offset:
+        # seek back to current position to read the next tag
+        print_("Seeking back to current offset = {}".format(current_offset))
+        binary_file.seek(current_offset)
+
+    return tag
+
+
+def get_tag_values_from_ifds(tag_num, ifds):
+    """
+    Return values of a tag, if found in ifds. Return None otherwise.
+    Assuming any tag exists only once in all ifds.
+    """
+    for key, ifd in ifds.items():
+        if tag_num in ifd.tags:
+            return ifd.tags[tag_num].values
+    return None
+
+
+def eight_bytes_to_fraction(eight_bytes, endian_sign, signed):
+    """
+    Convert 8-byte array into a Fraction. Take care of endian and sign.
+    """
+    if signed:
+        num = struct.unpack(endian_sign + "l", eight_bytes[:4])[0]
+        den = struct.unpack(endian_sign + "l", eight_bytes[4:])[0]
+    else:
+        num = struct.unpack(endian_sign + "L", eight_bytes[:4])[0]
+        den = struct.unpack(endian_sign + "L", eight_bytes[4:])[0]
+    den = den if den != 0 else 1
+    return Fraction(num, den)

PolyuColor/raw_prc_pipeline/fs.py

@@ -0,0 +1,43 @@
+import cv2
+import numpy as np
+
+
+def perform_flash(source, a=5, target=-1, perform_gamma_correction=True):
+    rows, cols, _ = source.shape
+
+    v = np.max(source, axis=2)
+    vd = np.copy(v)
+    vd[vd == 0] = 1e-9
+    result = source / (a * np.exp(np.mean(np.log(vd))) + np.tile(np.expand_dims(vd, axis=2), (1, 1, 3)))
+
+    if perform_gamma_correction:
+        result **= 1.0 / 2.2
+
+    if target >= 0:
+        result *= target / np.mean((0.299 * result[:, :, 2] + 0.587 * result[:, :, 1] + 0.114 * result[:, :, 0]))
+    else:
+        result *= 255.0 / np.max(result)
+
+    return result
+
+
+def perform_storm(source, a=5, target=-1, kernels=(1, 4, 16, 64, 256), perform_gamma_correction=True):
+    rows, cols, _ = source.shape
+
+    v = np.max(source, axis=2)
+    vd = np.copy(v)
+    vd[vd == 0] = 1e-9
+    lv = np.log(vd)
+    result = sum([source / np.tile(
+        np.expand_dims(a * np.exp(cv2.boxFilter(lv, -1, (int(min(rows // kernel, cols // kernel)),) * 2)) + vd, axis=2),
+        (1, 1, 3)) for kernel in kernels])
+
+    if perform_gamma_correction:
+        result **= 1.0 / 2.2
+
+    if target >= 0:
+        result *= target / np.mean((0.299 * result[:, :, 2] + 0.587 * result[:, :, 1] + 0.114 * result[:, :, 0]))
+    else:
+        result *= 255.0 / np.max(result)
+
+    return result

PolyuColor/raw_prc_pipeline/grey_pixels.py

@@ -0,0 +1,102 @@
+"""
+unofficial implementation of paper
+    "Efficient Illuminant Estimation for Color Constancy Using Grey Pixels"
+"""
+
+import cv2
+import numpy as np
+
+_params = {
+    'patch_size': 3,
+    'blur_kernel': 7,
+    "top_n": 0.1,
+    'eps': 1e-6,
+    'threshold': 1e-1
+}
+
+
+def _compute_local_std(log_img: np.ndarray) -> np.ndarray:
+    mean = cv2.blur(log_img, (_params['patch_size'], _params['patch_size']),
+                    borderType=cv2.BORDER_REPLICATE)
+    sq_mean = cv2.blur(log_img**2,
+                       (_params['patch_size'], _params['patch_size']),
+                       borderType=cv2.BORDER_REPLICATE)
+    tmp = sq_mean - mean**2
+    tmp[tmp < 0] = 0
+    std_dev = np.sqrt(tmp)
+    return std_dev
+
+
+def _compute_local_derive_gaussian(img: np.ndarray,
+                                   kernel_half_size=2,
+                                   sigma: float = .5):
+    x, y = np.meshgrid(np.arange(-kernel_half_size, kernel_half_size + 1),
+                       np.arange(-kernel_half_size, kernel_half_size + 1))
+    ssq = sigma**2
+    kernel = -x * np.exp(-(x**2 + y**2) / (2 * ssq)) / (np.pi * ssq)
+
+    ans = cv2.filter2D(img, -1, kernel, borderType=cv2.BORDER_REPLICATE)
+    ans = np.abs(ans)
+    return ans
+
+
+def _compute_pixel_std(img: np.ndarray) -> np.ndarray:
+    mean = np.mean(img, axis=-1)
+    sq_mean = np.mean(img**2, axis=-1)
+    tmp = sq_mean - mean**2
+    tmp[tmp < 0] = 0
+    std_dev = np.sqrt(tmp)
+    return std_dev
+
+
+def _compute_grey_index_map(img: np.ndarray, method='std') -> np.ndarray:
+    mask = np.any(img < 2e-2, axis=-1) | np.any(img > 1 - _params['threshold'],
+                                                axis=-1)
+    img = img * 65535 + 1
+    log_img = np.log(img) + _params['eps']
+    if method == 'std':
+        iim_map = _compute_local_std(log_img)
+    else:
+        iim_map = _compute_local_derive_gaussian(log_img)
+    mask |= np.all(iim_map < _params['eps'], axis=-1)
+
+    Ds = _compute_pixel_std(iim_map)
+    Ds /= (iim_map.mean(axis=-1) + _params['eps'])
+
+    l_value = img.mean(axis=-1)
+
+    Ps = Ds / l_value
+
+    Ps /= (Ps.max() + _params['eps'])
+
+    Ps[mask] = Ps.max()
+
+    grey_index_map = cv2.blur(Ps,
+                              (_params['blur_kernel'], _params['blur_kernel']))
+
+    return grey_index_map, mask
+
+
+def grey_pixels(img: np.ndarray) -> np.ndarray:
+    h, w, c = img.shape
+    img = img.reshape(-1, c)
+    pixel_num = int(h * w * _params['top_n'] / 100)
+
+    grey_index_map, mask = _compute_grey_index_map(img, method='std')
+    valid_num = np.sum(~mask)
+    if valid_num < pixel_num:
+        return np.array([1., 1., 1.])
+    grey_index_map = np.ravel(grey_index_map)
+    indexes = np.argsort(grey_index_map)[:pixel_num]
+
+    candidates: np.ndarray = img[indexes]
+    r_avg, g_avg, b_avg = candidates.mean(axis=0)
+
+    r_avg /= (g_avg + _params['eps'])
+    b_avg /= (g_avg + _params['eps'])
+    if r_avg < 0.2 or r_avg > 5 or b_avg < 0.2 or b_avg > 5:
+        r_avg = 1.
+        b_avg = 1.
+
+    res = np.array([r_avg, 1., b_avg])
+    return res

PolyuColor/raw_prc_pipeline/lsc.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+from utils import *
+
+
+def simple_lsc(raw: np.ndarray, shading: np.ndarray) -> np.ndarray:
+    """
+    Simple LSC algorithm.
+
+    :param raw: raw image
+    :param shading: shading image
+    :param dark_th: threshold for detecting dark image
+    :return: LSC-corrected image
+    """
+    rggb_calibrated = pack_raw(shading)
+    rggb_raw = pack_raw(raw)
+    rggb_raw /= rggb_calibrated
+    ret = depack_raw(rggb_raw)
+    ret = np.clip(ret, 0, 1)
+
+    return ret

PolyuColor/raw_prc_pipeline/model.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+
+
+class UNetSeeInDark(nn.Module):
+    def __init__(self, in_channels=4, out_channels=4):
+        super(UNetSeeInDark, self).__init__()
+        # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.conv1_1 = nn.Conv2d(in_channels,
+                                 32,
+                                 kernel_size=3,
+                                 stride=1,
+                                 padding=1)
+        self.conv1_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
+        self.pool1 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv2_1 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.conv2_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+        self.pool2 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv3_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+        self.conv3_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+        self.pool3 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv4_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+        self.conv4_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.pool4 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv5_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
+        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+
+        self.upv6 = nn.ConvTranspose2d(512, 256, 2, stride=2)
+        self.conv6_1 = nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1)
+        self.conv6_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+
+        self.upv7 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.conv7_1 = nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1)
+        self.conv7_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+
+        self.upv8 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.conv8_1 = nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1)
+        self.conv8_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+
+        self.upv9 = nn.ConvTranspose2d(64, 32, 2, stride=2)
+        self.conv9_1 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1)
+        self.conv9_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
+
+        self.conv10_1 = nn.Conv2d(32, out_channels, kernel_size=1, stride=1)
+
+    def forward(self, x):
+        conv1 = self.lrelu(self.conv1_1(x))
+        conv1 = self.lrelu(self.conv1_2(conv1))
+        pool1 = self.pool1(conv1)
+
+        conv2 = self.lrelu(self.conv2_1(pool1))
+        conv2 = self.lrelu(self.conv2_2(conv2))
+        pool2 = self.pool1(conv2)
+
+        conv3 = self.lrelu(self.conv3_1(pool2))
+        conv3 = self.lrelu(self.conv3_2(conv3))
+        pool3 = self.pool1(conv3)
+
+        conv4 = self.lrelu(self.conv4_1(pool3))
+        conv4 = self.lrelu(self.conv4_2(conv4))
+        pool4 = self.pool1(conv4)
+
+        conv5 = self.lrelu(self.conv5_1(pool4))
+        conv5 = self.lrelu(self.conv5_2(conv5))
+
+        up6 = self.upv6(conv5)
+        up6 = torch.cat([up6, conv4], 1)
+        conv6 = self.lrelu(self.conv6_1(up6))
+        conv6 = self.lrelu(self.conv6_2(conv6))
+
+        up7 = self.upv7(conv6)
+        up7 = torch.cat([up7, conv3], 1)
+        conv7 = self.lrelu(self.conv7_1(up7))
+        conv7 = self.lrelu(self.conv7_2(conv7))
+
+        up8 = self.upv8(conv7)
+        up8 = torch.cat([up8, conv2], 1)
+        conv8 = self.lrelu(self.conv8_1(up8))
+        conv8 = self.lrelu(self.conv8_2(conv8))
+
+        up9 = self.upv9(conv8)
+        up9 = torch.cat([up9, conv1], 1)
+        conv9 = self.lrelu(self.conv9_1(up9))
+        conv9 = self.lrelu(self.conv9_2(conv9))
+
+        conv10 = self.conv10_1(conv9)
+        # out = nn.functional.pixel_shuffle(conv10, 2)
+        out = conv10
+        return out
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                m.weight.data.normal_(0.0, 0.02)
+                if m.bias is not None:
+                    m.bias.data.normal_(0.0, 0.02)
+            if isinstance(m, nn.ConvTranspose2d):
+                m.weight.data.normal_(0.0, 0.02)
+
+    def lrelu(self, x):
+        outt = torch.max(0.2 * x, x)
+        return outt

PolyuColor/raw_prc_pipeline/pipeline.py

@@ -0,0 +1,331 @@
+"""
+Demo raw processing pipeline and pipeline executor.
+"""
+
+import time
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+
+from raw_prc_pipeline.pipeline_utils import *
+from utils import *
+
+from .contrast_enhancement import contrast_enhancement
+from .grey_pixels import *
+from .lsc import *
+from .model import *
+from .sharpening import *
+from .tone_mapping import *
+
+
+class ModelEncaupsulation(nn.Module):
+    def __init__(self, model):
+        super(ModelEncaupsulation, self).__init__()
+        self.model = model()
+
+    def forward(self, x):
+        return self.model(x)
+
+
+model = ModelEncaupsulation(UNetSeeInDark)
+parameters = torch.load('resources/sid_fp32_best.ckpt')
+model.float()
+model.load_state_dict(parameters['state_dict'])
+model.eval()
+model = model.cuda()
+shading_grid = cv2.imread('resources/average_shading.png',
+                          cv2.IMREAD_UNCHANGED).astype(np.float32) / 65535.0
+
+
+class RawProcessingPipelineDemo:
+    """
+    Demonstration pipeline of raw image processing.
+
+    This pipeline is a baseline pipeline to process raw image.
+    The public methods of this class are successive steps of raw image processing pipeline.
+    The declaration order of the public methods must correspond to the order in which these methods (steps) are supposed to be called when processing raw image.
+
+    It is assumed that each public method has 2 parameters:
+    raw_img : ndarray
+        Array with images data.
+    img_meta : Dict
+        Some metadata of image.
+    
+    Also each such public method must return an image (ndarray) as the result of processing.
+    """
+    def __init__(self,
+                 illumination_estimation='',
+                 denoise_flg=True,
+                 tone_mapping='Flash',
+                 out_landscape_width=None,
+                 out_landscape_height=None,
+                 color_matrix=[
+                     1.06835938, -0.29882812, -0.14257812, -0.43164062,
+                     1.35546875, 0.05078125, -0.1015625, 0.24414062, 0.5859375
+                 ]):
+        """
+        RawProcessingPipelineDemo __init__ method.
+
+        Parameters
+        ----------
+        illumination_estimation : str, optional
+            Options for illumination estimation algorithms: '', 'gw', 'wp', 'sog', 'iwp', by default ''.
+        denoise_flg : bool, optional
+            Denoising flag, by default True.
+            If True, resulted images will be denoised with some predefined parameters.
+        tone_mapping : str, optional
+            Options for tone mapping methods, defined in function `apply_tone_map` from `pipeline_utils` module.
+            By default 'Flash'.
+        out_landscape_width : int, optional
+            The width of output image (when orientation is landscape). If None, the image resize will not be performed.
+            By default None.
+        out_landscape_height : int, optional
+            The height of output image (when orientation is landscape). If None, the image resize will not be performed.
+            By default None.
+        color_matrix : list, optional
+            Avg color tranformation matrix. If None, average color transformation matrix of Huawei Mate 40 Pro is used.
+        """
+
+        self.params = locals()
+        del self.params['self']
+        self.current_step = 0
+        self.shading_grid = None
+        self.lut3d = None
+        self.k_th = 2.5e-3
+        self.k_max = 0.02750327847
+        self.k_min = 2.32350645e-06
+
+    # Linearization not handled.
+    def linearize_raw(self, raw_img, img_meta):
+        self.current_step += 1
+
+        return raw_img
+
+    def normalize(self, linearized_raw, img_meta):
+        self.start_time = time.perf_counter()
+        ret = normalize(linearized_raw, img_meta['black_level'],
+                        img_meta['white_level'])
+        return ret
+
+    def luma_shading_correction(self, normalized, img_meta):
+        if img_meta['noise_profile'][0] >= self.k_th:
+            lsc_image = normalized
+        else:
+            lsc_image = simple_lsc(normalized, shading_grid)
+        return lsc_image
+
+    def pack_raw_image(self, normalized, img_meta):
+        ret = pack_raw(torch.from_numpy(normalized)).numpy()
+        return ret
+
+    def denoise(self, packed_raw, img_meta):
+        packed_raw = torch.from_numpy(packed_raw)
+        packed_raw = packed_raw.permute(2, 0, 1).unsqueeze(0).cuda()
+        packed_raw = resize_rggb(packed_raw,
+                                 target_height=1536,
+                                 target_width=2048)
+        with torch.no_grad():
+            denoised_packed_raw = model(packed_raw.float()).float()
+            denoised_packed_raw = denoised_packed_raw.squeeze(0).permute(
+                1, 2, 0).cpu()
+        denoised_raw = depack_raw(denoised_packed_raw)
+        denoised_raw = np.clip(denoised_raw.numpy(), 0, 1)
+        return denoised_raw
+
+    def demosaic(self, denoised, img_meta):
+        ret = simple_demosaic(denoised, img_meta['cfa_pattern'])
+        ret = np.clip(ret, 0, 1)
+        return ret
+
+    def auto_white_balance(self, demosaic, img_meta):
+        wb_params = patch_based_white_balance(demosaic)
+        white_balanced = white_balance(demosaic, wb_params)
+
+        return white_balanced
+
+    def xyz_transform(self, white_balanced, img_meta):
+        # in case of absence of color matrix we use mean color matrix
+        if "color_matrix_1" not in img_meta.keys():
+            img_meta["color_matrix_1"] = self.params["color_matrix"]
+            img_meta["color_matrix_2"] = self.params["color_matrix"]
+        ret = apply_color_space_transform(white_balanced,
+                                          img_meta['color_matrix_1'],
+                                          img_meta['color_matrix_2'])
+        return ret
+
+    def srgb_transform(self, xyz, img_meta):
+        ret = transform_xyz_to_srgb(xyz)
+        return ret
+
+    def ltm(self, srgb, img_meta):
+        k = img_meta['noise_profile'][0]
+        scale_ratio = k * 20000
+        scale_ratio = np.clip(scale_ratio, 5, 500)
+        first_k = self.k_th / 2
+        second_k = self.k_th
+        third_k = self.k_th * 2
+        forth_k = self.k_th * 3
+        if k < first_k:
+            s = 0.8
+        elif k < second_k:
+            s = 0.75
+        elif k < third_k:
+            s = 0.7
+        elif k < forth_k:
+            s = 0.6
+        else:
+            s = 0.5
+        ret = local_tone_mapping(srgb, scale_ratio=scale_ratio, mode=1, s=s)
+
+        return ret
+
+    def autocontrast(self, srgb, img_meta):
+        k = img_meta['noise_profile'][0]
+        ret = contrast_enhancement(srgb, k, self.k_th)
+
+        return ret
+
+    def resize(self, srgb, img_meta):
+        h, w, c = srgb.shape
+        target_height = 768
+        target_width = 1024
+
+        if self.params['out_landscape_width'] is not None and self.params[
+                'out_landscape_height'] is not None:
+            target_height = self.params['out_landscape_height']
+            target_width = self.params['out_landscape_width']
+        if h != self.params['out_landscape_height'] or w != self.params[
+                'out_landscape_width']:
+            srgb = cv2.resize(srgb, (target_width, target_height),
+                              interpolation=cv2.INTER_CUBIC)
+        return srgb
+
+    def to_uint8(self, srgb, img_meta):
+        self.current_step = 0
+        srgb = np.clip(srgb, 0, 1)
+        return (srgb * 255).astype(np.uint8)
+
+    def fix_orientation(self, img, img_meta):
+        ret = fix_orientation(img, img_meta['orientation'])
+        return ret
+
+
+class PipelineExecutor:
+    """
+    Pipeline executor class.
+
+    This class can be used to successively execute the steps of some image pipeline class (for example `RawProcessingPipelineDemo`).
+    The declaration order of the public methods of pipeline class must correspond to the order in which these methods (steps) are supposed to be called when processing image.
+
+    It is assumed that each public method of the pipeline class has 2 parameters:
+    raw_img : ndarray
+        Array with images data.
+    img_meta : Dict
+        Some meta data of image.
+    
+    Also each such public method must return an image (ndarray) as the result of processing.
+    """
+    def __init__(self,
+                 img,
+                 img_meta,
+                 img_name,
+                 pipeline_obj,
+                 first_stage=None,
+                 last_stage=None,
+                 save_dir="debug_output"):
+        """
+        PipelineExecutor __init__ method.
+
+        Parameters
+        ----------
+        img : ndarray
+            Image that should be processed by pipeline.
+        img_meta : Dict
+            Some image metadata.
+        pipeline_obj : pipeline object
+            Some pipeline object such as RawProcessingPipelineDemo.
+        first_stage : str, optional
+            The name of first public method of pipeline object that should be called by PipelineExecutor.
+            If None, the first public method from defined in pipeline object will be considered as `first_stage` method.
+            By default None.
+        last_stage : str, optional
+            The name of last public method of pipeline object that should be called by PipelineExecutor.
+            If None, the last public method from defined in pipeline object will be considered as `last_stage` method.
+            By default None.
+        """
+        self.pipeline_obj = pipeline_obj
+        self.img_name = img_name
+        self.save_dir = os.path.join(save_dir, img_name)
+        self.pipeline_obj.save_dir = self.save_dir
+        self.stages_dict = self._init_stages()
+        self.stages_names, self.stages = list(self.stages_dict.keys()), list(
+            self.stages_dict.values())
+
+        if first_stage is None:
+            self.next_stage_indx = 0
+        else:
+            assert first_stage in self.stages_names, f"Invalid first_stage={first_stage}. Try use the following stages: {self.stages_names}"
+            self.next_stage_indx = self.stages_names.index(first_stage)
+
+        if last_stage is None:
+            self.last_stage_indx = len(self.stages_names) - 1
+        else:
+            assert last_stage in self.stages_names, f"Invalid last_stage={last_stage}. Try use the following stages: {self.stages_names}"
+            self.last_stage_indx = self.stages_names.index(last_stage)
+            if self.next_stage_indx > self.last_stage_indx:
+                print(
+                    f'Warning: the specified first_stage={first_stage} follows the specified last_stage={last_stage}, so using __call__ no image processing will be done.'
+                )
+
+        self.current_image = img
+        self.img_meta = img_meta
+
+    def _init_stages(self):
+        stages = {
+            func: getattr(self.pipeline_obj, func)
+            for func in self.pipeline_obj.__class__.__dict__
+            if callable(getattr(self.pipeline_obj, func))
+            and not func.startswith("_")
+        }
+        return stages
+
+    @property
+    def next_stage(self):
+        if self.next_stage_indx < len(self.stages):
+            return self.stages_names[self.next_stage_indx]
+        else:
+            return None
+
+    @property
+    def last_stage(self):
+        return self.stages_names[self.last_stage_indx]
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.next_stage_indx < len(self.stages):
+            stage_func = self.stages[self.next_stage_indx]
+            self.current_image = stage_func(self.current_image, self.img_meta)
+            self.next_stage_indx += 1
+            return self.current_image
+        else:
+            raise StopIteration
+
+    def __call__(self):
+        """
+        PipelineExecutor __call__ method.
+
+        This method will sequentially execute the methods defined in the pipeline object from the `first_stage` to the `last_stage` inclusive.
+
+        Returns
+        -------
+        ndarray
+            Resulted processed raw image.
+        """
+        for current_image in self:
+            if self.next_stage_indx > self.last_stage_indx:
+                return current_image
+        return self.current_image

PolyuColor/raw_prc_pipeline/pipeline_utils.py

@@ -0,0 +1,575 @@
+"""
+Camera pipeline utilities.
+"""
+
+import os
+from fractions import Fraction
+
+import cv2
+import exifread
+import numpy as np
+import rawpy
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+# from exifread import Ratio
+from exifread.utils import Ratio
+from PIL import Image, ImageOps
+from scipy.io import loadmat
+from skimage.restoration import denoise_bilateral
+from skimage.transform import resize as skimage_resize
+
+from raw_prc_pipeline.exif_utils import get_tag_values_from_ifds, parse_exif
+from raw_prc_pipeline.fs import perform_flash, perform_storm
+
+EPS = 1e-9
+
+
+def get_visible_raw_image(image_path):
+    raw_image = rawpy.imread(image_path).raw_image_visible.copy()
+    # raw_image = rawpy.imread(image_path).raw_image.copy()
+    return raw_image
+
+
+def get_image_tags(image_path):
+    with open(image_path, 'rb') as f:
+        tags = exifread.process_file(f)
+    return tags
+
+
+def get_image_ifds(image_path):
+    ifds = parse_exif(image_path, verbose=False)
+    return ifds
+
+
+def get_metadata(image_path):
+    metadata = {}
+    tags = get_image_tags(image_path)
+    ifds = get_image_ifds(image_path)
+    metadata['linearization_table'] = get_linearization_table(tags, ifds)
+    metadata['black_level'] = get_black_level(tags, ifds)
+    metadata['white_level'] = get_white_level(tags, ifds)
+    metadata['cfa_pattern'] = get_cfa_pattern(tags, ifds)
+    metadata['as_shot_neutral'] = get_as_shot_neutral(tags, ifds)
+    color_matrix_1, color_matrix_2 = get_color_matrices(tags, ifds)
+    metadata['color_matrix_1'] = color_matrix_1
+    metadata['color_matrix_2'] = color_matrix_2
+    metadata['orientation'] = get_orientation(tags, ifds)
+    # isn't used
+    metadata['noise_profile'] = get_noise_profile(tags, ifds)
+    # ...
+    # fall back to default values, if necessary
+    if metadata['black_level'] is None:
+        metadata['black_level'] = 0
+        print("Black level is None; using 0.")
+    if metadata['white_level'] is None:
+        metadata['white_level'] = 2**16
+        print("White level is None; using 2 ** 16.")
+    if metadata['cfa_pattern'] is None:
+        metadata['cfa_pattern'] = [0, 1, 1, 2]
+        print("CFAPattern is None; using [0, 1, 1, 2] (RGGB)")
+    if metadata['as_shot_neutral'] is None:
+        metadata['as_shot_neutral'] = [1, 1, 1]
+        print("AsShotNeutral is None; using [1, 1, 1]")
+    if metadata['color_matrix_1'] is None:
+        metadata['color_matrix_1'] = [1] * 9
+        print("ColorMatrix1 is None; using [1, 1, 1, 1, 1, 1, 1, 1, 1]")
+    if metadata['color_matrix_2'] is None:
+        metadata['color_matrix_2'] = [1] * 9
+        print("ColorMatrix2 is None; using [1, 1, 1, 1, 1, 1, 1, 1, 1]")
+    if metadata['orientation'] is None:
+        metadata['orientation'] = 0
+        print("Orientation is None; using 0.")
+    # ...
+    return metadata
+
+
+def get_linearization_table(tags, ifds):
+    possible_keys = [
+        'Image Tag 0xC618', 'Image Tag 50712', 'LinearizationTable',
+        'Image LinearizationTable'
+    ]
+    return get_values(tags, possible_keys)
+
+
+def get_black_level(tags, ifds):
+    possible_keys = [
+        'Image Tag 0xC61A', 'Image Tag 50714', 'BlackLevel', 'Image BlackLevel'
+    ]
+    vals = get_values(tags, possible_keys)
+    if vals is None:
+        # print("Black level not found in exifread tags. Searching IFDs.")
+        vals = get_tag_values_from_ifds(50714, ifds)
+    return vals
+
+
+def get_white_level(tags, ifds):
+    possible_keys = [
+        'Image Tag 0xC61D', 'Image Tag 50717', 'WhiteLevel', 'Image WhiteLevel'
+    ]
+    vals = get_values(tags, possible_keys)
+    if vals is None:
+        # print("White level not found in exifread tags. Searching IFDs.")
+        vals = get_tag_values_from_ifds(50717, ifds)
+    return vals
+
+
+def get_cfa_pattern(tags, ifds):
+    possible_keys = ['CFAPattern', 'Image CFAPattern']
+    vals = get_values(tags, possible_keys)
+    if vals is None:
+        # print("CFAPattern not found in exifread tags. Searching IFDs.")
+        vals = get_tag_values_from_ifds(33422, ifds)
+    return vals
+
+
+def get_as_shot_neutral(tags, ifds):
+    possible_keys = [
+        'Image Tag 0xC628', 'Image Tag 50728', 'AsShotNeutral',
+        'Image AsShotNeutral'
+    ]
+    return get_values(tags, possible_keys)
+
+
+def get_color_matrices(tags, ifds):
+    possible_keys_1 = [
+        'Image Tag 0xC621', 'Image Tag 50721', 'ColorMatrix1',
+        'Image ColorMatrix1'
+    ]
+    color_matrix_1 = get_values(tags, possible_keys_1)
+    possible_keys_2 = [
+        'Image Tag 0xC622', 'Image Tag 50722', 'ColorMatrix2',
+        'Image ColorMatrix2'
+    ]
+    color_matrix_2 = get_values(tags, possible_keys_2)
+    #print(f'Color matrix 1:{color_matrix_1}')
+    #print(f'Color matrix 2:{color_matrix_2}')
+    #print(np.sum(np.abs(np.array(color_matrix_1) - np.array(color_matrix_2))))
+    return color_matrix_1, color_matrix_2
+
+
+def get_orientation(tags, ifds):
+    possible_tags = ['Orientation', 'Image Orientation']
+    return get_values(tags, possible_tags)
+
+
+def get_noise_profile(tags, ifds):
+    possible_keys = [
+        'Image Tag 0xC761', 'Image Tag 51041', 'NoiseProfile',
+        'Image NoiseProfile'
+    ]
+    vals = get_values(tags, possible_keys)
+    if vals is None:
+        # print("Noise profile not found in exifread tags. Searching IFDs.")
+        vals = get_tag_values_from_ifds(51041, ifds)
+    return vals
+
+
+def get_values(tags, possible_keys):
+    values = None
+    for key in possible_keys:
+        if key in tags.keys():
+            values = tags[key].values
+    return values
+
+
+def normalize(raw_image, black_level, white_level):
+    if type(black_level) is list and len(black_level) == 1:
+        black_level = float(black_level[0])
+    if type(white_level) is list and len(white_level) == 1:
+        white_level = float(white_level[0])
+    black_level_mask = black_level
+    if type(black_level) is list and len(black_level) == 4:
+        if type(black_level[0]) is Ratio:
+            black_level = ratios2floats(black_level)
+        if type(black_level[0]) is Fraction:
+            black_level = fractions2floats(black_level)
+        black_level_mask = np.zeros(raw_image.shape)
+        idx2by2 = [[0, 0], [0, 1], [1, 0], [1, 1]]
+        step2 = 2
+        for i, idx in enumerate(idx2by2):
+            black_level_mask[idx[0]::step2, idx[1]::step2] = black_level[i]
+    normalized_image = raw_image.astype(np.float32) - black_level_mask
+    # if some values were smaller than black level
+    normalized_image[normalized_image < 0] = 0
+    normalized_image = normalized_image / (white_level - black_level_mask)
+    return normalized_image
+
+
+def ratios2floats(ratios):
+    floats = []
+    for ratio in ratios:
+        floats.append(float(ratio.num) / ratio.den)
+    return floats
+
+
+def fractions2floats(fractions):
+    floats = []
+    for fraction in fractions:
+        floats.append(float(fraction.numerator) / fraction.denominator)
+    return floats
+
+
+def illumination_parameters_estimation(current_image,
+                                       illumination_estimation_option):
+    ie_method = illumination_estimation_option.lower()
+    if ie_method == "gw":
+        mask = np.any(current_image < 1e-2, axis=-1) | np.any(
+            current_image > 1 - 5e-2, axis=-1)
+        if np.sum(~mask) == 0:
+            return np.array([1, 1, 1])
+        valid_pixels = current_image[~mask]
+        ie = np.mean(valid_pixels, axis=0)
+        ie /= ie[1]
+        ie[ie < 1e-1] = 1
+        return ie
+    elif ie_method == "sog":
+        sog_p = 4.
+        ie = np.mean(current_image**sog_p, axis=(0, 1))**(1 / sog_p)
+        ie /= ie[1]
+        return ie
+    elif ie_method == "wp":
+        ie = np.max(current_image, axis=(0, 1))
+        ie /= ie[1]
+        return ie
+    elif ie_method == "iwp":
+        samples_count = 20
+        sample_size = 20
+        rows, cols = current_image.shape[:2]
+        data = np.reshape(current_image, (rows * cols, 3))
+        maxima = np.zeros((samples_count, 3))
+        for i in range(samples_count):
+            maxima[i, :] = np.max(data[
+                np.random.randint(low=0, high=rows *
+                                  cols, size=(sample_size)), :],
+                                  axis=0)
+        ie = np.mean(maxima, axis=0)
+        ie /= ie[1]
+        return ie
+    else:
+        raise ValueError(
+            'Bad illumination_estimation_option value! Use the following options: "gw", "wp", "sog", "iwp"'
+        )
+
+
+def sample_acceptable_white_point(lower=0.35, upper=0.5, samples_count=10):
+    x = np.linspace(lower, upper, samples_count)
+    a, b, c = 2.14325171, 7.12239676, 0.10934688
+    y = a * np.exp(-b * x) + c
+    ret = np.stack([x, y], axis=1)
+    return ret
+
+
+def patch_based_white_balance(img: np.ndarray, split_ratio=2) -> np.ndarray:
+    h, w = img.shape[:2]
+    patch_size = (h // split_ratio, w // split_ratio)
+    patches = [
+        img[i:i + patch_size[0], j:j + patch_size[1], :]
+        for i in range(0, h, patch_size[0])
+        for j in range(0, w, patch_size[1])
+    ]
+    white_points = []
+    for patch in patches:
+        white_point = illumination_parameters_estimation(patch, "gw")
+        white_points.append(white_point)
+    white_points.append(illumination_parameters_estimation(img, "gw"))
+    white_points = np.array(white_points)
+
+    white_point = white_point_regularization(white_points)
+    return white_point
+
+
+def white_point_regularization(white_points: np.ndarray,
+                               radius=0.125) -> np.ndarray:
+    """
+    Regularize the white point vector to avoid numerical instability.
+    """
+    centers = np.array([[0.339, 1, 0.361], [0.367, 1,
+                                            0.289], [0.398, 1, 0.237],
+                        [0.464, 1, 0.198], [0.39, 1, 0.29], [0.480, 1, 0.187],
+                        [0.535, 1, 0.165], [0.582, 1, 0.135]])
+    center_weights = np.zeros(len(centers), dtype=np.float32)
+    mini_dist = float("inf")
+    global_cloest_center = None
+    for i, wp in enumerate(white_points):
+        distances = np.linalg.norm(np.abs(centers - wp), axis=1)
+        for i, distance in enumerate(distances):
+            if distance <= radius:
+                center_weights[i] += 1. / (distance + 1e-6)
+            if distance < mini_dist:
+                mini_dist = distance
+                global_cloest_center = centers[i]
+    if center_weights.sum() == 0:
+        white_point = global_cloest_center
+        return white_point
+    center_weights /= center_weights.sum()
+    weighted_white_points = centers * center_weights[:, np.newaxis]
+    white_point = weighted_white_points.sum(axis=0)
+    return white_point
+
+
+def white_balance(demosaic_img, as_shot_neutral):
+    if type(as_shot_neutral[0]) is Ratio:
+        as_shot_neutral = ratios2floats(as_shot_neutral)
+
+    as_shot_neutral = np.asarray(as_shot_neutral)
+    # transform vector into matrix
+    if as_shot_neutral.shape == (3, ):
+        as_shot_neutral = np.diag(1. / as_shot_neutral)
+
+    assert as_shot_neutral.shape == (3, 3)
+
+    white_balanced_image = np.dot(demosaic_img, as_shot_neutral.T)
+    white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)
+
+    return white_balanced_image
+
+
+def simple_demosaic(img, cfa_pattern):
+    raw_colors = np.asarray(cfa_pattern).reshape((2, 2))
+    demosaiced_image = np.zeros((img.shape[0] // 2, img.shape[1] // 2, 3))
+    for i in range(2):
+        for j in range(2):
+            ch = raw_colors[i, j]
+            if ch == 1:
+                demosaiced_image[:, :, ch] += img[i::2, j::2] / 2
+            else:
+                demosaiced_image[:, :, ch] = img[i::2, j::2]
+    return demosaiced_image
+
+
+def denoise_image(demosaiced_image):
+    multichannel = False
+    if len(demosaiced_image.shape) == 3:
+        multichannel = True
+    current_image = denoise_bilateral(demosaiced_image,
+                                      sigma_color=None,
+                                      sigma_spatial=1.,
+                                      channel_axis=-1,
+                                      mode='reflect',
+                                      multichannel=multichannel)
+    return current_image
+
+
+def apply_color_space_transform(demosaiced_image, color_matrix_1,
+                                color_matrix_2):
+    if isinstance(color_matrix_1[0], Fraction):
+        color_matrix_1 = fractions2floats(color_matrix_1)
+    if isinstance(color_matrix_2[0], Fraction):
+        color_matrix_2 = fractions2floats(color_matrix_2)
+    xyz2cam1 = np.reshape(np.asarray(color_matrix_1), (3, 3))
+    xyz2cam2 = np.reshape(np.asarray(color_matrix_2), (3, 3))
+    # normalize rows (needed?)
+    xyz2cam1 = xyz2cam1 / np.sum(xyz2cam1, axis=1, keepdims=True)
+    xyz2cam2 = xyz2cam2 / np.sum(xyz2cam1, axis=1, keepdims=True)
+    # inverse
+    cam2xyz1 = np.linalg.inv(xyz2cam1)
+    cam2xyz2 = np.linalg.inv(xyz2cam2)
+    # for now, use one matrix  # TODO: interpolate btween both
+    # simplified matrix multiplication
+    xyz_image = cam2xyz1[np.newaxis,
+                         np.newaxis, :, :] * demosaiced_image[:, :,
+                                                              np.newaxis, :]
+    xyz_image = np.sum(xyz_image, axis=-1)
+    xyz_image = np.clip(xyz_image, 0.0, 1.0)
+    return xyz_image
+
+
+def transform_xyz_to_srgb(xyz_image):
+    # srgb2xyz = np.array([[0.4124564, 0.3575761, 0.1804375],
+    #                      [0.2126729, 0.7151522, 0.0721750],
+    #                      [0.0193339, 0.1191920, 0.9503041]])
+
+    # xyz2srgb = np.linalg.inv(srgb2xyz)
+
+    xyz2srgb = np.array([[3.2404542, -1.5371385, -0.4985314],
+                         [-0.9692660, 1.8760108, 0.0415560],
+                         [0.0556434, -0.2040259, 1.0572252]])
+
+    # normalize rows (needed?)
+    xyz2srgb = xyz2srgb / np.sum(xyz2srgb, axis=-1, keepdims=True)
+
+    srgb_image = xyz2srgb[np.newaxis,
+                          np.newaxis, :, :] * xyz_image[:, :, np.newaxis, :]
+    srgb_image = np.sum(srgb_image, axis=-1)
+    srgb_image = np.clip(srgb_image, 0.0, 1.0)
+    return srgb_image
+
+
+def reverse_orientation(image, orientation):
+    # 1 = Horizontal(normal)
+    # 2 = Mirror horizontal
+    # 3 = Rotate 180
+    # 4 = Mirror vertical
+    # 5 = Mirror horizontal and rotate 270 CW
+    # 6 = Rotate 90 CW
+    # 7 = Mirror horizontal and rotate 90 CW
+    # 8 = Rotate 270 CW
+    rev_orientations = np.array([1, 2, 3, 4, 5, 8, 7, 6])
+    return fix_orientation(image, rev_orientations[orientation - 1])
+
+
+def apply_gamma(x, gamma=1.5):
+    return x**(1.0 / gamma)
+    # x = x.copy()
+    # idx = x <= 0.0031308
+    # x[idx] *= 12.92
+    # x[idx == False] = (x[idx == False] ** (1.0 / 2.4)) * 1.055 - 0.055
+    # return x
+
+
+def apply_tone_map(x, tone_mapping='Base'):
+    if tone_mapping == 'Flash':
+        return perform_flash(x, perform_gamma_correction=0) / 255.
+    elif tone_mapping == 'Storm':
+        return perform_storm(x, perform_gamma_correction=0) / 255.
+    elif tone_mapping == 'Drago':
+        tonemap = cv2.createTonemapDrago()
+        return tonemap.process(x.astype(np.float32))
+    elif tone_mapping == 'Mantiuk':
+        tonemap = cv2.createTonemapMantiuk()
+        return tonemap.process(x.astype(np.float32))
+    elif tone_mapping == 'Reinhard':
+        tonemap = cv2.createTonemapReinhard()
+        return tonemap.process(x.astype(np.float32))
+    elif tone_mapping == 'Linear':
+        return np.clip(x / np.sort(x.flatten())[-50000], 0, 1)
+    elif tone_mapping == 'Base':
+        # return 3 * x ** 2 - 2 * x ** 3
+        # tone_curve = loadmat('tone_curve.mat')
+        tone_curve = loadmat(
+            os.path.join(os.path.dirname(os.path.realpath(__file__)),
+                         'tone_curve.mat'))
+        tone_curve = tone_curve['tc']
+        x = np.round(x * (len(tone_curve) - 1)).astype(int)
+        tone_mapped_image = np.squeeze(tone_curve[x])
+        return tone_mapped_image
+    else:
+        raise ValueError(
+            'Bad tone_mapping option value! Use the following options: "Base", "Flash", "Storm", "Linear", "Drago", "Mantiuk", "Reinhard"'
+        )
+
+
+def autocontrast(output_image, cutoff_prcnt=2, preserve_tone=False):
+    if preserve_tone:
+        min_val, max_val = np.percentile(output_image,
+                                         [cutoff_prcnt, 100 - cutoff_prcnt])
+        output_image = (output_image - min_val) / (max_val - min_val)
+    else:
+        channels = [None] * 3
+        for ch in range(3):
+            min_val, max_val = np.percentile(
+                output_image[..., ch], [cutoff_prcnt, 100 - cutoff_prcnt])
+            channels[ch] = (output_image[..., ch] - min_val) / (max_val -
+                                                                min_val)
+        output_image = np.dstack(channels)
+    output_image = np.clip(output_image, 0, 1)
+    return output_image
+
+
+def autocontrast_using_pil(img, cutoff=(2, 0.2)):
+    img_uint8 = np.clip(255 * img, 0, 255).astype(np.uint8)
+    img_pil = Image.fromarray(img_uint8)
+    img_pil = ImageOps.autocontrast(img_pil, cutoff=cutoff, preserve_tone=True)
+    output_image = np.array(img_pil).astype(np.float32) / 255
+    return output_image
+
+
+def raw_rgb_to_cct(rawRgb, xyz2cam1, xyz2cam2):
+    """Convert raw-RGB triplet to corresponding correlated color temperature (CCT)"""
+    pass
+    # pxyz = [.5, 1, .5]
+    # loss = 1e10
+    # k = 1
+    # while loss > 1e-4:
+    #     cct = XyzToCct(pxyz)
+    #     xyz = RawRgbToXyz(rawRgb, cct, xyz2cam1, xyz2cam2)
+    #     loss = norm(xyz - pxyz)
+    #     pxyz = xyz
+    #     fprintf('k = %d, loss = %f\n', [k, loss])
+    #     k = k + 1
+    # end
+    # temp = cct
+
+
+def resize_using_skimage(img, width=1296, height=864):
+    out_shape = (height, width) + img.shape[2:]
+    if img.shape == out_shape:
+        return img
+    out_img = skimage_resize(img,
+                             out_shape,
+                             preserve_range=True,
+                             anti_aliasing=True)
+    out_img = out_img.astype(np.uint8)
+    return out_img
+
+
+def resize_using_pil(img, width=1296, height=864):
+    img_pil = Image.fromarray(img)
+    out_size = (width, height)
+    if img_pil.size == out_size:
+        return img
+    out_img = img_pil.resize(out_size, Image.ANTIALIAS)
+    out_img = np.array(out_img)
+    return out_img
+
+
+def fix_orientation(image, orientation):
+    # 1 = Horizontal(normal)
+    # 2 = Mirror horizontal
+    # 3 = Rotate 180
+    # 4 = Mirror vertical
+    # 5 = Mirror horizontal and rotate 270 CW
+    # 6 = Rotate 90 CW
+    # 7 = Mirror horizontal and rotate 90 CW
+    # 8 = Rotate 270 CW
+    orientation_dict = [
+        "Horizontal (normal)", "Mirror horizontal", "Rotate 180",
+        "Mirror vertical", "Mirror horizontal and rotate 270 CW",
+        "Rotate 90 CW", "Mirror horizontal and rotate 90 CW", "Rotate 270 CW"
+    ]
+    orientation_dict = {v: k for k, v in enumerate(orientation_dict)}
+    orientation = orientation_dict[orientation] + 1
+
+    if orientation == 1:
+        pass
+    elif orientation == 2:
+        image = cv2.flip(image, 0)
+    elif orientation == 3:
+        image = cv2.rotate(image, cv2.ROTATE_180)
+    elif orientation == 4:
+        image = cv2.flip(image, 1)
+    elif orientation == 5:
+        image = cv2.flip(image, 0)
+        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+    elif orientation == 6:
+        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
+    elif orientation == 7:
+        image = cv2.flip(image, 0)
+        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
+    elif orientation == 8:
+        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+
+    return image
+
+
+def compute_lsc_gain(rggb):
+    gains = []
+    for ch in range(4):
+        channel_max = rggb[:, :, ch].max()
+        gain = rggb[:, :, ch] / (channel_max * 0.8)
+        gains.append(gain.clip(0., 1.0))
+    return gains
+
+
+def resize_rggb(rggb_raw: torch.Tensor, target_height=768, target_width=1024):
+    height, width = rggb_raw.shape[-2:]
+    target_size = (target_width,
+                   target_height) if height > width else (target_width, 1024)
+    resized_raw = TF.resize(rggb_raw,
+                            size=(target_size),
+                            interpolation=TF.InterpolationMode.BICUBIC,
+                            antialias=True)
+    resized_raw = torch.clamp(resized_raw, 0, 1)
+
+    return resized_raw

PolyuColor/raw_prc_pipeline/sharpening.py

@@ -0,0 +1,14 @@
+import cv2
+import numpy as np
+
+
+def sharpen_image_with_unsharp_masking(image, sigma=1.0, alpha=1.0):
+    """
+    sharp operation
+    """
+    image = (image * 65535).astype(np.uint16)
+    blurred = cv2.GaussianBlur(image, (0, 0), sigmaX=sigma, sigmaY=sigma)
+    sharpened = cv2.addWeighted(image, 1 + alpha, blurred, -alpha, 0.15)
+    sharpened_image = np.clip(sharpened / 65535.0, 0, 1)
+
+    return sharpened_image

PolyuColor/raw_prc_pipeline/tone_mapping.py

@@ -0,0 +1,150 @@
+import numpy as np
+import cv2
+from utils import *
+
+__all__ = ['local_tone_mapping', 'lmhe_global_tone_mapping']
+
+
+def _filterGaussianWindow(img: np.ndarray, window):
+    w = max(np.round(window), 3)
+    if w % 2 == 0:
+        w += 1
+    img_blur = cv2.GaussianBlur(img, (w, w),
+                                0,
+                                borderType=cv2.BORDER_REPLICATE)
+    return img_blur
+
+
+def _remove_specials(img: np.ndarray, replace_value=1) -> np.ndarray:
+    mask = np.isinf(img) | np.isnan(img)
+    img[mask] = replace_value
+    return img
+
+
+def _ashikhmin_filtering(image: np.ndarray, Ashikhmin_sMax=5) -> tuple:
+    r, c = image.shape
+    threshold = 0.5
+
+    Lfiltered = np.zeros((r, c, Ashikhmin_sMax), dtype=image.dtype)
+    LC = np.zeros((r, c, Ashikhmin_sMax), dtype=image.dtype)
+    for i in range(Ashikhmin_sMax):
+        Lfiltered[:, :, i] = _filterGaussianWindow(image, i + 1)
+        LC[:, :, i] = _remove_specials(
+            np.abs(Lfiltered[:, :, i] -
+                   _filterGaussianWindow(image, (i + 1) * 2)) /
+            Lfiltered[:, :, i])
+
+    L_adapt = -np.ones_like(image)
+    for i in range(Ashikhmin_sMax):
+        LC_i = LC[:, :, i]
+        mask = LC_i < threshold
+        L_adapt[mask] = Lfiltered[:, :, i][mask]
+
+    mask = L_adapt < 0
+    L_adapt[mask] = Lfiltered[:, :, -1][mask]
+    L_detail = _remove_specials(image / L_adapt)
+    L_detail = np.clip(L_detail, 0, None)
+
+    return L_adapt, L_detail
+
+
+def _tvi_ashikhmin(img: np.ndarray) -> np.ndarray:
+    Lout = np.zeros_like(img, dtype=img.dtype)
+    mask = img < 0.0034
+    Lout[mask] = img[mask] / 0.0014
+
+    mask = (img >= 0.0034) & (img < 1)
+    Lout[mask] = 2.4483 + np.log(img[mask] / 0.0034) / 0.4027
+
+    mask = (img >= 1) & (img < 7.2444)
+    Lout[mask] = 16.5630 + (img[mask] - 1) / 0.4027
+
+    mask = img >= 7.2444
+    Lout[mask] = 32.0693 + np.log(img[mask] / 7.2444) / 0.0556
+
+    return Lout
+
+
+def local_tone_mapping(image: np.ndarray,
+                       scale_ratio: float = 50,
+                       s=0.7,
+                       mode=1) -> np.ndarray:
+    """
+    Local tone mapping function.
+
+    Parameters:
+    image: np.ndarray
+        Input image with shape (h, w, 3), range [0, 1], color space sRGB.
+    scale_ratio: float
+        Scale ratio of the input image to the output image. luminance 1 equals to 10,000 cd/m^2.
+    s: float
+        s factor of the local tone mapping function, control the staturation of the image.
+    mode: int
+        Gain map appling mode.
+
+    Returns:
+    img_out: np.ndarray
+        Output image with shape (h, w, 3)
+    """
+    image = image / scale_ratio
+
+    r = image[:, :, 0]
+    g = image[:, :, 1]
+    b = image[:, :, 2]
+
+    lumin = 0.2126729 * r + 0.7151522 * g + 0.0721750 * b
+    ld_max = 100
+    L, Ldetail = _ashikhmin_filtering(lumin)
+    maxL = L.max()
+    minL = L.min()
+    maxL_TVI = _tvi_ashikhmin(maxL)
+    minL_TVI = _tvi_ashikhmin(minL)
+    Ld = (ld_max / 100) * (_tvi_ashikhmin(L) - minL_TVI) / (maxL_TVI -
+                                                            minL_TVI)
+    new_lumin = Ld * Ldetail
+    lumin[lumin <= 0] = 1
+    if mode == 1:
+        img_out = (new_lumin[:, :, np.newaxis] *
+                   ((image / lumin[:, :, np.newaxis])**s))
+    else:
+        img_out = new_lumin[:, :, np.newaxis] * (
+            (image / lumin[:, :, np.newaxis] - 1) * s + 1)
+    img_out = _remove_specials(img_out)
+    img_out = np.clip(img_out, 0, 1)
+    return img_out
+
+
+def compute_y(img: np.ndarray) -> np.ndarray:
+    y = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2]
+    return y
+
+
+def lmhe_global_tone_mapping(img: np.ndarray,
+                             mu: float = 7,
+                             bit_depth=10,
+                             s=0.7) -> np.ndarray:
+    """
+    Log based modified histogram equalization
+    Args:
+        img: input image, range [0, 1]
+        mu: parameter for lmhe
+        bit_depth: bit depth of tone mapping curve
+        protect_ratio: protect ratio of the image, the protected area will not be tone mapped
+    Returns:
+        tone mapped image
+    """
+    y = compute_y(img)
+    bit_counts = 2**bit_depth
+    hist, bins = np.histogram(y.ravel(), bit_counts, [0, 1])
+    m = np.log(hist * hist.max() * (10**(-mu)) + 1) / (np.log(hist.max()**2 *
+                                                              (10**(-mu)) + 1))
+    cdf = np.cumsum(m)
+    cdf_m = np.ma.masked_equal(cdf, 0)
+    cdf_m = (cdf_m - cdf_m.min()) / (cdf_m.max() - cdf_m.min())
+    cdf = np.ma.filled(cdf_m, 0)
+    y_new = np.interp(y.ravel(), bins[:-1], cdf)
+    y_new = y_new.reshape(img.shape[0], img.shape[1])
+    y[y == 0] = 1
+    img_out = y_new[:, :, np.newaxis] * (
+        (img / y[:, :, np.newaxis] - 1) * s + 1)
+    return img_out

PolyuColor/requirements.txt

@@ -0,0 +1,8 @@
+ExifRead==3.0.0
+numpy==1.23.4
+opencv_python==4.8.1.78
+Pillow==10.2.0
+rawpy==0.19.0
+scikit_image==0.19.3
+scipy==1.12.0
+tqdm==4.64.1

PolyuColor/resources/sid_fp32_best.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70f38d04429c00b43637137c8210c074518f0c2a8ac33e729b1d31a7f3d8265a
+size 31057848

PolyuColor/run.py

@@ -0,0 +1,142 @@
+import argparse
+import os
+import time
+import warnings
+from pathlib import Path
+
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+from raw_prc_pipeline import (expected_img_ext, expected_landscape_img_height,
+                              expected_landscape_img_width)
+from raw_prc_pipeline.pipeline import (PipelineExecutor,
+                                       RawProcessingPipelineDemo)
+from utils import fraction_from_json, json_read
+
+warnings.filterwarnings("ignore")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=
+        'Demo script for processing PNG images with given metadata files.')
+    parser.add_argument(
+        '-p',
+        '--png_dir',
+        type=str,
+        default='data',
+        help='Path of the directory containing PNG images with metadata files')
+    parser.add_argument(
+        '-o',
+        '--out_dir',
+        type=str,
+        default=None,
+        help=
+        'Path to the directory where processed images will be saved. Images will be saved in JPG format.'
+    )
+    parser.add_argument(
+        '-ie',
+        '--illumination_estimation',
+        type=str,
+        default='gw',
+        help=
+        'Options for illumination estimation algorithms: "gw", "wp", "sog", "iwp".'
+    )
+    parser.add_argument(
+        '-tm',
+        '--tone_mapping',
+        type=str,
+        default='Storm',
+        help=
+        'Options for tone mapping algorithms: "Base", "Flash", "Storm", "Linear", "Drago", "Mantiuk", "Reinhard".'
+    )
+    parser.add_argument(
+        '-n',
+        '--denoising_flg',
+        action='store_false',
+        help=
+        'Denoising flag. By default resulted images will be denoised with some default parameters.'
+    )
+    parser.add_argument('-m',
+                        '--camera_matrix',
+                        type=float,
+                        nargs=9,
+                        default=[
+                            1.06835938, -0.29882812, -0.14257812, -0.43164062,
+                            1.35546875, 0.05078125, -0.1015625, 0.24414062,
+                            0.5859375
+                        ],
+                        help='Mean color matrix of Hauwei Mate 40 Pro')
+    args = parser.parse_args()
+
+    if args.out_dir is None:
+        args.out_dir = args.png_dir
+
+    return args
+
+
+class PNGProcessingDemo:
+    def __init__(self, ie_method, tone_mapping, denoising_flg, camera_matrix,
+                 save_dir):
+        self.camera_matrix = camera_matrix
+        self.save_dir = save_dir
+        self.pipeline_demo = RawProcessingPipelineDemo(
+            illumination_estimation=ie_method,
+            denoise_flg=denoising_flg,
+            tone_mapping=tone_mapping,
+            out_landscape_height=expected_landscape_img_height,
+            out_landscape_width=expected_landscape_img_width)
+        self.process_times = []
+
+    def __call__(self, png_path: Path, out_path: Path):
+        # parse raw img
+        raw_image = cv2.imread(str(png_path), cv2.IMREAD_UNCHANGED)
+        # parse metadata
+        metadata = json_read(png_path.with_suffix('.json'),
+                             object_hook=fraction_from_json)
+        start_time = time.perf_counter()
+        # executing img pipelinex
+        pipeline_exec = PipelineExecutor(raw_image,
+                                         metadata,
+                                         os.path.basename(
+                                             str(png_path)).split('.')[0],
+                                         self.pipeline_demo,
+                                         save_dir=self.save_dir)
+        # process img
+        output_image = pipeline_exec()
+        end_time = time.perf_counter()
+        self.process_times.append(end_time - start_time)
+
+        # save results
+        output_image = cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(str(out_path), output_image,
+                    [cv2.IMWRITE_JPEG_QUALITY, 100])
+
+
+def main(png_dir, out_dir, illumination_estimation, tone_mapping,
+         denoising_flg, camera_matrix):
+    png_dir = Path(png_dir)
+    out_dir = Path(out_dir)
+    out_dir.mkdir(exist_ok=True)
+
+    png_paths = list(png_dir.glob('*.png'))
+    out_paths = [
+        out_dir / png_path.with_suffix(expected_img_ext).name
+        for png_path in png_paths
+    ]
+
+    png_processor = PNGProcessingDemo(illumination_estimation, tone_mapping,
+                                      denoising_flg, camera_matrix,
+                                      str(out_dir))
+
+    for png_path, out_path in tqdm(zip(png_paths, out_paths),
+                                   total=len(png_paths)):
+        png_processor(png_path, out_path)
+    print("Average processing time: {:.2f}s".format(
+        np.mean(png_processor.process_times)))
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(**vars(args))

PolyuColor/run.sh

@@ -0,0 +1,2 @@
+#!/usr/bin/env bash
+python run.py -p data

PolyuColor/utils/__init__.py

@@ -0,0 +1,36 @@
+from fractions import Fraction
+from pathlib import Path
+from json import JSONEncoder
+from .utils import *
+from .image_utils import *
+
+def rmtree(path: Path):
+    if path.is_file():
+        path.unlink()
+    else:
+        for ch in path.iterdir():
+            rmtree(ch)
+        path.rmdir()
+
+
+def safe_save(fpath, data, save_fun, rewrite=False, error_msg='File {fpath} exists! To rewite it use `--rewrite` flag', **kwargs):
+    if not fpath.is_file() or rewrite:
+        save_fun(str(fpath), data, **kwargs)
+    else:
+        raise FileExistsError(error_msg.format(fpath=fpath))
+
+
+class FractionJSONEncoder(JSONEncoder):
+    def default(self, o):
+        if isinstance(o, Fraction):
+            return {'Fraction': [o.numerator, o.denominator]}
+        else:
+            return o.__dict__
+
+
+def fraction_from_json(json_object):
+    if 'Fraction' in json_object:
+        return Fraction(*json_object['Fraction'])
+    return json_object
+
+

PolyuColor/utils/image_utils.py

@@ -0,0 +1,69 @@
+import cv2
+import numpy as np
+import torch
+
+
+def save_img(img, name, gamma=False):
+    if gamma:
+        img = np.power(img, 1/2.2)
+    img = np.clip(img, 0, 1)
+    img = (img * 65535).astype(np.uint16)
+    if img.ndim == 3:
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    cv2.imwrite(name, img)
+
+
+def compute_y(img: np.ndarray) -> np.ndarray:
+    y = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2]
+    return y
+
+
+def compute_raw_y(img: np.ndarray) -> np.ndarray:
+    g1 = img[..., 1]
+    g2 = img[..., 2]
+    ret = (g1 + g2) / 2
+    return ret
+
+
+def pack_raw(im):
+    # pack Bayer image to 4 channels
+    if isinstance(im, torch.Tensor):
+        im = torch.unsqueeze(im, dim=-1)
+        img_shape = im.shape
+        H = img_shape[0]
+        W = img_shape[1]
+
+        out = torch.cat((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :],
+                         im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]),
+                        dim=-1)
+    elif isinstance(im, np.ndarray):
+        im = np.expand_dims(im, axis=-1)
+        img_shape = im.shape
+        H = img_shape[0]
+        W = img_shape[1]
+
+        out = np.concatenate((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :],
+                              im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]),
+                             axis=-1)
+    return out
+
+
+def depack_raw(im):
+    # unpack 4 channels to Bayer image
+    img_shape = im.shape
+    H = img_shape[0]
+    W = img_shape[1]
+    if isinstance(im, torch.Tensor):
+        output = torch.zeros((H * 2, W * 2), dtype=im.dtype)
+    elif isinstance(im, np.ndarray):
+        output = np.zeros((H * 2, W * 2), dtype=im.dtype)
+    img_shape = output.shape
+    H = img_shape[0]
+    W = img_shape[1]
+
+    output[0:H:2, 0:W:2] = im[:, :, 0]
+    output[0:H:2, 1:W:2] = im[:, :, 1]
+    output[1:H:2, 1:W:2] = im[:, :, 2]
+    output[1:H:2, 0:W:2] = im[:, :, 3]
+
+    return output

PolyuColor/utils/utils.py

@@ -0,0 +1,56 @@
+from PIL import Image
+import json
+import os
+
+def json_read(fname, **kwargs):
+    with open(fname) as j:
+        data = json.load(j, **kwargs)
+    return data
+
+
+def json_save(fname, data, indent_len=4, **kwargs):
+    with open(fname, "w") as f:
+        s = json.dumps(data, sort_keys=True, ensure_ascii=False,
+                       indent=" " * indent_len, **kwargs)
+        f.write(s)
+
+
+def process_wb_from_txt(txt_path):
+    with open(txt_path, 'r') as fh:
+        txt = [line.rstrip().split() for line in fh]
+
+    txt = [[float(k) for k in row] for row in txt]
+
+    assert len(txt) in [1, 3]
+
+    if len(txt) == 1:
+        # wb vector
+        txt = txt[0]
+
+    return txt
+
+
+def process_ids_from_txt(txt_path):
+    with open(txt_path, 'r') as fh:
+        temp = fh.read().splitlines()
+    return temp
+
+
+def save_txt(p, s):
+    with open(p, 'w') as text_file:
+        text_file.write(s)
+
+
+def downscale_jpg(img_path, new_shape, quality_perc=100):
+    img = Image.open(img_path)
+    if (img.size[0], img.size[1]) != new_shape:
+        new_img = img.resize(new_shape, Image.ANTIALIAS)
+        new_img.save(img_path[:-len('.jpg')] + '.jpg',
+                     'JPEG', quality=quality_perc)
+
+
+def rename_img(img_path):
+    if img_path.lower().endswith('jpeg'):
+        os.rename(img_path, img_path[:-len('jpeg')] + 'jpg')
+    else:
+        os.rename(img_path, img_path[:-len('JPG')] + 'jpg')