MiAlgo/utils.py

import numpy as np
import cv2
import cv2 as cv

import os
import json
from fractions import Fraction
from exifread.utils import Ratio
from scipy.io import loadmat
from PIL import Image, ImageOps
import torch
from copy import deepcopy


def json_read(fname, **kwargs):
    with open(fname) as j:
        data = json.load(j, **kwargs)
    return data

def fraction_from_json(json_object):
    if 'Fraction' in json_object:
        return Fraction(*json_object['Fraction'])
    return json_object

def fractions2floats(fractions):
    floats = []
    for fraction in fractions:
        floats.append(float(fraction.numerator) / fraction.denominator)
    return floats

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

def normalize(raw_image, black_level, white_level, bl_fix=0):
    if type(black_level) is list and len(black_level) == 1:
        black_level = float(black_level[0])
    if type(white_level) is list and len(white_level) == 1:
        white_level = float(white_level[0])
    black_level_mask = black_level
    if type(black_level) is list and len(black_level) == 4:
        if type(black_level[0]) is Ratio:
            black_level = ratios2floats(black_level)
        if type(black_level[0]) is Fraction:
            black_level = fractions2floats(black_level)
        black_level_mask = np.zeros(raw_image.shape)
        idx2by2 = [[0, 0], [0, 1], [1, 0], [1, 1]]
        step2 = 2
        for i, idx in enumerate(idx2by2):
            black_level_mask[idx[0]::step2, idx[1]::step2] = black_level[i]

    if bl_fix:
        black_level_mask = black_level_mask + bl_fix

    normalized_image = (raw_image.astype(np.float32) - black_level_mask) / (white_level - black_level_mask)
    normalized_image = np.clip(normalized_image, 0, 1)
    return normalized_image

def bayer_to_offsets(bayer_pattern):
    """
    Transform bayer pattern to offsets in order 'RGrBGb'
    n.b. Support 'RGrBGb' bayer pattern only.
    Args:
        bayer_pattern: string, e.g. 'rggb'. Must be one of 'rggb', 'grbg', 'gbrg', 'bggr'

    Returns:
        offsets: packed raw image with 4 channels
    """
    bayer_pattern = bayer_pattern.lower()
    assert bayer_pattern in ['rggb', 'grbg', 'gbrg', 'bggr'], 'WRONG BAYER PATTERN!'

    if bayer_pattern == 'rggb':
        offsets = [[0,0],[0,1],[1,1],[1,0]]
    elif bayer_pattern == 'grbg':
        offsets = [[0,1],[0,0],[1,0],[1,1]]
    elif bayer_pattern == 'gbrg':
        offsets = [[1,0],[0,0],[0,1],[1,1]]
    else: #bayer_pattern == 'bggr':
        offsets = [[1,1],[0,1],[0,0],[1,0]]

    return offsets
    
def pack_raw_to_4ch(rawim, offsets):
    """
    Pack raw to h/2 x w/2 x 4n with order "RGrBGb..." RGBG RGBG RGBG
    n.b. Support ordinary bayer pattern only.
    Args:
        rawim: numpy.ndarray in shape (h, w, ...)
        bayer_pattern: string, e.g. "rggb". Must be one of "rggb", "grbg", "gbrg", "bggr"

    Returns:
        out: packed raw image with 4n channels
    """


    if rawim.ndim == 2:
        rawim = np.expand_dims(rawim, axis=-1)
        rawim_pack = np.concatenate((rawim[offsets[0][0]::2, offsets[0][1]::2],
                                rawim[offsets[1][0]::2, offsets[1][1]::2],
                                rawim[offsets[2][0]::2, offsets[2][1]::2],
                                rawim[offsets[3][0]::2, offsets[3][1]::2]), axis=-1)
    elif rawim.ndim ==3:
        frame_num = rawim.shape[2]
        rawim_pack = np.zeros((int(rawim.shape[0]/2), int(rawim.shape[1]/2), rawim.shape[2] * 4))
        for i in range(frame_num):
            rawim_temp = rawim[:,:,i]
            rawim_temp = np.expand_dims(rawim_temp, axis=-1)
            rawim_temp_pack = np.concatenate((rawim_temp[offsets[0][0]::2, offsets[0][1]::2],
                                              rawim_temp[offsets[1][0]::2, offsets[1][1]::2],
                                              rawim_temp[offsets[2][0]::2, offsets[2][1]::2],
                                              rawim_temp[offsets[3][0]::2, offsets[3][1]::2]), axis=-1)

            rawim_pack[:, :, i * 4:(i + 1) * 4] = rawim_temp_pack

    return rawim_pack



