Code associated with our paper "Graphics2RAW: Mapping Computer Graphics Images to Sensor RAW Images". The paper has been accepted to the International Conference on Computer Vision (ICCV'23).
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.
"""

import os
import pickle
import argparse
from utils.img_utils import *
import scipy.io
from data_preparation.data_generator_illum_est import get_gt_illum_by_fname


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--graphics_path', type=str,
                        help='path to graphics images',
                        default='illum_est_expts/synthia/SYNTHIA_RAND_CVPR16/RGB'
                        )
    parser.add_argument('--train_val_set', type=str,
                        help='only use these images for training and validation',
                        default='assets/split_files/illum_est/synthia_train_val_list.p'
                        )
    parser.add_argument('--save_path', type=str,
                        help='path to exr images',
                        default='illum_est_expts/data/SamsungNX2000/ours/'
                        )
    parser.add_argument('--target_camera', type=str,
                        help='use camera name to select CM1 and CM2',
                        default='SamsungNX2000'
                        )
    parser.add_argument('--cm_file', type=str,
                        help='path to NUS CM1 and CM2',
                        default='assets/container_dngs/NUS_CST_mats.p'
                        )
    parser.add_argument('--mat_file', type=str,
                        help='mat file with gt illuminations',
                        default='illum_est_expts/nus_metadata/nus_outdoor_gt_illum_mats/SamsungNX2000_gt.mat'
                        )
    parser.add_argument('--use_train_val_illums_only', action='store_true',
                        help='sample illuminants from the convex hull built from training and validation illuminants only'
                        )
    parser.add_argument('--split_file', type=str,
                        help='path to file specifying which filenames belong to which split, example:path/to/SamsungNX2000_train_valid_test_split_fns.p')
    parser.add_argument('--max_illums', type=int,
                        help='Maximum number of illuminants used to build the convex hull. None for no max limit',
                        default=None)
    parser.add_argument('--no_safe_invert', action='store_true',
                        help='Do not use safe invert for highlight regions, setting this flag will make the saturated regions have a color cast',
                        )
    parser.add_argument('--rgb_gain', action='store_true',
                        help='Apply random global gain to images',
                        )
    parser.add_argument('--rgb_gain_mean', type=float, default=1.0,
                        help='Mean value for rgb gain, used only when rgb_gain is set to True.',
                        )

    args = parser.parse_args()

    print(args)

    return args


if __name__ == '__main__':

    RAND_SEED = 101
    np.random.seed(seed=RAND_SEED)

    args = parse_args()

    train_val_set = pickle.load(open(args.train_val_set, 'rb'))

    metadata = pickle.load(open(args.cm_file, 'rb'))[args.target_camera]
    cmD65 = metadata['cm_D65']
    cmA = metadata['cm_A']

    # load illuminants
    if args.use_train_val_illums_only:
        # Use only training and validation illuminants to build the convex hull
        gt_illum_by_fname = get_gt_illum_by_fname(args.mat_file)
        splits = pickle.load(open(args.split_file, 'rb'))['real_split']
        split_fns_train = splits['train']
        split_fns_valid = splits['valid']
        split_fns = np.append(split_fns_train, split_fns_valid)
        print('Number of training and validation illums: ', len(split_fns))

        fnames = [os.path.basename(f)[:-4] for f in split_fns]
        gt_illum = [gt_illum_by_fname[fname] for fname in fnames]
        gt_illum = np.concatenate([gt_illum])  # all 150 train and validation illums

        # sample within the training and validation illums
        if args.max_illums is not None:
            seed = 0
            np.random.seed(seed)
            rand_idx = np.random.choice(len(gt_illum), args.max_illums, replace=False)
            gt_illum = gt_illum[rand_idx, :]
            np.random.seed(RAND_SEED)  # reset the seed
        print('Number of illums used to build the convex hull: ', len(gt_illum))

    else:
        gt_illum = scipy.io.loadmat(args.mat_file)
        gt_illum = gt_illum['groundtruth_illuminants']
        gt_illum[:, 0], gt_illum[:, 1], gt_illum[:, 2] = get_illum_normalized_by_g(gt_illum)
    illum_mean = np.mean(gt_illum, 0)
    illum_cov = np.cov(np.transpose(gt_illum))

    xyz2srgb_mat = get_xyz_to_srgb_mat()

    gt_illum_array = []
    file_name_array = []

    for i, filename in enumerate(train_val_set):
        savename = filename
        print(i, filename)

        graphics_img = cv2.imread(os.path.join(args.graphics_path, filename), -1)
        graphics_img = graphics_img[:, :, ::-1]
        graphics_img = np.array(graphics_img).astype(np.float32) / 255.0  # SYNTHIA images are 8 bit
        graphics_img = np.clip(graphics_img, 0, 1)

        # De-gamma
        graphics_img = graphics_img ** 2.2

        # Sample an illuminant
        while True:
            wb_vec = np.random.multivariate_normal(illum_mean, illum_cov, 1).squeeze()
            if in_hull(np.expand_dims(wb_vec[[0, 2]], axis=0), gt_illum[:, [0, 2]]):
                break
        gt_illum_array.append(wb_vec)
        print(wb_vec)

        # Compute CST matrix
        cst_mat = get_cst_matrix(cmD65, cmA, wb_vec)

        # sRGB to CIE XYZ (Invert XYZ2sRGB)
        raw_est = apply_combined_mat(graphics_img, np.linalg.inv(xyz2srgb_mat))

        # CIE XYZ to device RGB (Invert CST)
        raw_est = apply_combined_mat(raw_est, np.linalg.inv(cst_mat))

        # Inverse digital gain (optional)
        rgb_gain = np.random.normal(loc=args.rgb_gain_mean, scale=0.1) if args.rgb_gain else 1.0

        # Invert WB
        if not args.no_safe_invert:
            raw_est = safe_invert_gains(raw_est, wb_vec, rgb_gain)
        else:
            wb_mat = get_wb_as_matrix(wb_vec)
            raw_est *= rgb_gain
            raw_est = apply_combined_mat(raw_est, np.linalg.inv(wb_mat))
        cv2.imwrite(
            os.path.join(args.save_path, savename[:-4] + '.png'),
            (raw_est[:, :, [2, 1, 0]] * 65535).astype(np.uint16))
        file_name_array.append(savename[:-4])

    gt_values = {'gt_illum': gt_illum_array, 'filenames': file_name_array}

    pickle.dump(gt_values, open(os.path.join(args.save_path, 'gt_illum.p'), "wb"))

    print('Done')

data_generation/invert_synthia_graphics2raw.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.
"""

import os
import pickle
import argparse
from utils.img_utils import *
from data_generation.unprocess import unprocess
import tensorflow as tf


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--graphics_path', type=str,
                        help='path to graphics images',
                        default='illum_est_expts/synthia/SYNTHIA_RAND_CVPR16/RGB'
                        )
    parser.add_argument('--train_val_set', type=str,
                        help='only use these images for training and validation',
                        default='assets/split_files/illum_est/synthia_train_val_list.p'
                        )
    parser.add_argument('--save_path', type=str,
                        help='path to exr images',
                        default='illum_est_expts/data/SamsungNX2000/upi'
                        )

    args = parser.parse_args()

    print(args)

    return args


if __name__ == '__main__':
    RAND_SEED = 101
    np.random.seed(RAND_SEED)
    tf.random.set_seed(RAND_SEED)

    args = parse_args()

    train_val_set = pickle.load(open(args.train_val_set, 'rb'))

    gt_illum_array = []
    file_name_array = []

    for i, filename in enumerate(train_val_set):
        savename = filename
        print(i, filename)

        graphics_img = cv2.imread(os.path.join(args.graphics_path, filename), -1)
        graphics_img = graphics_img[:, :, ::-1]
        graphics_img = np.array(graphics_img).astype(np.float32) / 255.0  # SYNTHIA images are 8 bit
        graphics_img = np.clip(graphics_img, 0, 1)

        graphics_img = tf.convert_to_tensor(graphics_img, dtype=tf.float32)

        raw_est, metadata = unprocess(graphics_img)
        raw_est = raw_est.numpy().astype(np.float32)
        for k, v in metadata.items():
            metadata[k] = v.numpy().astype(np.float32)

        wb_vec = np.array([1 / metadata['red_gain'], 1, 1 / metadata['blue_gain']])

        gt_illum_array.append(wb_vec)
        print(wb_vec)

        cv2.imwrite(
            os.path.join(args.save_path, savename[:-4] + '.png'),
            (raw_est[:, :, [2, 1, 0]] * 65535).astype(np.uint16))
        file_name_array.append(savename[:-4])

    gt_values = {'gt_illum': gt_illum_array, 'filenames': file_name_array}

    pickle.dump(gt_values, open(os.path.join(args.save_path, 'gt_illum.p'), "wb"))

    print('Done')

data_generation/invert_synthia_upi.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.
"""

import cv2
from utils.img_utils import update_hex_image,update_wb_values
from binascii import hexlify, unhexlify
import argparse
from glob import glob
import os
from pipeline.pipeline_utils import get_metadata
import copy

"""
Package clean RAW image (averaged from 30 ISO 50 frames) into a DNG. 
"""


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dng_folder_path', type=str,
                        help='path to day-to-night nighttime dataset iso50 dngs',
                        )
    parser.add_argument('--raw_folder_path', type=str,
                        help='path to day-to-night nighttime averaged clean raw images',
                        )
    parser.add_argument('--save_path', type=str,
                        help='path to save to',
                        default='neural_isp_expts/data/clean_raw_dngs'
                        )
    parser.add_argument('--container_dng_path', type=str,
                        help='container dng',
                        default='assets/container_dngs/container_dng_S20_FE_main_rectilinear_OFF_noise_OFF.dng'
                        )
    parser.add_argument('--wb_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=50816
                        )
    parser.add_argument('--image_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=59288
                        )

    args = parser.parse_args()

    print(args)

    return args


if __name__ == '__main__':

    args = parse_args()
    assert os.path.exists(args.save_path), f'{args.save_path} does not exist!'

    dnglist = sorted(glob(os.path.join(args.dng_folder_path,'*.dng')))
    rawlist = sorted(glob(os.path.join(args.raw_folder_path,'*.png')))

    with open(args.container_dng_path, "rb") as fn:
        myhexc = hexlify(fn.read())
    myhexc = bytearray(myhexc)

    for i in range(len(dnglist)):
        print(os.path.basename(dnglist[i]),os.path.basename(rawlist[i]))
        bayer = cv2.imread(rawlist[i],-1)
        metadata = get_metadata(dnglist[i])

        myhex = copy.deepcopy(myhexc)

        wb_dng = metadata['as_shot_neutral']

        myhex = update_wb_values(myhex, wb_dng, args.wb_start)
        myhex = update_hex_image(myhex,bayer,args.image_start)

        db = unhexlify(myhex)

        savepath = os.path.join(args.save_path,os.path.basename(dnglist[i]))
        with open(savepath,"wb") as fb:
            fb.write(db)

    print('Done')
data_generation/package_clean_nighttime_raw_to_dng.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Package graphics exr image to DNG file
"""

import os
import argparse
from glob import glob
from binascii import unhexlify
import scipy.io

from pipeline.pipeline_utils import get_metadata
from utils.img_utils import *
import copy
import pickle

RAND_SEED = 101
np.random.seed(RAND_SEED)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--exr_folder_path', type=str,
                        help='path to exr images')  # expected directory structure: exr_folder_path/<scene name>/<GT_xx_xx>/xxx.exr
    parser.add_argument('--save_path', type=str,
                        help='path to save to',
                        default='neural_isp_expts/data/graphics_dngs_graphics2raw')
    parser.add_argument('--container_dng_path', type=str,
                        help='path to container dng',
                        default='assets/container_dngs/container_dng_S20_FE_main_rectilinear_OFF_gain_OFF_noise_OFF_cam_calib_OFF.dng'
                        )
    parser.add_argument('--wb_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=50816
                        )
    parser.add_argument('--image_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=59288
                        )
    parser.add_argument('--train_val_set', type=str,
                        help='create dngs only for these images',
                        default='assets/split_files/graphics2raw_train_val_list.p'
                        )
    parser.add_argument('--no_safe_invert', action='store_true',
                        help='Do not use safe invert for highlight regions, setting this flag will make the saturated regions have a color cast',
                        )
    parser.add_argument('--rgb_gain', action='store_true',
                        help='Apply random global gain to images',
                        )
    parser.add_argument('--rgb_gain_mean', type=float, default=0.8,
                        help='Mean value for rgb gain, used only when rgb_gain is set to True.',
                        )
    parser.add_argument('--max_illums', type=int, default=None,
                        help='Maximum number of illuminants used to build the convex hull. None for no max limit',
                        )
    parser.add_argument('--illum_seed', type=int, default=None,
                        help='Random seed for sampling illuminants used to build the convex hull.',
                        )
    args = parser.parse_args()
    print(args)

    return args


