Spaces:
Runtime error
Runtime error
| import os, argparse, cv2, random | |
| import numpy as np | |
| from skimage import exposure | |
| def guideFilter(I, p, winSize, eps): | |
| mean_I = cv2.blur(I, winSize) | |
| mean_p = cv2.blur(p, winSize) | |
| mean_II = cv2.blur(I * I, winSize) | |
| mean_Ip = cv2.blur(I * p, winSize) | |
| var_I = mean_II - mean_I * mean_I | |
| cov_Ip = mean_Ip - mean_I * mean_p | |
| a = cov_Ip / (var_I + eps) | |
| b = mean_p - a * mean_I | |
| mean_a = cv2.blur(a, winSize) | |
| mean_b = cv2.blur(b, winSize) | |
| q = mean_a * I + mean_b | |
| return q | |
| def syn_low(img, light, img_gray, light_max=3, | |
| light_min=2, noise_max=0.08, noise_min=0.03): | |
| light = guideFilter(light, img_gray,(3,3),0.01)[:, :, np.newaxis] | |
| n = np.random.uniform(noise_min, noise_max) | |
| R = img / (light + 1e-7) | |
| L = (light + 1e-7) ** np.random.uniform(light_min, light_max) | |
| return np.clip(R * L + np.random.normal(0, n, img.shape), 0, 1) | |
| def syn_haze(img, depth, beta_max=2.0, beta_min=1.0, A_max=0.9, A_min=0.6, | |
| color_max=0, color_min=0): | |
| beta = np.random.rand(1) * (beta_max - beta_min) + beta_min | |
| t = np.exp(-np.minimum(1 - cv2.blur(depth,(22,22)),0.7) * beta) | |
| A = np.random.rand(1) * (A_max - A_min) + A_min | |
| A_random = np.random.rand(3) * (color_max - color_min) + color_min | |
| A = A + A_random | |
| return np.clip(img * t + A * (1 - t), 0, 1) | |
| def syn_data(hq_file, light_file, depth_file, rain_file, snow_file, out_file, | |
| low, haze, rain, snow): | |
| file_list = os.listdir(hq_file) | |
| rain_list = os.listdir(rain_file) | |
| snow_list = os.listdir(snow_file) | |
| num_rain = random.sample(range(0,len(rain_list)),len(rain_list)) | |
| num_snow = random.sample(range(0,len(snow_list)),len(snow_list)) | |
| for i in range(1, len(file_list)): | |
| img = cv2.imread(hq_file+file_list[i]) | |
| w, h, _ = img.shape | |
| light = cv2.cvtColor(cv2.imread(light_file + file_list[i]), cv2.COLOR_RGB2GRAY) / 255.0 | |
| depth = cv2.imread(depth_file + file_list[i]) / 255.0 | |
| rain_mask = cv2.imread(rain_file + rain_list[num_rain[i]]) / 255.0 | |
| rain_mask = cv2.resize(rain_mask,(h,w)) | |
| snow_mask = cv2.imread(snow_file + snow_list[num_snow[i]]) / 255.0 | |
| snow_mask = cv2.resize(snow_mask, (h, w)) | |
| img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)/ 255.0 | |
| lq = img.copy()/255.0 | |
| color_dis = 1 | |
| if low: | |
| lq = syn_low(lq, light, img_gray) | |
| if rain: | |
| lq = lq+rain_mask | |
| if snow: | |
| lq = lq*(1-snow_mask)+color_dis*snow_mask | |
| if haze: | |
| lq = syn_haze(lq, depth) | |
| # out = np.concatenate((lq*255.0,img),1) | |
| out = lq*255.0 | |
| cv2.imwrite(out_file + file_list[i], out) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| # load model | |
| parser.add_argument("--hq-file", type=str, default = './data/clear/') | |
| parser.add_argument("--light-file", type=str, default = './data/light_map/') | |
| parser.add_argument("--depth-file", type=str, default = './data/depth_map/') | |
| parser.add_argument("--rain-file", type=str, default = './data/rain_mask/') | |
| parser.add_argument("--snow-file", type=str, default = './data/snow_mask/') | |
| parser.add_argument("--out-file", type=str, default = './out/') | |
| parser.add_argument("--low", action='store_true') | |
| parser.add_argument("--haze", action='store_true') | |
| parser.add_argument("--rain", action='store_true') | |
| parser.add_argument("--snow", action='store_true') | |
| args = parser.parse_args() | |
| syn_data(args.hq_file, args.light_file, args.depth_file, args.rain_file, | |
| args.snow_file, args.out_file, args.low, args.haze, args.rain, args.snow) |