Spaces:
Running
Running
File size: 1,970 Bytes
499e141 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import rawpy
import numpy as np
import glob, os
import colour_demosaicing
import imageio
import argparse
from PIL import Image as PILImage
import scipy.io as scio
parser = argparse.ArgumentParser(description="data preprocess")
parser.add_argument("--camera", type=str, default="NIKON_D700", help="Camera Name")
parser.add_argument("--Bayer_Pattern", type=str, default="RGGB", help="Bayer Pattern of RAW")
parser.add_argument("--JPEG_Quality", type=int, default=90, help="Jpeg Quality of the ground truth.")
args = parser.parse_args()
camera_name = args.camera
Bayer_Pattern = args.Bayer_Pattern
JPEG_Quality = args.JPEG_Quality
dng_path = sorted(glob.glob('/mnt/nvme2n1/hyz/data/' + camera_name + '/DNG/*.cr2'))
rgb_target_path = '/mnt/nvme2n1/hyz/data/'+ camera_name + '/RGB/'
raw_input_path = '/mnt/nvme2n1/hyz/data/' + camera_name + '/RAW/'
if not os.path.isdir(rgb_target_path):
os.mkdir(rgb_target_path)
if not os.path.isdir(raw_input_path):
os.mkdir(raw_input_path)
def flip(raw_img, flip):
if flip == 3:
raw_img = np.rot90(raw_img, k=2)
elif flip == 5:
raw_img = np.rot90(raw_img, k=1)
elif flip == 6:
raw_img = np.rot90(raw_img, k=3)
else:
pass
return raw_img
for path in dng_path:
print("Start Processing %s" % os.path.basename(path))
raw = rawpy.imread(path)
file_name = path.split('/')[-1].split('.')[0]
im = raw.postprocess(use_camera_wb=True,no_auto_bright=True)
flip_val = raw.sizes.flip
cwb = raw.camera_whitebalance
raw_img = raw.raw_image_visible
if camera_name == 'Canon_EOS_5D':
raw_img = np.maximum(raw_img - 127.0, 0)
de_raw = colour_demosaicing.demosaicing_CFA_Bayer_bilinear(raw_img, Bayer_Pattern)
de_raw = flip(de_raw, flip_val)
rgb_img = PILImage.fromarray(im).save(rgb_target_path + file_name + '.jpg', quality = JPEG_Quality, subsampling = 1)
np.savez(raw_input_path + file_name + '.npz', raw=de_raw, wb=cwb)
|