def apply_color_space_transform(demosaiced_image, color_matrix_1, color_matrix_2):
    if isinstance(color_matrix_1[0], Fraction):
        color_matrix_1 = fractions2floats(color_matrix_1)
    if isinstance(color_matrix_2[0], Fraction):
        color_matrix_2 = fractions2floats(color_matrix_2)
    xyz2cam1 = np.reshape(np.asarray(color_matrix_1), (3, 3))
    # normalize rows (needed?)
    xyz2cam1 = xyz2cam1 / np.sum(xyz2cam1, axis=1, keepdims=True)
    # inverse
    cam2xyz1 = np.linalg.inv(xyz2cam1)
    # for now, use one matrix  # TODO: interpolate btween both
    # simplified matrix multiplication
    
    xyz_image = cam2xyz1[np.newaxis, np.newaxis, :, :] * \
        demosaiced_image[:, :, np.newaxis, :]
    xyz_image = np.sum(xyz_image, axis=-1)
    xyz_image = np.clip(xyz_image, 0.0, 1.0)
    return xyz_image




def transform_xyz_to_srgb(xyz_image):
    xyz2srgb = np.array([[3.0799, -1.5372, -0.5428],
                         [-0.9212, 1.8760, 0.0452],
                         [0.0529, -0.2040, 1.1512]])   
    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 apply_gamma(x):
    # return x ** (1.0 / 2.4)
    x = x.copy()
    gray = cv2.cvtColor(x.astype(np.float32), cv2.COLOR_BGR2GRAY)
    idx = gray <= 0.0031308
    idx = np.stack([idx, idx, idx], axis=-1)
    x[idx] *= 12.92
    x[idx == False] = (x[idx == False] ** (1.0 / 2.4)) * 1.055 - 0.055
    return x


def fix_orientation(image, orientation):
    type1 = "Horizontal (normal)"
    type2 = "Mirror horizontal"
    type3 = "Rotate 180"
    type4 = "Mirror vertical"
    type5 = "Mirror horizontal and rotate 270 CW"
    type6 = "Rotate 90 CW"
    type7 = "Mirror horizontal and rotate 90 CW"
    type8 = "Rotate 270 CW"

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

    if orientation == type1:
        pass
    elif orientation == type2:
        image = cv.flip(image, 0)
    elif orientation == type3:
        image = cv.rotate(image, cv.ROTATE_180)
    elif orientation == type4:
        image = cv.flip(image, 1)
    elif orientation == type5:
        image = cv.flip(image, 0)
        image = cv.rotate(image, cv.ROTATE_90_COUNTERCLOCKWISE)
    elif orientation == type6:
        image = cv.rotate(image, cv.ROTATE_90_CLOCKWISE)
    elif orientation == type7:
        image = cv.flip(image, 0)
        image = cv.rotate(image, cv.ROTATE_90_CLOCKWISE)
    elif orientation == type8:
        image = cv.rotate(image, cv.ROTATE_90_COUNTERCLOCKWISE)
    else:
        raise

    return image


def apply_tone_map(x, tone_mapping='Base'):
    if tone_mapping == 'Flash':
        return perform_flash(x, perform_gamma_correction=0)/255.
    elif tone_mapping == 'Storm':
        return perform_storm(x, perform_gamma_correction=0)/255.
    elif tone_mapping == 'Drago':
        tonemap = cv2.createTonemapDrago()
        return tonemap.process(x.astype(np.float32))
    elif tone_mapping == 'Mantiuk':
        tonemap = cv2.createTonemapMantiuk()
        return tonemap.process(x.astype(np.float32))
    elif tone_mapping == 'Reinhard':
        tonemap = cv2.createTonemapReinhard()
        return tonemap.process(x.astype(np.float32))
    elif tone_mapping == 'Linear':
        return np.clip(x/np.sort(x.flatten())[-50000], 0, 1)
    elif tone_mapping == 'Base':
        tone_curve = loadmat(os.path.join(os.path.dirname(
            os.path.realpath(__file__)), 'tone_curve.mat'))
        tone_curve = tone_curve['tc']
        x = np.round(x * (len(tone_curve) - 1)).astype(int)
        tone_mapped_image = np.squeeze(tone_curve[x])
        return tone_mapped_image
    else:
        raise ValueError(
            'Bad tone_mapping option value! Use the following options: "Base", "Flash", "Storm", "Linear", "Drago", "Mantiuk", "Reinhard"')