if __name__ == '__main__':

    args = parse_args()
    assert os.path.exists(args.save_path), f'{args.save_path} does not exist!'

    exrlist = []

    allsubfol = [f.path for f in os.scandir(args.exr_folder_path) if f.is_dir()]
    for subfol in allsubfol:
        allsubsubfol = [f.path for f in os.scandir(os.path.join(args.exr_folder_path, subfol)) if f.is_dir()]
        for subsubfol in allsubsubfol:
            exrlistfol = sorted(glob(os.path.join(args.exr_folder_path, subfol, subsubfol, '*.exr')))
            exrlist.append(exrlistfol)

    train_val_set = pickle.load(open(args.train_val_set, "rb"))

    metadata = get_metadata(args.container_dng_path)
    metadata = get_extra_tags(args.container_dng_path, metadata)

    # IMPORTANT: for S20FE, color_matrix_1 corresponds to D65,
    # color_matrix_2 corresponds to Standard Light A
    # The order may be different for different cameras, need to check beforehand!
    cmD65 = metadata['color_matrix_1']
    cmA = metadata['color_matrix_2']

    # Illumination dictionary
    # load nighttime illuminants
    gt_illum = scipy.io.loadmat('./data_generation/night_dict_v2.mat')
    gt_illum = gt_illum['gt_illum']
    gt_illum = gt_illum[0:45, :]
    # discard outliers
    indd = gt_illum[:, 0] < 1
    gt_illum = gt_illum[indd, :]

    if args.max_illums is not None:
        # Sample N illums from 39 illums
        np.random.seed(args.illum_seed)
        rand_idx = np.random.choice(len(gt_illum), args.max_illums, replace=False)
        gt_illum = gt_illum[rand_idx, :]
        np.random.seed(RAND_SEED)  # reset the seed

    print(gt_illum.shape[0], 'num illums')

    gt_illum[:, 0], gt_illum[:, 1], gt_illum[:, 2] = get_illum_normalized_by_g(gt_illum)

    illum_mean = np.mean(gt_illum, 0)
    illum_cov = np.cov(np.transpose(gt_illum))

    # for S20FE main camera
    white_level = 1023
    black_level = 64

    w = 4032
    h = 3024

    with open(args.container_dng_path, "rb") as fn:
        myhexc = hexlify(fn.read())
    myhexc = bytearray(myhexc)

    xyz2srgb_mat = get_xyz_to_srgb_mat()

    for i in range(len(exrlist)):
        pathnames = exrlist[i][0].split('/')
        savename = pathnames[-3] + '_' + pathnames[-1]

        if savename[:-4] in train_val_set:
            print(exrlist[i])

            exr_img = cv2.imread(exrlist[i][0], cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
            exr_img = exr_img[:, :, ::-1]

            # Resize to the shape of S20 FE RAW image: 3024 x 4032
            if exr_img.shape[0] < h and exr_img.shape[1] < w:
                exr_img = cv2.resize(exr_img, dsize=(w, h))
            elif exr_img.shape[0] < h and exr_img.shape[1] > w:
                exr_img = exr_img[:, 0:w]
                exr_img = cv2.resize(exr_img, dsize=(w, h))
            elif exr_img.shape[0] > h and exr_img.shape[1] < w:
                exr_img = exr_img[0:h, :]
                exr_img = cv2.resize(exr_img, dsize=(w, h))
            else:
                exr_img = exr_img[0:h, 0:w]

            exr_img = np.clip(exr_img, 0, 1)

            # Sample an illuminant
            while True:
                wb_vec = np.random.multivariate_normal(illum_mean, illum_cov, 1).squeeze()
                if in_hull(np.expand_dims(wb_vec[[0, 2]], axis=0), gt_illum[:, [0, 2]]):
                    break
            print(wb_vec)

            # Compute CST matrix
            cst_mat = get_cst_matrix(cmD65, cmA, wb_vec)

            # sRGB to CIE XYZ (Invert XYZ2sRGB)
            raw_est = apply_combined_mat(exr_img, np.linalg.inv(xyz2srgb_mat))

            # CIE XYZ to device RGB (Invert CST)
            raw_est = apply_combined_mat(raw_est, np.linalg.inv(cst_mat))

            # Inverse digital gain (optional)
            rgb_gain = np.random.normal(loc=args.rgb_gain_mean, scale=0.1) if args.rgb_gain else 1.0

            # Invert WB
            if not args.no_safe_invert:
                raw_est = safe_invert_gains(raw_est, wb_vec, rgb_gain)
            else:
                wb_mat = get_wb_as_matrix(wb_vec)
                raw_est *= rgb_gain
                raw_est = apply_combined_mat(raw_est, np.linalg.inv(wb_mat))

            wb_dng = wb_vec

            # Mosaic
            raw_est = RGB2bayer(raw_est)  # Bayer for S20 FE: G R B G

            # Denormalize
            raw_est = raw_est * (white_level - black_level) + black_level
            raw_est[raw_est < 0] = 0
            raw_est[raw_est > white_level] = white_level

            # Save to DNG
            myhex = copy.deepcopy(myhexc)
            myhex = update_wb_values(myhex, wb_dng, args.wb_start)
            myhex = update_hex_image(myhex, raw_est, args.image_start)

            db = unhexlify(myhex)

            savepath = os.path.join(args.save_path, savename[:-4] + '.dng')
            with open(savepath, "wb") as fb:
                fb.write(db)

    print('Done')

data_generation/package_exr_to_dng_graphics2raw.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Package graphics exr image to DNG file
"""

import os
import argparse
from glob import glob
from binascii import unhexlify

import tensorflow as tf

from utils.img_utils import *
import copy
from data_generation.unprocess import unprocess
import pickle


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--exr_folder_path', type=str,
                        help='path to exr images')  # expected directory structure exr_folder_path/<scene name>/<GT_xx_xx>/xxx.exr
    parser.add_argument('--save_path', type=str,
                        help='path to save to',
                        default='./neural_isp_expts/data/graphics_dngs_upi')
    parser.add_argument('--container_dng_path', type=str,
                        help='path to save to',
                        default='assets/container_dngs/container_dng_S20_FE_main_rectilinear_OFF_gain_OFF_noise_OFF_cam_calib_OFF.dng'
                        )
    parser.add_argument('--wb_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=50816
                        )
    parser.add_argument('--image_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=59288
                        )
    parser.add_argument('--colormatrix1_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=50240
                        )
    parser.add_argument('--colormatrix2_start', type=int,
                        help='magic value for S20 FE main camera',
                        default=50384
                        )
    parser.add_argument('--train_val_set', type=str,
                        help='create dngs only for these images',
                        default='assets/split_files/graphics2raw_train_val_list.p'
                        )

    args = parser.parse_args()

    print(args)

    return args


if __name__ == '__main__':
    RAND_SEED = 101
    np.random.seed(RAND_SEED)
    tf.random.set_seed(RAND_SEED)

    args = parse_args()
    assert os.path.exists(args.save_path), f'{args.save_path} does not exist!'

    exrlist = []

    allsubfol = [f.path for f in os.scandir(args.exr_folder_path) if f.is_dir()]
    for subfol in allsubfol:
        allsubsubfol = [f.path for f in os.scandir(os.path.join(args.exr_folder_path, subfol)) if f.is_dir()]
        for subsubfol in allsubsubfol:
            exrlistfol = sorted(glob(os.path.join(args.exr_folder_path, subfol, subsubfol, '*.exr')))
            exrlist.append(exrlistfol)

    train_val_set = pickle.load(open(args.train_val_set, "rb"))

    # for S20 FE main camera
    white_level = 1023
    black_level = 64

    w = 4032
    h = 3024

    with open(args.container_dng_path, "rb") as fn:
        myhexc = hexlify(fn.read())
    myhexc = bytearray(myhexc)

    for i in range(len(exrlist)):
        pathnames = exrlist[i][0].split('/')
        savename = pathnames[-3] + '_' + pathnames[-1]
        
        
        if savename[:-4] in train_val_set:
            print(exrlist[i])

            exr_img = cv2.imread(exrlist[i][0], cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)  # .astype('float32')
            exr_img = exr_img[:, :, ::-1]  # .exr image is BGR, change it to RGB for processing

            # resize to 3024 x 4032
            if exr_img.shape[0] < h and exr_img.shape[1] < w:
                exr_img = cv2.resize(exr_img, dsize=(w, h))
            elif exr_img.shape[0] < h and exr_img.shape[1] > w:
                exr_img = exr_img[:, 0:w]
                exr_img = cv2.resize(exr_img, dsize=(w, h))
            elif exr_img.shape[0] > h and exr_img.shape[1] < w:
                exr_img = exr_img[0:h, :]
                exr_img = cv2.resize(exr_img, dsize=(w, h))
            else:
                exr_img = exr_img[0:h, 0:w]

            exr_img = np.clip(exr_img, 0, 1)
            exr_img = exr_img ** (1/2.2)
            exr_img = tf.convert_to_tensor(exr_img, dtype=tf.float32)

            raw_est, metadata = unprocess(exr_img)
            raw_est = raw_est.numpy().astype(np.float32)
            for k, v in metadata.items():
                metadata[k] = v.numpy().astype(np.float32)

            wb_dng = np.array([1 / metadata['red_gain'], 1, 1 / metadata['blue_gain']])

            # Bayer for S20 FE
            # GR
            # BG
            raw_est = RGB2bayer(raw_est)

            # denormalize
            raw_est = raw_est * (white_level - black_level) + black_level
            raw_est[raw_est < 0] = 0
            raw_est[raw_est > white_level] = white_level

            colormatrix1= metadata['xyz2cam'].flatten().astype(np.float64)

            myhex = copy.deepcopy(myhexc)
            myhex = update_wb_values(myhex, wb_dng, args.wb_start)
            myhex = update_hex_image(myhex, raw_est, args.image_start)
            myhex = update_colormatrix1_values(myhex, colormatrix1, args.colormatrix1_start)
            myhex = update_colormatrix1_values(myhex, colormatrix1, args.colormatrix2_start) # repeat for cm2

            db = unhexlify(myhex)

            savepath = os.path.join(args.save_path, savename[:-4] + '.dng')
            with open(savepath, "wb") as fb:                
                fb.write(db)

    print('Done')

data_generation/package_exr_to_dng_upi.py
SamsungLabs/graphics2raw
"""
Copy over the code from https://github.com/timothybrooks/unprocessing/blob/master/unprocess.py
and make modifications according to docs/upi.md
"""

import tensorflow as tf
import pickle


def unprocess(image):
    pass
data_generation/unprocess.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import glob
import cv2
import numpy as np

from pipeline.pipeline_utils import demosaic
from pipeline.pipeline import run_pipeline
from pipeline.raw_utils import stack_rggb_channels, RGGB2Bayer
from os.path import join
from noise_profiler.image_synthesizer import synthesize_noisy_image_v2


np.random.seed(101)


def get_all_patches(img, patch_size=48, stride=48):
    h, w, = img.shape[:2]
    patches = []
    # extract patches
    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            x = img[i:i + patch_size, j:j + patch_size, ...]
            patches.append(x)
    return patches


def normalize_s20fe(img):
    white_level = 1023.0
    black_level = 64.0
    img = (img - black_level) / (white_level - black_level)
    return img


def denormalize_s20fe(img):
    white_level = 1023.0
    black_level = 64.0
    img = img * (white_level - black_level) + black_level
    return img


def post_process_stacked_images_s20fe(img, cfa_pattern):
    """
    Process an image that went through stack_bayer_norm_gamma_s20fe
    to demosaiced, gamma-ed RAW RGB
    :param img:
    :param cfa_pattern:
    :return:
    """
    img = (img ** 2.2) * 1023.0  # [0, 1023]
    img = np.clip(img, 64, 1023)  # [64, 1023]
    img = (img - 64) / (1023 - 64)  # [0, 1] normalization with black level subtraction
    img = RGGB2Bayer(img)
    img = demosaic(img, cfa_pattern, alg_type='EA')
    img = img ** (1 / 2.2)
    return img


def post_process_stacked_images_s20fe_to_srgb(img, metadata):
    """
    Process an image that went through stack_bayer_norm_gamma_s20fe
    to sRGB
    :param img:
    :param metadata:
    :return:
    """
    img = (img ** 2.2) * 1023.0  # [0, 1023]
    img = np.clip(img, 64, 1023)  # [64, 1023]
    img = RGGB2Bayer(img)

    stages = [
        'raw',
        'normal',
        'lens_shading_correction',
        'white_balance',
        'demosaic',
        'xyz',
        'srgb',
        'fix_orient',
        'gamma',
    ]

    params = {
        'input_stage': 'raw',
        'output_stage': 'gamma',
        'demosaicer': 'EA',
    }

    img = run_pipeline(img, params=params, metadata=metadata, stages=stages)
    return img


def linearize_stacked_images_s20fe(img):
    """
    Process an image that went through stack_bayer_norm_gamma_s20fe
    to stacked, normalized, linear RAW RGB in the range of [0, 1]
    with black level subtraction
    :param img: a RAW image that went through stack_bayer_norm_gamma_s20fe
    :return:
    """
    img = img ** 2.2
    img = img * 1023.0  # [0, 1023], no black level subtraction
    img = normalize_s20fe(img)  # [0, 1], with black level subtraction
    return img


def stack_bayer_norm_gamma_s20fe(img, clip_bot=False):
    """

    :param img: not normalized bayer image, [black_level, white_level]
    :param clip_bot: clip minimum value; maximum value is always clipped to 1023
    :return:
    """
    img = img.astype(np.float32)
    img = np.clip(img, 0.0, 1023.0)
    if clip_bot:
        img = np.clip(img, 64.0, 1023.0)
    img = img / 1023.0
    img = stack_rggb_channels(img)  # do not pass in bayer_pattern to not change the stacking order
    img = img ** (1 / 2.2)
    return img


def process_image(file_name, data_type='input', noise_model=None, iso=100):
    """
    Assumptions:
    Camera: S20FE
    h == 3024 and w == 4032

    :param iso: 1600 or 3200
    :param noise_model:  hg noise model
    :param data_type:
    input_graphics_raw: clean_raw -> noise -> stacked -> clip to 0~1023 -> gamma
    target_graphics_raw: clean_raw -> stacked -> clip to 64~1023 -> gamma
    input_real: noisy_raw -> stacked -> clip to 0~1023 -> gamma
    target_real: clean_raw -> stacked -> clip to 64~1023 -> gamma
    :return:
    """
    img = cv2.imread(file_name, cv2.IMREAD_UNCHANGED)  # bayer image
    h, w = img.shape[:2]
    assert h == 3024 and w == 4032, 'Wrong height and width!'

    # Generate input patches
    if data_type == 'input_graphics_raw':
        # Pre-generated noise
        img = synthesize_noisy_image_v2(img, model=noise_model['noise_model'],
                                        dst_iso=iso, min_val=0,
                                        max_val=1023,
                                        iso2b1_interp_splines=noise_model['iso2b1_interp_splines'],
                                        iso2b2_interp_splines=noise_model['iso2b2_interp_splines'])

        img_ = stack_bayer_norm_gamma_s20fe(img, clip_bot=False)

    # Generate target patches
    elif data_type == 'input_real':
        img_ = stack_bayer_norm_gamma_s20fe(img, clip_bot=False)
    else:
        # Covered cases: target_graphics_raw, target_real
        # img: linear, [black_level, white_level], bayer
        img_ = stack_bayer_norm_gamma_s20fe(img, clip_bot=True)

    img_ = img_.astype(np.float32)
    return img_


def datagenerator_raw(data_dir='dummy_dataset/train', batch_size=128, patch_size=48, stride=48, verbose=True,
                      data_type='graphics_raw', noise_model=None, iso=3200):
    if 'graphics' in data_type:
        input_list = sorted(glob.glob(join(data_dir, 'clean_raw', '*.png')))
        target_list = input_list
    elif 'real' in data_type:
        input_list = sorted(glob.glob(join(data_dir, 'noisy_raw', '*.png')))
        target_list = sorted(glob.glob(join(data_dir, 'clean_raw', '*.png')))
    else:
        raise Exception('Unexpected data type')

    # initialize
    patches_input = []
    patches_target = []
    # generate patches
    for i in range(len(input_list)):
        img_input = process_image(input_list[i], data_type=f'input_{data_type}', noise_model=noise_model, iso=iso)
        img_target = process_image(target_list[i], data_type=f'target_{data_type}')
        img_pair = np.dstack([img_input, img_target])
        patches = get_all_patches(img_pair, patch_size, stride)
        patches = np.array(patches)
        assert patches[0].shape[-1] == 8  # both input and target are stacked
        split = 4
        patch_input = patches[..., :split]
        patch_target = patches[..., split:]
        patches_input.append(patch_input)
        patches_target.append(patch_target)
        if verbose:
            print(str(i + 1) + '/' + str(len(input_list)) + ' is done')

    patches_input = np.concatenate(patches_input)
    patches_target = np.concatenate(patches_target)

    discard_n = len(patches_input) - len(patches_input) // batch_size * batch_size

    patches_input = np.delete(patches_input, range(discard_n), axis=0)
    patches_target = np.delete(patches_target, range(discard_n), axis=0)

    print('^_^-raw training data finished-^_^')
    return patches_input, patches_target


data_preparation/data_generator_denoise.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import glob
import cv2
import numpy as np
import pickle
import os
import scipy
from utils.img_utils import get_illum_normalized_by_g


def shuffle_files(fnames):
    """
    Shuffle graphics images to cover all scenes
    graphics_names: sorted graphics image names
    """
    np.random.seed(101)
    fnames = np.array(fnames)
    indices = np.arange(len(fnames))
    np.random.shuffle(indices)
    fnames = fnames[indices]
    return fnames


def get_all_patches(img, patch_size=48, stride=48):
    h, w, = img.shape[:2]
    patches = []
    # extract patches
    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            x = img[i:i + patch_size, j:j + patch_size, ...]
            patches.append(x)
    return patches


def gen_patches_illum_est(file_name, patch_size=48, stride=48):
    """
    Crop one image into patches
    :return:
    """
    img = np.array(cv2.imread(file_name, cv2.IMREAD_UNCHANGED))[:, :, ::-1].astype(np.float32)  # img is now in rgb
    img = img / 65535.0
    img = np.clip(img, 0, 1)
    patches = get_all_patches(img, patch_size, stride)
    return patches


def get_gt_illum_by_fname(illum_file):
    gt_illum_by_fname = {}
    if illum_file.endswith('.p'):
        content = pickle.load(open(illum_file, 'rb'))
        gt_illums = content['gt_illum']  # already normalized by g
        filenames = content['filenames']

    elif illum_file.endswith('.mat'):
        content = scipy.io.loadmat(illum_file)
        gt_illums = content['groundtruth_illuminants']
        gt_illums[:, 0], gt_illums[:, 1], gt_illums[:, 2] = get_illum_normalized_by_g(gt_illums)
        filenames = [name[0][0] + '_st4' for name in content['all_image_names']]  # array([array(['SamsungNX2000_0001'])) -> 'SamsungNX2000_0001_st4'
    else:
        raise Exception('Unsupported gt illum file type.')

    for illum, fname in zip(gt_illums, filenames):
        gt_illum_by_fname[fname] = illum
    return gt_illum_by_fname


def datagenerator_illum_est(dataset_dir='', illum_file='', split_file='', split='train',
                            batch_size=128, patch_size=48, stride=48, debug=False):
    file_list = sorted(glob.glob(os.path.join(dataset_dir, '*.png')))  # get name list of all .png files
    file_list = shuffle_files(file_list)  # shuffle to sample from all graphics scenes
    gt_illum_by_fname = get_gt_illum_by_fname(illum_file)

    split_indices = pickle.load(open(split_file, 'rb'))[split]
    file_list = file_list[split_indices]

    if debug:
        # For debugging only, check if images in the split are expected
        mysplit = 'valid' if split == 'val' else split
        split_fns_fp = split_file[:-5] + 'fns.p' 
        split_type = 'graphics_split' if dataset_dir.split('/')[-1] != 'real' else 'real_split'
        split_fns = pickle.load(open(split_fns_fp, 'rb'))[split_type][mysplit]
        for myf, gtf in zip(file_list, split_fns):
            assert os.path.basename(myf) == gtf, f'Panic! {os.path.basename(myf)}, {gtf}'
        print(f'Loaded images are correct.')
        # end of debugging

    in_patches = []
    gt_illums = []

    # generate patches
    for file in file_list:
        patches = gen_patches_illum_est(file, patch_size, stride)
        fname = os.path.basename(file)[:-4]
        gt_illum = gt_illum_by_fname[fname]
        in_patches.append(patches)
        gt_illums.append([gt_illum] * len(patches))

    in_patches = np.concatenate(in_patches)
    gt_illums = np.concatenate(gt_illums)
    discard_n = len(in_patches) - len(in_patches) // batch_size * batch_size
    in_patches = np.delete(in_patches, range(discard_n), axis=0)
    gt_illums = np.delete(gt_illums, range(discard_n), axis=0)

    print(f'^_^-{split} data finished-^_^')
    return in_patches, gt_illums


data_preparation/data_generator_illum_est.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


import glob
import cv2
import numpy as np
import pickle

from pipeline.pipeline import run_pipeline


aug_times = 1


def data_aug(img, mode=0, is_batch=False):
    """
    :param is_batch:
    When is_batch == True: img (b, h, w, c)
    When is_batch == False: img (h, w, c)
    """
    flipud_axis = 1 if is_batch else 0
    rot_axes = (1, 2) if is_batch else (0, 1)

    if mode == 0:
        return img
    elif mode == 1:
        return np.flip(img, axis=flipud_axis)
    elif mode == 2:
        return np.rot90(img, axes=rot_axes)
    elif mode == 3:
        return np.flip(np.rot90(img, axes=rot_axes), axis=flipud_axis)
    elif mode == 4:
        return np.rot90(img, k=2, axes=rot_axes)
    elif mode == 5:
        return np.flip(np.rot90(img, k=2, axes=rot_axes), axis=flipud_axis)
    elif mode == 6:
        return np.rot90(img, k=3, axes=rot_axes)
    elif mode == 7:
        return np.flip(np.rot90(img, k=3, axes=rot_axes), axis=flipud_axis)


def get_all_patches(img, patch_size=48, stride=48):
    h, w = img.shape[:2]
    patches = []
    # extract patches
    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            x = img[i:i + patch_size, j:j + patch_size, ...]
            # data aug
            for k in range(0, aug_times):
                x_aug = data_aug(x, mode=0)  # np.random.randint(0,8))
                patches.append(x_aug)
    return patches


def gen_patches_sRGB(file_name, patch_size=48, stride=48):
    # read image
    img = cv2.imread(file_name, cv2.IMREAD_UNCHANGED)
    img = np.array(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), dtype=np.float32)
    patches = get_all_patches(img, patch_size, stride)
    return patches


def gen_patches_raw(file_name, meta_name, wb_illum='avg', patch_size=48, stride=48, is_PS_sRGB=True):
    """
    :param is_PS_sRGB: whether target sRGB is produced by PS
    :return:
    """

    meta_data_org = pickle.load(open(meta_name, "rb"))

    if wb_illum == 'avg':
        meta_data_org['as_shot_neutral'] = meta_data_org['avg_night_illuminant']  # modify as_shot_neutral

    img = cv2.imread(file_name, cv2.IMREAD_UNCHANGED)

    if is_PS_sRGB:
        params = {
            'save_as': 'png',  # options: 'jpg', 'png', 'tif', etc.
            'white_balancer': 'default',  # options: default, or self-defined module
            'demosaicer': '',  # options: '' for simple interpolation,
            #          'EA' for edge-aware,
            #          'VNG' for variable number of gradients,
            #          'menon2007' for Menon's algorithm
            'tone_curve': 'simple-s-curve',  # options: 'simple-s-curve', 'default', or self-defined module
            'output_stage': 'default_cropping',
        }

        stages = ['raw', 'active_area_cropping', 'linearization', 'normal', 'lens_shading_correction',
                'white_balance', 'demosaic', 'default_cropping']

    else:
        params = {
            'save_as': 'png',  # options: 'jpg', 'png', 'tif', etc.
            'white_balancer': 'default',  # options: default, or self-defined module
            'demosaicer': '',  # options: '' for simple interpolation,
            #          'EA' for edge-aware,
            #          'VNG' for variable number of gradients,
            #          'menon2007' for Menon's algorithm
            'tone_curve': 'simple-s-curve',  # options: 'simple-s-curve', 'default', or self-defined module
            'output_stage': 'demosaic',
        }

        stages = ['raw', 'active_area_cropping', 'linearization', 'normal', 'white_balance',
                  'demosaic']

    img = run_pipeline(img, params=params, metadata=meta_data_org, stages=stages)
    img = img ** (1 / 2.2)

    patches = get_all_patches(img, patch_size, stride)
    return patches


def datagenerator_sRGB(data_dir='dummy_dataset/train/clean', batch_size=128, patch_size=48, stride=48, verbose=False,
                       data_mode='train'):
    file_list = sorted(glob.glob(data_dir + '/*.png'))  # get name list of all .png files
    # initialize
    all_patches = []
    # generate patches
    for i in range(len(file_list)):
        patches = gen_patches_sRGB(file_list[i], patch_size=patch_size, stride=stride)
        all_patches.append(patches)
        if verbose:
            print(str(i + 1) + '/' + str(len(file_list)) + ' is done')
    all_patches = np.concatenate(all_patches)
    discard_n = len(all_patches) - len(all_patches) // batch_size * batch_size
    all_patches = np.delete(all_patches, range(discard_n), axis=0)

    print(f'^_^-sRGB {data_mode} data finished-^_^')
    return all_patches


def datagenerator_raw(data_dir='dummy_dataset/train/clean_raw', meta_dir='dummy_dataset/train/metadata_raw',
                      wb_illum='avg', batch_size=128, patch_size=48, stride=48, is_PS_sRGB=False, verbose=True,
                      data_mode='train'):

    file_list = sorted(glob.glob(data_dir + '/*.png'))  # get name list of all .png files

    meta_list = sorted(glob.glob(meta_dir + '/*.p'))  # get name list of all .p files

    # initialize
    all_patches = []
    # generate patches
    for i in range(len(file_list)):
        patches = gen_patches_raw(file_list[i], meta_list[i], wb_illum=wb_illum, patch_size=patch_size, stride=stride,
                                  is_PS_sRGB=is_PS_sRGB)
        all_patches.append(patches)
        if verbose:
            print(str(i + 1) + '/' + str(len(file_list)) + ' is done')
    all_patches = np.concatenate(all_patches)
    discard_n = len(all_patches) - len(all_patches) // batch_size * batch_size
    all_patches = np.delete(all_patches, range(discard_n), axis=0)

    print(f'^_^-raw {data_mode} data finished-^_^')
    return all_patches


if __name__ == '__main__':
    batch_size, patch_size, stride = 128, 48, 48
    data_sRGB = datagenerator_sRGB(data_dir='real_night/train/clean', batch_size=batch_size, patch_size=patch_size,
                                   stride=stride)
    data_raw = datagenerator_raw(data_dir='real_night/train/clean_raw', meta_dir='dummy_dataset/train/metadata_raw',
                                 wb_illum='avg', batch_size=batch_size, patch_size=patch_size, stride=stride)
    print('Done!')

data_preparation/data_generator_neural_isp.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

# no need to run this code separately

import os
import torch
from torchvision import transforms
import data_preparation.data_generator_denoise as dg
from noise_profiler.image_synthesizer import load_noise_model

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def numpy2tensor(ims, preload_on_cuda):
    ims = torch.from_numpy(ims)
    ims = ims.permute(0, 3, 1, 2)
    if preload_on_cuda:
        ims = ims.to(device)
    return ims


class DatasetRAW(object):
    def __init__(self, root, batch_size, patch_size, stride, preload_on_cuda, iso=3200, data_type='graphics_raw'):
        print('Inside raw data generator')

        noise_model_path = './noise_profiler/h-gauss-s20-v1'
        noise_model, iso2b1_interp_splines, iso2b2_interp_splines = load_noise_model(path=noise_model_path)
        noise_model_obj = {
            'noise_model': noise_model,
            'iso2b1_interp_splines': iso2b1_interp_splines,
            'iso2b2_interp_splines': iso2b2_interp_splines,
        }

        input_ims, target_ims = dg.datagenerator_raw(data_dir=os.path.join(root),
                                                     batch_size=batch_size,
                                                     patch_size=patch_size, stride=stride,
                                                     data_type=data_type,
                                                     noise_model=noise_model_obj, iso=iso)
        print('Number of patches ' + str(input_ims.shape[0]))
        print('Outside raw data generator \n')

        self.input_ims = numpy2tensor(input_ims, preload_on_cuda)
        self.target_ims = numpy2tensor(target_ims, preload_on_cuda)

        assert patch_size % 2 == 0  # crop stride must be even to preserve the bayer pattern

    def __getitem__(self, idx):
        # load images
        img = self.input_ims[idx]  # c, h, w
        target = self.target_ims[idx]
        return img, target

    def __len__(self):
        return len(self.input_ims)


if __name__ == '__main__':
    from utils.torch_utils import torchvision_visualize_raw
    data_dir = 'graphics_dataset'
    batch_size, patch_size, stride = 16, 64, 64
    preload_on_cuda = False

    image_datasets = {
        x: DatasetRAW(os.path.join(data_dir, x), batch_size, patch_size, stride, preload_on_cuda)
        for x in ['val']}

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
                                                  shuffle=True, num_workers=0)
                   for x in ['val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['val']}

    # Get a batch of training data
    inputs, targets = next(iter(dataloaders['val']))  # b, c, h, w;

    img_grid = torchvision_visualize_raw([inputs, targets])
    img_grid = transforms.ToPILImage()(img_grid.cpu())
    img_grid.save("debug_inputs_targets_grid.png")
    print('Done!')

data_preparation/dataset_denoise.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import torch
import torchvision
from torchvision import transforms
import data_preparation.data_generator_illum_est as dg

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class DatasetIllumEst(object):
    def __init__(self, dataset_dir, illum_file, split_file, split, batch_size, patch_size, stride, on_cuda):

        print('Inside illum est data generator')
        in_patches, gt_illums = dg.datagenerator_illum_est(dataset_dir=dataset_dir,
                                                           illum_file=illum_file,
                                                           split_file=split_file,
                                                           split=split,
                                                           batch_size=batch_size,
                                                           patch_size=patch_size,
                                                           stride=stride)

        in_patches = torch.from_numpy(in_patches)
        in_patches = in_patches.permute(0, 3, 1, 2)  # b, c, h, w
        gt_illums = torch.from_numpy(gt_illums)  # b, 3
        if on_cuda:
            in_patches = in_patches.to(device)
            gt_illums = gt_illums.to(device)

        print('Number of patches ' + str(in_patches.shape[0]))
        print('Number of illums ' + str(gt_illums.shape[0]))
        print('Outside illum est data generator \n')
        self.in_patches = in_patches
        self.gt_illums = gt_illums

    def __getitem__(self, idx):
        # load images
        img = self.in_patches[idx]
        gt_illum = self.gt_illums[idx]

        return img, gt_illum

    def __len__(self):
        return len(self.in_patches)


if __name__ == '__main__':
    image_datasets = {
        x: DatasetIllumEst(dataset_dir='illum_est_expts/data/ours',
                           illum_file='illum_est_expts/data/ours/gt_illum.p',
                           split_file='illum_est_expts/data/SamsungNX2000_train_valid_test_split_idx.p',
                           split=x,
                           batch_size=32,
                           patch_size=48,
                           stride=48,
                           on_cuda=True)
        for x in ['val']}

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=32,
                                                  shuffle=True, num_workers=0)
                   for x in ['val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['val']}

    # Get a batch of training data
    inputs, targets = next(iter(dataloaders['val']))

    inputs = inputs ** (1 / 2.2)
    target_ims = targets[..., None, None] * torch.ones((32, 3, 48, 48))  # 32, 3, 1, 1 * 32, 3, 48, 48
    # Make a grid from batch
    targets_grid = torchvision.utils.make_grid(target_ims)
    inputs_grid = torchvision.utils.make_grid(inputs)

    targets_grid = transforms.ToPILImage()(targets_grid.cpu())
    inputs_grid = transforms.ToPILImage()(inputs_grid.cpu())

    targets_grid.save("debug_targets_grid.png")
    inputs_grid.save("debug_inputs_grid.png")

    print('Done!')

data_preparation/dataset_illum_est.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


import os
import torch
import torchvision
from torchvision import transforms
import data_preparation.data_generator_neural_isp as dg

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class DatasetRAW(object):
    def __init__(self, root, batch_size, patch_size, stride, wb_illum, on_cuda, which_input, is_PS_sRGB,
                 data_mode='train'):

        print('Inside raw data generator')
        rawimgs = dg.datagenerator_raw(data_dir=os.path.join(root, which_input),
                                       meta_dir=os.path.join(root, 'metadata_raw'),
                                       wb_illum=wb_illum, batch_size=batch_size,
                                       patch_size=patch_size, stride=stride,
                                       is_PS_sRGB=is_PS_sRGB, data_mode=data_mode)

        rawimgs = torch.from_numpy(rawimgs)
        rawimgs = rawimgs.permute(0, 3, 1, 2)
        if on_cuda:
            rawimgs = rawimgs.to(device)

        print('Number of patches ' + str(rawimgs.shape[0]))
        print('Outside raw data generator \n')
        self.rawimgs = rawimgs

        print('Inside sRGB data generator')
        srgbimgs = dg.datagenerator_sRGB(data_dir=os.path.join(root, 'clean'), batch_size=batch_size,
                                         patch_size=patch_size, stride=stride, data_mode=data_mode)
        print('Number of patches ' + str(srgbimgs.shape[0]))
        print('Outside sRGB data generator \n')
        srgbimgs = torch.from_numpy(srgbimgs)
        srgbimgs = srgbimgs.permute(0, 3, 1, 2)
        srgbimgs = srgbimgs / 255.0
        if on_cuda:
            srgbimgs = srgbimgs.to(device)
        self.srgbimgs = srgbimgs

    def __getitem__(self, idx):
        # load images
        img = self.rawimgs[idx]  # c, h, w
        target = self.srgbimgs[idx]
        return img, target

    def __len__(self):
        return len(self.rawimgs)


if __name__ == '__main__':
    data_dir = 'graphics_dataset'
    batch_size, patch_size, stride = 16, 64, 64
    wb_illum = 'asn'
    on_cuda = False
    which_input = 'clean_raw'

    image_datasets = {
        x: DatasetRAW(os.path.join(data_dir, x), batch_size, patch_size, stride, wb_illum, on_cuda, which_input,
                      is_PS_sRGB=True, data_mode=x)
        for x in ['val']}

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
                                                  shuffle=True, num_workers=0)
                   for x in ['val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['val']}

    # Get a batch of training data
    inputs, targets = next(iter(dataloaders['val']))

    # Make a grid from batch
    targets_grid = torchvision.utils.make_grid(targets)
    inputs_grid = torchvision.utils.make_grid(inputs)

    targets_grid = transforms.ToPILImage()(targets_grid.cpu())
    inputs_grid = transforms.ToPILImage()(inputs_grid.cpu())

    targets_grid.save("debug_targets_grid.png")
    inputs_grid.save("debug_inputs_grid.png")

    print('Done!')

data_preparation/dataset_neural_isp.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Input folder: contains DNGs
Output folders:
1) train : clean_raw, metadata_raw
2) val : clean_raw, metadata_raw
"""


import argparse
import cv2
import pickle
from utils.general_utils import check_dir
import os
from glob import glob
from shutil import rmtree
from pipeline.pipeline_utils import get_metadata, get_visible_raw_image


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--dng_dir', type=str, help='dng dir')
    parser.add_argument('--val_names', default='Living,Playground,Rustic', type=str, help='name of test dataset')
    parser.add_argument('--save_dir', default='graphics_dataset', type=str, help='save dir')
    args = parser.parse_args()

    if os.path.isdir(args.save_dir): rmtree(args.save_dir)
      
    print(args.dng_dir)

    # create directories
    check_dir(args.save_dir)
    for fol in ['train', 'val']:  check_dir(os.path.join(args.save_dir, fol))
    for fol in ['train', 'val']:
        subfols = ['clean_raw', 'metadata_raw']
        for subfol in subfols: check_dir(os.path.join(args.save_dir, fol, subfol))

    allfiles = [os.path.basename(x) for x in sorted(glob(os.path.join(args.dng_dir, '*.dng')))]

    valnames = [item for item in args.val_names.split(',')]

    index = {'train': [i for i in range(len(allfiles)) if allfiles[i].split('_')[0] not in valnames],
             'val': [i for i in range(len(allfiles)) if allfiles[i].split('_')[0] in valnames]}

    input_dir_dng = args.dng_dir

    for fol in ['train', 'val']:
        for ind in index[fol]:
            for subfol in ['clean_raw']:
                cleanrawimg = get_visible_raw_image(os.path.join(input_dir_dng, allfiles[ind]))
                destination = os.path.join(args.save_dir, fol, subfol, allfiles[ind][:-4] + '.png')
                cv2.imwrite(destination, cleanrawimg)

            for subfol in ['metadata_raw']:
                metadata = get_metadata(os.path.join(input_dir_dng, allfiles[ind]))
                pickle.dump(metadata, open(os.path.join(args.save_dir, fol, subfol, allfiles[ind][:-4] + '.p'), "wb"))

data_preparation/initial_data_prep_denoise.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Input is two folders: one containing DNGs and other containing tiffs
Output folders: (to match current day-to-night code)
1) train : clean, clean_raw, metadata_raw, noisy_raw
2) val : clean, clean_raw, metadata_raw, noisy_raw
"""

import argparse
import cv2
import pickle
from utils.general_utils import check_dir
import os
from glob import glob
from shutil import copyfile, rmtree
from pipeline.pipeline_utils import get_metadata, get_visible_raw_image

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--dng_dir', type=str, help='dng dir')
    parser.add_argument('--tif_dir', type=str, help='tif dir')
    parser.add_argument('--val_names', default='Living,Playground,Rustic', type=str, help='name of test dataset')
    parser.add_argument('--save_dir', default='graphics_dataset', type=str, help='save dir')
    parser.add_argument('--noisy', action='store_true', help='noisy dataset')

    args = parser.parse_args()

    if os.path.isdir(args.save_dir): rmtree(args.save_dir)
      
    print(args.dng_dir)
    print(args.tif_dir)

    # create directories
    check_dir(args.save_dir)
    for fol in ['train', 'val']:  check_dir(os.path.join(args.save_dir, fol))
    for fol in ['train', 'val']:
        subfols = ['clean_raw', 'clean', 'metadata_raw']
        if args.noisy:
            subfols.append('noisy_raw')
        for subfol in subfols: check_dir(os.path.join(args.save_dir, fol, subfol))

    allfiles = [os.path.basename(x) for x in sorted(glob(os.path.join(args.dng_dir, '*.dng')))]

    valnames = [item for item in args.val_names.split(',')]

    index = {'train': [i for i in range(len(allfiles)) if allfiles[i].split('_')[0] not in valnames],
             'val': [i for i in range(len(allfiles)) if allfiles[i].split('_')[0] in valnames]}

    input_dir_dng = args.dng_dir
    input_dir_tif = args.tif_dir

    for fol in ['train', 'val']:
        for ind in index[fol]:
            for subfol in ['clean']:
                cleanimg = cv2.imread(os.path.join(input_dir_tif, allfiles[ind][:-4] + '.tif'), cv2.IMREAD_UNCHANGED)
                destination = os.path.join(args.save_dir, fol, subfol, allfiles[ind][:-4] + '.png')
                cv2.imwrite(destination, cleanimg)

            for subfol in ['clean_raw']:
                cleanrawimg = get_visible_raw_image(os.path.join(input_dir_dng, allfiles[ind]))
                destination = os.path.join(args.save_dir, fol, subfol, allfiles[ind][:-4] + '.png')
                cv2.imwrite(destination, cleanrawimg)
        
            if args.noisy:
                for subfol in ['noisy_raw']:
                    source = os.path.join(input_dir_dng, 'noisy_raw', allfiles[ind][:-4] + '.png')
                    destination = os.path.join(args.save_dir, fol, subfol, allfiles[ind][:-4] + '.png')
                    copyfile(source,destination)

            for subfol in ['metadata_raw']:
                metadata = get_metadata(os.path.join(input_dir_dng, allfiles[ind]))
                pickle.dump(metadata, open(os.path.join(args.save_dir, fol, subfol, allfiles[ind][:-4] + '.p'), "wb"))

data_preparation/initial_data_prep_neural_isp.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Convert PS tif to png
"""

