import numpy as np

 import struct

 from PIL import Image

 import rawpy

 import glob

 import os

 def conv(v):

     s = 0

     for i in range(len(v)):

         s += i * v[i]

         v[i] = s

     return v

 def diff(v):

     n = len(v)

     v_diff = np.zeros([len(v)])

     for i in range(n-1):

         v_diff[n-1-i] = v[n-1-i] - v[n-1-i-1]

     return v_diff

 def gray_ps(rgb):

     return np.power(np.power(rgb[:,:,0], 2.2) * 0.2973 + np.power(rgb[:,:,1], 2.2) * 0.6274 + np.power(rgb[:,:,2], 2.2) * 0.0753, 1/2.2) + 1e-7

 def HDR(x, curve_ratio):

     gray_scale = np.expand_dims(gray_ps(x), axis=-1)

     gray_scale_new = np.power(gray_scale, curve_ratio)

     return np.minimum(x * gray_scale_new / gray_scale, 1.0)

 def gray_world_balance(x):

     mean_ = np.maximum(np.mean(x, axis=(0, 1)), 1e-6)

     ratio_ = mean_.mean() / mean_

     return np.minimum(x * ratio_, 1.0)

 def show_bbf(path):

     print('====== %s =====' % path)

     max_level = 1023

     black_level = 64

     height = 3024

     width = 4032

     if path.endswith('dng') or path.endswith('DNG'):

         raw = rawpy.imread(path)

         im = raw.raw_image_visible.astype(np.float32)

         res_path = str.replace(path, '.dng', '.jpg')

     elif path.endswith('raw') or path.endswith('RAW'):

         raw = open(path, 'rb').read()

         raw = struct.unpack('H'*int(len(raw)/2), raw)

         im = np.float32(raw)

         res_path = str.replace(path, '.raw', '.jpg')

     else:

         assert False

     im = im.reshape(height, width)

     im = np.maximum(im - black_level, 0) / (max_level - black_level)

     # AMPLIFICATION

     # n = 256

     # std = 0.05

     # sample_rate = 1

     # extreme_dark_ratio = 13 * std / 0.05

     # light_scene_ratio = 4.0 * std / 0.05

     # light_threshold = 0.04

     # decay_in_light_ratio = 0.3

     # light_radius = 5

     # total = height * width / (sample_rate*sample_rate) / 4

     #

     # bins = np.arange(n + 1) / (n - 1)

     # hists = [None] * 5

     # hists[0], _ = np.histogram(im[0:1512:sample_rate, 0:2016:sample_rate], bins)

     # hists[1], _ = np.histogram(im[0:1512:sample_rate, 2016:4032:sample_rate], bins)

     # hists[2], _ = np.histogram(im[1512:3024:sample_rate, 0:2016:sample_rate], bins)

     # hists[3], _ = np.histogram(im[1512:3024:sample_rate, 2016:4032:sample_rate], bins)

     # hists[4] = (hists[0] + hists[1] + hists[2] + hists[3])/4

     # convs = [None] * 5

     # min_conv = 255

     # max_conv = 0

     # final_ratio = 1.0

     # is_dark = False

     #

     # for i in range(5):

     #     hists[i] = hists[i] / total

     #     convs[i] = conv(hists[i][0:n])

     #     print(convs[i][-1])

     #     if convs[i][-1]<min_conv:

     #         min_conv = convs[i][-1]

     #     if convs[i][-1]>max_conv:

     #         max_conv = convs[i][-1]

     #

     # print('min=%.6f, max=%.6f' % (min_conv, max_conv))

     #

     # hist_conv = convs[4]

     # hist_diff = diff(hist_conv)

     # ratio = std / (hist_conv[-1] / n)

     # print("Normal ratio=%.6f" % ratio)

     # if ratio < 1:

     #     print("Daylight scene found!")

     #     final_ratio = 1.0

     # # check if exists high contrast scene

     # if min_conv < 3 and max_conv/(min_conv + 1e-7) > 2.2:

     #     print('high contrast scene detected!')

     #     final_ratio = min(2.0, ratio)

     # else:

     #     if max(hist_diff[-light_radius:]) > light_threshold:

     #         print('Light Found')

     #         if ratio > extreme_dark_ratio:

     #             final_ratio = ratio * decay_in_light_ratio

     #             print('Extreme Dark')

     #         else:

     #             final_ratio = min(light_scene_ratio, ratio)

     #     else:

     #         if ratio > extreme_dark_ratio:

     #             print('Extreme Dark')

     #             is_dark = True

     #             final_ratio = extreme_dark_ratio*extreme_dark_ratio*extreme_dark_ratio/(ratio*ratio)

     #             if 4 < final_ratio < 6:

     #                 final_ratio = 4

     #         else:

     #             final_ratio = ratio

     #     if ratio < 1.0:

     #         final_ratio = 1.0

     # if 5 < final_ratio < 12 and not is_dark:

     #     final_ratio *= 0.7

     # im = np.minimum(im * final_ratio, 1.0)

     # print('ratio=%.6f' % final_ratio)

     im = np.expand_dims(im, axis=2)

     H = im.shape[0]

     W = im.shape[1]

     out = np.concatenate((

         im[0:H:2, 1:W:2, :],

         (im[0:H:2, 0:W:2, :] + im[1:H:2, 1:W:2, :])/2.0,

         im[1:H:2, 0:W:2, :]),

         axis=2)

     if path.endswith('dng') or path.endswith('DNG'):

         wb = np.array(raw.camera_whitebalance[:3], np.float32)

         wb = wb / wb[1]

         out = np.minimum(out * wb, 1.0)

     else:

         out = gray_world_balance(out)

     out = np.minimum(out * 0.2 / out[:, :, 1].mean(), 1.0)

     out = HDR(out, 0.35)

     Image.fromarray(np.uint8(out*255)).save(res_path)

 def dng2raw(path):

     raw = rawpy.imread(path)

     im = raw.raw_image_visible.astype(np.ushort)

     h, w = im.shape

     res_path = str.replace(path, '.dng', '.raw')

     with open(res_path, 'wb')as fp:

         for i in range(h):

             for j in range(w):

                 a = struct.pack('<H', im[i, j])

                 fp.write(a)

 def decode_sony(path):

     if os.path.split(path)[-1].split('.')[-1] != 'ARW':

         print('Error: image type not match!')

         exit(1)

     im_ = rawpy.imread(path)

     data_ = im_.raw_image_visible.astype(np.float32)

     h_, w_ = data_.shape

     save_path_ = str.replace(path, '.ARW', '.JPG')

     data_ = np.maximum(data_ - 512, 0) / (16383 - 512)

     data_ = np.expand_dims(data_, axis=2)

     data_ = np.concatenate((

         data_[0:h_:2, 0:w_:2, :],

         (data_[0:h_:2, 1:w_:2, :] + data_[1:h_:2, 0:w_:2, :]) / 2.0,

         data_[1:h_:2, 1:w_:2, :]),

         axis=2)

     # white balance

     wb = np.array(im_.camera_whitebalance[:3], np.float32)

     wb = wb / wb[1]

     data_ = np.minimum(data_ * wb, 1.0)

     data_ = np.minimum(data_ * 0.2 / data_[:, :, 1].mean(), 1.0)

     data_ = HDR(data_, 0.35)

     Image.fromarray(np.uint8(data_ * 255)).save(save_path_)

 def decode_sony_rawpy(path):

     if os.path.split(path)[-1].split('.')[-1] != 'ARW':

         print('Error: image type not match!')

         exit(1)

     im_ = rawpy.imread(path)

     save_path_ = str.replace(path, '.ARW', '.JPG')

     Image.fromarray(im_.postprocess(use_camera_wb=True)).save(save_path_)

 if __name__ == '__main__':

     # files = glob.glob('D:/data/report_dng/*.dng')

     # for i in range(len(files)):

     #     show_bbf(files[i])

     # files = glob.glob('D:/data/Sony/long/00057_00_10s.ARW')

     # im = rawpy.imread(files[0])

     # Image.fromarray(im.postprocess(use_camera_wb=True)).show()

     # files = glob.glob('D:/data/Sony/long/*.ARW')

     # for i in range(len(files)):

     #     decode_sony_rawpy(files[i])

     # files = glob.glob('D:/data/Sony/long/RAW/30s/*.ARW')

     # for i in range(len(files)):

     #     decode_sony(files[i])

     # files = glob.glob('D:/data/Sony/dataset/illu_detect/day/*.dng')

     # for i in range(len(files)):

     #     dng2raw(files[i])

     # files = glob.glob('C:/Users/Administrator/Desktop/ILLU/*.dng')

     files = glob.glob('D:/data/LightOnOff/*.dng')

     for i in range(len(files)):

         # dng2raw(files[i])

         show_bbf(files[i])