def process_all(images_path, depth_path, output_path):
img_names = my_utils.os_listdir(images_path)
depth_names = my_utils.os_listdir(depth_path)
beta = 0
pbar = tqdm.tqdm(total=len(img_names))
for name_file, depth_file in zip(img_names, depth_names):
pbar.update(1)
img = cv2.imread(os.path.join(images_path, name_file))
gray_img = cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY)
# divided by 256 to convert it into metres
original_depth = cv2.imread(os.path.join(depth_path, depth_file),
cv2.IMREAD_UNCHANGED) / 256
smooth_depth = improve_depth(gray_img, original_depth, threshold=beta)
np.save(os.path.join(output_path, name_file), smooth_depth)
def load_streak_database(self):
'''
Function to load and store the texture maps.
Streaks are stored in a list.
'''
if not os.path.exists(self.streaks_path):
print("No existing path for streak database (", self.streaks_path,
")")
exit(-1)
tmp = []
norm_coeff_path = self.norm_coeff_path
norm_coeffs = {}
with open(norm_coeff_path, 'r') as file:
lines = file.readlines()
for line in lines:
if line[:2] == 'cv':
coeff = int(line[2:])
continue
norm_coeffs.update({
coeff: [float(v) for v in line.split('\n')[0].split(' ')[:-1]]
})
for file_name in my_utils.os_listdir(self.streaks_path):
name = os.path.splitext(file_name)[0]
coeff, osc = name.split('_')
if len(coeff) == 3:
coeff = int(coeff[-1:])
else:
coeff = int(coeff[-2:])
osc = int(osc[-1:])
drop_image = cv2.imread(os.path.join(self.streaks_path, file_name),
cv2.IMREAD_ANYDEPTH)
drop_image = cv2.cvtColor(drop_image, cv2.COLOR_GRAY2BGR)
drop_image_norm = ((255.0 * norm_coeffs[coeff][osc] * drop_image) /
65535.0).astype(np.uint8)
tmp.append(drop_image_norm)
self.ratio = np.append(self.ratio,
tmp[-1].shape[1] / tmp[-1].shape[0])
self.ratio = np.unique(self.ratio)
self.streaks_light = np.array(tmp)
def run(self):
process_t0 = time.time()
folders_num = len(self.images)
# case for any number of sequences and supported rain intensities
for folder_idx, sequence in enumerate(self.sequences):
folder_t0 = time.time()
print('\nSequence: ' + sequence)
sim_num = len(self.particles[sequence])
depth_folder = self.depth[sequence]
for sim_idx, sim_weather in enumerate(self.weather):
weather, fallrate = sim_weather["weather"], sim_weather[
"fallrate"]
out_seq_dir = os.path.join(self.output_root, sequence)
out_dir = os.path.join(out_seq_dir, weather,
'{}mm'.format(fallrate))
sim_file = self.particles[sequence][sim_idx]
# Resolve output path
path_exists = os.path.exists(out_dir)
if path_exists:
if self.conflict_strategy == "skip":
pass
elif self.conflict_strategy == "overwrite":
pass
elif self.conflict_strategy == "rename_folder":
out_dir_, out_shift = out_dir, 0
while os.path.exists(out_dir_ +
'_copy%05d' % out_shift):
out_shift += 1
out_dir = out_dir_ + '_copy%05d' % out_shift
else:
raise NotImplementedError
# Create directory
os.makedirs(out_dir, exist_ok=True)
# Default fog-like rain parameters
fog_params = {
"rain_intensity": fallrate,
"focal": self.focal,
"f_number": self.f_number,
"angle": 90,
"exposure": self.exposure,
"camera_gain": self.camera_gain
}
if "nuscenes" in self.dataset:
files = self.images[sequence]
depth_files = self.depth[sequence]
depth_file = depth_files[0]
assert depth_file.endswith(
".npy"
), "nuscenes processing only works with .npy for depth"
# depth = np.load(depth_file)
if "gan" in self.dataset:
# HARDCODED since these values are not known in nuscenes_gan
imW, imH = (1600, 900)
else:
img = cv2.imread(files[0])
imH, imW = img.shape[0:2]
else:
files = natsorted(
np.array([
os.path.join(self.images[sequence], picture)
for picture in my_utils.os_listdir(
self.images[sequence])
]))
depth_files = natsorted(
np.array([
os.path.join(depth_folder, depth)
for depth in my_utils.os_listdir(depth_folder)
]))
im = files[0]
if im.endswith(".png"):
imH, imW = cv2.imread(im).shape[0:2]
elif im.endswith(".npy"):
imH, imW = np.load(im).shape[0:2]
else:
raise Exception("Invalid extension", im)
imH = imH // self.settings["render_scale"]
imW = imW // self.settings["render_scale"]
if self.camera_gain:
fog_params["camera_gain"] = self.camera_gain
print('Simulation: rain {}mm/hr'.format(fallrate))
# loading StreaksDBManager, RainRenderer, FOVComputation and EnvMapGenerator
self.db = DBManager(streaks_path_xml=sim_file,
streaks_path=self.texture,
norm_coeff_path=self.norm_coeff)
self.renderer = RainRenderer(focal=self.focal,
f_number=self.f_number,
focus_plane=6,
radius=10,
fov=165)
self.fov_comp = FovComputation(camera=np.array([0, 0, 0]))
map_generator = EnvironmentMapGenerator(self.focal, imW, imH)
# loading fog class
FOG = add_attenuation.FogRain(**fog_params)
# loading calib files for frames
calib_files = self.calib[sequence]
# loading streaks from Streaks Database
self.db.load_streak_database()
# creating drops based on the simulator file
self.db.load_streaks_from_xml(self.dataset,
self.settings, [imW, imH],
use_pickle=False,
verbose=self.verbose)
env_map_input = self.env_type if self.irrad_type == 'garg' else self.env_type + '_' + self.irrad_type
frame_render_dict = list(self.db.streaks_simulator.values())
f_start, f_end, f_step = self.frame_start, self.frame_end, self.frame_step
f_end = len(files) if f_end is None else min(f_end, len(files))
if self.frames:
# prone to go "boom", so we clip and remove 'wrong' ids
idx = np.unique(np.clip(self.frames, 0,
f_end - 1)).tolist()
else:
idx = list(range(f_start, f_end, f_step)) # to make it
f_num = len(idx)
sim_t0 = time.time()
print("{} images".format(len(idx)))
frames_exist_nb = 0
for f_idx, i in enumerate(idx):
image_file = files[i]
depth_file = depth_files[i]
# Compute frame index (independent of starting frame) to allow deterministic / reproducible rendering
if self.dataset == 'nuscenes':
# It could be useful for other dataset, but, for the moment, lets consign this little gem of a
# code to nuscenes
# If f_end was not supplied, we could see if the number of file is not equal to the number of
# simulated drop... if so, lets remap them
render_ix = np.linspace(0,
len(frame_render_dict),
len(files),
endpoint=False,
dtype=int)
f_name_idx = render_ix[i]
else:
f_name_idx = i
assert os.path.exists(
image_file), "Image file {} does not exist".format(
image_file)
assert os.path.exists(
depth_file), "Depth file {} does not exist".format(
depth_file)
# Ensure deterministic behavior
np.random.seed(f_name_idx)
frame_t0 = time.time()
frame = frame_render_dict[f_name_idx %
len(frame_render_dict)]
file_name = os.path.split(image_file)[-1]
out_rainy_path = os.path.join(
out_dir, 'rainy_image',
'{}.png'.format(file_name[:-4]))
out_rainy_mask_path = os.path.join(
out_dir, 'rain_mask', '{}.png'.format(file_name[:-4]))
out_env_path = os.path.join(
out_seq_dir, 'envmap', '{}.png'.format(file_name[:-4]))
frame_exists = os.path.exists(
out_rainy_path) or os.path.exists(out_rainy_mask_path)
if frame_exists:
if self.conflict_strategy == "skip":
frames_exist_nb += 1
continue
elif self.conflict_strategy == "overwrite":
pass
else:
raise NotImplementedError
# TODO adding functions of depth weighting for other dataset
if USE_DEPTH_WEIGHTING == 1 and calib_files is not None:
# Compute the drop depth map to allow weighting envmap from drop FOV
depth_drop_evaluator = DropDepthMap(
filename=calib_files[0])
# flush should happens after a while
if self.verbose:
sys.stdout.write('\r' + my_utils.process_eta_str(
process_t0, folder_idx, folders_num, folder_t0,
sim_idx, sim_num, sim_t0, f_idx, f_num, frame_t0) +
' ')
# two copies of bg because one is used for rain drop calculation
# and the other is changed after adding each drop
bg = cv2.imread(image_file) / 255.0
if self.settings["render_scale"] != 1:
bg = cv2.resize(
bg,
(int(bg.shape[1] // self.settings["render_scale"]),
int(bg.shape[0] //
self.settings["render_scale"])))
if FOG_ATT == 1:
# Depth map is used in weighting the luminance effect of the environment map on a single drop
if depth_file.endswith(".png"):
depth = cv2.imread(depth_file,
cv2.IMREAD_UNCHANGED)
if depth is None:
print('Missing/Corrupted depth data (%s)' %
depth_file)
continue
depth = depth.astype(np.float32) / 256.
elif depth_file.endswith(".npy"):
depth = np.load(depth_file)
else:
raise Exception("Invalid extension")
# Apply depth and render scale
depthHW = np.array([
int((depth.shape[0] * self.settings["depth_scale"])
// self.settings["render_scale"]),
int((depth.shape[1] * self.settings["depth_scale"])
// self.settings["render_scale"])
])
if not np.all(depth.shape[:2] == depthHW):
depth = cv2.resize(depth, (depthHW[1], depthHW[0]))
assert (np.all(depth.shape[:2] <= bg.shape[:2])
), "Depth cannot be larger than the image"
# Strategy to apply if RGB and Depth size mismatch
if not np.all(depth.shape[:2] == bg.shape[:2]):
# print("\nDepth {} size ({},{}) differs from image ({},{}). Will assume depth is crop centered.".format(image_file, depth.shape[0], depth.shape[1], bg.shape[0], bg.shape[1]))
bg = my_utils.crop_center(bg, depth.shape[0],
depth.shape[1])
else:
# In that case, no need for depth, but it's still used down in the code, for more less nothing
depth = np.zeros((bg.shape[1], bg.shape[0]), np.float)
rainy_bg = FOG.fog_rain_layer(bg, depth)
# rain layer is the image of the rendered rain drops blended with the background
rain_layer = np.zeros((bg.shape[0], bg.shape[1], 4),
np.float64)
# rainy_mask keeps track of the pixels of the background that have already been occupied by rain.
# This is used in cases of occluding or partially occluded drops
rainy_mask = np.zeros((bg.shape[0], bg.shape[1]),
np.float64)
rainy_saturation_mask = np.zeros(
(bg.shape[0], bg.shape[1], 3), np.float64)
# Environment map of the frame using (Christopher Cameron, 2005):
# http://www.cs.cmu.edu/afs/andrew/scs/cs/15-463/f05/pub/www/projects/fproj/cmcamero/report.pdf
if 'ours' in env_map_input:
# print('\nGenerating environment map')
self.BGR_env_map = map_generator.generate_map(rainy_bg)
elif 'pano' in env_map_input:
# print('\nLoading Environment Pano')
self.BGR_env_map = cv2.imread(
os.path.join('../data', 'panos',
file_name)) / 255.0
else:
raise NotImplementedError
self.env_map_xyY = my_utils.convert_rgb_to_xyY(
self.BGR_env_map[..., ::-1])
self.env_map_xyY[np.isnan(self.env_map_xyY)] = 0
self.solid_angle_map = solid_angle.get_solid_angles(
self.BGR_env_map)
# Render only streaks inside the frame
streak_dict = frame.streaks
streak_dict = {
keys: values
for keys, values in streak_dict.items()
if 1 <= values.max_width < max(imH, imW)
and 1 <= values.length < max(imH, imW) and (
(0 <= values.image_position_start[0] < imW
and 0 <= values.image_position_start[1] < imH) or
(0 <= values.image_position_end[0] < imW
and 0 <= values.image_position_end[1] < imH))
}
if USE_DEPTH_WEIGHTING == 1:
xyz_coord = depth_drop_evaluator.get_world_points(
depth)
assert streak_dict.__len__() <= 2 ** 16, \
"Assert that the number of drops doesn't overpass the uint16 rain_mask capacity"
streak_list = list(streak_dict.values())
drop_num = len(streak_list)
drop_process_t0 = time.time()
for drop_idx, drop_dict in enumerate(streak_list):
# Returns the rainy image, rainy_mask, drop, blended drop and the starting coord of
# the drop in image
rainy_bg, rainy_mask, rainy_saturation_mask, \
drop, blended_drop, minC = self.compute_drop(bg, drop_dict, rainy_bg,
rainy_mask, rainy_saturation_mask)
if blended_drop is not None:
rain_layer = self.renderer.make_rain_layer(
drop, blended_drop, rain_layer, rainy_mask,
minC)
else:
print(
"Trace: rain drop {} in sequence {} in image {} ({})"
.format(drop_idx, sequence, f_idx, f_name_idx))
# Compute progress
avg_drop_time = (time.time() -
drop_process_t0) / (drop_idx + 1)
if self.verbose or drop_idx == 0:
sys.stdout.write('\r' + my_utils.process_eta_str(
process_t0, folder_idx, folders_num, folder_t0,
sim_idx, sim_num, sim_t0, f_idx, f_num,
frame_t0, drop_idx, drop_num) + '\t\t' +
'%.1fms /drop' %
(1000. * avg_drop_time) +
' ')
# Create all output directories
os.makedirs(os.path.dirname(out_rainy_path), exist_ok=True)
os.makedirs(os.path.dirname(out_rainy_mask_path),
exist_ok=True)
if self.save_envmap:
os.makedirs(os.path.dirname(out_env_path),
exist_ok=True)
# mean contrast adjusted image
rainy_bg_mean = np.mean(rainy_bg)
bg_mean = np.mean(bg)
difference_mean = rainy_bg_mean - bg_mean
rainy_bg_copy = rainy_bg - difference_mean
plt.imsave(out_rainy_path,
np.clip(rainy_bg_copy[..., ::-1], 0, 1))
plt.imsave(out_rainy_mask_path, rainy_mask)
if self.save_envmap:
plt.imsave(out_env_path, self.BGR_env_map[..., ::-1])
if frames_exist_nb > 0:
print("Skipped {}/{} already existing renderings".format(
frames_exist_nb, len(idx)))
print("\n\nEnd of the simulation")