import argparse
import cv2
from utils.general_utils import check_dir
import os
from glob import glob
from shutil import rmtree

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--tif_dir', default='neural_isp_expts/data/clean_srgb_tiffs/', type=str, help='tif dir')
    parser.add_argument('--save_dir', default='real_dataset/clean', type=str, help='save dir')

    args = parser.parse_args()

    # remove png files already inside
    if os.path.isdir(args.save_dir): rmtree(args.save_dir)

    # create directory again
    check_dir(args.save_dir)

    allfiles = sorted(glob(os.path.join(args.tif_dir, '*.tif')))

    for fil in allfiles:
        cleanimg = cv2.imread(os.path.join(fil), cv2.IMREAD_UNCHANGED)
        destination = os.path.join(args.save_dir, os.path.basename(fil)[:-4] + '.png')
        cv2.imwrite(destination, cleanimg)


data_preparation/initial_data_prep_tif2png.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from shutil import copyfile, rmtree
import os
from glob import glob
import argparse
import cv2
import pickle

from pipeline.pipeline_utils import get_metadata, get_visible_raw_image
from utils.general_utils import check_dir


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--base_path', default='real_dataset', type=str,
                        help='base address')
    parser.add_argument('--save_path', default='real_dataset_k_fold', type=str,
                        help='output address')
    parser.add_argument('--which_fold', default=0, type=int, help='which fold for testing [0,1,2]')
    parser.add_argument('--kfold_indices', default='assets/split_files/day2night_k_fold_indices.p', type=str,
                        help='load saved k-fold indices')
    parser.add_argument('--with_noise', default=0, type=int, help='for noisy case')
    parser.add_argument('--only_iso_3200', action='store_true',
                        help='True: only run on iso_3200.')
    parser.add_argument('--only_iso_1600', action='store_true',
                        help='True: only run on iso_1600.')
    parser.add_argument('--test_on_val', default=0, type=int, help='for testing on validation set')

    args = parser.parse_args()

    if args.only_iso_3200:
        assert not args.only_iso_1600
    elif args.only_iso_1600:
        assert not args.only_iso_3200
    # if they are both false, mix iso 1600 and 3200 equally
    print(args)

    return args


def split_func(args, folder_path, k_fold_indices):
    save_path = args.save_path

    train_index = k_fold_indices['train_index_all'][args.which_fold]
    val_index = k_fold_indices['val_index_all'][args.which_fold]
    test_index = k_fold_indices['test_index_all'][args.which_fold]

    if not args.only_iso_3200 and not args.only_iso_1600:
        if folder_path == 'iso_1600':
            index = {'train': train_index[0::2], 'val': val_index[0::2]}
        elif folder_path == 'iso_3200':
            index = {'train': train_index[1::2], 'val': val_index[1::2]}
        else:
            index = {'train': train_index, 'val': val_index}
    else:
        index = {'train': train_index, 'val': val_index}

    input_dir = args.base_path

    print(index['train'])
    print('...')
    print(index['val'])
    print('...')
    print(test_index)
    print('...')
    print(args.which_fold, len(index['train']), len(index['val']), len(test_index))

    # create directories if they don't exist
    check_dir(save_path)
    for fol in ['train', 'val', 'test']:  check_dir(os.path.join(save_path, fol))
    for fol in ['train', 'val']:
        for subfol in ['clean_raw', 'noisy_raw', 'clean', 'metadata_raw']: check_dir(os.path.join(save_path, fol, subfol))
    for fol in ['test']:
        for subfol in ['clean_raw', 'clean', 'dng']: check_dir(os.path.join(save_path, fol, subfol))
    check_dir(os.path.join(save_path, 'test', 'dng', folder_path))
    if args.test_on_val:
        check_dir(os.path.join(save_path, 'val', 'dng', folder_path))

    allfiles = [os.path.basename(x) for x in sorted(glob(os.path.join(input_dir, 'clean', '*.png')))]

    for fol in ['train', 'val']:
        for ind in index[fol]:
            for subfol in ['clean_raw', 'clean']:
                source = os.path.join(input_dir, subfol, allfiles[ind])
                destination = os.path.join(save_path, fol, subfol, allfiles[ind])
                copyfile(source, destination)

            for subfol in ['noisy_raw']:
                rawimg = get_visible_raw_image(os.path.join(input_dir, 'dng', folder_path, allfiles[ind][:-4] + '.dng'))
                destination = os.path.join(save_path, fol, subfol, allfiles[ind])
                cv2.imwrite(destination, rawimg)

            for subfol in ['metadata_raw']:
                metadata = get_metadata(os.path.join(input_dir, 'dng', folder_path, allfiles[ind][:-4] + '.dng'))
                pickle.dump(metadata, open(os.path.join(save_path, fol, subfol, allfiles[ind][:-4] + '.p'), "wb"))

            if fol == 'val' and args.test_on_val:
                print('saving val dngs')
                source = os.path.join(input_dir, 'dng', folder_path, allfiles[ind][:-4] + '.dng')
                destination = os.path.join(save_path, fol, 'dng', folder_path, allfiles[ind][:-4] + '.dng')
                copyfile(source, destination)

    for fol in ['test']:
        clean_raw_dir = os.listdir(os.path.join(save_path, fol, 'clean_raw'))
        clean_dirs_empty = len(clean_raw_dir) == 0
        for ind in test_index:
            if clean_dirs_empty:
                for subfol in ['clean_raw', 'clean']:
                    source = os.path.join(input_dir, subfol, allfiles[ind])
                    destination = os.path.join(save_path, fol, subfol, allfiles[ind])
                    copyfile(source, destination)

            for subfol in ['dng']:
                source = os.path.join(input_dir, subfol, folder_path, allfiles[ind][:-4] + '.dng')
                destination = os.path.join(save_path, fol, subfol, folder_path, allfiles[ind][:-4] + '.dng')
                copyfile(source, destination)


if __name__ == "__main__":
    args = parse_args()
    k_fold_indices = pickle.load(open(args.kfold_indices, "rb"))

    if os.path.isdir(args.save_path): rmtree(args.save_path)

    if args.with_noise:
        if args.only_iso_3200:
            split_func(args, 'iso_3200', k_fold_indices)
        elif args.only_iso_1600:
            split_func(args, 'iso_1600', k_fold_indices)
        else:
            split_func(args, 'iso_1600', k_fold_indices)
            split_func(args, 'iso_3200', k_fold_indices)
    else:
        split_func(args, 'iso_50', k_fold_indices)

    print('Done!')

data_preparation/k_fold_split_data.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import time
import shutil
from jobs.denoise_graphics2raw_iso3200 import prep_environment, movedata, get_parser, train, test

if __name__ == '__main__':
    prep_environment()
    parser = get_parser()
    args = parser.parse_args()

    iso = 1600
    exp_name = f'denoise_graphics2raw_iso{iso}'

    input_dir = 'graphics_dngs_graphics2raw'
    movedata(input_dir)

    num_runs = args.num_runs
    for i in range(num_runs):
        timestamp = int(time.time())
        train(exp_name, timestamp, iso)
        test(exp_name, timestamp, iso)
        print(f'-----------Experiment: {exp_name}_{timestamp}-----------')

        shutil.rmtree('results')

jobs/denoise_graphics2raw_iso1600.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
import time
import argparse
import shutil
from jobs.job_utils import prep_environment, get_day2night_data

DATA_ROOT_DIR = 'neural_isp_expts'  # uses the same RAW data as neural ISP
EXPT_ROOT_DIR = 'denoising_expts'


def movedata(input_dir):
    print('Preparing data...')
    get_day2night_data()
    os.system(f'python3 -m data_preparation.initial_data_prep_denoise --dng_dir {DATA_ROOT_DIR}/data/{input_dir};')


def train(exp_name, timestamp, iso):
    folder_name = f'{exp_name}_{timestamp}'
    args_str = f"--savefoldername {folder_name} \
            --exp_dir {EXPT_ROOT_DIR}/expts/ \
            --lr 0.0001 \
            --num_epochs 100 \
            --milestones 90 \
            --patch_size 128 \
            --batch_size 32 \
            --loss l1 \
            --restormer_dim 8 \
            --data_type graphics_raw \
            --iso {iso} \
            --model_save_freq 400"  # not saving intermediate models

    print('Start training...')
    print('args_str = {}'.format(args_str))
    os.system('python3 train_denoise.py {} '.format(args_str))


def test(exp_name, timestamp, iso):
    folder_name = f'{exp_name}_{timestamp}'
    print('Start testing...')
    args_str = f"--model_dir {folder_name} \
            --exp_dir {EXPT_ROOT_DIR}/expts/ \
            --set_name iso_{iso} \
            --restormer_dim 8"

    print('Start testing...')
    print('args_str = {}'.format(args_str))
    os.system('python3 test_denoise.py {} '.format(args_str))

    target_dir = f'{EXPT_ROOT_DIR}/expts/{folder_name}/results/'
    os.makedirs(target_dir, exist_ok=True)
    os.system(f'cp results/* {target_dir} -r -v;')


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_runs', '-n', default=1, type=int, help='Number of runs.')
    return parser


if __name__ == '__main__':
    prep_environment()
    parser = get_parser()
    args = parser.parse_args()
    
    iso = 3200
    exp_name = f'denoise_graphics2raw_iso{iso}'

    input_dir = 'graphics_dngs_graphics2raw'
    movedata(input_dir)

    num_runs = args.num_runs
    for i in range(num_runs):
        timestamp = int(time.time())
        train(exp_name, timestamp, iso)
        test(exp_name, timestamp, iso)
        print(f'-----------Experiment: {exp_name}_{timestamp}-----------')

        shutil.rmtree('results')

jobs/denoise_graphics2raw_iso3200.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import time
from jobs.denoise_real_iso3200 import get_parser, movedata, train, test
from jobs.job_utils import prep_environment
import shutil


if __name__ == '__main__':
    prep_environment()
    parser = get_parser()
    args = parser.parse_args()

    fold = args.fold
    iso = 1600
    exp_name = f'denoise_real_iso{iso}_f{fold}'

    movedata(fold)

    num_runs = args.num_runs
    for i in range(num_runs):
        timestamp = int(time.time())
        train(exp_name, timestamp, args.restormer_dim, args.epochs)
        test(exp_name, timestamp, args.restormer_dim)
        print(f'-----------Experiment: {exp_name}_{timestamp}-----------')

        shutil.rmtree('results')

jobs/denoise_real_iso1600.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
import time
import argparse
from jobs.job_utils import prep_environment, get_day2night_data
from jobs.denoise_graphics2raw_iso3200 import EXPT_ROOT_DIR
import shutil


def movedata(fold, iso):
    print('Preparing data...')
    get_day2night_data()
    os.system(f'python3 -m data_preparation.k_fold_split_data --which_fold {fold} --with_noise 1 --only_iso_{iso};')


def train(exp_name, timestamp, iso):
    folder_name = f'{exp_name}_{timestamp}'
    args_str = f"--savefoldername {folder_name} \
            --exp_dir {EXPT_ROOT_DIR}/expts/ \
            --data_dir real_dataset_k_fold \
            --lr 0.0001 \
            --num_epochs 100 \
            --milestones 90 \
            --patch_size 128 \
            --batch_size 32 \
            --loss l1 \
            --restormer_dim 8 \
            --data_type real \
            --iso {iso} \
            --model_save_freq 400"  # not saving intermediate models

    print('Start training...')
    print('args_str = {}'.format(args_str))
    os.system('python3 train_denoise.py {} '.format(args_str))


def test(exp_name, timestamp, iso):
    folder_name = f'{exp_name}_{timestamp}'
    print('Start testing...')
    args_str = f"--model_dir {folder_name} \
            --exp_dir {EXPT_ROOT_DIR}/expts/ \
            --set_dir real_dataset_k_fold/test \
            --set_name iso_{iso} \
            --restormer_dim 8"

    print('Start testing...')
    print('args_str = {}'.format(args_str))
    os.system('python3 test_denoise.py {} '.format(args_str))

    target_dir = f'{EXPT_ROOT_DIR}/expts/{folder_name}/results/'
    os.makedirs(target_dir, exist_ok=True)
    os.system(f'cp results/* {target_dir} -r -v;')


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--fold', '-f', default=0, type=int, help='Which fold of the real data to train and test.')
    parser.add_argument('--num_runs', '-n', default=1, type=int, help='Number of runs.')
    return parser


if __name__ == '__main__':
    prep_environment()
    parser = get_parser()
    args = parser.parse_args()
    fold = args.fold

    iso = 3200
    exp_name = f'denoise_real_iso{iso}_f{fold}'

    movedata(fold, iso)

    num_runs = args.num_runs
    for i in range(num_runs):
        timestamp = int(time.time())
        train(exp_name, timestamp, iso)
        test(exp_name, timestamp, iso)
        print(f'-----------Experiment: {exp_name}_{timestamp}-----------')

        shutil.rmtree('results')

jobs/denoise_real_iso3200.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import time
import shutil
from jobs.denoise_graphics2raw_iso3200 import train, test, get_parser, prep_environment, movedata


if __name__ == '__main__':
    prep_environment()
    parser = get_parser()
    args = parser.parse_args()

    iso = 1600
    exp_name = f'denoise_upi_iso{iso}'
    input_dir = 'graphics_dngs_upi'
    movedata(input_dir)
    
    num_runs = args.num_runs
    for i in range(num_runs):
        timestamp = int(time.time())
        train(exp_name, timestamp, iso)
        test(exp_name, timestamp, iso)
        print(f'-----------Experiment: {exp_name}_{timestamp}-----------')

        shutil.rmtree('results')

jobs/denoise_upi_iso1600.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import time
import shutil
from jobs.denoise_graphics2raw_iso3200 import train, test, get_parser, prep_environment, movedata


if __name__ == '__main__':
    prep_environment()
    parser = get_parser()
    args = parser.parse_args()

    iso = 3200
    exp_name = f'denoise_upi_iso{iso}'
    input_dir = 'graphics_dngs_upi'
    movedata(input_dir)

    num_runs = args.num_runs
    for i in range(num_runs):
        timestamp = int(time.time())
        train(exp_name, timestamp, iso)
        test(exp_name, timestamp, iso)
        print(f'-----------Experiment: {exp_name}_{timestamp}-----------')

        shutil.rmtree('results')

jobs/denoise_upi_iso3200.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


import os
from jobs.job_utils import prep_environment
"""
Requirements:
illum_est_expts/data/<camera>/real
illum_est_expts/data/<camera>/real/<camera>_gt.mat
illum_est_expts/nus_metadata/
illum_est_expts/nus_metadata/nus_outdoor_gt_illum_mats
illum_est_expts/synthia/SYNTHIA_RAND_CVPR16/RGB
"""


def main():
    graphics_path = 'illum_est_expts/synthia/SYNTHIA_RAND_CVPR16/RGB'
    
    cameras = [
        'Canon1DsMkIII', 'Canon600D', 'FujifilmXM1', 'NikonD5200', 'OlympusEPL6', 
        'PanasonicGX1',
        'SamsungNX2000', 
        'SonyA57',
        'NikonD40'
    ]

    method = 'ours'

    for camera in cameras:
        camera_dir = os.path.join('illum_est_expts/data', camera)
        save_path = os.path.join(camera_dir, method)

        os.makedirs(save_path, exist_ok=True)

        mat_file = os.path.join('illum_est_expts/nus_metadata', 'nus_outdoor_gt_illum_mats',
                                f'{camera}_gt.mat')
        split_file = f'assets/split_files/illum_est/{camera}_train_valid_test_split_fns.p'

        assert os.path.exists(mat_file)
        assert os.path.exists(split_file)
        
        args_str = f"--save_path {save_path} \
                --target_camera {camera} \
                --mat_file {mat_file} \
                --use_train_val_illums_only \
                --split_file {split_file} \
                --graphics_path {graphics_path} \
                --rgb_gain \
                --rgb_gain_mean 0.8"
                # do safe invert
                # max illums None

        os.system('python3 -m data_generation.invert_synthia_graphics2raw {} '.format(args_str))


if __name__ == '__main__':
    prep_environment()
    main()






jobs/generate_dataset_illum_est_graphics2raw.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


import os
from jobs.job_utils import prep_environment


def main():
    cameras = ['SamsungNX2000']

    graphics_path = 'illum_est_expts/synthia/SYNTHIA_RAND_CVPR16/RGB'

    method = 'upi'

    for camera in cameras:
        save_path = os.path.join('illum_est_expts/data', camera, method)
        os.makedirs(save_path, exist_ok=True)

        args_str = f"--save_path {save_path} \
                     --graphics_path {graphics_path}"

        os.system('python3 -m data_generation.invert_synthia_upi {} '.format(args_str))


if __name__ == '__main__':
    prep_environment()
    main()






jobs/generate_dataset_illum_est_upi.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


import os
from jobs.job_utils import prep_environment


def main():
    save_path = 'neural_isp_expts/data/graphics_dngs_graphics2raw'
    os.makedirs(save_path, exist_ok=True)

    exr_folder_path = 'path/to/graphics2RAW-dataset/'

    args_str = f"--save_path {save_path} \
            --exr_folder_path {exr_folder_path} \
            --rgb_gain \
            --rgb_gain_mean 0.8"
    # do apply safe invert

    os.system('python3 -m data_generation.package_exr_to_dng_graphics2raw {} '.format(args_str))


if __name__ == '__main__':
    prep_environment()
    main()






jobs/generate_dataset_isp_denoise_graphics2raw.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
from jobs.job_utils import prep_environment


def main():
    save_path = 'neural_isp_expts/data/graphics_dngs_upi'
    os.makedirs(save_path, exist_ok=True)

    exr_folder_path = 'path/to/graphics2RAW-dataset/'

    args_str = f"--save_path {save_path} \
                --exr_folder_path {exr_folder_path}"

    os.system('python3 -m data_generation.package_exr_to_dng_upi {}'.format(args_str))


if __name__ == '__main__':
    prep_environment()
    main()






jobs/generate_dataset_isp_denoise_upi.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
import time
from jobs.job_utils import prep_environment
import argparse


ROOT_DIR = 'illum_est_expts'

def train(camera, method, exp_name, dataset_dir):
    illum_fname = f'{camera}_gt.mat' if method == 'real' else 'gt_illum.p'

    args_str = f"--dataset-dir {os.path.join(dataset_dir, camera, method)} \
            --savefoldername {exp_name} \
            --illum_file {os.path.join(dataset_dir, camera, method, illum_fname)} \
            --split_file assets/split_files/illum_est/{camera}_train_valid_test_split_idx.p \
            --exp_dir {ROOT_DIR}/expts \
            --on-cuda \
            --model_save_freq 3000 \
            --num-epochs 2000"
            # not saving intermediate models

    print('Start training...')
    print('args_str = {}'.format(args_str))
    os.system('python3 train_illum_est.py {} '.format(args_str))


def test(camera, exp_name, dataset_dir):
    print('Start testing...')
    args_str = f"--dataset-dir {os.path.join(dataset_dir, camera, 'real')} \
            --exp_name {exp_name} \
            --illum_file {os.path.join(dataset_dir, camera, 'real', f'{camera}_gt.mat')} \
            --split_file assets/split_files/illum_est/{camera}_train_valid_test_split_idx.p \
            --exp_dir {ROOT_DIR}/expts"

    print('Start testing...')
    print('args_str = {}'.format(args_str))
    os.system('python3 test_illum_est.py {} '.format(args_str))


if __name__ == '__main__':
    prep_environment()
    parser = argparse.ArgumentParser()
    parser.add_argument('--cameras', '-c', default='SamsungNX2000,SonyA57', type=str, help='Which cameras to run on. Comma-separated string.')
    parser.add_argument('--methods', '-m', default='ours,real', type=str, help='Which methods to run on. Comma-separated string.')
    parser.add_argument('--num_runs', '-n', default=5, type=int, help='Number of repeated runs.')
    args = parser.parse_args()

    exp_id = 'illum_est'
    cameras = [] if args.cameras == '' else args.cameras.split(',')
    methods = [] if args.methods == '' else args.methods.split(',')

    for method in methods:
        for camera in cameras:
            for i in range(args.num_runs):
                timestamp = int(time.time())
                exp_name = f'{camera}_{exp_id}_{method}_{timestamp}'

                dataset_dir = f'{ROOT_DIR}/data'
                train(camera, method, exp_name, dataset_dir)
                test(camera, exp_name, dataset_dir)

jobs/illum_est.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import time
from jobs.job_utils import prep_environment
from jobs.illum_est import train, test, ROOT_DIR
import argparse


if __name__ == '__main__':
    prep_environment()
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_runs', '-n', default=5, type=int, help='Number of repeated runs.')
    args = parser.parse_args()

    method = 'upi'
    cameras = [
        'Canon1DsMkIII',
        'Canon600D',
        'FujifilmXM1',
        'NikonD5200',
        'OlympusEPL6',
        'PanasonicGX1',
        'SamsungNX2000',
        'SonyA57',
        'NikonD40'
    ]
    exp_id = 'illum_est'
    dataset_dir = f'{ROOT_DIR}/data'

    for i in range(args.num_runs):
        timestamp = int(time.time())
        exp_name = f'{exp_id}_{method}_{timestamp}'

        train('SamsungNX2000', method, exp_name, dataset_dir)

        for camera in cameras:
            test(camera, exp_name, dataset_dir)
            print(f'UPI testing done: {camera}')

jobs/illum_est_upi.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os


def prep_environment():
    os.system('pip install -r requirements.txt')


def get_day2night_data():
    # copy subdirectories in night_real into ./real_dataset 
    assert os.path.exists('real_dataset')
    assert os.path.exists('real_dataset/clean')
    assert os.path.exists('real_dataset/clean_raw')
    assert os.path.exists('real_dataset/dng')

jobs/job_utils.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
import time
import argparse
from jobs.job_utils import prep_environment, get_day2night_data


ROOT_DIR = 'neural_isp_expts'

def movedata(input_dir, target_dir):
    print('Preparing data...')
    get_day2night_data()
    
    tif_dir = f'{ROOT_DIR}/data/clean_srgb_tiffs'
    os.system(f'python3 -m data_preparation.initial_data_prep_tif2png --tif_dir {tif_dir};')
    os.system(f'python3 -m data_preparation.initial_data_prep_neural_isp --dng_dir {ROOT_DIR}/data/{input_dir} --tif_dir {ROOT_DIR}/data/{target_dir};')


def train(input_type, exp_name, timestamp):
    args_str = f"--which-input {input_type} \
            --savefoldername {exp_name}_{input_type}_{timestamp} \
            --exp_dir {ROOT_DIR}/expts/ \
            --on-cuda \
            --model_save_freq 600 \
            --num-epochs 250"

    print('Start training...')
    print('args_str = {}'.format(args_str))
    os.system('python3 train_neural_isp.py {} '.format(args_str))


def test(input_type, exp_name, timestamp):
    print('Start testing...')
    folder_name = f'{exp_name}_{input_type}_{timestamp}'
    args_str = f"--exp_dir {ROOT_DIR}/expts/ \
                 --model_dir {folder_name}"

    print('Start testing...')
    print('args_str = {}'.format(args_str))
    os.system('python3 test_neural_isp.py {} '.format(args_str))

    target_dir = f'{ROOT_DIR}/expts/{folder_name}/results/'
    os.makedirs(target_dir, exist_ok=True)
    os.system(f'cp results/* {target_dir} -r -v;')


if __name__ == '__main__':
    prep_environment()
    parser = argparse.ArgumentParser()
    parser.add_argument('--timestamp', '-t', default=None, type=str, help='Time stamp to distinguish between runs.')
    parser.add_argument('--eval', '-e', action='store_true', help='Evaluation only, no training.')
    args = parser.parse_args()
    timestamp = args.timestamp if args.timestamp else int(time.time())

    input_type = 'clean_raw'
    exp_name = 'neural_isp_graphics2raw'

    input_dir = 'graphics_dngs_graphics2raw'
    target_dir = 'graphics_srgb_graphics2raw'
    movedata(input_dir, target_dir)
    if not args.eval:
        train(input_type, exp_name, timestamp)
    test(input_type, exp_name, timestamp)





jobs/neural_isp_graphics2raw.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
import time
import argparse
import shutil
from jobs.job_utils import prep_environment, get_day2night_data
from jobs.neural_isp_graphics2raw import ROOT_DIR


def movedata(fold):
    print('Preparing data...')
    get_day2night_data()
    
    tif_dir = f'{ROOT_DIR}/data/clean_srgb_tiffs'
    os.system(f'python3 -m data_preparation.initial_data_prep_tif2png --tif_dir {tif_dir};')
    os.system(f'python3 -m data_preparation.k_fold_split_data --which_fold {fold} --with_noise 0;')


def train(input_type, folder_name):
    args_str = f"--which-input {input_type} \
            --savefoldername {folder_name} \
            --exp_dir {ROOT_DIR}/expts/ \
            --on-cuda \
            --data-dir real_dataset_k_fold \
            --model_save_freq 600 \
            --num-epochs 250"

    print('Start training...')
    print('args_str = {}'.format(args_str))
    os.system('python3 train_neural_isp.py {} '.format(args_str))


def test(folder_name):
    args_str = f"--exp_dir {ROOT_DIR}/expts/ \
            --model_dir {folder_name} \
            --set_dir real_dataset_k_fold/test"

    print('Start testing...')
    print('args_str = {}'.format(args_str))
    os.system('python3 test_neural_isp.py {} '.format(args_str))

    target_dir = f'{ROOT_DIR}/expts/{folder_name}/results/'
    os.makedirs(target_dir, exist_ok=True)
    os.system(f'cp results/* {target_dir} -r -v;')


if __name__ == '__main__':
    prep_environment()

    parser = argparse.ArgumentParser()
    parser.add_argument('--fold', '-f', default=0, type=int, help='Which fold of the real data to train and test.')
    args = parser.parse_args()
    fold = args.fold

    input_type = 'clean_raw'
    exp_name = 'neural_isp_real'
    num_runs = 2
    for i in range(num_runs):
        timestamp = int(time.time())
        folder_name = f'{exp_name}_{input_type}_{fold}_{timestamp}'
        movedata(fold)
        train(input_type, folder_name)
        test(folder_name)
        print(f'-----------Experiment: {folder_name}-----------')

        shutil.rmtree('real_dataset')
        shutil.rmtree('real_dataset_k_fold')
        shutil.rmtree('results')

jobs/neural_isp_real.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import time
from jobs.job_utils import prep_environment
from jobs.neural_isp_graphics2raw import movedata, train, test
import argparse


if __name__ == '__main__':
    prep_environment()
    parser = argparse.ArgumentParser()
    parser.add_argument('--timestamp', '-t', default=None, type=str, help='Time stamp to distinguish between runs.')
    parser.add_argument('--eval', '-e', action='store_true', help='Evaluation only, no training.')
    args = parser.parse_args()
    timestamp = args.timestamp if args.timestamp else int(time.time())

    input_type = 'clean_raw'
    exp_name = 'neural_isp_upi'

    input_dir = 'graphics_dngs_upi'
    target_dir = 'graphics_srgb_upi'
    movedata(input_dir, target_dir)
    if not args.eval:
        train(input_type, exp_name, timestamp)
    test(input_type, exp_name, timestamp)





