Spaces:
Running
Running
File size: 2,725 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
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)
# HistEQ start
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")
# HistEQ end
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)
|