|
import glob |
|
import rawpy |
|
import cv2 |
|
import os |
|
import numpy as np |
|
import colour_demosaicing |
|
from tqdm import tqdm |
|
|
|
|
|
def process_raw(args, path, w_new, h_new): |
|
raw = rawpy.imread(str(path)).raw_image_visible |
|
if "_00200_" in str(path) or "_00100_" in str(path): |
|
raw = np.clip(raw.astype("float32") - 512, 0, 65535) |
|
else: |
|
raw = np.clip(raw.astype("float32") - 2048, 0, 65535) |
|
img = colour_demosaicing.demosaicing_CFA_Bayer_bilinear(raw, "RGGB").astype( |
|
"float32" |
|
) |
|
img = np.clip(img, 0, 16383) |
|
|
|
|
|
if args.histeq: |
|
img2 = np.zeros_like(img) |
|
for i in range(3): |
|
hist, bins = np.histogram(img[..., i].flatten(), 16384, [0, 16384]) |
|
cdf = hist.cumsum() |
|
cdf_normalized = cdf * float(hist.max()) / cdf.max() |
|
cdf_m = np.ma.masked_equal(cdf, 0) |
|
cdf_m = (cdf_m - cdf_m.min()) * 16383 / (cdf_m.max() - cdf_m.min()) |
|
cdf = np.ma.filled(cdf_m, 0).astype("uint16") |
|
img2[..., i] = cdf[img[..., i].astype("int16")] |
|
img[..., i] = img2[..., i].astype("float32") |
|
|
|
|
|
m = img.mean() |
|
d = np.abs(img - img.mean()).mean() |
|
img = (img - m + 2 * d) / 4 / d * 255 |
|
image = np.clip(img, 0, 255) |
|
|
|
image = cv2.resize( |
|
image.astype("float32"), (w_new, h_new), interpolation=cv2.INTER_AREA |
|
) |
|
|
|
if args.histeq: |
|
path = str(path) |
|
os.makedirs( |
|
"/".join(path.split("/")[:-2] + [path.split("/")[-2] + "-npy"]), |
|
exist_ok=True, |
|
) |
|
np.save( |
|
"/".join( |
|
path.split("/")[:-2] |
|
+ [path.split("/")[-2] + "-npy"] |
|
+ [path.split("/")[-1].replace("cr2", "npy")] |
|
), |
|
image, |
|
) |
|
else: |
|
path = str(path) |
|
os.makedirs( |
|
"/".join(path.split("/")[:-2] + [path.split("/")[-2] + "-npy-nohisteq"]), |
|
exist_ok=True, |
|
) |
|
np.save( |
|
"/".join( |
|
path.split("/")[:-2] |
|
+ [path.split("/")[-2] + "-npy-nohisteq"] |
|
+ [path.split("/")[-1].replace("cr2", "npy")] |
|
), |
|
image, |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
import argparse |
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--H", type=int, default=int(640)) |
|
parser.add_argument("--W", type=int, default=int(960)) |
|
parser.add_argument("--histeq", action="store_true") |
|
parser.add_argument("--dataset_dir", type=str, default="/data/hyz/MID/") |
|
args = parser.parse_args() |
|
|
|
path_ls = glob.glob(args.dataset_dir + "/*/pair*/?????/*") |
|
for path in tqdm(path_ls): |
|
process_raw(args, path, args.W, args.H) |
|
|