jobs/neural_isp_upi.py
SamsungLabs/graphics2raw
"""
Implementation of
Convolutional Mean: A Simple Convolutional Neural Network for Illuminant Estimation

@INPROCEEDINGS{convmean,
  title={Convolutional Mean: {A} Simple Convolutional Neural Network for Illuminant Estimation},
  author={Han Gong},
  booktitle={BMVC},
  year={2019},
}
"""

import torch.nn as nn


class IllumEstNet(nn.Module):

    def __init__(self, in_channels=3, out_channels=1, num_filters=7):
        super(IllumEstNet, self).__init__()

        # conv1
        self.conv1_1 = nn.Conv2d(in_channels, num_filters, (3, 3), padding=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.relu1_1 = nn.ReLU(inplace=True)

        # conv2
        self.conv2_1 = nn.Conv2d(num_filters, num_filters * 2, (3, 3), padding=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.relu2_1 = nn.ReLU(inplace=True)

        # conv3
        self.conv3_1 = nn.Conv2d(num_filters * 2, out_channels, (1, 1), padding=0)  # 1x1 conv
        self.relu3_1 = nn.ReLU(inplace=True)

        # per channel global average pooling
        self.aap = nn.AdaptiveAvgPool2d((1, 1))  # b, c, 1, 1

    def forward(self, x):
        # x: b, 3, p, p
        x = self.conv1_1(x)  # b, 7, p, p
        x = self.pool1(x)  # b, 7, p/2, p/2
        x = self.relu1_1(x)

        x = self.conv2_1(x)  # b, 14, p/2, p/2
        x = self.pool2(x)  # b, 14, p/4, p/4  16 16
        x = self.relu2_1(x)

        x = self.conv3_1(x)  # b, 3, p/4, p/4
        x = self.relu3_1(x)
        x = self.aap(x)  # b, 3, 1, 1
        x = x.squeeze(dim=-1).squeeze(dim=-1)  # b, 3
        return x

model_archs/cnn.py
SamsungLabs/graphics2raw
"""
Copy over the code from https://github.com/swz30/Restormer/blob/main/basicsr/models/archs/restormer_arch.py as is
"""

import torch.nn as nn


class Restormer(nn.Module):
    def __init__(self):
        super(Restormer, self).__init__()
        pass

model_archs/restormer.py
SamsungLabs/graphics2raw
# From: https://github.com/mateuszbuda/brain-segmentation-pytorch/blob/master/unet.py
# MIT License: https://github.com/mateuszbuda/brain-segmentation-pytorch/blob/master/LICENSE

from collections import OrderedDict

import torch
import torch.nn as nn


class UNet(nn.Module):

    def __init__(self, in_channels=3, out_channels=1, init_features=32):
        super(UNet, self).__init__()

        features = init_features
        self.encoder1 = UNet._block(in_channels, features, name="enc1")
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder2 = UNet._block(features, features * 2, name="enc2")
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder3 = UNet._block(features * 2, features * 4, name="enc3")
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder4 = UNet._block(features * 4, features * 8, name="enc4")
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.bottleneck = UNet._block(features * 8, features * 16, name="bottleneck")

        self.upconv4 = nn.ConvTranspose2d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.decoder4 = UNet._block((features * 8) * 2, features * 8, name="dec4")
        self.upconv3 = nn.ConvTranspose2d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.decoder3 = UNet._block((features * 4) * 2, features * 4, name="dec3")
        self.upconv2 = nn.ConvTranspose2d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.decoder2 = UNet._block((features * 2) * 2, features * 2, name="dec2")
        self.upconv1 = nn.ConvTranspose2d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = UNet._block(features * 2, features, name="dec1")

        self.conv = nn.Conv2d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )

    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return torch.sigmoid(self.conv(dec1))

    @staticmethod
    def _block(in_channels, features, name):
        return nn.Sequential(
            OrderedDict(
                [
                    (
                        name + "conv1",
                        nn.Conv2d(
                            in_channels=in_channels,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm1", nn.BatchNorm2d(num_features=features)),
                    (name + "relu1", nn.ReLU(inplace=True)),
                    (
                        name + "conv2",
                        nn.Conv2d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm2", nn.BatchNorm2d(num_features=features)),
                    (name + "relu2", nn.ReLU(inplace=True)),
                ]
            )
        )
model_archs/unet.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""
import os
import numpy as np

from noise_profiler.img_utils import brightness_transfer
from noise_profiler.spatial_correlation.utils import sample_cov


def synthesize_noisy_image(src_image, metadata, model, src_iso, dst_iso, min_val, max_val, dst_image=None, debug=False,
                           per_channel=False, fix_exp=None, match_bright_gt=True, spat_var=False, cross_corr=False,
                           cov_mat_fn=None, iso2b1_interp_splines=None, iso2b2_interp_splines=None):
    """
    Synthesize a high-ISO image with `dst_iso` from a lowISO one `src_image` with `src_iso`.
    Using a heteroscedastic Gaussian noise model `model`.
    :param src_image: Clean (ISO-50) image.
    :param metadata: Image metadata.
    :param model: Noise model.
    :param src_iso: ISO of `src_image`.
    :param dst_iso: ISO of noisy image to be synthesized.
    :param min_val: Minimum image intensity.
    :param max_val: Maximum image intensity.
    :param dst_image: If not None, match the brightness with provided destination image `dst_image`.
    :param debug: Whether to perform some debugging steps.
    :param per_channel: Whether to apply model per channel.
    :param fix_exp: Whether to fix image exposure before and after noise synthesis.
    :param match_bright_gt: Optionally, match the brightness with provided source image `src_image`.
    :param spat_var: Simulate spatial variations in noise.
    :param cross_corr: Simulate spatial cross correlations of noise.
    :param cov_mat_fn: Filename of covariance matrix used to simulate spatial correlations.
    :param iso2b1_interp_splines: Interpolation/extrapolation splines for shot noise (beta1)
    :param iso2b2_interp_splines: Interpolation/extrapolation splines for read noise (beta2)
    :return: Synthesized noisy image, and optionally, a another copy of the image with brightness matching `dst_image`.
    """
    # make a copy
    image = src_image.copy().astype(np.float32)

    # brightness transfer using a target image
    if dst_image is not None:
        image_adjusted = brightness_transfer(image, dst_image)
    else:
        image_adjusted = None

    # apply lens shading correction
    # image_lsc = apply_lens_gain(image, metadata['lens_gain_resized'], metadata['cfa_pattern'],
    #                             operation='multiply')
    # image_lsc_clip = np.clip(image_lsc, min_val, max_val)

    # TODO: account for noise already in image
    # src_params = model[src_iso]

    if fix_exp is not None:
        # update image exposure to match target ISO before synthesizing noise, then re-scale it back later
        image_fix_exp = np.round(image.astype(np.float32) * fix_exp).astype(np.uint32)
    else:
        image_fix_exp = None

    # get noise params (shot, read), per channel
    dst_params = get_noise_params(model, iso2b1_interp_splines, iso2b2_interp_splines, dst_iso)

    # compute noise variance, std. dev.
    if per_channel:
        dst_var = np.zeros(shape=image.shape)
        bp_idx = [[0, 0], [0, 1], [1, 0], [1, 1]]
        for ch in range(4):
            i0 = bp_idx[ch][0]
            j0 = bp_idx[ch][1]
            if fix_exp is not None:
                dst_var[i0::2, j0::2] = image_fix_exp[i0::2, j0::2] * dst_params[ch, 0] + dst_params[ch, 1]
            else:
                dst_var[i0::2, j0::2] = image[i0::2, j0::2] * dst_params[ch, 0] + dst_params[ch, 1]
    else:
        dst_var = image * dst_params[0] + dst_params[1]

    # simulate variance of noise variance
    if spat_var:
        dst_var[dst_var < 0] = 0
        dst_var += np.random.normal(loc=0, scale=1, size=image.shape) * np.sqrt(dst_var)

    dst_var[dst_var < 0] = 0

    # std. dev.
    dst_std = np.sqrt(dst_var)

    if cross_corr:
        # Simplex noise, with cross correlations
        # noise = generate_simplex_noise_template_tiles(image.shape[0], image.shape[1], scale=2)
        # sample noise from estimated covariance matrix
        cov_mat = np.load(cov_mat_fn)
        noise = sample_cov(cov_mat, image.shape[0], image.shape[1])
        noise = noise[0]
    else:
        # Normal Gaussian noise
        noise = np.random.normal(loc=0, scale=1, size=image.shape)

    # scale by heteroscedastic standard deviation
    noise *= dst_std

    if fix_exp is not None:
        noise /= fix_exp

    # invert lens shading correction for noise only
    # noise_inv_lsc = apply_lens_gain(noise, metadata['lens_gain_resized'], metadata['cfa_pattern'],
    #                                 operation='divide')

    # add noise
    noisy_image = (image + noise).astype(image.dtype)

    # invert lens shading correction for noisy image
    # noisy_image = apply_lens_gain(noisy_image, metadata['lens_gain_resized'], metadata['cfa_pattern'],
    #                               operation='divide')

    # defective pixels
    noisy_image = synthesize_defective_pixels(noisy_image)

    # match brightness of source image
    if match_bright_gt:
        noisy_image = brightness_transfer(noisy_image, src_image)

    # clip
    noisy_image = np.clip(noisy_image, min_val, max_val)

    # apply noise on image with adjusted brightness
    noisy_image_adjusted = None
    if dst_image is not None:
        if per_channel:
            dst_var_adj = np.zeros(shape=image_adjusted.shape)
            bp_idx = [[0, 0], [0, 1], [1, 0], [1, 1]]
            for ch in range(4):
                i0 = bp_idx[ch][0]
                j0 = bp_idx[ch][1]
                dst_var_adj[i0::2, j0::2] = image_adjusted[i0::2, j0::2] * dst_params[ch, 0] + dst_params[ch, 1]
        else:
            dst_var_adj = image_adjusted * dst_params[0] + dst_params[1]

        # simulate variance of noise variance
        if spat_var:
            dst_var_adj[dst_var_adj < 0] = 0
            dst_var_adj += np.random.normal(loc=0, scale=1, size=image.shape) * np.sqrt(dst_var_adj)

        dst_var_adj[dst_var_adj < 0] = 0

        dst_std_adj = np.sqrt(dst_var_adj)
        noise_adj = np.random.normal(loc=0, scale=1, size=image_adjusted.shape)
        noise_adj *= dst_std_adj
        noisy_image_adjusted = image_adjusted + noise_adj
        noisy_image_adjusted = np.clip(noisy_image_adjusted, min_val, max_val)

    # results
    synth_res = dict()
    synth_res['noisy_image'] = noisy_image
    synth_res['noisy_image_adjusted'] = noisy_image_adjusted
    return synth_res


def synthesize_noisy_image_v2(src_image, model, metadata=None, src_iso=None, min_val=None, max_val=None, dst_iso=None,
                              dst_image=None, per_channel=True, fix_exp=None, match_bright_gt=True, spat_var=False,
                              cross_corr=False, cov_mat_fn=None, iso2b1_interp_splines=None,
                              iso2b2_interp_splines=None,
                              debug=False, black_diff=0):
    """
    Synthesize a noisy image from `src_image` using a heteroscedastic Gaussian noise model `model`.
    :param src_image: Clean/Semi-clean image.
    :param metadata: Image metadata.
    :param model: Noise model.
    :param src_iso: ISO of `src_image`.
    :param dst_iso: ISO of noisy image to be synthesized.
    :param min_val: Minimum image intensity.
    :param max_val: Maximum image intensity.
    :param dst_image: If not None, match the brightness with provided destination image `dst_image`.
    :param debug: Whether to perform some debugging steps.
    :param per_channel: Whether to apply model per channel.
    :param fix_exp: Whether to fix image exposure before and after noise synthesis.
    :param match_bright_gt: Optionally, match the brightness with provided source image `src_image`.
    :param spat_var: Simulate spatial variations in noise.
    :param cross_corr: Simulate spatial cross correlations of noise.
    :param cov_mat_fn: Filename of covariance matrix used to simulate spatial correlations.
    :param iso2b1_interp_splines: Interpolation/extrapolation splines for shot noise (beta1)
    :param iso2b2_interp_splines: Interpolation/extrapolation splines for read noise (beta2)
    :param black_diff: Difference in black level from the original image and the synthesized noisy image.
        (e.g., original_black - synthetic_black).
    :return: Synthesized noisy image, and optionally, a another copy of the image with brightness matching `dst_image`.
    """
    # make a copy
    image = src_image.copy().astype(np.float32)

    if fix_exp is not None:
        # update image exposure to match target ISO before synthesizing noise, then re-scale it back later
        image_fix_exp = np.round(image.astype(np.float) * fix_exp).astype(np.uint32)
    else:
        image_fix_exp = None

    # if target ISO is not specified, select a random value
    if dst_iso is None:
        dst_iso = np.random.randint(50, 3201)

    if iso2b1_interp_splines is None or iso2b2_interp_splines is None:
        iso2b1_interp_splines = model['iso2b1_interp_splines']
        iso2b2_interp_splines = model['iso2b2_interp_splines']

    # get noise params (shot, read), per channel
    if dst_iso in model:
        dst_params = model[dst_iso]
    else:
        dst_params = np.zeros((4, 2))
        for c in range(4):
            dst_params[c, 0] = iso2b1_interp_splines[c](dst_iso)
            dst_params[c, 1] = iso2b2_interp_splines[c](dst_iso)

    # compute noise variance, std. dev.
    if per_channel:
        dst_var = np.zeros(shape=image.shape)
        bp_idx = [[0, 0], [0, 1], [1, 0], [1, 1]]
        for ch in range(4):
            i0 = bp_idx[ch][0]
            j0 = bp_idx[ch][1]
            if fix_exp is not None:
                dst_var[i0::2, j0::2] = image_fix_exp[i0::2, j0::2] * dst_params[ch, 0] + dst_params[ch, 1]
            else:
                # Fix: account for difference in black level between original image and synthetic noisy image
                dst_var[i0::2, j0::2] = (image[i0::2, j0::2] + black_diff) * dst_params[ch, 0] + dst_params[ch, 1]
    else:
        dst_var = image * dst_params[0] + dst_params[1]

    # simulate variance of noise variance
    if spat_var:
        dst_var[dst_var < 0] = 0
        dst_var += np.random.normal(loc=0, scale=1, size=image.shape) * np.sqrt(dst_var)

    dst_var[dst_var < 0] = 0

    # std. dev.
    dst_std = np.sqrt(dst_var)

    # Normal Gaussian noise
    noise = np.random.normal(loc=0, scale=1, size=image.shape)

    # scale by heteroscedastic standard deviation
    noise *= dst_std

    if fix_exp is not None:
        noise /= fix_exp

    # add noise
    noisy_image = (image + noise).astype(image.dtype)

    # clip
    noisy_image = np.clip(noisy_image, min_val, max_val)

    return noisy_image


def synthesize_defective_pixels(image, min_def_perc=0.001, max_def_perc=0.01):
    def_pix_perc = np.random.uniform(min_def_perc, max_def_perc)
    n_def_pix = int(def_pix_perc * image.shape[0] * image.shape[1])
    y_def = np.random.uniform(0, image.shape[0], size=n_def_pix).astype(np.int32)
    x_def = np.random.uniform(0, image.shape[1], size=n_def_pix).astype(np.int32)
    def_multiplier = 1.1
    def_std = .05
    image[y_def, x_def] *= np.random.normal(def_multiplier, def_std, size=n_def_pix)
    return image


def load_noise_model(path):
    """
    Load noise model.
    :param path: Path to noise model: either a directory to a non-mixture model or a file path to a mixture model.
    :return: Corresponding noise parameter data.
    """
    if os.path.isdir(path):
        noise_model_path = os.path.join(path, 'model_params.npy')
        iso2b1_interp_splines_fn = os.path.join(path, 'iso2b1_interp_splines.npy')
        iso2b2_interp_splines_fn = os.path.join(path, 'iso2b2_interp_splines.npy')
        return load_model(noise_model_path, iso2b1_interp_splines_fn, iso2b2_interp_splines_fn)
    else:
        return load_noise_mixture_model(path)


def load_model(noise_model_path, iso2b1_interp_splines_fn, iso2b2_interp_splines_fn):
    """
    Load noise model files.
    :param noise_model_path: Path to model directory.
    :param iso2b1_interp_splines_fn: Shot noise interpolation spline file name.
    :param iso2b2_interp_splines_fn: Read noise interpolation spline file name.
    :return: Noise discrete parameters, shot noise interpolation spline, read noise interpolation spline.
    """
    model_arr = np.load(noise_model_path)
    model_ = dict()
    for im_idx in range(model_arr.shape[0] // 4):
        model_[model_arr[im_idx * 4, 0]] = model_arr[im_idx * 4:(im_idx + 1) * 4, 2:4]
    # load interpolation/extrapolation splines for noise parameters
    iso2b1_interp_splines = np.load(iso2b1_interp_splines_fn, allow_pickle=True)
    iso2b2_interp_splines = np.load(iso2b2_interp_splines_fn, allow_pickle=True)
    return model_, iso2b1_interp_splines, iso2b2_interp_splines


def load_noise_mixture_model(model_path):
    """
    Load a noise mixture model.
    :param model_path: Path to model file.
    :return: Noise mixture model (dictionary).
    """
    return np.load(model_path, allow_pickle=True)[()]


def get_noise_params(model, iso2b1_interp_splines, iso2b2_interp_splines, dst_iso):
    """
    Get noise parameters for a given ISO from a noise model.
    :param model: Noise model (single or mixture).
    :param iso2b1_interp_splines: Shot noise interpolation spline.
    :param iso2b2_interp_splines: Read noise interpolation spline.
    :param dst_iso: Target ISO.
    :return: Noise parameter for the given ISO: read and shot noise per channel.
    """
    if 'shot' in model:  # this is a mixture model
        noise_params = get_noise_params_mixture_model(model, dst_iso)
    else:  # non-mixture model
        noise_params = get_noise_params_interp(iso2b1_interp_splines, iso2b2_interp_splines, dst_iso)
    return noise_params


def get_noise_params_interp(iso2b1_interp_splines, iso2b2_interp_splines, dst_iso):
    """
    Get noise parameters (shot, read) per channel, for a given ISO.
    :param iso2b1_interp_splines: Shot noise interpolation spline.
    :param iso2b2_interp_splines: Read noise interpolation spline.
    :param dst_iso: Target ISO level.
    :return: Noise parameters.
    """
    noise_params = np.zeros((4, 2))
    for c in range(4):
        noise_params[c, 0] = iso2b1_interp_splines[c](dst_iso)
        noise_params[c, 1] = iso2b2_interp_splines[c](dst_iso)
    return noise_params


def get_noise_params_mixture_model(model, dst_iso):
    """
    Get noise parameters from a noise mixture model for a given ISO.
    :param model: Noise mixture model.
    :param dst_iso: Target ISO.
    :return: Noise parameter for the given ISO: read and shot noise per channel.
    """
    n_componets = len(model['shot'])
    mixture_idx = np.random.choice(n_componets, size=1, replace=True, p=model['weight'])[0]
    return get_noise_params_interp(model['shot'][mixture_idx], model['read'][mixture_idx], dst_iso)

noise_profiler/image_synthesizer.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import numpy as np


def brightness_transfer(image, dst_image, white_level=1023):
    subsample = 16
    src_image_sample = image[::subsample, ::subsample].flatten()
    dst_image_sample = dst_image[::subsample, ::subsample].flatten()
    bright_func = np.polyfit(src_image_sample, dst_image_sample, 1)
    image_adjusted = image * bright_func[0] + bright_func[1]
    image_adjusted = np.clip(image_adjusted, 0, white_level)
    return image_adjusted


def brightness_transfer_v1(image, dst_image, white_level=1023):
    mean1 = np.mean(image)
    mean2 = np.mean(dst_image)
    std1 = np.std(image)
    std2 = np.std(dst_image)
    std1 = max(std1, 1e-8)
    std2 = max(std2, 1e-8)
    image_adjusted = (image - mean1) / std1 * std2 + mean2
    image_adjusted = np.clip(image_adjusted, 0, white_level)
    return image_adjusted


def brightness_transfer_v11(image, dst_image, white_level=1023):
    mean1 = np.mean(image)
    mean2 = np.mean(dst_image)
    image_adjusted = image - mean1 + mean2
    image_adjusted = np.clip(image_adjusted, 0, white_level)
    return image_adjusted


def brightness_transfer_v2(image, dst_image, white_level=1023):
    subsample = 4
    src_image_sample = image[::subsample, ::subsample].flatten()
    dst_image_sample = dst_image[::subsample, ::subsample].flatten()
    bright_func = np.polyfit(src_image_sample, dst_image_sample, 1)
    image_adjusted = image * bright_func[0] + bright_func[1]
    image_adjusted = np.clip(image_adjusted, 0, white_level)
    return image_adjusted


def standardize(x):
    mean = np.mean(x)
    std = np.std(x)
    return (x - mean) / std


def normalize(x):
    min_ = np.min(x)
    max_ = np.max(x)
    return (x - min_) / (max_ - min_)

noise_profiler/img_utils.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.
"""
import cv2
import numpy as np
import matplotlib.pyplot as plt

from noise_profiler.img_utils import normalize


def sample_norm(h_, w_):
    return np.random.normal(0, 1, (h_, w_))


def sample_cov(cov_mat_, h_, w_):
    n_var = 4
    sample_cov_ = np.random.multivariate_normal([0] * n_var, cov_mat_, h_ * w_ // n_var)
    sample_cov_image_ = np.zeros((h_, w_))
    sample_cov_image_[0::2, 0::2] = sample_cov_[:, 0].reshape((h_ // 2, w_ // 2))
    sample_cov_image_[0::2, 1::2] = sample_cov_[:, 1].reshape((h_ // 2, w_ // 2))
    sample_cov_image_[1::2, 0::2] = sample_cov_[:, 2].reshape((h_ // 2, w_ // 2))
    sample_cov_image_[1::2, 1::2] = sample_cov_[:, 3].reshape((h_ // 2, w_ // 2))
    return sample_cov_image_, sample_cov_


# save images
def save_image(im, save_fn, sc=4):
    cv2.imwrite(save_fn, cv2.resize((normalize(im) * 255).astype(np.uint8), dsize=(im.shape[1] * sc, im.shape[0] * sc),
                                    interpolation=cv2.INTER_NEAREST))


# plot covariance matrix as a heat map
def plot_cov_mat(cov_mat_, save_fn):
    fig = plt.figure()
    plt.imshow(cov_mat_)
    ax = plt.gca()
    plt.colorbar()
    plt.clim(0, 1)
    for i in range(4):
        for j in range(4):
            ax.annotate("{:.4f}".format(cov_mat_[i, j]), xy=(i - .25, j))
    plt.tick_params(axis='both', which='both', bottom=False, top=False, left=False, right=False, labelbottom=True)
    plt.xticks(np.arange(0, 4), labels=['Gr', 'R', 'B', 'Gb'])
    plt.yticks(np.arange(0, 4), labels=['Gr', 'R', 'B', 'Gb'])
    plt.tight_layout()
    # plt.show()
    plt.savefig(save_fn)
    plt.close(fig)

noise_profiler/spatial_correlation/utils.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.


Utility function for working with correlated color temperatures.
"""

import numpy as np
import sys


def dot(x, y):
    return np.sum(x * y, axis=-1)


def norm(x):
    return np.sqrt(dot(x, x))


def raw_rgb_to_cct(raw_rgb, xyz2cam1, xyz2cam2):
    """Convert raw-RGB triplet to corresponding correlated color temperature (CCT)"""
    # pxyz = [.5, 1, .5]
    pxyz = [.3, .3, .3]
    # cct = xyz2cct(pxyz)
    loss = 1e10
    k = 1
    cct = 6500  # default
    max_iter = 10000
    iter = 0
    while loss > 1e-4 and iter < max_iter:
        cct = xyz2cct(pxyz)
        xyz = raw_rgb_to_xyz(raw_rgb, cct, xyz2cam1, xyz2cam2)
        loss = norm(xyz - pxyz)
        # loss = ang_err(xyz, pxyz)
        pxyz = xyz
        k = k + 1
        iter += 1
        if(iter == max_iter):
          print('Max iter reached')
    return cct


def raw_rgb_to_xyz(raw_rgb, temp, xyz2cam1, xyz2cam2):
    # TODO: what if temperature > 6500?
    # RawRgbToXyz Convert raw-RGB triplet to corresponding XYZ
    cct1 = 6500  # D65, DNG code = 21
    cct2 = 2500  # A, DNG code = 17
    cct1inv = 1 / cct1
    cct2inv = 1 / cct2
    tempinv = 1 / temp
    g = (tempinv - cct2inv) / (cct1inv - cct2inv)
    h = 1 - g
    # if g < 0:
    #     g = 0
    # if h < 0:
    #     h = 0
    # if h > 1:
    #     h = 1
    xyz2cam = g * xyz2cam1 + h * xyz2cam2
    xyz = np.matmul(np.linalg.inv(xyz2cam), np.transpose(raw_rgb))

    # xyz = xyz / xyz(2)
    # xyz_ = [xyz(1) / sum(xyz), xyz(2) / sum(xyz), 1]
    # xyz = xyz_

    return xyz


# /* LERP(a,b,c) = linear interpolation macro, is 'a' when c == 0.0 and 'b' when c == 1.0 */
def lerp(a, b, c):
    return (b - a) * c + a


def xyz2cct(xyz):
    #     % /*
    # %  *      Name:   XYZtoCorColorTemp.c
    # %  *
    # %  *      Author: Bruce Justin Lindbloom
    # %  *
    # %  *      Copyright (c) 2003 Bruce Justin Lindbloom. All rights reserved.
    # %  *
    # %  *      Input:  xyz = pointer to the input array of X, Y and Z color components (in that order).
    # %  *              temp = pointer to where the computed correlated color temperature should be placed.
    # %  *
    # %  *      Output: *temp = correlated color temperature, if successful.
    # %  *                    = unchanged if unsuccessful.
    # %  *
    # %  *      Return: 0 if successful, else -1.
    # %  *
    # %  *      Description:
    # %  *              This is an implementation of Robertson's method of computing the correlated color
    # %  *              temperature of an XYZ color. It can compute correlated color temperatures in the
    # %  *              range [1666.7K, infinity].
    # %  *
    # %  *      Reference:
    # %  *              "Color Science: Concepts and Methods, Quantitative Data and Formulae", Second Edition,
    # %  *              Gunter Wyszecki and W. S. Stiles, John Wiley & Sons, 1982, pp. 227, 228.
    # %  */
    #

    rt = [  # /* reciprocal temperature (K) */
        sys.float_info.min, 10.0e-6, 20.0e-6, 30.0e-6, 40.0e-6, 50.0e-6,
        60.0e-6, 70.0e-6, 80.0e-6, 90.0e-6, 100.0e-6, 125.0e-6,
        150.0e-6, 175.0e-6, 200.0e-6, 225.0e-6, 250.0e-6, 275.0e-6,
        300.0e-6, 325.0e-6, 350.0e-6, 375.0e-6, 400.0e-6, 425.0e-6,
        450.0e-6, 475.0e-6, 500.0e-6, 525.0e-6, 550.0e-6, 575.0e-6,
        600.0e-6
    ]

    uvt = [
        [0.18006, 0.26352, -0.24341],
        [0.18066, 0.26589, -0.25479],
        [0.18133, 0.26846, -0.26876],
        [0.18208, 0.27119, -0.28539],
        [0.18293, 0.27407, -0.30470],
        [0.18388, 0.27709, -0.32675],
        [0.18494, 0.28021, -0.35156],
        [0.18611, 0.28342, -0.37915],
        [0.18740, 0.28668, -0.40955],
        [0.18880, 0.28997, -0.44278],
        [0.19032, 0.29326, -0.47888],
        [0.19462, 0.30141, -0.58204],
        [0.19962, 0.30921, -0.70471],
        [0.20525, 0.31647, -0.84901],
        [0.21142, 0.32312, -1.0182],
        [0.21807, 0.32909, -1.2168],
        [0.22511, 0.33439, -1.4512],
        [0.23247, 0.33904, -1.7298],
        [0.24010, 0.34308, -2.0637],
        [0.24792, 0.34655, -2.4681],
        # /* Note: 0.24792 is a corrected value for the error found in W&S as 0.24702 */
        [0.25591, 0.34951, -2.9641],
        [0.26400, 0.35200, -3.5814],
        [0.27218, 0.35407, -4.3633],
        [0.28039, 0.35577, -5.3762],
        [0.28863, 0.35714, -6.7262],
        [0.29685, 0.35823, -8.5955],
        [0.30505, 0.35907, -11.324],
        [0.31320, 0.35968, -15.628],
        [0.32129, 0.36011, -23.325],
        [0.32931, 0.36038, -40.770],
        [0.33724, 0.36051, -116.45]
    ]

    # us = 0
    # vs = 0
    # p = 0
    di = 0
    # dm = 0
    i = 0
    if (xyz[0] < 1.0e-20) and (xyz[1] < 1.0e-20) and (xyz[2] < 1.0e-20):
        return -1  # /* protect against possible divide-by-zero failure */
    us = (4.0 * xyz[0]) / (xyz[0] + 15.0 * xyz[1] + 3.0 * xyz[2])
    vs = (6.0 * xyz[1]) / (xyz[0] + 15.0 * xyz[1] + 3.0 * xyz[2])
    dm = 0.0
    for i in range(31):
        di = (vs - uvt[i][1]) - uvt[i][2] * (us - uvt[i][0])
        if (i > 0) and (((di < 0.0) and (dm >= 0.0)) or ((di >= 0.0) and (dm < 0.0))):
            break  # /* found lines bounding (us, vs) : i-1 and i */
        dm = di

    if i == 31:
        # /* bad XYZ input, color temp would be less than minimum of 1666.7 degrees, or too far towards blue */
        return -1
    di = di / np.sqrt(1.0 + uvt[i][2] * uvt[i][2])
    dm = dm / np.sqrt(1.0 + uvt[i - 1][2] * uvt[i - 1][2])
    p = dm / (dm - di)  # /* p = interpolation parameter, 0.0 : i-1, 1.0 : i */
    p = 1.0 / (lerp(rt[i - 1], rt[i], p))
    cct = p
    return cct  # /* success */


def interpolate_cst(xyz2cam1, xyz2cam2, temp):
    # RawRgbToXyz Convert raw-RGB triplet to corresponding XYZ
    cct1 = 6500  # D65, DNG code = 21
    cct2 = 2500  # A, DNG code = 17
    cct1inv = 1 / cct1
    cct2inv = 1 / cct2
    tempinv = 1 / temp
    g = (tempinv - cct2inv) / (cct1inv - cct2inv)
    h = 1 - g
    if g < 0:
        g = 0
    if h < 0:
        h = 0
    if h > 1:
        h = 1
    xyz2cam = g * xyz2cam1 + h * xyz2cam2
    return xyz2cam

pipeline/cct_utils.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

Utility functions for handling DNG opcode lists.
"""
import struct
import numpy as np
from .exif_utils import get_tag_values_from_ifds


class Opcode:
    def __init__(self, id_, dng_spec_ver, option_bits, size_bytes, data):
        self.id = id_
        self.dng_spec_ver = dng_spec_ver
        self.size_bytes = size_bytes
        self.option_bits = option_bits
        self.data = data


def parse_opcode_lists(ifds):
    # OpcodeList1, 51008, 0xC740
    # Applied to raw image as read directly form file

    # OpcodeList2, 51009, 0xC741
    # Applied to raw image after being mapped to linear reference values
    # That is, after linearization, black level subtraction, normalization, and clipping

    # OpcodeList3, 51022, 0xC74E
    # Applied to raw image after being demosaiced

    opcode_list_tag_nums = [51008, 51009, 51022]
    opcode_lists = {}
    for i, tag_num in enumerate(opcode_list_tag_nums):
        opcode_list_ = get_tag_values_from_ifds(tag_num, ifds)
        if opcode_list_ is not None:
            opcode_list_ = bytearray(opcode_list_)
            opcodes = parse_opcodes(opcode_list_)
            opcode_lists.update({tag_num: opcodes})
        else:
            pass

    return opcode_lists


def parse_opcodes(opcode_list):
    """
    Parse a byte array representing an opcode list.
    :param opcode_list: An opcode list as a byte array.
    :return: Opcode lists as a dictionary.
    """
    # opcode lists are always stored in big endian
    endian_sign = ">"

    # opcode IDs
    # 9: GainMap
    # 1: Rectilinear Warp

    # clip to
    # [0, 2^32 - 1] for OpcodeList1
    # [0, 2^16 - 1] for OpcodeList2
    # [0, 1] for OpcodeList3

    i = 0
    num_opcodes = struct.unpack(endian_sign + "I", opcode_list[i:i + 4])[0]
    i += 4

    opcodes = {}
    for j in range(num_opcodes):
        opcode_id_ = struct.unpack(endian_sign + "I", opcode_list[i:i + 4])[0]
        i += 4
        dng_spec_ver = [struct.unpack(endian_sign + "B", opcode_list[i + k:i + k + 1])[0] for k in range(4)]
        i += 4
        option_bits = struct.unpack(endian_sign + "I", opcode_list[i:i + 4])[0]
        i += 4

        # option bits
        if option_bits & 1 == 1:  # optional/unknown
            pass
        elif option_bits & 2 == 2:  # can be skipped for "preview quality", needed for "full quality"
            pass
        else:
            pass

        opcode_size_bytes = struct.unpack(endian_sign + "I", opcode_list[i:i + 4])[0]
        i += 4

        opcode_data = opcode_list[i:i + opcode_size_bytes]
        i += opcode_size_bytes

        # GainMap (lens shading correction map)
        if opcode_id_ == 9:
            opcode_gain_map_data = parse_opcode_gain_map(opcode_data)
            opcode_data = opcode_gain_map_data
            # change opcode_id from 9 to 9.0, 9.1, 9.2, ... to handle different lsc maps for different channels
            opcode_id_ = opcode_id_ + j / 10.
        # WarpRectilinear
        elif opcode_id_ == 1:
            opcode_rect_warp_data = parse_opcode_rect_warp(opcode_data)
            opcode_data = opcode_rect_warp_data
        # FixBadPixelsList
        elif opcode_id_ == 5:
            bad_pixels_list = parse_bad_pixels_list(opcode_data)
            opcode_data = bad_pixels_list

        # set opcode object
        opcode = Opcode(id_=opcode_id_, dng_spec_ver=dng_spec_ver, option_bits=option_bits,
                        size_bytes=opcode_size_bytes,
                        data=opcode_data)
        opcodes.update({opcode_id_: opcode})

    return opcodes


def parse_bad_pixels_list(opcode_data):
    endian_sign = ">"  # big
    opcode_dict = {'bad_points': [], 'bad_rects': []}

    i = 0
    opcode_dict['bayer_phase'] = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
    i += 4

    opcode_dict['bad_point_count'] = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
    i += 4

    opcode_dict['bad_rect_count'] = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
    i += 4

    for j in range(opcode_dict['bad_point_count']):
        bad_point_row = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
        i += 4
        bad_point_col = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
        i += 4
        opcode_dict['bad_points'].append((bad_point_row, bad_point_col))

    for j in range(opcode_dict['bad_rect_count']):
        bad_point_top = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
        i += 4
        bad_point_left = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
        i += 4
        bad_point_bot = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
        i += 4
        bad_point_right = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
        i += 4
        opcode_dict['bad_points'].append((bad_point_top, bad_point_left, bad_point_bot, bad_point_right))

    return opcode_dict


def parse_opcode_rect_warp(opcode_data):
    endian_sign = ">"  # big
    opcode_dict = {}
    '''
    opcode_dict = {
        'N': 3,
        'coefficient_set': [
            {
                'k_r0': 1,
                'k_r1': 0,
                'k_r2': 0,
                'k_r3': 0,
                'k_t0': 0,
                'k_t1': 0,              
             },
            ...
        ],
        'cx': 0.5,
        'cy': 0.5
    }
    '''
    i = 0
    num_planes = struct.unpack(endian_sign + "L", opcode_data[i:i + 4])[0]
    i += 4

    opcode_dict['N'] = num_planes
    opcode_dict['coefficient_set'] = []

    for j in range(num_planes):
        keys = ['k_r0', 'k_r1', 'k_r2', 'k_r3', 'k_t0', 'k_t1']
        coefficient_set = {}
        for key in keys:
            coefficient_set[key] = struct.unpack(endian_sign + "d", opcode_data[i:i + 8])[0]
            i += 8
        opcode_dict['coefficient_set'].append(coefficient_set)

    opcode_dict['cx'] = struct.unpack(endian_sign + "d", opcode_data[i:i + 8])[0]
    i += 8
    opcode_dict['cy'] = struct.unpack(endian_sign + "d", opcode_data[i:i + 8])[0]

    return opcode_dict


def parse_opcode_gain_map(opcode_data):
    endian_sign = ">"  # big
    opcode_dict = {}
    keys = ['top', 'left', 'bottom', 'right', 'plane', 'planes', 'row_pitch', 'col_pitch', 'map_points_v',
            'map_points_h', 'map_spacing_v', 'map_spacing_h', 'map_origin_v', 'map_origin_h', 'map_planes', 'map_gain']
    dtypes = ['L'] * 10 + ['d'] * 4 + ['L'] + ['f']
    dtype_sizes = [4] * 10 + [8] * 4 + [4] * 2  # data type size in bytes
    counts = [1] * 15 + [0]  # 0 count means variable count, depending on map_points_v and map_points_h
    # values = []

    i = 0
    for k in range(len(keys)):
        if counts[k] == 0:  # map_gain
            counts[k] = opcode_dict['map_points_v'] * opcode_dict['map_points_h']

        if counts[k] == 1:
            vals = struct.unpack(endian_sign + dtypes[k], opcode_data[i:i + dtype_sizes[k]])[0]
            i += dtype_sizes[k]
        else:
            vals = []
            for j in range(counts[k]):
                vals.append(struct.unpack(endian_sign + dtypes[k], opcode_data[i:i + dtype_sizes[k]])[0])
                i += dtype_sizes[k]

        opcode_dict[keys[k]] = vals

    opcode_dict['map_gain_2d'] = np.reshape(opcode_dict['map_gain'],
                                            (opcode_dict['map_points_v'], opcode_dict['map_points_h']))

    return opcode_dict

pipeline/dng_opcode.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


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'),
}

pipeline/exif_data_formats.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""


import struct
from fractions import Fraction
from .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)

pipeline/exif_utils.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from .pipeline_utils import get_visible_raw_image, get_metadata, normalize, white_balance, demosaic, \
    denoise, apply_color_space_transform, transform_xyz_to_srgb, apply_gamma, apply_tone_map, fix_orientation, \
    fix_missing_params, lens_shading_correction, active_area_cropping, lens_distortion_correction, default_cropping, \
    resize, apply_local_tone_map, transform_xyz_to_prophoto, transform_prophoto_to_srgb


def run_module(image, module, built_in_function, built_in_args):
    if type(module) is list and len(module) == 2:
        image = module[0](image, **module[1])
    elif type(module) is str:
        image = built_in_function(image, **built_in_args)
    else:
        raise ValueError('Invalid input module.')
    return image


def linearization_stage(current_image, metadata):
    linearization_table = metadata['linearization_table']
    if linearization_table is not None:
        print('Linearization table found. Not handled.')
        # TODO
    return current_image


def lens_shading_correction_stage(current_image, metadata, clip):
    gain_map_opcode = None
    if 'opcode_lists' in metadata:
        if 51009 in metadata['opcode_lists']:
            opcode_list_2 = metadata['opcode_lists'][51009]
            gain_map_opcode = [opcode_list_2[opcode_id] for opcode_id in opcode_list_2.keys() if 9 <= opcode_id <= 9.3]

    lsc_map = None
    if 'lsc_map' in metadata:
        lsc_map = metadata['lsc_map']

    if gain_map_opcode is not None:
        current_image = lens_shading_correction(current_image, gain_map_opcode=gain_map_opcode,
                                                bayer_pattern=metadata['cfa_pattern'], clip=clip)
    elif lsc_map is not None:
        current_image = lens_shading_correction(current_image, gain_map_opcode=None,
                                                bayer_pattern=metadata['cfa_pattern'], gain_map=metadata['lsc_map'],
                                                clip=clip)
    return current_image


def lens_distortion_correction_stage(current_image, metadata, clip):
    if 'opcode_lists' in metadata:
        if 51022 in metadata['opcode_lists']:
            opcode_list_3 = metadata['opcode_lists'][51022]
            rect_warp_opcode = opcode_list_3[1]
            current_image = lens_distortion_correction(current_image, rect_warp_opcode=rect_warp_opcode,
                                                       clip=clip)
    return current_image


def run_pipeline(image_or_path, params=None, metadata=None, stages=None, clip=True):
    if type(image_or_path) == str:
        image_path = image_or_path
        # raw image data
        raw_image = get_visible_raw_image(image_path)
        # metadata
        metadata = get_metadata(image_path)
    else:
        raw_image = image_or_path.copy()
        # must provide metadata
        if metadata is None:
            raise ValueError("Must provide metadata when providing image data in first argument.")

    # take a deep copy of params as it will be modified below
    params = params.copy()

    # fill any missing parameters with default values
    params = fix_missing_params(params)

    '''
    Function performed at each stage. Follows this format:
    * {'stage_name': [function_name, function_params]} 
    * Assumes the function takes in `current_image` as the first parameter. 
    '''
    operation_by_stage = {
        'active_area_cropping': [active_area_cropping, {'active_area': metadata['active_area']}],
        'default_cropping': [default_cropping, {'default_crop_origin': metadata['default_crop_origin'],
                                                'default_crop_size': metadata['default_crop_size']}],
        'linearization': [linearization_stage, {'metadata': metadata}],
        'normal': [normalize, {
            'black_level': metadata['black_level'],
            'white_level': metadata['white_level'],
            'black_level_delta_h': metadata['black_level_delta_h'],
            'black_level_delta_v': metadata['black_level_delta_v'],
            'clip': clip}],
        'lens_shading_correction': [lens_shading_correction_stage, {'metadata': metadata, 'clip': clip}],
        'white_balance': [run_module, {
            'module': params['white_balancer'],
            'built_in_function': white_balance,
            'built_in_args': {
                'alg_type': params['white_balancer'],
                'as_shot_neutral': metadata['as_shot_neutral'],
                'cfa_pattern': metadata['cfa_pattern'],
                'clip': clip
            }
        }],
        'demosaic': [run_module, {
            'module': params['demosaicer'],
            'built_in_function': demosaic,
            'built_in_args': {
                'cfa_pattern': metadata['cfa_pattern'],
                'output_channel_order': 'RGB',
                'alg_type': params['demosaicer'],
            }
        }],
        'lens_distortion_correction': [lens_distortion_correction_stage, {'metadata': metadata, 'clip': clip}],
        'denoise': [run_module, {
            'module': params['denoiser'],
            'built_in_function': denoise,
            'built_in_args': {
                'alg_type': params['denoiser'],
                'cfa_pattern': metadata['cfa_pattern']
            }
        }],
        'xyz': [apply_color_space_transform, {
            'color_matrix_1': metadata['color_matrix_1'],
            'color_matrix_2': metadata['color_matrix_2'],
            'illuminant': metadata['as_shot_neutral']
        }],
        'prophoto': [transform_xyz_to_prophoto, {}],
        'srgb': [transform_prophoto_to_srgb, {}],
        'xyz2srgb': [transform_xyz_to_srgb, {}],
        'fix_orient': [fix_orientation, {'orientation': metadata['orientation']}],
        'gamma': [apply_gamma, {}],
        'tone': [run_module, {
            'module': params['tone_curve'],
            'built_in_function': apply_tone_map,
            'built_in_args': {
                'tone_curve': params['tone_curve']
            }
        }],
        'resize': [resize, {'target_size': (raw_image.shape[1], raw_image.shape[0])}],
        'local_tone_mapping': [run_module, {
            'module': params['local_tone_mapping'],
            'built_in_function': apply_local_tone_map,
            'built_in_args': {
                'alg_type': params['local_tone_mapping'],
                'channel_order': 'RGB',
                'clahe_clip_limit': params['clahe_clip_limit'],
                'clahe_grid_size': params['clahe_grid_size']
            }
        }],
    }

    if not stages:
        stages = ['raw', 'active_area_cropping', 'linearization', 'normal', 'lens_shading_correction', 'white_balance',
                  'demosaic', 'lens_distortion_correction', 'denoise', 'xyz', 'prophoto', 'srgb', 'fix_orient', 'gamma',
                  'tone', 'local_tone_mapping', 'default_cropping', 'resize']

    input_stage = params['input_stage']
    output_stage = params['output_stage']
    if input_stage not in stages \
            or output_stage not in stages \
            or stages.index(input_stage) > stages.index(output_stage):
        raise ValueError('Invalid input/output stage: input_stage = {}, output_stage = {}'.format(input_stage,
                                                                                                  output_stage))
    # Handle transforming directly from xyz to srgb
    srgb_idx = stages.index('srgb') if 'srgb' in stages else -1
    if srgb_idx > 1 and stages[srgb_idx - 1] == 'xyz':
        stages[srgb_idx] = 'xyz2srgb'

    input_idx = stages.index(input_stage)
    output_idx = stages.index(output_stage)
    current_image = raw_image

    for stage in stages[input_idx + 1:output_idx + 1]:
        operation = operation_by_stage[stage]
        current_image = operation[0](current_image, **operation[1])
    return current_image

pipeline/pipeline.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.


Camera pipeline utilities.
"""
import math
import os
import cv2
import numpy as np
import exifread
import rawpy
from .cct_utils import raw_rgb_to_cct, interpolate_cst
from fractions import Fraction
from exifread.utils import Ratio
from scipy.io import loadmat
from colour_demosaicing import demosaicing_CFA_Bayer_Menon2007
from .exif_utils import parse_exif, get_tag_values_from_ifds
from .dng_opcode import parse_opcode_lists
from .raw_utils import stack_rgb_channels


def get_visible_raw_image(image_path):
    raw_image = rawpy.imread(image_path).raw_image_visible.copy()
    return raw_image


def load_raw_burst(image_paths):
    """loads bayer data from raw dngs"""
    images = []
    for i in range(len(image_paths)):
        image = get_visible_raw_image(image_paths[i])
        images.append(image.astype(np.float32))

    images = np.stack(images)

    print(f'burst of shape {list(images.shape)} loaded')
    return images


def save_image_stage(image, out_path, stage, save_as, bit_depth=8, stage_id=''):
    if bit_depth != 8 and save_as == 'jpg':
        save_as = 'png'

    if os.path.isdir(out_path):
        output_image_path = os.path.join(out_path, 'merged_{}_{}.{}'.format(stage_id, stage, save_as))
    else:
        out_path_no_ext, _ = os.path.splitext(out_path)
        output_image_path = '{}_{}_{}.{}'.format(out_path_no_ext, stage_id, stage, save_as)

    output_image = (image * (2 ** bit_depth - 1)).astype('uint' + str(bit_depth))

    if len(output_image.shape) > 2:
        output_image = output_image[:, :, ::-1]
    if save_as == 'jpg':
        cv2.imwrite(output_image_path, output_image, [cv2.IMWRITE_JPEG_QUALITY, 100])
    else:
        cv2.imwrite(output_image_path, output_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['active_area'] = get_active_area(tags, ifds)
    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)
    metadata['noise_profile'] = get_noise_profile(tags, ifds)
    metadata['iso'] = get_iso(ifds)
    metadata['exposure_time'] = get_exposure_time(ifds)
    metadata['default_crop_origin'] = get_default_crop_origin(ifds)
    metadata['default_crop_size'] = get_default_crop_size(ifds)
    metadata['black_level_delta_h'] = get_black_level_delta_h(ifds)
    metadata['black_level_delta_v'] = get_black_level_delta_v(ifds)
    metadata['make'] = get_make(ifds)
    metadata['model'] = get_model(ifds)
    # ...

    # opcode lists
    metadata['opcode_lists'] = parse_opcode_lists(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_active_area(tags, ifds):
    possible_keys = ['Image Tag 0xC68D', 'Image Tag 50829', 'ActiveArea', 'Image ActiveArea']
    vals = get_values(tags, possible_keys)
    if vals is None:
        vals = get_tag_values_from_ifds(50829, ifds)
    return vals


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)
    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_iso(ifds):
    # 0x8827	34855
    return get_tag_values_from_ifds(34855, ifds)


def get_exposure_time(ifds):
    # 0x829a	33434
    exposure_time = get_tag_values_from_ifds(33434, ifds)[0]
    return float(exposure_time.numerator) / float(exposure_time.denominator)


def get_default_crop_origin(ifds):
    return get_tag_values_from_ifds(50719, ifds)


def get_default_crop_size(ifds):
    return get_tag_values_from_ifds(50720, ifds)


def get_black_level_delta_h(ifds):
    bldh = get_tag_values_from_ifds(50715, ifds)
    if bldh and type(bldh[0]) is Fraction:
        bldh = [float(bldh[i].numerator) / bldh[i].denominator for i in range(len(bldh))]
    return bldh


def get_black_level_delta_v(ifds):
    bldv = get_tag_values_from_ifds(50716, ifds)
    if bldv and type(bldv[0]) is Fraction:
        bldv = [float(bldv[i].numerator) / bldv[i].denominator for i in range(len(bldv))]
    return bldv


def get_make(ifds):
    make = get_tag_values_from_ifds(271, ifds)
    if make:
        make = b''.join(make[:-1]).decode()
    return make


def get_model(ifds):
    model = get_tag_values_from_ifds(272, ifds)
    if model:
        model = b''.join(model[:-1]).decode()
    return model


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 active_area_cropping(image, active_area):
    if active_area is not None and active_area != [0, 0, image.shape[0], image.shape[1]]:
        image = image[active_area[0]: active_area[2], active_area[1]:active_area[3]]

    return image


def default_cropping(image, default_crop_origin, default_crop_size):
    if default_crop_origin is not None and default_crop_size is not None:
        if type(default_crop_origin[0]) is Fraction:
            default_crop_origin = [float(x.numerator) / float(x.denominator) for x in default_crop_origin]
        if type(default_crop_size[0]) is Fraction:
            default_crop_size = [float(x.numerator) / float(x.denominator) for x in default_crop_size]
        if np.any([x != int(x) for x in default_crop_size]):
            raise ValueError('Default crop size is not integer, default_crop_size = {}'.format(default_crop_size))

        # when default_crop_origin and default_crop_size come in (H,W) rather than (W,H), flip the order
        if (default_crop_size[0] < default_crop_size[1] and image.shape[0] < image.shape[1]) or \
                (default_crop_size[0] > default_crop_size[1] and image.shape[0] > image.shape[1]):
            default_crop_size.reverse()
            default_crop_origin.reverse()

        # check if any elements in default crop origin or default crop size is not an integer
        if np.any([x != int(x) for x in default_crop_origin]):
            xs, ys = np.meshgrid(np.arange(default_crop_size[0]) + default_crop_origin[0],
                                 np.arange(default_crop_size[1]) + default_crop_origin[1])
            xs = xs.astype(np.float32)
            ys = ys.astype(np.float32)
            image = cv2.remap(image, xs, ys, cv2.INTER_LINEAR)
        else:
            image = image[int(default_crop_origin[1]):int(default_crop_origin[1] + default_crop_size[1]),
                    int(default_crop_origin[0]):int(default_crop_origin[0] + default_crop_size[0]), ...]
    return image


def resize(image, target_size):
    """
    target_size: (width, height)
    """
    return cv2.resize(image, dsize=target_size, interpolation=cv2.INTER_LINEAR)


def normalize(raw_image, black_level, white_level, black_level_delta_h=None, black_level_delta_v=None, clip=True):
    h, w = raw_image.shape[:2]

    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)
        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]

    # add black level delta h, v
    if black_level_delta_h:
        black_level_mask += np.tile(np.array(black_level_delta_h)[np.newaxis, :], [h, 1])
    if black_level_delta_v:
        black_level_mask += np.tile(np.array(black_level_delta_v)[:, np.newaxis], [1, w])

    normalized_image = raw_image.astype(np.float32) - black_level_mask
    # if some values were smaller than black level
    if clip:
        normalized_image[normalized_image < 0] = 0
    normalized_image = normalized_image / (white_level - black_level_mask)
    return normalized_image


def denormalize(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)
        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]
    denormalized_image = raw_image.astype(np.float32) * (white_level - black_level_mask) + black_level_mask
    denormalized_image[denormalized_image < black_level] = black_level
    denormalized_image[denormalized_image > white_level] = white_level
    return denormalized_image


def ratios2floats(ratios):
    floats = []
    for ratio in ratios:
        floats.append(float(ratio.num) / ratio.den)
    return floats


def calc_gw(image, cfa_pattern):
    if len(image.shape) != 3:
        # bayer image
        image = stack_rgb_channels(image, cfa_pattern)

    gw_wb_val = np.reshape(image, (-1, 3)).mean(axis=0)
    illuminant = gw_wb_val / gw_wb_val[1]
    return illuminant


def white_balance(normalized_image, as_shot_neutral, cfa_pattern, clip=True, alg_type='default'):
    """
    :param normalized_image:
    :param as_shot_neutral: required when alg_type=='default', not required when alg_type=='gw'
    :param cfa_pattern: required when image is bayer, not required wh
    :param clip:
    :param alg_type:
    :return:
    """
    if alg_type == 'gw':
        illuminant = calc_gw(normalized_image, cfa_pattern)
    else:
        if type(as_shot_neutral[0]) is Ratio:
            as_shot_neutral = ratios2floats(as_shot_neutral)
        illuminant = as_shot_neutral

    if len(normalized_image.shape) == 3:
        return white_balance_rgb(normalized_image, illuminant, clip)

    idx2by2 = [[0, 0], [0, 1], [1, 0], [1, 1]]
    step2 = 2
    white_balanced_image = np.zeros(normalized_image.shape)
    for i, idx in enumerate(idx2by2):
        idx_y = idx[0]
        idx_x = idx[1]
        white_balanced_image[idx_y::step2, idx_x::step2] = \
            normalized_image[idx_y::step2, idx_x::step2] / illuminant[cfa_pattern[i]]
    if clip:
        white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)
    return white_balanced_image


def white_balance_rgb(image_rgb, illuminant, clip=True):
    illuminant_ = illuminant.copy()
    if type(illuminant_[0]) is Ratio:
        illuminant_ = ratios2floats(illuminant_)
    illuminant_ = np.reshape(illuminant_, (1, 1, 3))
    white_balanced_image = image_rgb / illuminant_
    if clip:
        white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)
    return white_balanced_image