def autocontrast_using_pil(img, cutoff=2):
    img_uint8 = np.clip(255*img, 0, 255).astype(np.uint8)
    img_pil = Image.fromarray(img_uint8)
    img_pil = ImageOps.autocontrast(img_pil, cutoff=cutoff)
    output_image = np.array(img_pil).astype(np.float32) / 255
    return output_image

def usm_sharp(img, weight=0.666, radius=10, threshold=10):
    if radius % 2 == 0:
        radius += 1
    blur = cv2.GaussianBlur(img, (radius, radius), 0)
    residual = img - blur
    mask = np.abs(residual) * 255 > threshold
    mask = mask.astype('float32')
    soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)

    sharp = img + weight * residual
    sharp = np.clip(sharp, 0, 1)
    return soft_mask * sharp + (1 - soft_mask) * img


def white_balance(demosaic_img, as_shot_neutral):
    if type(as_shot_neutral[0]) is Ratio:
        as_shot_neutral = ratios2floats(as_shot_neutral)

    as_shot_neutral = np.asarray(as_shot_neutral)
    # transform vector into matrix
    if as_shot_neutral.shape == (3,):
        as_shot_neutral = np.diag(1./as_shot_neutral)

    assert as_shot_neutral.shape == (3, 3)
    white_balanced_image = np.dot(demosaic_img, as_shot_neutral.T)
    white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)

    return white_balanced_image


def get_net(NET, checkpoint_path, device):
    net = NET()
    load_net = torch.load(checkpoint_path, map_location="cpu")

    try:
        load_net = load_net['params']
    except:
        pass

    net.load_state_dict(load_net, strict=True)
    net = net.to(device)
    net = net.eval()
    return net



def pyrblend(img1, img2, mask):
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    mask = mask.astype(np.float32)
    
    # 图像向下取样, 构造高斯金字塔: [原图，下取样1次，下取样2次，下取样3次，下取样4次]
    levels = 5  # 高斯金字塔层数
    gaussPyr1, gaussPyr2 = [img1], [img2]  # 原始图像为高斯金字塔第 0 层
    for i in range(1, levels):  # 高斯金字塔共 5 层: 0,1,2,3,4
        gaussPyr1.append(cv2.pyrDown(gaussPyr1[i-1]))  # 计算第 i 层高斯金字塔
        gaussPyr2.append(cv2.pyrDown(gaussPyr2[i-1]))

    # 图像向上取样, 构造拉普拉斯金字塔 [第1层残差，第2层残差，第3层残差，第4层残差]
    lapPyr1, lapPyr2 = [], []  # 从最顶层开始恢复
    for i in range(levels-1):  # 拉普拉斯金字塔有 4 层: 0,1,2,3
        lapPyr1.append(gaussPyr1[i] - cv2.pyrUp(gaussPyr1[i+1]))
        lapPyr2.append(gaussPyr2[i] - cv2.pyrUp(gaussPyr2[i+1]))

    # 拉普拉斯金字塔左右拼接
    lapStack = []
    for i in range(levels-1):  # 拉普拉斯金字塔共 4 层: 0,1,2,3
        rows, cols, channel = lapPyr1[i].shape
        lmask = cv2.resize(mask, (cols, rows), interpolation=cv2.INTER_AREA)
        lmask = lmask[:, :, np.newaxis]
        
        splicing = lapPyr1[i] * (1-lmask) + lapPyr2[i]*lmask
        lapStack.append(splicing)

    # 由拼接后的Laplace金字塔恢复原图像
    rows, cols, channel = gaussPyr1[-1].shape  # 高斯金字塔顶层 G4:(32,32)
    lmask = cv2.resize(mask, (cols, rows), interpolation=cv2.INTER_AREA)
    lmask = lmask[:, :, np.newaxis]
    stackG4 = gaussPyr1[-1] * (1-lmask) + gaussPyr2[-1]*lmask # 拼接高斯金字塔顶层
    stackG3 = lapStack[3] + cv2.pyrUp(stackG4)  # stackG3:(64,64)
    stackG2 = lapStack[2] + cv2.pyrUp(stackG3)  # stackG2:(128,128)
    stackG1 = lapStack[1] + cv2.pyrUp(stackG2)  # stackG1:(256,256)
    stackG0 = lapStack[0] + cv2.pyrUp(stackG1)  # stackG0:(512,512)

    return stackG0