def lens_distortion_correction(image, rect_warp_opcode, clip=True):
    # TODO This function does not work as expected on images tested, need to
    #  investigate whether it's caused by inaccurate metadata or the algorithm.
    # TODO currently we assume there is only one coeff set, need to extend this to 3
    result_image = np.zeros_like(image)

    num_planes = rect_warp_opcode.data['N']
    coeff_set = rect_warp_opcode.data['coefficient_set'][0]
    k_r0 = coeff_set['k_r0']
    k_r1 = coeff_set['k_r1']
    k_r2 = coeff_set['k_r2']
    k_r3 = coeff_set['k_r3']
    k_t0 = coeff_set['k_t0']
    k_t1 = coeff_set['k_t1']
    cx = rect_warp_opcode.data['cx']
    cy = rect_warp_opcode.data['cy']

    x0 = y0 = 0
    x1 = image.shape[1]
    y1 = image.shape[0]
    cx = x0 + cx * (x1 - x0)
    cy = y0 + cy * (y1 - y0)
    mx = max(abs(x0 - cx), abs(x1 - cx))
    my = max(abs(y0 - cy), abs(y1 - cy))
    m = math.sqrt(mx ** 2 + my ** 2)

    for y, row in enumerate(image):
        for x, col in enumerate(row):
            dx = (x - cx) / m
            dy = (y - cy) / m
            r = math.sqrt(dx ** 2 + dy ** 2)
            f = k_r0 + k_r1 * (r ** 2) + k_r2 * (r ** 4) + k_r3 * (r ** 6)
            dxr = f * dx
            dyr = f * dy
            dxt = k_t0 * 2 * dx * dy + k_t1 * (r ** 2 + 2 * (dx ** 2))
            dyt = k_t1 * 2 * dx * dy + k_t0 * (r ** 2 + 2 * (dy ** 2))
            x_ = cx + m * (dxr + dxt)
            y_ = cy + m * (dyr + dyt)
            for p in range(len(col)):
                # nearest neighbor

                # y_ = int(round(y_))
                # x_ = int(round(x_))
                #
                # y_ = max(y_, 0)
                # x_ = max(x_, 0)
                #
                # y_ = min(y_, image.shape[0] - 1)
                # x_ = min(x_, image.shape[1] - 1)
                #
                # result_image[y, x, p] = image[y_, x_, p]

                # bilinear interpolation

                x_ratio = x_ - math.floor(x_)
                p_y0 = x_ratio * image[math.floor(y_), math.ceil(x_), p] + (1 - x_ratio) * image[
                    math.floor(y_), math.floor(x_), p]
                p_y1 = x_ratio * image[math.ceil(y_), math.ceil(x_), p] + (1 - x_ratio) * image[
                    math.ceil(y_), math.floor(x_), p]
                y_ratio = y_ - math.floor(y_)
                pixel = y_ratio * p_y1 + (1 - y_ratio) * p_y0
                result_image[y, x, p] = pixel

    if clip:
        result_image = np.clip(result_image, 0.0, 1.0)
    return result_image


def lens_shading_correction(raw_image, gain_map_opcode, bayer_pattern, gain_map=None, clip=True):
    """
    Apply lens shading correction map.
    :param raw_image: Input normalized (in [0, 1]) raw image.
    :param gain_map_opcode: Gain map opcode.
    :param bayer_pattern: Bayer pattern (RGGB, GRBG, ...).
    :param gain_map: Optional gain map to replace gain_map_opcode. 1 or 4 channels in order: R, Gr, Gb, and B.
    :param clip: Whether to clip result image to [0, 1].
    :return: Image with gain map applied; lens shading corrected.
    """
    if gain_map is not None:
        # resize gain map, make it 4 channels, if needed
        gain_map = cv2.resize(gain_map, dsize=(raw_image.shape[1] // 2, raw_image.shape[0] // 2),
                              interpolation=cv2.INTER_LINEAR)
    elif len(gain_map_opcode) == 1:
        gain_map = gain_map_opcode[0].data['map_gain_2d']
        # resize gain map, make it 4 channels, if needed
        gain_map = cv2.resize(gain_map, dsize=(raw_image.shape[1] // 2, raw_image.shape[0] // 2),
                              interpolation=cv2.INTER_LINEAR)
        gain_map = np.tile(gain_map[..., np.newaxis], [1, 1, 4])
    else:
        # The 4 gain maps should come in the order of R G G B (verified by cfa_pattern and the top/left values)
        resized_gain_maps = [
            cv2.resize(opcode_data.data['map_gain_2d'], dsize=(raw_image.shape[1] // 2, raw_image.shape[0] // 2),
                       interpolation=cv2.INTER_LINEAR) for opcode_data in gain_map_opcode]
        gain_map = np.dstack(resized_gain_maps)

    # if gain_map_opcode:
    #     # TODO: consider other parameters
    #
    #     top = gain_map_opcode.data['top']
    #     left = gain_map_opcode.data['left']
    #     bottom = gain_map_opcode.data['bottom']
    #     right = gain_map_opcode.data['right']
    #     rp = gain_map_opcode.data['row_pitch']
    #     cp = gain_map_opcode.data['col_pitch']
    #
    #     gm_w = right - left
    #     gm_h = bottom - top

        # gain_map = cv2.resize(gain_map, dsize=(gm_w, gm_h), interpolation=cv2.INTER_LINEAR)

        # TODO
        # if top > 0:
        #     pass
        # elif left > 0:
        #     left_col = gain_map[:, 0:1]
        #     rep_left_col = np.tile(left_col, [1, left])
        #     gain_map = np.concatenate([rep_left_col, gain_map], axis=1)
        # elif bottom < raw_image.shape[0]:
        #     pass
        # elif right < raw_image.shape[1]:
        #     pass

    result_image = raw_image.copy()

    # one channel
    # result_image[::rp, ::cp] *= gain_map[::rp, ::cp]

    # per bayer channel
    upper_left_idx = [[0, 0], [0, 1], [1, 0], [1, 1]]
    bayer_pattern_idx = np.array(bayer_pattern)
    # blue channel index --> 3
    bayer_pattern_idx[bayer_pattern_idx == 2] = 3
    # second green channel index --> 2
    if bayer_pattern_idx[3] == 1:
        bayer_pattern_idx[3] = 2
    else:
        bayer_pattern_idx[2] = 2
    for c in range(4):
        i0 = upper_left_idx[c][0]
        j0 = upper_left_idx[c][1]
        result_image[i0::2, j0::2] *= gain_map[:, :, bayer_pattern_idx[c]]

    if clip:
        result_image = np.clip(result_image, 0.0, 1.0)

    return result_image


def get_opencv_demsaic_flag(cfa_pattern, output_channel_order, alg_type='VNG'):
    # using opencv edge-aware demosaicing
    if alg_type != '':
        alg_type = '_' + alg_type
    if output_channel_order == 'BGR':
        if cfa_pattern == [0, 1, 1, 2]:  # RGGB
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_BG2BGR' + alg_type)
        elif cfa_pattern == [2, 1, 1, 0]:  # BGGR
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_RG2BGR' + alg_type)
        elif cfa_pattern == [1, 0, 2, 1]:  # GRBG
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_GB2BGR' + alg_type)
        elif cfa_pattern == [1, 2, 0, 1]:  # GBRG
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_GR2BGR' + alg_type)
        else:
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_BG2BGR' + alg_type)
            print("CFA pattern not identified.")
    else:  # RGB
        if cfa_pattern == [0, 1, 1, 2]:  # RGGB
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_BG2RGB' + alg_type)
        elif cfa_pattern == [2, 1, 1, 0]:  # BGGR
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_RG2RGB' + alg_type)
        elif cfa_pattern == [1, 0, 2, 1]:  # GRBG
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_GB2RGB' + alg_type)
        elif cfa_pattern == [1, 2, 0, 1]:  # GBRG
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_GR2RGB' + alg_type)
        else:
            opencv_demosaic_flag = eval('cv2.COLOR_BAYER_BG2RGB' + alg_type)
            print("CFA pattern not identified.")
    return opencv_demosaic_flag


def denoise(image, alg_type='fgs', cfa_pattern=None):
    if alg_type == 'fgs':
        guide = (image * 255).astype(dtype=np.uint8)
        denoised_image = cv2.ximgproc.fastGlobalSmootherFilter(guide, image, 100, 0.75)
    else:
        denoised_image = image

    return denoised_image


def demosaic(bayer_image, cfa_pattern, output_channel_order='RGB', alg_type='VNG', clip=True):
    """
    Demosaic a Bayer image.
    :param bayer_image: Image in Bayer format, single channel.
    :param cfa_pattern: Bayer/CFA pattern.
    :param output_channel_order: Either RGB or BGR.
    :param alg_type: algorithm type. options: '', 'EA' for edge-aware, 'VNG' for variable number of gradients
    :param clip: Whether to clip values to [0, 1].
    :return: Demosaiced image.
    """
    if alg_type == 'VNG':
        max_val = 255
        wb_image = (bayer_image * max_val).astype(dtype=np.uint8)
    else:
        max_val = 16383
        wb_image = (bayer_image * max_val).astype(dtype=np.uint16)

    if alg_type in ['', 'EA', 'VNG']:
        opencv_demosaic_flag = get_opencv_demsaic_flag(cfa_pattern, output_channel_order, alg_type=alg_type)
        demosaiced_image = cv2.cvtColor(wb_image, opencv_demosaic_flag)
    elif alg_type == 'menon2007':
        cfa_pattern_str = "".join(["RGB"[i] for i in cfa_pattern])
        demosaiced_image = demosaicing_CFA_Bayer_Menon2007(wb_image, pattern=cfa_pattern_str)
    else:
        raise ValueError('Unsupported demosaicing algorithm, alg_type = {}'.format(alg_type))

    demosaiced_image = demosaiced_image.astype(dtype=np.float32) / max_val

    if clip:
        demosaiced_image = np.clip(demosaiced_image, 0, 1)

    return demosaiced_image


def apply_color_space_transform(image, color_matrix_1, color_matrix_2, illuminant=None, interpolate_csts=True):
    if type(color_matrix_1[0]) is Ratio:
        color_matrix_1 = ratios2floats(color_matrix_1)
    if type(color_matrix_2[0]) is Ratio:
        color_matrix_2 = ratios2floats(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(xyz2cam2, axis=1, keepdims=True)

    if interpolate_csts and illuminant is not None:
        # interpolate between CSTs based on illuminant
        cct = raw_rgb_to_cct(illuminant, xyz2cam1, xyz2cam2)
        xyz2cam_interp = interpolate_cst(xyz2cam1, xyz2cam2, cct)
        xyz2cam_interp = xyz2cam_interp / np.sum(xyz2cam_interp, axis=1, keepdims=True)
        cam2xyz_interp = np.linalg.inv(xyz2cam_interp)
        xyz_image = cam2xyz_interp[np.newaxis, np.newaxis, :, :] * image[:, :, np.newaxis, :]
    else:
        # for now, use one matrix
        # simplified matrix multiplication
        # inverse
        cam2xyz1 = np.linalg.inv(xyz2cam1)
        xyz_image = cam2xyz1[np.newaxis, np.newaxis, :, :] * 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_prophoto(xyz_image):
    xyz2pp = np.array([[1.3459433, -0.2556075, -0.0511118],
                       [-0.5445989, 1.5081673, 0.0205351],
                       [0.0000000, 0.0000000, 1.2118128]])
    xyz2pp = xyz2pp / np.sum(xyz2pp, axis=-1, keepdims=True)
    pp_image = np.matmul(xyz_image, xyz2pp.T)
    pp_image = np.clip(pp_image, 0.0, 1.0)
    return pp_image


def transform_prophoto_to_srgb(pp_image):
    # Solved from ppn2srgbn = xyz2srgbn @ xyz2ppn.inv
    pp2srgb = np.array([[1.84770997, -0.53597631, -0.31173367],
                        [-0.25511275, 1.24992018, 0.00519257],
                        [-0.0164456, -0.14912261, 1.1655682]])
    srgb_image = np.matmul(pp_image, pp2srgb.T)
    srgb_image = np.clip(srgb_image, 0.0, 1.0)
    return srgb_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 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

    if type(orientation) is list:
        orientation = orientation[0]

    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 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):
    return x ** (1.0 / 2.2)


def apply_tone_map(x, tone_curve='simple-s-curve'):
    if tone_curve == 'simple-s-curve':
        tone_mapped_image = 3 * x ** 2 - 2 * x ** 3
    else:
        # 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


def apply_local_tone_map(x, alg_type='clahe', channel_order='RGB', clahe_clip_limit=1.0, clahe_grid_size=(8, 8)):
    if alg_type == 'clahe':
        to_ycrcb_flag = eval('cv2.COLOR_{}2YCR_CB'.format(channel_order))
        from_ycrcb_flag = eval('cv2.COLOR_YCR_CB2{}'.format(channel_order))

        max_val = 65535
        x = (x * max_val).astype(dtype=np.uint16)

        y_cr_cb = cv2.cvtColor(x, to_ycrcb_flag)
        clahe = cv2.createCLAHE(clipLimit=clahe_clip_limit, tileGridSize=clahe_grid_size)
        x = np.dstack((clahe.apply(y_cr_cb[:, :, 0]), y_cr_cb[:, :, 1], y_cr_cb[:, :, 2]))
        x = cv2.cvtColor(x, from_ycrcb_flag)

        x = x.astype(dtype=np.float32) / max_val
    else:
        raise ValueError('Unsupported local tone mapping algorithm, alg_type = {}'.format(alg_type))
    return x


def process_to_save(image, out_dtype='uint8', channel_order='bgr'):
    """
    Process an RGB image to be saved with OpenCV.
    :param image: Input image.
    :param out_dtype: Target data type (e.g., 'uint8', 'uint16', ...).
    :param channel_order: Output channel order (e.g., 'bgr' for OpenCV, ...).
    :return: Processed image in the target data type and channel order.
    """
    in_dtype = str(image.dtype)

    if in_dtype != out_dtype:

        # normalize with source data type
        if in_dtype == 'uint8':
            image = image.astype('float32') / 255.0
        elif in_dtype == 'uint16':
            image = image.astype('float32') / 65535.0
        else:
            pass  # assuming float

        # quantize with target data type
        if out_dtype == 'uint8':
            max_val = 255
        elif out_dtype == 'uint16':
            max_val = 65535
        else:
            max_val = 255  # default
        image = (image * max_val).astype(out_dtype)

    # rearrange channel order, if needed
    if channel_order == 'bgr':
        image = image[:, :, ::-1]

    return image


def fix_missing_params(params):
    """
    Fix params dictionary by filling missing parameters with default values.
    :param params: Input params dictionary.
    :return: Fixed params dictionary.
    """
    params_fixed = params.copy()
    default_params = {
        'input_stage': 'raw',
        'output_stage': 'tone',
        'save_as': 'jpg',
        'white_balancer': 'default',
        'demosaicer': '',
        'denoiser': 'fgs',
        'tone_curve': 'simple-s-curve',
        'local_tone_mapping': 'clahe',  # options: 'clahe', or self-defined module
        'clahe_clip_limit': 1.0,
        'clahe_grid_size': (8, 8),
    }
    for key in default_params.keys():
        if key not in params_fixed:
            params_fixed[key] = default_params[key]
    return params_fixed

pipeline/pipeline_utils.py
SamsungLabs/graphics2raw
"""
Author(s):
Abdelrahman Abdelhamed (a.abdelhamed@samsung.com)

Copyright (c) 2023 Samsung Electronics Co., Ltd.

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import numpy as np

def RGGB2Bayer(im, _cfa_pattern=[0,1,1,2]):
    # add 1 to the second-green and the blue channel (e.g., [0, 1, 1, 2] will be [0, 1, 2, 3])
    _cfa_pattern_arr=np.asarray(_cfa_pattern)
    _cfa_pattern_arr[_cfa_pattern_arr == 2] += 1
    _cfa_pattern_arr[2:][_cfa_pattern_arr[2:] == 1] += 1
    # convert RGGB stacked image to one channel Bayer
    bayer = np.zeros((im.shape[0] * 2, im.shape[1] * 2))
    bayer[0::2, 0::2] = im[:, :, _cfa_pattern_arr[0]]
    bayer[0::2, 1::2] = im[:, :, _cfa_pattern_arr[1]]
    bayer[1::2, 0::2] = im[:, :, _cfa_pattern_arr[2]]
    bayer[1::2, 1::2] = im[:, :, _cfa_pattern_arr[3]]
    return bayer

def stack_rggb_channels(raw_image, bayer_pattern=None):
    """
    Stack the four channels of a CFA/Bayer raw image along a third dimension.
    """
    if bayer_pattern is None:
        bayer_pattern = [0, 1, 1, 2]
    height, width = raw_image.shape
    channels = []
    pattern = np.array(bayer_pattern)
    # add 1 to the second-green and the blue channel (e.g., [0, 1, 1, 2] will be [0, 1, 2, 3])
    pattern[pattern == 2] += 1
    pattern[2:][pattern[2:] == 1] += 1
    idx = [[0, 0], [0, 1], [1, 0], [1, 1]]
    for c in pattern:
        raw_image_c = raw_image[idx[c][0]:height:2, idx[c][1]:width:2].copy()
        channels.append(raw_image_c)

    # special case: channels re-ordered to [B G2 G1 R] instead of [R G1 G2 B] when
    # bayer_pattern==[G B R G]; need to flip it back.
    if bayer_pattern == [1, 2, 0, 1]:
        channels.reverse()
    channels = np.stack(channels, axis=-1)
    return channels


def stack_rgb_channels(raw_image, bayer_pattern):
    """
    Stack the four channels in a CFA/Bayer image into 3 RGB channels, averaging the two G channels.
    """
    rggb = stack_rggb_channels(raw_image, bayer_pattern)
    rgb = np.zeros((rggb.shape[0], rggb.shape[1], 3), dtype=np.float32)
    rgb[:, :, 0] = rggb[:, :, 0]
    rgb[:, :, 1] = (rggb[:, :, 1] + rggb[:, :, 2]) / 2.0
    rgb[:, :, 2] = rggb[:, :, 3]
    return rgb


def rggb_to_rgb(image_4ch, bayer_pattern):
    bayer_pattern = list(bayer_pattern)
    g1_idx = bayer_pattern.index(1)
    g2_idx = 3 if g1_idx == 0 else 2
    r_idx = bayer_pattern.index(0)
    b_idx = bayer_pattern.index(2)
    g = np.mean([image_4ch[:, :, g1_idx], image_4ch[:, :, g2_idx]], axis=0)
    rgb = np.stack([image_4ch[:, :, r_idx], g, image_4ch[:, :, b_idx]], axis=-1)
    return rgb

pipeline/raw_utils.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from model_archs.restormer import Restormer
from utils.general_utils import save_args, get_git_info, str2int_arr
import argparse
import os, time, datetime
import numpy as np
from skimage.measure import compare_psnr, compare_ssim
import cv2
import torch
from pipeline.pipeline_utils import get_metadata, get_visible_raw_image
from utils.general_utils import check_dir
from data_preparation.data_generator_denoise import post_process_stacked_images_s20fe, stack_bayer_norm_gamma_s20fe, \
    linearize_stacked_images_s20fe, post_process_stacked_images_s20fe_to_srgb
import glob


def to_tensor(img):
    img = torch.from_numpy(img.astype(np.float32))
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    return img


def from_tensor(img):
    img = img.permute(0, 2, 3, 1)
    img = img.cpu().detach().numpy()
    return np.squeeze(img)


def log(*args, **kwargs):
    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), *args, **kwargs)


def save_result(result, path):
    path = path if path.find('.') != -1 else path + '.png'
    ext = os.path.splitext(path)[-1]
    if ext in ('.txt', '.dlm'):
        np.savetxt(path, result, fmt='%2.4f')
    else:
        cv2.imwrite(path, result)


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--set_dir', default='real_dataset', type=str, help='directory of test dataset')
    parser.add_argument('--set_name', default='iso_50', type=str, help='name of test dataset')
    parser.add_argument('--model_dir', default='no_wb_model', type=str, help='directory of the model')
    parser.add_argument('--model_name', default='bestmodel.pt', type=str, help='name of the model')
    parser.add_argument('--restormer_dim', default=8, type=int, help='Restormer dim.')
    parser.add_argument('--result_dir', default='results', type=str, help='directory of results')
    parser.add_argument('--save_num', default=0, type=int, help='number of images to save')
    parser.add_argument('--save_visual', action='store_true', help='save save_num number of images')
    parser.add_argument('--save_visual_only', action='store_true', help='only save visual results')
    parser.add_argument('--save_fns', default='', type=str,
                        help='path to a file specifying which images to generate visual results for. '
                             'Used only when save_visual_only is True.')

    parser.add_argument('--exp_dir', default='./', type=str, help='directory to save experiment data to')
    args = parser.parse_args()

    to_save = []
    if args.save_visual_only:
        with open(args.save_fns, 'r') as f:
            lines = f.readlines()
            lines = [l.strip() for l in lines]
        to_save = lines

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = Restormer(inp_channels=4,
                      out_channels=4,
                      dim=args.restormer_dim,
                      num_blocks=[4, 6, 6, 8],
                      num_refinement_blocks=4,
                      heads=[1, 2, 4, 8],
                      ffn_expansion_factor=2.66,
                      bias=False,
                      LayerNorm_type='BiasFree',
                      dual_pixel_task=False)

    model.load_state_dict(torch.load(os.path.join(args.exp_dir, args.model_dir, 'models', args.model_name)))
    model = model.to(device)
    model.eval()  # Set model to evaluate mode
    print('model loaded')

    ### File IO
    set_name = args.set_name
    fullsavepath = os.path.join(args.result_dir,
                                os.path.basename(args.set_dir) + '_' + set_name + '_' + os.path.basename(
                                    args.model_dir) + '_' + args.model_name[:-3])
    check_dir(args.result_dir)
    check_dir(fullsavepath)
    save_args(args, fullsavepath)

    with open(os.path.join(fullsavepath, 'test_info.txt'), 'w') as f:
        f.write(get_git_info())

    psnrs = []
    ssims = []

    if not args.save_visual_only:
        fsrgb = open(os.path.join(fullsavepath, "results.txt"), "w")

    ### Inference
    test_dir = args.set_dir
    for c, dng_path in enumerate(sorted(glob.glob(os.path.join(args.set_dir, 'dng', set_name, '*.dng')))):
        im = os.path.basename(dng_path)
        if args.save_visual_only:
            tag = im[:-4]
            if tag not in to_save:
                print('Skipping ', tag)
                continue

        print('Processing ' + str(c + 1))
        if set_name == 'iso_50':
            x = cv2.imread(os.path.join(args.set_dir, 'clean_raw', im[:-4] + '.png'), cv2.IMREAD_UNCHANGED)
        else:
            x = get_visible_raw_image(os.path.join(args.set_dir, 'dng', set_name, im))

        meta_path = os.path.join(args.set_dir, 'dng', set_name, im)
        meta_data = get_metadata(meta_path)

        x_gamma = stack_bayer_norm_gamma_s20fe(x, clip_bot=False)
        x_gamma = to_tensor(x_gamma)

        y = np.array(cv2.imread(os.path.join(args.set_dir, 'clean_raw', im[:-4] + '.png'), cv2.IMREAD_UNCHANGED))
        y_gamma = stack_bayer_norm_gamma_s20fe(y, clip_bot=True)
        y_gamma = np.clip(y_gamma, 0, 1)

        x_gamma = x_gamma.to(device)
        start_time = time.time()
        with torch.no_grad():
            pred_gamma = model(x_gamma)  # inference
        elapsed_time = time.time() - start_time
        pred_gamma = pred_gamma.cpu()

        pred_gamma = from_tensor(pred_gamma)
        pred_gamma = np.clip(np.squeeze(pred_gamma), 0, 1)

        y_lin = linearize_stacked_images_s20fe(y_gamma)
        pred_lin = linearize_stacked_images_s20fe(pred_gamma)
        psnr_x = compare_psnr(y_lin, pred_lin, data_range=1)
        ssim_x = compare_ssim(y_lin, pred_lin, multichannel=True, data_range=1)
        psnrs.append(psnr_x)
        ssims.append(ssim_x)

        log('{0:10s} \n PSNR = {1:2.2f}dB, SSIM = {2:1.4f}, Time = {3:2.4f} seconds'.format(im, psnr_x,
                                                                                            ssim_x, elapsed_time))
        if not args.save_visual_only:
            fsrgb.write('{0:10s} : PSNR = {1:2.2f}dB, SSIM = {2:1.4f}, Time = {3:2.4f} seconds \n'.format(im, psnr_x,
                                                                                                          ssim_x,
                                                                                                          elapsed_time))

        if (args.save_visual and c < args.save_num) or args.save_visual_only:
            x_gamma = from_tensor(x_gamma).squeeze()  # drop the batch dimension, assume bs = 1
            # 4-channel images
            y_gamma_ = (255 * post_process_stacked_images_s20fe(y_gamma, meta_data['cfa_pattern'])).astype('uint8')
            pred_gamma_ = (255 * post_process_stacked_images_s20fe(pred_gamma, meta_data['cfa_pattern'])).astype('uint8')
            x_gamma_ = (255 * post_process_stacked_images_s20fe(x_gamma, meta_data['cfa_pattern'])).astype('uint8')

            name, ext = os.path.splitext(im)
            save_result(x_gamma_[:, :, [2, 1, 0]],
                        path=os.path.join(fullsavepath, name + '_input.png')) 
            save_result(pred_gamma_[:, :, [2, 1, 0]],
                        path=os.path.join(fullsavepath, name + f'_output_psnr{np.round(psnr_x, 4)}_ssim{np.round(ssim_x, 4)}.png')) 
            save_result(y_gamma_[:, :, [2, 1, 0]],
                        path=os.path.join(fullsavepath, name + '_target.png')) 
            
            y_srgb = (255 * post_process_stacked_images_s20fe_to_srgb(y_gamma, meta_data)).astype('uint8')
            pred_srgb = (255 * post_process_stacked_images_s20fe_to_srgb(pred_gamma, meta_data)).astype('uint8')
            x_srgb = (255 * post_process_stacked_images_s20fe_to_srgb(x_gamma, meta_data)).astype('uint8')

            save_result(x_srgb[:, :, [2, 1, 0]],
                        path=os.path.join(fullsavepath, name + '_input_srgb.png')) 
            save_result(pred_srgb[:, :, [2, 1, 0]],
                        path=os.path.join(fullsavepath, name + f'_output_srgb_psnr{np.round(psnr_x, 4)}_ssim{np.round(ssim_x, 4)}.png')) 
            save_result(y_srgb[:, :, [2, 1, 0]],
                        path=os.path.join(fullsavepath, name + '_target_srgb.png'))  

    psnr_avg = np.mean(psnrs)
    ssim_avg = np.mean(ssims)
    psnrs.append(psnr_avg)
    ssims.append(ssim_avg)

    print()
    log('Dataset: {0:10s} \n Avg. PSNR = {1:2.4f}dB, Avg. SSIM = {2:1.4f}'.format(set_name, psnr_avg, ssim_avg))

    if not args.save_visual_only:
        fsrgb.write(
            '\nDataset: {0:10s} \n Avg. PSNR = {1:2.4f}dB, Avg. SSIM = {2:1.4f}'.format(set_name, psnr_avg, ssim_avg))
        fsrgb.close()

test_denoise.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from model_archs.cnn import IllumEstNet
from utils.general_utils import save_args
import argparse
import os, time, datetime
import numpy as np
import cv2
import torch
from utils.general_utils import check_dir
from data_preparation.data_generator_illum_est import shuffle_files, get_gt_illum_by_fname
import glob
import pickle
from utils.img_utils import compute_ang_error


def to_tensor(img):
    img = torch.from_numpy(img.astype(np.float32))
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    return img


def from_tensor(img):
    img = img.permute(0, 2, 3, 1)
    img = img.cpu().detach().numpy()
    return np.squeeze(img)


def log(*args, **kwargs):
    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), *args, **kwargs)


def save_results(ang_err_array, out_fn=None, method=''):
    ang_err_array = np.array(ang_err_array)
    err_arr_sorted = np.sort(ang_err_array)
    fourth = int(np.round(ang_err_array.shape[0] / 4.0))
    best25_mean_err = np.round(np.mean(err_arr_sorted[:fourth]), 4)
    worst25_mean_err = np.round(np.mean(err_arr_sorted[-fourth:]), 4)
    arr_len = ang_err_array.shape[0]
    mean_err = np.round(ang_err_array.mean(), 4)
    median_err = np.round(np.median(ang_err_array), 4)

    print(arr_len, mean_err, median_err, best25_mean_err, worst25_mean_err)

    if out_fn:
        if not os.path.exists(out_fn):
            f = open(out_fn, 'w')
            f.write('exp / angular_errs, num_test, mean, median, best 25%, worst 25%\n')
        else:
            f = open(out_fn, 'a')
        f.write(f'{method}, {arr_len}, {mean_err}, {median_err}, {best25_mean_err}, {worst25_mean_err}\n')
        f.close()


if __name__ == '__main__':

    """
    Assumed dataset directory structure:
    dataset_root_dir
        camera
            method

    Assumed experiment directory structure:
    exp_dir
        exp_name
            models
            results
                bestmodel
    """

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset-dir', default='illum_est_expts/data/ours', type=str, help='folder of png images')
    parser.add_argument(
        '--illum_file', default='illum_est_expts/data/ours/gt_illum.p', type=str,
        help='path to split indices')
    parser.add_argument(
        '--split_file', default='illum_est_expts/data/SamsungNX2000_train_valid_test_split_idx.p', type=str,
        help='path to split indices')
    parser.add_argument('--exp_name', default='illum_est_expt', type=str, help='experiment name, relative to exp_dir')
    parser.add_argument('--model_name', default='bestmodel.pt', type=str, help='name of the model')
    parser.add_argument('--num_filters', type=int, default=7, help='number of filters for CNN layers ')
    parser.add_argument('--exp_dir', default='./', type=str, help='directory to save experiment data to')
    parser.add_argument('--debug', action='store_true', help='debug mode')
    args = parser.parse_args()

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    exp_name = args.exp_name

    model = IllumEstNet(in_channels=3, out_channels=3, num_filters=args.num_filters)

    model.load_state_dict(torch.load(os.path.join(args.exp_dir, exp_name, 'models', args.model_name)))
    model = model.to(device)
    model.eval()  # Set model to evaluate mode
    print('model loaded')

    full_save_path = os.path.join(args.exp_dir, exp_name, 'results', args.model_name[:-3])
    check_dir(full_save_path)
    save_args(args, full_save_path)

    ang_errs = []
    ssims = []

    camera, method = args.dataset_dir.split('/')[-2:]  # assume dataset_dir points to .../camera/method
    fname_result = f'results_{camera}.txt'
    f_result = open(os.path.join(full_save_path, fname_result), "w")

    # Data loading
    file_list = sorted(glob.glob(os.path.join(args.dataset_dir, '*.png')))  # get name list of all .png files
    file_list = shuffle_files(file_list)  # shuffle to sample from all scenes
    gt_illum_by_fname = get_gt_illum_by_fname(args.illum_file)

    split_indices = pickle.load(open(args.split_file, 'rb'))['test']
    file_list = file_list[split_indices]

    if args.debug:
        # For debugging only, check if images in the split are expected
        split_fns_fp = args.split_file[:-5] + 'fns.p'
        split_type = 'graphics_split' if args.dataset_dir.split('/')[-1] != 'real' else 'real_split'
        split_fns = pickle.load(open(split_fns_fp, 'rb'))[split_type]['test']
        for myf, gtf in zip(file_list, split_fns):
            assert os.path.basename(myf) == gtf, f'Panic! {os.path.basename(myf)}, {gtf}'
        print(f'Loaded images are correct.')
        # end of debugging

    for file in file_list:
        img = np.array(cv2.imread(file, cv2.IMREAD_UNCHANGED))[:, :, ::-1].astype(np.float32)  # img is now rgb
        img = img / 65535.0
        img = np.clip(img, 0, 1)

        fname = os.path.basename(file)[:-4]
        gt_illum = gt_illum_by_fname[fname]

        x = to_tensor(img)
        x = x.to(device)
        start_time = time.time()
        with torch.no_grad():
            y_ = model(x)  # inference
        elapsed_time = time.time() - start_time
        y_ = y_.cpu().numpy().squeeze()
        ang_err = compute_ang_error(y_, gt_illum)
        ang_errs.append(ang_err)

        log('{0:10s} \n ang_err = {1:2.2f} deg, Time = {2:2.4f} seconds, Pred: {3}, GT: {4}'.format(fname, ang_err,
                                                                                                    elapsed_time,
                                                                                                    y_, gt_illum))
        f_result.write('{0:10s} : ang_err = {1:2.2f} deg, '
                       'Time = {2:2.4f} seconds, Pred: {3}, GT: {4} \n'.format(fname,
                                                                               ang_err, elapsed_time, y_, gt_illum))

    ang_err_avg = np.mean(ang_errs)

    print()
    log('Dataset: {0:10s} \n Avg. Ang Err = {1:2.4f} deg'.format(f'{camera}_{method}', ang_err_avg))

    f_result.write('\nDataset: {0:10s} \n Avg. Ang Err = {1:2.4f} deg'.format(f'{camera}_{method}', ang_err_avg))
    f_result.close()

    aggr_result_fp = os.path.join(args.exp_dir, 'results.csv')
    index = exp_name.find(method)
    result_header = f'{camera}_{exp_name}' if camera not in exp_name else exp_name
    save_results(np.array(ang_errs), aggr_result_fp, method=result_header)

test_illum_est.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from model_archs.unet import UNet
from utils.general_utils import save_args
import argparse
import os, time, datetime
import numpy as np
from skimage.measure import compare_psnr, compare_ssim
import cv2
import torch
from pipeline.pipeline import run_pipeline
from pipeline.pipeline_utils import get_metadata, get_visible_raw_image
from utils.general_utils import check_dir
from utils.img_utils import calc_deltaE_rgb


def to_tensor(img):
    img = torch.from_numpy(img.astype(np.float32))
    img = img.unsqueeze(0).permute(0, 3, 1, 2)
    return img


def from_tensor(img):
    img = img.permute(0, 2, 3, 1)
    img = img.cpu().detach().numpy()
    return np.squeeze(img)


def log(*args, **kwargs):
    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), *args, **kwargs)


def save_result(result, path):
    path = path if path.find('.') != -1 else path + '.png'
    ext = os.path.splitext(path)[-1]
    if ext in ('.txt', '.dlm'):
        np.savetxt(path, result, fmt='%2.4f')
    else:
        cv2.imwrite(path, result)


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--set_dir', default='real_dataset', type=str, help='directory of test dataset')
    parser.add_argument('--model_dir', default='neural_isp_model', type=str, help='directory of the model')
    parser.add_argument('--model_name', default='bestmodel.pt', type=str, help='name of the model')
    parser.add_argument('--result_dir', default='results', type=str, help='directory of results')
    parser.add_argument('--save_result', action='store_true', help='save image')
    parser.add_argument('--save_num', default=0, type=int, help='number of images to save')
    parser.add_argument('--is-PS-sRGB', default=True, action='store_true', help='False: our in-house ISP sRGB like in day-to-night experiments, True: Photoshop sRGB')
    parser.add_argument('--num_filters', type=int, default=32, help='number of filters for UNet layers ')
    parser.add_argument('--exp_dir', default='./', type=str, help='directory to save experiment data to')
    parser.add_argument('--save_visual_only', action='store_true', help='True: only save visual results; save_result '
                                                                        'and save_num will not be used.')
    parser.add_argument('--save_fns', default='', type=str,
                        help='path to a file specifying which images to generate visual results for. '
                             'Used only when save_visual_only is True.')
    args = parser.parse_args()

    to_save = []
    if args.save_visual_only:
        with open(args.save_fns, 'r') as f:
            lines = f.readlines()
            lines = [l.strip() for l in lines]
        to_save = lines

    if args.is_PS_sRGB:
        params = {
            'save_as': 'png',  # options: 'jpg', 'png', 'tif', etc.
            'white_balancer': 'default',  # options: default, or self-defined module
            'demosaicer': '',  # options: '' for simple interpolation,
            #          'EA' for edge-aware,
            #          'VNG' for variable number of gradients,
            #          'menon2007' for Menon's algorithm
            'tone_curve': 'simple-s-curve',  # options: 'simple-s-curve', 'default', or self-defined module
            'output_stage': 'default_cropping',
        }

        stages = ['raw', 'active_area_cropping', 'linearization', 'normal', 'lens_shading_correction',
                'white_balance', 'demosaic', 'default_cropping']
    else:
        params = {
            'save_as': 'png',  # options: 'jpg', 'png', 'tif', etc.
            'white_balancer': 'default',  # options: default, or self-defined module
            'demosaicer': '',  # options: '' for simple interpolation,
            #          'EA' for edge-aware,
            #          'VNG' for variable number of gradients,
            #          'menon2007' for Menon's algorithm
            'tone_curve': 'simple-s-curve',  # options: 'simple-s-curve', 'default', or self-defined module
            'output_stage': 'demosaic',
        }

        stages = ['raw', 'active_area_cropping', 'linearization', 'normal', 'white_balance',
                  'demosaic']

    set_names_list = ['iso_50']

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = UNet(in_channels=3, out_channels=3, init_features=args.num_filters)
    
    model_path = os.path.join(args.exp_dir, args.model_dir, 'models', args.model_name)

    model.load_state_dict(torch.load(model_path))
    model = model.to(device)
    model.eval()  # Set model to evaluate mode
    print('model loaded')

    for folder_index, set_cur in enumerate(set_names_list):
        print(set_cur)

        fullsavepath = os.path.join(args.result_dir,
                                   os.path.basename(args.set_dir) + '_' + set_cur + '_' + os.path.basename(args.model_dir) + '_' + args.model_name[:-3])

        check_dir(args.result_dir)
        check_dir(fullsavepath)
        save_args(args, fullsavepath)

        psnrs = []
        ssims = []
        deltaEs = []

        if not args.save_visual_only:
            fsrgb = open(os.path.join(fullsavepath, "results.txt"), "w")

        for c, im in enumerate(sorted(os.listdir(os.path.join(args.set_dir, 'dng', set_cur)))):
            if args.save_visual_only:
                tag = im[:-4]
                if tag not in to_save:
                    continue

            print('Processing ' + str(c + 1) + ' ' + os.path.basename(im))
            if set_cur == 'iso_50':
                x = cv2.imread(os.path.join(args.set_dir, 'clean_raw', im[:-4]+'.png'), cv2.IMREAD_UNCHANGED)
            else:
                x = get_visible_raw_image(os.path.join(args.set_dir, 'dng', set_cur, im))

            meta_path = os.path.join(args.set_dir, 'dng', set_cur, im)
            meta_data_org = get_metadata(meta_path)

            x = run_pipeline(x, params=params, metadata=meta_data_org, stages=stages)

            x = x ** (1 / 2.2)

            x = to_tensor(x)

            y = cv2.imread(os.path.join(args.set_dir, 'clean', im[:-4]+'.png'), cv2.IMREAD_UNCHANGED)
            y = np.array(cv2.cvtColor(y, cv2.COLOR_BGR2RGB), dtype=np.float32)
            if y.shape[0] > y.shape[1]:
                y = np.rot90(y)
            y = y.astype('uint8')

            x = x.to(device)
            start_time = time.time()
            with torch.no_grad():
                y_ = model(x)  # inference
            elapsed_time = time.time() - start_time
            y_ = y_.cpu()

            y_ = (y_).permute(0, 2, 3, 1).numpy()
            y_ = np.clip(np.squeeze(y_), 0, 1)
            y_ = (255 * y_).astype('uint8')

            psnr_x_ = compare_psnr(y, y_, data_range=255)
            ssim_x_ = compare_ssim(y, y_, multichannel=True, data_range=255)
            deltaE_x_ = calc_deltaE_rgb(y, y_)
            psnrs.append(psnr_x_)
            ssims.append(ssim_x_)
            deltaEs.append(deltaE_x_)

            if (args.save_result and c < args.save_num) or args.save_visual_only:
                name, ext = os.path.splitext(im)
                x_ = x.permute(0, 2, 3, 1).cpu().numpy().squeeze()  # drop the batch dimension, assume bs = 1
                x_ = (255 * x_).astype('uint8')
                save_result(x_[:, :, [2, 1, 0]],
                            path=os.path.join(fullsavepath, name + '_input.png'))
                save_result(y_[:, :, [2, 1, 0]],
                            path=os.path.join(fullsavepath, name + f'_output_psnr{np.round(psnr_x_, 4)}_ssim{np.round(ssim_x_, 4)}.png'))
                save_result(y[:, :, [2, 1, 0]],
                            path=os.path.join(fullsavepath, name + '_target.png'))

            log('{0:10s} \n PSNR = {1:2.2f}dB, SSIM = {2:1.4f}, DeltaE = {3:2.4f} Time = {4:2.4f} seconds'.format(im, psnr_x_,
                                                                                            ssim_x_, deltaE_x_, elapsed_time))
            if not args.save_visual_only:
                fsrgb.write('{0:10s} : PSNR = {1:2.2f}dB, SSIM = {2:1.4f}, DeltaE = {3:2.4f} Time = {4:2.4f} seconds \n'.format(im, psnr_x_,
                                                                                                     ssim_x_, deltaE_x_, elapsed_time))

        psnr_avg = np.mean(psnrs)
        ssim_avg = np.mean(ssims)
        deltaE_avg = np.mean(deltaEs)
        psnrs.append(psnr_avg)
        ssims.append(ssim_avg)

        print()
        log('Dataset: {0:10s} \n Avg. PSNR = {1:2.4f}dB, Avg. SSIM = {2:1.4f}, Avg. DeltaE = {3:2.4f}'.format(set_cur, psnr_avg, ssim_avg, deltaE_avg))

        if not args.save_visual_only:
            fsrgb.write('\nDataset: {0:10s} \n Avg. PSNR = {1:2.4f}dB, Avg. SSIM = {2:1.4f}, Avg. DeltaE = {3:2.4f}'.format(set_cur, psnr_avg, ssim_avg, deltaE_avg))
            fsrgb.close()

test_neural_isp.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from model_archs.restormer import Restormer
from utils.general_utils import save_args, get_git_info
from utils.torch_utils import torchvision_visualize_raw
import torch
from torch.nn.modules.loss import _Loss
import torch.optim as optim
from data_preparation.dataset_denoise import DatasetRAW
from torch.optim import lr_scheduler
import time
import os
import argparse

from torch.utils.tensorboard import SummaryWriter


def parse_args():
    parser = argparse.ArgumentParser(description='Day-to-night train')
    parser.add_argument(
        '--data_dir', default='graphics_dataset', type=str, help='folder of training and validation images')
    parser.add_argument(
        '--savefoldername', default='models', type=str, help='folder to save trained models to')
    parser.add_argument(
        '--exp_dir', default='./', type=str, help='directory to save experiment data to')
    parser.add_argument(
        '--patch_size', type=int, default=64, help='patch size')
    parser.add_argument(
        '--stride', type=int, default=64, help='stride when cropping patches')
    parser.add_argument(
        '--batch_size', type=int, default=128, help='batch size')
    parser.add_argument(
        '--lr', type=float, default=0.001, help='learning rate')
    parser.add_argument(
        '--milestones', default='400', type=str, help='milestones as comma separated string')
    parser.add_argument(
        '--scheduler', default='MultiStepLR', type=str, help='MultiStepLR')
    parser.add_argument(
        '--num_epochs', type=int, default=500, help='number of epochs')
    parser.add_argument(
        '--loss', type=str, default='sum_squared_error', help='loss function')
    parser.add_argument(
        '--tboard_freq', type=int, default=200, help='frequency of writing to tensorboard')
    parser.add_argument('--preload_on_cuda', default=False, action='store_true',
                        help='False: load each batch on cuda, True: load all data directly on cuda')
    parser.add_argument('--iso', default=3200, type=int, help='target ISO')
    parser.add_argument('--data_type', default='graphics_raw', type=str,
                        help='Type of experiment. '
                             'graphics_raw: graphics data inverted to raw by various methods; '
                             'real: real S20FE captures.')
    parser.add_argument('--restormer_dim', default=8, type=int, help='Restormer dim.')
    parser.add_argument('--model_save_freq', type=int, default=100, help='save model per model_save_freq epochs')
    args = parser.parse_args()

    print(args)

    return args


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def mypsnr(img1, img2):
    mse = torch.mean(((img1 * 255.0).floor() - (img2 * 255.0).floor()) ** 2, dim=[1, 2, 3])
    mse[torch.nonzero((mse == 0), as_tuple=True)] = 0.05
    psnrout = torch.mean(20 * torch.log10(255.0 / torch.sqrt(mse)))
    return psnrout


class sum_squared_error(_Loss):
    """
    Definition: sum_squared_error = 1/2 * nn.MSELoss(reduction = 'sum')
    The backward is defined as: input-target
    """

    def __init__(self, size_average=None, reduce=None, reduction='sum'):
        super(sum_squared_error, self).__init__(size_average, reduce, reduction)

    def forward(self, input, target):
        return torch.nn.functional.mse_loss(input, target, size_average=None, reduce=None, reduction='sum').div_(2)


def main(args):
    milestones = [item for item in args.milestones.split(',')]
    for i in range(len(milestones)):
        milestones[i] = int(milestones[i])

    savefoldername = args.savefoldername
    exp_dir = args.exp_dir
    tb_dir = os.path.join(exp_dir, savefoldername, 'tensorboard')

    os.makedirs(tb_dir, exist_ok=True)
    writers = [SummaryWriter(os.path.join(tb_dir, savefoldername))]

    modsavepath = os.path.join(exp_dir, savefoldername, 'models')
    if not (os.path.exists(modsavepath) and os.path.isdir(modsavepath)):
        os.makedirs(modsavepath)

    save_args(args, modsavepath)

    image_datasets = {
        x: DatasetRAW(os.path.join(args.data_dir, x), args.batch_size, args.patch_size, args.stride,
                      args.preload_on_cuda,
                      iso=args.iso, data_type=args.data_type)
        for x in ['train', 'val']}

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=args.batch_size,
                                                  shuffle=True, num_workers=0)
                   for x in ['train', 'val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}

    epoch_loss = {x: 0.0 for x in ['train', 'val']}
    epoch_psnr = {x: 0.0 for x in ['train', 'val']}

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = Restormer(inp_channels=4,
                      out_channels=4,
                      dim=args.restormer_dim,
                      num_blocks=[4, 6, 6, 8],
                      num_refinement_blocks=4,
                      heads=[1, 2, 4, 8],
                      ffn_expansion_factor=2.66,
                      bias=False,
                      LayerNorm_type='BiasFree',
                      dual_pixel_task=False)

    num_params = count_parameters(model)

    model = model.to(device)
    if args.loss == 'sum_squared_error':
        criterion = sum_squared_error()
    elif args.loss == 'l1':
        criterion = torch.nn.L1Loss()
    else:
        raise Exception('Unrecognized loss function: ', args.loss)
        
    params = model.parameters()
    optimizer = optim.Adam(params, lr=args.lr)
    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.1)

    # training loop starts here
    since = time.time()

    best_loss = 10 ** 6
    best_psnr = 0.0
    best_epoch = 0

    for epoch in range(args.num_epochs):
        print('Epoch {}/{}'.format(epoch, args.num_epochs - 1))
        print('-' * 10)

        running_loss_tboard = 0.0
        running_psnr_tboard = 0.0

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_psnr = 0.0

            # Iterate over data.
            for i, (inputs, targets) in enumerate(dataloaders[phase]):

                if not args.preload_on_cuda:
                    inputs = inputs.to(device)
                    targets = targets.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):

                    outputs = model(inputs)
                    loss = criterion(outputs, targets)
                    psnrout = mypsnr(outputs, targets)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                        running_loss_tboard += loss.item()
                        running_psnr_tboard += psnrout.item()
                        if (i + 1) % args.tboard_freq == 0:  # every tboard_freq mini-batches...

                            # ...log the running loss
                            for writer in writers:
                                writer.add_scalar('iter_loss',
                                                  running_loss_tboard / args.tboard_freq,
                                                  epoch * len(dataloaders[phase]) + i)

                                writer.add_scalar('iter_psnr',
                                                  running_psnr_tboard / args.tboard_freq,
                                                  epoch * len(dataloaders[phase]) + i)

                            running_loss_tboard = 0.0
                            running_psnr_tboard = 0.0

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_psnr += psnrout.item() * inputs.size(0)

            if phase == 'train':
                scheduler.step()

            epoch_loss[phase] = running_loss / dataset_sizes[phase]
            epoch_psnr[phase] = running_psnr / dataset_sizes[phase]

            if phase == 'val':
                # ...log the running loss
                for writer in writers:
                    writer.add_scalars('epoch_loss',
                                       {'train': epoch_loss['train'], 'val': epoch_loss['val']},
                                       epoch + 1)

                    writer.add_scalars('epoch_psnr',
                                       {'train': epoch_psnr['train'], 'val': epoch_psnr['val']},
                                       epoch + 1)
                    if (epoch + 1) % (args.num_epochs // 3) == 0:
                        img_grid = torchvision_visualize_raw([inputs, outputs, targets])
                        writer.add_image('val_epoch_' + str(epoch), img_grid)

            print('{} Loss: {:.6f} PSNR: {:.4f}'.format(
                phase, epoch_loss[phase], epoch_psnr[phase]))

            # save the model
            if phase == 'val' and epoch_loss[phase] < best_loss:
                best_loss = epoch_loss[phase]
                best_psnr = epoch_psnr[phase]
                best_epoch = epoch
                torch.save(model.state_dict(), os.path.join(modsavepath, 'bestmodel.pt'))

        # save the model
        if (epoch + 1) % args.model_save_freq == 0:
            torch.save(model.state_dict(), os.path.join(modsavepath, f'model_ep_{epoch + 1}.pt'))

        print()

    time_elapsed = time.time() - since
    train_str = 'Training complete in {:.0f}m {:.0f}s\n'.format(
        time_elapsed // 60, time_elapsed % 60)
    train_str += 'Best val loss: {:4f}\n'.format(best_loss)
    train_str += 'Best val psnr: {:4f}\n'.format(best_psnr)
    train_str += 'Best epoch: {}\n'.format(best_epoch)
    train_str += f'Dataset lengths: {dataset_sizes}\n'
    train_str += f'Number of model trainable parameters: {num_params}\n'
    train_str += f'{get_git_info()}\n'
    with open(os.path.join(modsavepath, 'train_info.txt'), 'w') as f:
        f.write(train_str)


if __name__ == "__main__":
    args = parse_args()
    main(args)

train_denoise.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from model_archs.cnn import IllumEstNet
from utils.general_utils import save_args
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from data_preparation.dataset_illum_est import DatasetIllumEst
from torch.optim import lr_scheduler
import time
import os
import argparse
from torch.utils.tensorboard import SummaryWriter


def parse_args():

    parser = argparse.ArgumentParser(description='Day-to-night train')
    parser.add_argument(
        '--dataset-dir', default='illum_est_expts/data/ours', type=str, help='folder of png images')
    parser.add_argument(
        '--illum_file', default='illum_est_expts/data/ours/gt_illum.p', type=str,
        help='path to split indices')
    parser.add_argument(
        '--split_file', default='illum_est_expts/data/SamsungNX2000_train_valid_test_split_idx.p', type=str,
        help='path to split indices')
    parser.add_argument(
        '--savefoldername', default='models', type=str, help='folder to save trained models to')
    parser.add_argument(
        '--exp_dir', default='./', type=str, help='directory to save experiment data to')
    parser.add_argument(
        '--patch-size', type=int, default=64, help='patch size')
    parser.add_argument(
        '--stride', type=int, default=64, help='stride when cropping patches')
    parser.add_argument(
        '--batch-size', type=int, default=128, help='batch size')
    parser.add_argument(
        '--lr', type=float, default=0.001, help='learning rate')
    parser.add_argument(
        '--milestones', default='', type=str, help='milestones as comma separated string')
    parser.add_argument(
        '--num-epochs', type=int, default=500, help='number of epochs')
    parser.add_argument(
        '--num_filters', type=int, default=7, help='number of filters for CNN layers ')
    parser.add_argument(
        '--tboard-freq', type=int, default=200, help='frequency of writing to tensorboard')
    parser.add_argument('--on-cuda', default=False, action='store_true', help='False: load each batch on cuda, True: load all data directly on cuda')
    parser.add_argument(
        '--model_save_freq', type=int, default=50, help='save model per model_save_freq epochs')
    parser.add_argument(
        '--tb_im_save_freq', type=int, default=500, help='save tensorboard images per tb_im_save_freq epochs')
    args = parser.parse_args()
    print(args)

    return args


EPS = 1e-9
PI = 22.0 / 7.0


def angular_loss(predicted, gt, shrink=True):
    predicted = predicted.reshape((-1, 3))
    gt = gt.reshape((-1, 3))
    cossim = torch.clamp(torch.sum(predicted * gt, dim=1) / (
            torch.norm(predicted, dim=1) * torch.norm(gt, dim=1) + EPS), -1, 1.)
    if shrink:
        angle = torch.acos(cossim * 0.9999999)
    else:
        angle = torch.acos(cossim)
    a_error = 180.0 / PI * angle
    a_error = torch.sum(a_error) / a_error.shape[0]
    return a_error


def count_parameters(model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)


def main(args):

    milestones = [int(item) for item in args.milestones.split(',')] if args.milestones else []

    savefoldername = args.savefoldername
    exp_dir = args.exp_dir
    tb_dir = os.path.join(exp_dir, savefoldername, 'tensorboard')

    os.makedirs(tb_dir, exist_ok=True)
    writers = [SummaryWriter(os.path.join(tb_dir, savefoldername))]

    modsavepath = os.path.join(exp_dir, savefoldername, 'models')
    if not (os.path.exists(modsavepath) and os.path.isdir(modsavepath)):
        os.makedirs(modsavepath)

    save_args(args, modsavepath)

    image_datasets = {
            x: DatasetIllumEst(dataset_dir=args.dataset_dir,
                               illum_file=args.illum_file,
                               split_file=args.split_file,
                               split=x,
                               batch_size=args.batch_size,
                               patch_size=args.patch_size,
                               stride=args.stride,
                               on_cuda=args.on_cuda)
            for x in ['train', 'val']}

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=args.batch_size,
                                                  shuffle=True, num_workers=0)
                   for x in ['train', 'val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}

    epoch_loss = {x: 0.0 for x in ['train', 'val']}
    epoch_ang_err = {x: 0.0 for x in ['train', 'val']}

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = IllumEstNet(in_channels=3, out_channels=3, num_filters=args.num_filters)
    num_params = count_parameters(model)
    model = model.to(device)

    criterion = nn.L1Loss()

    params = model.parameters()
    optimizer = optim.Adam(params, lr=args.lr)
    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.1)

    # training loop starts here
    since = time.time()

    best_loss = 10 ** 6
    best_ang_err = 0.0
    best_epoch = 0

    for epoch in range(args.num_epochs):
        print('Epoch {}/{}'.format(epoch, args.num_epochs - 1))
        print('-' * 10)

        running_loss_tboard = 0.0
        running_ang_err_tboard = 0.0

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_ang_err = 0.0

            # Iterate over data.
            for i, (inputs, targets) in enumerate(dataloaders[phase]):

                if not args.on_cuda:
                    inputs = inputs.to(device)
                    targets = targets.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):

                    outputs = model(inputs)
                    loss = criterion(outputs, targets)
                    ang_err_out = angular_loss(outputs, targets)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                        running_loss_tboard += loss.item()
                        running_ang_err_tboard += ang_err_out.item()
                        if (i+1) % args.tboard_freq == 0:  # every tboard_freq mini-batches...

                            # ...log the running loss
                            for writer in writers:
                                writer.add_scalar('loss',
                                                  running_loss_tboard / args.tboard_freq,
                                                  epoch * len(dataloaders[phase]) + i)

                                writer.add_scalar('ang_err',
                                                  running_ang_err_tboard / args.tboard_freq,
                                                  epoch * len(dataloaders[phase]) + i)

                            running_loss_tboard = 0.0
                            running_ang_err_tboard = 0.0

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_ang_err += ang_err_out.item() * inputs.size(0)

            if phase == 'train':
                scheduler.step()

            epoch_loss[phase] = running_loss / dataset_sizes[phase]
            epoch_ang_err[phase] = running_ang_err / dataset_sizes[phase]

            if phase == 'val':
                outputs_im = outputs[..., None, None] * torch.ones((args.batch_size, 3, args.patch_size,
                                                                    args.patch_size)).to(device)
                target_ims = targets[..., None, None] * torch.ones((args.batch_size, 3, args.patch_size,
                                                                    args.patch_size)).to(device)

                img_grid = torchvision.utils.make_grid(torch.cat((inputs[:, 0:3, :, :], outputs_im,
                                                                  target_ims), 2), normalize=True, range=(0, 1))
                # ...log the running loss
                for writer in writers:
                    writer.add_scalars('loss',
                                       {'train': epoch_loss['train'], 'val': epoch_loss['val']},
                                       (epoch + 1) * len(dataloaders['train']))

                    writer.add_scalars('ang_err',
                                       {'train': epoch_ang_err['train'], 'val': epoch_ang_err['val']},
                                       (epoch + 1) * len(dataloaders['train']))
                    if (epoch + 1) % args.tb_im_save_freq == 0:
                        writer.add_image('val_epoch_' + str(epoch), img_grid)

            print('{} Loss: {:.6f} Ang Err: {:.4f}'.format(
                phase, epoch_loss[phase], epoch_ang_err[phase]))

            # save the model
            if phase == 'val' and epoch_loss[phase] < best_loss:
                best_loss = epoch_loss[phase]
                best_ang_err = epoch_ang_err[phase]
                best_epoch = epoch
                torch.save(model.state_dict(), os.path.join(modsavepath, 'bestmodel.pt'))

        # save the model
        if (epoch + 1) % args.model_save_freq == 0:
            torch.save(model.state_dict(), os.path.join(modsavepath, f'model_ep_{epoch + 1}.pt'))

        print()

    time_elapsed = time.time() - since
    train_str = 'Training complete in {:.0f}m {:.0f}s\n'.format(
        time_elapsed // 60, time_elapsed % 60)
    train_str += 'Best val loss: {:4f}\n'.format(best_loss)
    train_str += 'Best val ang err: {:4f}\n'.format(best_ang_err)
    train_str += 'Best epoch: {}\n'.format(best_epoch)
    train_str += f'Number of model trainable parameters: {num_params}\n'
    with open(os.path.join(modsavepath, 'train_info.txt'), 'w') as f:
        f.write(train_str)


if __name__ == "__main__":

    args = parse_args()
    main(args)

train_illum_est.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

from model_archs.unet import UNet
from utils.general_utils import save_args
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from data_preparation.dataset_neural_isp import DatasetRAW
from torch.optim import lr_scheduler
import time
import os
import argparse

from torch.utils.tensorboard import SummaryWriter


def parse_args():

    parser = argparse.ArgumentParser(description='Day-to-night train')
    parser.add_argument(
        '--data-dir', default='graphics_dataset', type=str, help='folder of training and validation images')
    parser.add_argument(
        '--savefoldername', default='models', type=str, help='folder to save trained models to')
    parser.add_argument(
        '--exp_dir', default='./', type=str, help='directory to save experiment data to')
    parser.add_argument(
        '--which-input', default='clean_raw', type=str, help='clean_raw or noisy_raw')
    parser.add_argument(
        '--wb-illum', default='asn', type=str, help='asn (as-shot-neutral) or avg (average)')
    parser.add_argument(
        '--patch-size', type=int, default=64, help='patch size')
    parser.add_argument(
        '--stride', type=int, default=64, help='stride when cropping patches')
    parser.add_argument(
        '--batch-size', type=int, default=128, help='batch size')
    parser.add_argument(
        '--lr', type=float, default=0.001, help='learning rate')
    parser.add_argument(
        '--milestones', default='400', type=str, help='milestones as comma separated string')
    parser.add_argument(
        '--num-epochs', type=int, default=500, help='number of epochs')
    parser.add_argument(
        '--num_filters', type=int, default=32, help='number of filters for UNet layers ')
    parser.add_argument(
        '--tboard-freq', type=int, default=200, help='frequency of writing to tensorboard')
    parser.add_argument('--on-cuda', default=False, action='store_true', help='False: load each batch on cuda, True: load all data directly on cuda')

    parser.add_argument('--is-PS-sRGB', default=True, action='store_true', help='False: our in-house ISP sRGB like in day-to-night experiments, True: Photoshop sRGB')
    parser.add_argument(
        '--model_save_freq', type=int, default=50, help='save model per model_save_freq epochs')
    args = parser.parse_args()
    print(args)

    return args


def mypsnr(img1, img2):
    mse = torch.mean(((img1 * 255.0).floor() - (img2 * 255.0).floor()) ** 2, dim=[1, 2, 3])
    mse[torch.nonzero((mse == 0), as_tuple=True)] = 0.05
    psnrout = torch.mean(20 * torch.log10(255.0 / torch.sqrt(mse)))
    return psnrout


def main(args):

    milestones = [item for item in args.milestones.split(',')]
    for i in range(len(milestones)):
        milestones[i] = int(milestones[i])

    savefoldername = args.savefoldername
    exp_dir = args.exp_dir
    tb_dir = os.path.join(exp_dir, savefoldername, 'tensorboard')

    os.makedirs(tb_dir, exist_ok=True)
    writers = [SummaryWriter(os.path.join(tb_dir, savefoldername))]

    modsavepath = os.path.join(exp_dir, savefoldername, 'models')
    if not (os.path.exists(modsavepath) and os.path.isdir(modsavepath)):
        os.makedirs(modsavepath)

    save_args(args, modsavepath)

    image_datasets = {
            x: DatasetRAW(os.path.join(args.data_dir, x), args.batch_size, args.patch_size, args.stride, args.wb_illum,
                          args.on_cuda, args.which_input, args.is_PS_sRGB, data_mode=x)
            for x in ['train', 'val']}

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=args.batch_size,
                                                  shuffle=True, num_workers=0)
                   for x in ['train', 'val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}

    epoch_loss = {x: 0.0 for x in ['train', 'val']}
    epoch_psnr = {x: 0.0 for x in ['train', 'val']}

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = UNet(in_channels=3, out_channels=3, init_features=args.num_filters)
    model = model.to(device)

    criterion = nn.L1Loss()

    params = model.parameters()
    optimizer = optim.Adam(params, lr=args.lr)
    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.1)

    # training loop starts here
    since = time.time()

    best_loss = 10 ** 6
    best_psnr = 0.0
    best_epoch = 0

    for epoch in range(args.num_epochs):
        print('Epoch {}/{}'.format(epoch, args.num_epochs - 1))
        print('-' * 10)

        running_loss_tboard = 0.0
        running_psnr_tboard = 0.0

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_psnr = 0.0

            # Iterate over data.
            for i, (inputs, targets) in enumerate(dataloaders[phase]):

                if not args.on_cuda:
                    inputs = inputs.to(device)
                    targets = targets.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):

                    outputs = model(inputs)
                    loss = criterion(outputs, targets)
                    # psnrout = mypsnr((torch.clip(outputs, 0, 1)), targets)
                    psnrout = mypsnr(outputs, targets)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                        running_loss_tboard += loss.item()
                        running_psnr_tboard += psnrout.item()
                        if (i+1) % args.tboard_freq == 0:  # every tboard_freq mini-batches...

                            # ...log the running loss
                            for writer in writers:
                                writer.add_scalar('iter_loss',
                                                  running_loss_tboard / args.tboard_freq,
                                                  epoch * len(dataloaders[phase]) + i)

                                writer.add_scalar('iter_psnr',
                                                  running_psnr_tboard / args.tboard_freq,
                                                  epoch * len(dataloaders[phase]) + i)

                            running_loss_tboard = 0.0
                            running_psnr_tboard = 0.0

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_psnr += psnrout.item() * inputs.size(0)
            if phase == 'train':
                scheduler.step()

            epoch_loss[phase] = running_loss / dataset_sizes[phase]
            epoch_psnr[phase] = running_psnr / dataset_sizes[phase]

            if phase == 'val':
                img_grid = torchvision.utils.make_grid(torch.cat((inputs[:, 0:3, :, :], outputs,
                                                                  targets), 2), normalize=True, range=(0, 1))
                # ...log the running loss
                for writer in writers:
                    writer.add_scalars('epoch_loss',
                                       {'train': epoch_loss['train'], 'val': epoch_loss['val']},
                                       epoch + 1)

                    writer.add_scalars('epoch_psnr',
                                       {'train': epoch_psnr['train'], 'val': epoch_psnr['val']},
                                       epoch + 1)
                    if (epoch + 1) % 1000 == 0:
                        writer.add_image('val_epoch_' + str(epoch), img_grid)

            print('{} Loss: {:.6f} PSNR: {:.4f}'.format(
                phase, epoch_loss[phase], epoch_psnr[phase]))

            # save the model
            if phase == 'val' and epoch_loss[phase] < best_loss:
                best_loss = epoch_loss[phase]
                best_psnr = epoch_psnr[phase]
                best_epoch = epoch
                torch.save(model.state_dict(), os.path.join(modsavepath, 'bestmodel.pt'))

        # save the model
        if (epoch + 1) % args.model_save_freq == 0:
            torch.save(model.state_dict(), os.path.join(modsavepath, f'model_ep_{epoch + 1}.pt'))

        print()

    time_elapsed = time.time() - since
    train_str = 'Training complete in {:.0f}m {:.0f}s\n'.format(
        time_elapsed // 60, time_elapsed % 60)
    train_str += 'Best val loss: {:4f}\n'.format(best_loss)
    train_str += 'Best val psnr: {:4f}\n'.format(best_psnr)
    train_str += 'Best epoch: {}\n'.format(best_epoch)
    with open(os.path.join(modsavepath, 'train_info.txt'), 'w') as f:
        f.write(train_str)


if __name__ == "__main__":

    args = parse_args()
    main(args)

train_neural_isp.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import os
import errno
import subprocess
import sys


def check_dir(_path):
    if not os.path.exists(_path):
        try:
            os.makedirs(_path)
        except OSError as exc:  # Guard against race condition
            if exc.errno != errno.EEXIST:
                raise


def save_args(args, save_dir):
    """
    Source: https://github.com/VITA-Group/EnlightenGAN/blob/master/options/base_options.py
    EnlightenGAN base_options.py
    """
    args = vars(args)
    file_name = os.path.join(save_dir, 'args.txt')
    with open(file_name, 'wt') as opt_file:
        opt_file.write('------------ Options -------------\n')
        for k, v in sorted(args.items()):
            opt_file.write('%s: %s\n' % (str(k), str(v)))
        opt_file.write('-------------- End ----------------\n')

        opt_file.write('\n------------------------------\n')
        opt_file.write('Shell command:\n')
        opt_file.write(get_command())


def get_git_revision_hash() -> str:
    """
    Source: https://stackoverflow.com/questions/14989858/get-the-current-git-hash-in-a-python-script
    :return:
    """
    return subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('ascii').strip()


def get_git_revision_branch() -> str:
    return subprocess.check_output(['git', 'branch']).decode('ascii').strip()


def get_git_info() -> str:
    current_hash = get_git_revision_hash()
    current_branch = get_git_revision_branch()
    git_info = f'Git Info:\nCurrent commit: {current_hash}\nBranches:\n {current_branch}'
    return git_info


def str2int_arr(arr):
    # Parse comma-splited integer array
    return [int(e) for e in arr.split(',')]


def get_command() -> str:
    return " ".join(sys.argv[:])

utils/general_utils.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Abhijith Punnappurath (abhijith.p@samsung.com)
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import numpy as np
import cv2
from skimage import color

from pipeline.pipeline_utils import get_image_tags, get_values, get_image_ifds, ratios2floats, raw_rgb_to_cct, \
    interpolate_cst
from exifread.utils import Ratio
from fractions import Fraction
from binascii import hexlify
from scipy.spatial import Delaunay


def compute_ang_error(source, target):
    source = np.reshape(source, [-1, 3]).astype(np.float32)
    target = np.reshape(target, [-1, 3]).astype(np.float32)
    source_norm = np.sqrt(np.sum(np.power(source, 2), 1))
    target_norm = np.sqrt(np.sum(np.power(target, 2), 1))
    norm = source_norm * target_norm
    L = np.shape(norm)[0]
    inds = norm != 0
    angles = np.sum(source[inds, :] * target[inds, :], 1) / norm[inds]
    angles[angles > 1] = 1
    f = np.arccos(angles)
    f[np.isnan(f)] = 0
    f = f * 180 / np.pi
    return sum(f) / L


def imresize(img, scale_percent):
    width = int(img.shape[1] * scale_percent / 100)
    height = int(img.shape[0] * scale_percent / 100)
    dim = (width, height)

    # resize image
    return cv2.resize(img, dim, interpolation=cv2.INTER_LINEAR)


def get_camera_calibration_matrices(tags, ifds):
    possible_keys_1 = ['Image Tag 0xC623', 'CameraCalibration1', 'Image CameraCalibration1']
    camera_calibration_matrix_1 = get_values(tags, possible_keys_1)
    possible_keys_2 = ['Image Tag 0xC624', 'CameraCalibration2', 'Image CameraCalibration2']
    camera_calibration_matrix_2 = get_values(tags, possible_keys_2)
    return camera_calibration_matrix_1, camera_calibration_matrix_2


def get_forward_matrices(tags, ifds):
    possible_keys_1 = ['Image Tag 0xC714', 'ForwardMatrix1', 'Image ForwardMatrix1']
    forward_matrix_1 = get_values(tags, possible_keys_1)
    possible_keys_2 = ['Image Tag 0xC715', 'ForwardMatrix2', 'Image ForwardMatrix2']
    forward_matrix_2 = get_values(tags, possible_keys_2)
    return forward_matrix_1, forward_matrix_2


def get_extra_tags(image_path, metadata):
    ## inputs are path to dng file, and current metadata from pipeline
    tags = get_image_tags(image_path)
    ifds = get_image_ifds(image_path)

    camera_calibration_matrix_1, camera_calibration_matrix_2 = get_camera_calibration_matrices(tags, ifds)
    metadata['camera_calibration_matrix_1'] = camera_calibration_matrix_1
    metadata['camera_calibration_matrix_2'] = camera_calibration_matrix_2

    forward_matrix_1, forward_matrix_2 = get_forward_matrices(tags, ifds)
    metadata['forward_matrix_1'] = forward_matrix_1
    metadata['forward_matrix_2'] = forward_matrix_2

    return metadata


def white_balance_3channel(normalized_image, as_shot_neutral, clip=True):
    white_balanced_image = normalized_image / as_shot_neutral
    if clip:
        white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)
    return white_balanced_image


def get_wb_as_matrix(as_shot_neutral):
    if type(as_shot_neutral[0]) is Ratio:
        as_shot_neutral = ratios2floats(as_shot_neutral)
    return np.diag(1 / np.asarray(as_shot_neutral))


def get_cst_matrix(color_matrix_1, color_matrix_2, illuminant):
    if type(color_matrix_1[0]) is Ratio:
        color_matrix_1 = ratios2floats(color_matrix_1)
    if type(color_matrix_2[0]) is Ratio:
        color_matrix_2 = ratios2floats(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(xyz2cam2, axis=1, keepdims=True)

    # interpolate between CSTs based on illuminant
    cct = raw_rgb_to_cct(illuminant, xyz2cam1, xyz2cam2)
    # print(cct)
    xyz2cam_interp = interpolate_cst(xyz2cam1, xyz2cam2, cct)
    xyz2cam_interp = xyz2cam_interp / np.sum(xyz2cam_interp, axis=1, keepdims=True)
    cam2xyz_interp = np.linalg.inv(xyz2cam_interp)
    return cam2xyz_interp


def get_xyz_to_srgb_mat():
    xyz2srgb = np.array([[3.2404542, -1.5371385, -0.4985314],
                         [-0.9692660, 1.8760108, 0.0415560],
                         [0.0556434, -0.2040259, 1.0572252]])

    xyz2srgb = xyz2srgb / np.sum(xyz2srgb, axis=-1, keepdims=True)

    return xyz2srgb


def apply_combined_mat(image, mat):
    outimg = mat[np.newaxis, np.newaxis, :, :] * image[:, :, np.newaxis, :]
    outimg = np.sum(outimg, axis=-1)
    outimg = np.clip(outimg, 0.0, 1.0)
    return outimg


def RGB2bayer(img_rgb):
    h, w = img_rgb.shape[:2]
    bayer = np.empty((h, w), dtype=np.float32)

    (R, G, B) = cv2.split(img_rgb)

    # G R B G
    bayer[0::2, 0::2] = G[0::2, 0::2]  # top left
    bayer[0::2, 1::2] = R[0::2, 1::2]  # top right
    bayer[1::2, 0::2] = B[1::2, 0::2]  # bottom left
    bayer[1::2, 1::2] = G[1::2, 1::2]  # bottom right

    return bayer


def wb_float_to_rat(wb_values):
    print('Debug: Max 16')
    wb_frac_R = Fraction(wb_values[0]).limit_denominator(max_denominator=2 ** 16 - 1)
    wb_frac_B = Fraction(wb_values[2]).limit_denominator(max_denominator=2 ** 16 - 1)
    numer = np.array([wb_frac_R.numerator, 1, wb_frac_B.numerator]).astype(np.uint32)
    denom = np.array([wb_frac_R.denominator, 1, wb_frac_B.denominator]).astype(np.uint32)
    return numer, denom


def update_wb_values(myhex, wb_values, wb_start):
    if not isinstance(wb_values[0], Ratio):
        numer, denom = wb_float_to_rat(wb_values)
    else:
        numer = np.array([wb_values[0].numerator, 1, wb_values[2].numerator]).astype(np.uint32)
        denom = np.array([wb_values[0].denominator, 1, wb_values[2].denominator]).astype(np.uint32)

    mystart = wb_start
    myend = mystart + 8
    for i in range(3):
        myhex[mystart:myend] = hexlify(numer[i].tobytes())  # byteswap()
        # myhex[mystart:myend] = hexlify(swap32(numer[i]).to_bytes(4, byteorder='big'))
        mystart = myend
        myend = mystart + 8

        myhex[mystart:myend] = hexlify(denom[i].tobytes())  # byteswap()
        # myhex[mystart:myend] = hexlify(swap32(denom[i]).to_bytes(4, byteorder='big'))
        mystart = myend
        myend = mystart + 8

    return myhex


def update_hex_image(myhex, bayer, image_start):
    myhex[image_start:image_start + bayer.size * 4] = hexlify(bayer.astype(np.uint16).flatten().tobytes())
    return myhex


def get_illum_normalized_by_g(illum_in_arr):
    return illum_in_arr[:, 0] / illum_in_arr[:, 1], illum_in_arr[:, 1] / illum_in_arr[:, 1], illum_in_arr[:,
                                                                                             2] / illum_in_arr[:, 1]


def in_hull(p, hull):
    """
    Test if points in `p` are in `hull`
    `p` should be a `NxK` coordinates of `N` points in `K` dimensions
    `hull` is either a scipy.spatial.Delaunay object or the `MxK` array of the
    coordinates of `M` points in `K`dimensions for which Delaunay triangulation
    will be computed
    """
    if not isinstance(hull, Delaunay):
        hull = Delaunay(hull)

    return hull.find_simplex(p) >= 0


def array_float_to_rat(array_values):
    numer = []
    denom = []
    for i in range(len(array_values)):
        frac = Fraction(array_values[i]).limit_denominator(max_denominator=2 ** 16 - 1)
        numer.append(np.array([frac.numerator]).astype(np.int32))
        denom.append(np.array([frac.denominator]).astype(np.int32))
    return numer, denom


def update_colormatrix1_values(myhex, colormatrix1, colormatrix1_start):
    numer, denom = array_float_to_rat(colormatrix1)

    mystart = colormatrix1_start
    myend = mystart + 8
    for i in range(9):
        myhex[mystart:myend] = hexlify(numer[i].tobytes())  # byteswap()
        # myhex[mystart:myend] = hexlify(swap32(numer[i]).to_bytes(4, byteorder='big'))
        mystart = myend
        myend = mystart + 8

        myhex[mystart:myend] = hexlify(denom[i].tobytes())  # byteswap()
        # myhex[mystart:myend] = hexlify(swap32(denom[i]).to_bytes(4, byteorder='big'))
        mystart = myend
        myend = mystart + 8

    return myhex


def calc_deltaE_rgb(source, target):
    """
    :param source: rgb, 8 bit
    :param target: rgb, 8 bit
    :return:
    """
    source = color.rgb2lab(source)
    target = color.rgb2lab(target)
    source = np.reshape(source, [-1, 3]).astype(np.float32)
    target = np.reshape(target, [-1, 3]).astype(np.float32)
    delta_e = np.sqrt(np.sum(np.power(source - target, 2), 1))
    delta_e = sum(delta_e) / (np.shape(delta_e)[0])
    return delta_e


def safe_invert_gains(image, wb_vec, rgb_gain=1.0):
    """
    Inverts gains while safely handling saturated pixels.
    Source: numpy version of
    https://github.com/timothybrooks/unprocessing/blob/master/unprocess.py#L92
    """
    gains = np.array(wb_vec) * rgb_gain
    gains = gains[np.newaxis, np.newaxis, :]

    # Prevents dimming of saturated pixels by smoothly masking gains near white.
    gray = np.mean(image, axis=-1, keepdims=True)
    inflection = 0.9

    mask = (np.maximum(gray - inflection, 0.0) / (1.0 - inflection)) ** 2.0
    safe_gains = np.maximum(mask + (1.0 - mask) * gains, gains)
    image = image * safe_gains
    image = np.clip(image, 0.0, 1.0)
    return image

utils/img_utils.py
SamsungLabs/graphics2raw
"""
Copyright (c) 2023 Samsung Electronics Co., Ltd.

Author(s):
Luxi Zhao (lucy.zhao@samsung.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

"""

import torchvision
import torch


def torchvision_visualize_raw(images) -> torch.Tensor:
    """
    :param images: a list of torch tensors. Each item is a 4 channel stacked RAW image
    with bayer pattern [1, 0, 2, 1]. Item shape: (b, 4, h, w)
    :return: a stacked grid
    """
    # stacked in order of [1,0,2,1] -> want to swap 1 and 0 -> [0,1,2]
    images = [im[:, 0:3, :, :][:, [1, 0, 2], ...] for im in images]
    img_grid = torchvision.utils.make_grid(torch.cat(images, 2), normalize=False)
    return img_grid

utils/torch_utils.py
SamsungLabs/graphics2raw